Added linearization code

lantern_scripts/extract_linearized.py

@@ -0,0 +1,179 @@
+import argparse
+import json
+import os
+
+import cv2
+import numpy as np
+
+# fonction that apply the correction on the colors on the input images
+
+
+def apply_correction(coefs, img, CORRECTION_CURVE_TYPE) -> np.array:
+    B, G, R = coefs
+    original_shape = img.shape
+    img = img.astype(np.float32)
+    img = img.reshape(-1, 3)
+    if CORRECTION_CURVE_TYPE == "deg2":
+        img[:, 0] = B[0] + img[:, 0] * B[1] + img[:, 0] ** 2 * B[2]
+        img[:, 1] = G[0] + img[:, 1] * G[1] + img[:, 1] ** 2 * G[2]
+        img[:, 2] = R[0] + img[:, 2] * R[1] + img[:, 2] ** 2 * R[2]
+    if CORRECTION_CURVE_TYPE == "deg3":
+        img[:, 0] = B[0] + img[:, 0] * B[1] + \
+            img[:, 0] ** 2 * B[2] + img[:, 0] ** 3 * B[3]
+        img[:, 1] = G[0] + img[:, 1] * G[1] + \
+            img[:, 1] ** 2 * G[2] + img[:, 0] ** 3 * G[3]
+        img[:, 2] = R[0] + img[:, 2] * R[1] + \
+            img[:, 2] ** 2 * R[2] + img[:, 0] ** 3 * R[3]
+    elif CORRECTION_CURVE_TYPE == "deg1":
+        img[:, 0] = B[0] + img[:, 0] * B[1]
+        img[:, 1] = G[0] + img[:, 1] * G[1]
+        img[:, 2] = R[0] + img[:, 2] * R[1]
+    elif CORRECTION_CURVE_TYPE == "deg2_no_offset":
+        img[:, 0] = img[:, 0] * B[0] + img[:, 0] ** 2 * B[1]
+        img[:, 1] = img[:, 1] * G[0] + img[:, 1] ** 2 * G[1]
+        img[:, 2] = img[:, 2] * R[0] + img[:, 2] ** 2 * R[1]
+    elif CORRECTION_CURVE_TYPE == "deg3_no_offset":
+        img[:, 0] = img[:, 0] * B[0] + \
+            img[:, 0] ** 2 * B[1] + img[:, 0] ** 3 * B[2]
+        img[:, 1] = img[:, 1] * G[0] + \
+            img[:, 1] ** 2 * G[1] + img[:, 1] ** 3 * G[2]
+        img[:, 2] = img[:, 2] * R[0] + \
+            img[:, 2] ** 2 * R[1] + img[:, 2] ** 3 * R[2]
+    elif CORRECTION_CURVE_TYPE == "deg1_no_offset":
+        img[:, 0] = img[:, 0] * B[0]
+        img[:, 1] = img[:, 1] * G[0]
+        img[:, 2] = img[:, 2] * R[0]
+    img = img.clip(0, 255)
+    img = img.reshape(original_shape)
+    return img
+
+
+def extract_frames(pathToVideo, fps, resize=False, width=960, height=480, jpg=False, frames_name=None, out_path=None):
+
+    if out_path is None:
+        frames_path = "frames"
+    else:
+        frames_path = os.path.join(out_path, "frames")
+
+    if not os.path.exists(frames_path):
+        os.mkdir(frames_path)
+
+    # extract frames from the path video
+    if frames_name is None:
+        os.system(
+            '''ffmpeg -i "{}" -s {}x{} -vf fps={} {}/out%d.png'''.format(pathToVideo, width, height, str(fps), frames_path))
+    else:
+        os.system(
+            '''ffmpeg -i "{}" -s {}x{} -vf fps={} {}/{}%d.png'''.format(pathToVideo, width, height, str(fps), frames_path, frames_name))
+
+    if resize:
+        imgs = os.listdir(frames_path)
+        for file in imgs:
+            path = os.path.join(frames_path, file)
+            frame = cv2.imread(path)
+            frame = cv2.resize(frame, (width, height))
+            cv2.imwrite(path, frame)
+
+    if jpg:
+        imgs = os.listdir(frames_path)
+        for file in imgs:
+            path = os.path.join(frames_path, file)
+            frame = cv2.imread(path)
+            if path.find(".png") != -1:
+                cv2.imwrite(path.replace(".png", ".jpg"), frame,
+                            [int(cv2.IMWRITE_JPEG_QUALITY), 100])
+                os.remove(path)
+            elif path.find(".JPG") != -1:
+                cv2.imwrite(path.replace(".JPG", ".jpg"), frame,
+                            [int(cv2.IMWRITE_JPEG_QUALITY), 100])
+                os.remove(path)
+
+
+def linearize_frames(out_path=None):
+    # inearize all the images in the frame folder
+    # parameter of the color curve. stored in a dict in a json file resulting from the macduff_modified.py code
+    with open('data.json') as f:
+        data = json.load(f)
+    CORRECTION_CURVE_TYPE = data["type"]
+
+    if out_path is None:
+        frames_path = "frames"
+        lin_path = "linearized"
+    else:
+        frames_path = os.path.join(out_path, "frames")
+        lin_path = os.path.join(out_path, "linearized")
+
+    if not os.path.exists(lin_path):
+        os.mkdir(lin_path)
+
+    # calling the apply correction fonction for as long as there is images in the frames folder
+
+    dirs = os.listdir(frames_path)
+    for file in dirs:
+        im = cv2.imread(os.path.join(frames_path, file))
+        im_corrected = apply_correction(
+            (data["B"], data["G"], data["R"]), im, CORRECTION_CURVE_TYPE)
+        cv2.imwrite(lin_path+"/out_lin"+file, im_corrected)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description='Extract and linearize images')
+    parser.add_argument('-e', '--extract', action='store_true',
+                        help='if you want to extract frame of video. If so, --path and --fps are required')
+    parser.add_argument('-l', '--linearize', action='store_true',
+                        help='if you want to linearize the images in the "frames" folder')
+    parser.add_argument('-p', '--path', type=str, metavar='',
+                        help="path to the video you want to extract frame from")
+    parser.add_argument('-o', '--out_path', type=str, metavar='',
+                        help="path to the output frames")
+    parser.add_argument('-f', '--fps', type=float, metavar='',
+                        help="number of frame per second you want to extract from the video")
+    parser.add_argument('-r', '--resize', action='store_true',
+                        help="resize the extracted frames")
+    parser.add_argument('--height', type=int, metavar='',
+                        help="height of extracted frames")
+    parser.add_argument('-w', '--width', type=int, metavar='',
+                        help="width of extracted frames")
+    parser.add_argument('--jpg', action='store_true',
+                        help='save the frames in jpg format')
+    args = parser.parse_args()
+
+    # ============== START MAIN =============
+    #
+    print("=====================")
+    print("        Start        ")
+    print("=====================")
+
+    # check up parameters
+    if args.extract:
+        if args.path == None:
+            print("--path is empty, don't forget to input a path")
+            exit()
+        if args.fps == None:
+            print("--fps is empty, don't forget to specify the fps")
+            exit()
+        if args.fps == 0:
+            print("--fps is 0, it has to be a float higher than 0")
+            exit()
+    if args.extract == False and args.linearize == False:
+        print("It did nothing because you didnt ask for neither extract (-e) nor linearize (-l)")
+
+    # ================================================================================================
+    # =======================variable to modify - Path to the video===================================
+    # ================================================================================================
+    pathToVideo = args.path
+    # ================================================================================================
+    # ================================================================================================
+    # ================================================================================================
+
+    if args.extract:
+        extract_frames(pathToVideo, args.fps, args.resize,
+                       args.width, args.height, args.jpg, None, args.out_path)
+
+    if args.linearize:
+        linearize_frames(args.out_path)
+
+    print("=====================")
+    print("      The end        ")
+    print("=====================")

lantern_scripts/z10-color-correct-data.json

@@ -0,0 +1 @@
+{"B": [46.82997512817383, 0.11160952597856522, 0.002294188365340233], "G": [41.67820739746094, 0.09581712633371353, 0.0028952741995453835], "R": [37.4744873046875, 0.22809027135372162, 0.002323441905900836], "type": "deg2"}