easy_dewarp_reference_frames.py

# Name:
#   easy_dewarp_reference_frames.py
#
# Author:
#   Hyuk Ko ([email protected])
#
# Copyright (C) 2019 Hyuk Ko. All rights reserved.
#
#
# 3DE4.script.name:    Easy Dewarp Reference Frames...
#
# 3DE4.script.version:    v1.0
#
# 3DE4.script.gui:    Main Window::Python
#
# 3DE4.script.comment:    Export Dewarped images with Warp4
#
# MODIFY THIS AT YOUR OWN RISK


import os
import sys
import math
import time
import string
import platform
import subprocess
import multiprocessing
import multiprocessing.pool
from functools import partial


def getLDmodelParameterList(model):
    l = []
    for p in range(tde4.getLDModelNoParameters(model)):
        l.append(tde4.getLDModelParameterName(model, p))
    return l


TDE4_PATH = tde4.get3DEInstallPath()
OS = platform.system()
if OS == "Windows":
    WARP4_PATH = r"{0}\bin\warp4.exe".format(TDE4_PATH)
    DIRECTORY_SEPARATOR = "\\"
else:
    WARP4_PATH = "{0}/bin/warp4".format(TDE4_PATH)
    DIRECTORY_SEPARATOR = "/"

## GUI
req = tde4.createCustomRequester()

tde4.addTextFieldWidget(req, 'overscan_value', 'Overscan Value', "1.0")


## Main
cameraList = tde4.getCameraList()
cameraListSize = len(cameraList)

ret = tde4.postCustomRequester(req, "Easy Dewarp Reference Frames...", 500, 0, "Ok", "Cancel")

if ret==1:

    cmd_list = []
    for i in range(0, cameraListSize):
        camera = cameraList[i]
        inFilePath = tde4.getCameraFrameFilepath(cameraList[i], 1)
        inFileBase, inFileExt = os.path.splitext(inFilePath)
        outFilePath = "{0}_dewarped{1}".format(inFileBase, inFileExt)

        w_px = tde4.getCameraImageWidth(camera)
        h_px = tde4.getCameraImageHeight(camera)
        os_w_px = str(int(int(float(tde4.getWidgetValue(req, "overscan_value")) * w_px) / 2.0)*2) # Always even number
        os_h_px = str(int(int(float(tde4.getWidgetValue(req, "overscan_value")) * h_px) / 2.0)*2)# Always even number

        lens = tde4.getCameraLens(camera)
        model = tde4.getLensLDModel(lens)
        focal = tde4.getLensFocalLength(lens)
        focus = tde4.getLensFocus(lens)
        w_fb_cm = tde4.getLensFBackWidth(lens)
        h_fb_cm = tde4.getLensFBackHeight(lens)
        lco_x_cm = tde4.getLensLensCenterX(lens)
        lco_y_cm = tde4.getLensLensCenterY(lens)
        pxa = tde4.getLensPixelAspect(lens)

        parameters = ""
        for para in getLDmodelParameterList(model):
            parameters += str(tde4.getLensLDAdjustableParameter(lens, para, focal, focus)) + " "

        cmd = r'{0} -in {1} -out {2} -action remove_distortion -model "{3}" -parameters {4} -pixel_aspect {5} -lco {6} {7} -filmback {8} {9} -focal_length {10} -overscan {11} {12}'.format(WARP4_PATH, inFilePath, outFilePath, model, parameters, pxa, lco_x_cm, lco_y_cm, w_fb_cm, h_fb_cm, focal, os_w_px, os_h_px)
        cmd_list.append(cmd)

    p = multiprocessing.pool.ThreadPool(multiprocessing.cpu_count()-1) # Used ThreadPool instead of Pool. OS compat reasons...
    start = time.time()
    tde4.postProgressRequesterAndContinue("Easy Dewarp Reference Frames...", "Dewarping Frame 1 out of {0}...".format(len(cmd_list)), len(cmd_list), "Stop")
    for i, returncode in enumerate(p.imap(partial(subprocess.call, shell=True), cmd_list)):
        tde4.updateProgressRequester(i+1,"Dewarping Frame {0} out of {1}...".format(i+1, len(cmd_list)))
        if returncode != 0:
            print("%d command failed: %d" % (i, returncode))
    tde4.unpostProgressRequester()
    print 'Time elapsed: %s' % (time.time() - start)