def undistort()

in preprocess.py [0:0]

• 54 lines of code
• 7 McCabe index (conditional complexity)

def undistort(args):

    input_folder = args.input
    output_folder = args.output
    undistortion_file = args.undistort_with_calibration_file
    if os.path.isdir(undistortion_file):
        undistortion_file = os.path.join(undistortion_file, "lens_distortion.json")

    import json

    with open(undistortion_file, "r") as undistortion_file:
        undistortion_parameters = json.load(undistortion_file)

    if os.path.normpath(input_folder) == os.path.normpath(output_folder):
        distorted_folder = os.path.join(output_folder, "distorted_images/")
        undistorted_folder = os.path.join(input_folder, "images/")
        # backup distorted images
        shutil.move(undistorted_folder, distorted_folder)
    else:
        distorted_folder = os.path.join(input_folder, "images/")
        undistorted_folder = os.path.join(output_folder, "images/")

    create_folder(undistorted_folder)

    mask_folder = undistorted_folder[:-1] + "_mask/"
    create_folder(mask_folder)

    # undistort images
    local_parallel_processes = 5
    distorted_images = [
        file for file in os.listdir(distorted_folder) if file[-4:] in [".png", ".jpg"]
    ]
    from multiprocessing import Pool

    with Pool(local_parallel_processes) as pool:
        pool.map(
            _undistort_image,
            [
                (
                    i,
                    distorted_images,
                    undistorted_folder,
                    distorted_folder,
                    undistortion_parameters,
                    mask_folder,
                )
                for i in range(len(distorted_images))
            ],
        )

    # store new intrinsics in file
    with open(
        os.path.join(output_folder, "undistorted_calibration.txt"), "w"
    ) as output_calibration:
        # see for indexing into matrix: https://docs.opencv.org/2.4/modules/imgproc/doc/geometric_transformations.html#undistort
        output_calibration.write(
            "color fx " + str(undistortion_parameters["newcameramtx"][0][0]) + "\n"
        )
        output_calibration.write(
            "color fy " + str(undistortion_parameters["newcameramtx"][1][1]) + "\n"
        )
        output_calibration.write(
            "color cx " + str(undistortion_parameters["newcameramtx"][0][2]) + "\n"
        )
        output_calibration.write(
            "color cy " + str(undistortion_parameters["newcameramtx"][1][2]) + "\n"
        )