Add Play control, Adjust argument list

pillar_detection.py

@@ -1,3 +1,6 @@
+import io
+from contextlib import redirect_stderr
+
 import os
 from collections import Counter
 import cv2
@@ -7,8 +10,9 @@ import subprocess
 import argparse
 
 parser = argparse.ArgumentParser()
-parser.add_argument("-i", "--input_dir", default="./")
+parser.add_argument("-i", "--input_dir", default="./", help="Input vid directory")
-parser.add_argument("-s", "--scale", default="960:540")
+parser.add_argument("-o", "--output_dir", default="./out", help="Output vid directory")
+parser.add_argument("-d", "--debug", default=False, action="store_true", help="for debugging. short output")
 args = parser.parse_args()
 
 def resizeCropInfo(key, cropPos, frame_width, frame_height):
@@ -23,35 +27,35 @@ def resizeCropInfo(key, cropPos, frame_width, frame_height):
     
     w, h, x, y = cropPos
     # Left (Expand width)
-    if key == 0x250000:
+    if key == ord('j') or key == ord('J'):
         x -= 1
         w += 2
     # Right (Shrink width)
-    elif key == 0x270000:
+    elif key == ord('l') or key == ord('L'):
         x += 1
         w -= 2
     # Up (Expand height)
-    elif key == 0x260000:
+    elif key == ord('i') or key == ord('I'):
         y -= 1
         h += 2
     # Down (Shrink height)
-    elif key == 0x280000:
+    elif key == ord('k') or key == ord('K'):
         y += 1
         h -= 2
     # Num 4 (Move left)
-    elif key == 52:
+    elif key == ord('a') or key == ord('A'):
         x -= 1
         w -= 1
     # Num 6 (Move right)
-    elif key == 54:
+    elif key == ord('d') or key == ord('D'):
         w += 1
         x += 1
     # Num 8 (Move upward)
-    elif key == 56:
+    elif key == ord('w') or key == ord('W'):
         h -= 1
         y -= 1
     # Num 2 (Move downward)
-    elif key == 50:
+    elif key == ord('s') or key == ord('S'):
         h += 1
         y += 1
     _cropPos = w, h, x, y
@@ -64,33 +68,28 @@ def resizeCropInfo(key, cropPos, frame_width, frame_height):
 
     return cropPos
 
-# TODO:
-# cropinfo 로그로 쓸수 있게 txt로 저장
-
-
 def main():
-    vid_paths = sorted(glob.glob(os.path.join(args.input_dir, '*.avi')))
+    vid_types = ('*.avi', '*.mov', '*.mp4', '*.mxf')
-    _scale = args.scale.split(":")
+    vid_grapped = []
-    if len(_scale) != 2 :
+    for types in vid_types :
-        print("Invalid target scale : ", args.scale)
+        vid_grapped += sorted(glob.glob(os.path.join(args.input_dir, types)))
-        return
 
-    target_width, target_height = [int(x) for x in _scale]
+    # _scale = args.scale.split(":")
+    # if len(_scale) != 2 :
+    #     print("Invalid target scale : ", args.scale)
+    #     return
 
-    if not ((target_width > 0) and (target_height > 0)):
+    # target_width, target_height = [int(x) for x in _scale]
-        print("Invalid target scale : ", args.scale)
-        return
-    else :
-        print("Input dir :", args.input_dir)
-        print("Target scale: ", args.scale)
 
-    for fpath in vid_paths:
+    #--------------------------------------------------------
+    f = io.StringIO()
+    #--------------------------------------------------------
+
+    for fpath in vid_grapped:
         vid_name = os.path.basename(fpath)[:-4]
-        # Debugging
-        fpath = "01. In the sky - 블랙비트.avi"
 
         print("File to detect crop: ", fpath)
-        p = subprocess.Popen(["ffmpeg", "-ss", "30", "-i", fpath, "-vf", "cropdetect=limit=38",
+        p = subprocess.Popen(["ffmpeg", "-ss", "10", "-i", fpath, "-vf", "cropdetect=limit=38",
                              "-vframes", "3000", "-f", "null", "out.null"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
         infos = p.stderr.read().decode('utf-8')
         # print (infos)
@@ -103,30 +102,57 @@ def main():
         crop = mostCommonCrop[0][0]
 
         cropPos = crop.split(":")
+
         # Delete "crop=" in list cropPos
         cropPos[0] = cropPos[0][5:]
         cropPos = [int(x) for x in cropPos]
 
-        # Capture vid, Get width, height
+        with redirect_stderr(f):
+            # Capture vid, Get Vid info
             capture = cv2.VideoCapture(fpath)
             frame_width = capture.get(cv2.CAP_PROP_FRAME_WIDTH)
             frame_height = capture.get(cv2.CAP_PROP_FRAME_HEIGHT)
+            fps = capture.get(cv2.CAP_PROP_FPS)
+            frame_pos = 0
+            frame_count = 0
     
-        # Play speed : 33ms / frame
+            target_width, target_height = (frame_width, frame_height)
+            if not ((target_width > 0) and (target_height > 0)):
+                print("Invalid target scale : ", target_width, target_height)
+                return
+            else :
+                print("Input dir :", args.input_dir)
+                print("Target scale: ", target_width, target_height)
+
+            key = None
             while True:
-            key = cv2.waitKeyEx(33)
+                key = cv2.waitKeyEx(10)
                 if key > 0 :
-                # 27 == ESC
+                    # 13 == Enter key ASCII Code
-                if key == 27:
+                    if key == 13:
                         break
+                    # Pause
+                    elif key == ord('p') or key == ord('P'):
+                        key = cv2.waitKeyEx()
+                    # Replay
+                    elif key == ord('r') or key == ord('R'):
+                        capture = cv2.VideoCapture(fpath)
+                        continue
+                    # Move Forward (10s)
+                    elif key == ord('e') or key == ord('E'):
+                        capture.set(cv2.CAP_PROP_POS_FRAMES, frame_pos + (fps * 10))
+                    # Move Backward (10s)
+                    elif key == ord('q') or key == ord('Q'):
+                        capture.set(cv2.CAP_PROP_POS_FRAMES, frame_pos - (fps * 10))
+
 
                     cropPos = resizeCropInfo(key, cropPos, frame_width, frame_height)
 
                 frame_pos = capture.get(cv2.CAP_PROP_POS_FRAMES)
                 frame_count = capture.get(cv2.CAP_PROP_FRAME_COUNT)
 
-            if(frame_pos == frame_count):
+                if(frame_pos >= frame_count):
-                capture.open(vid)
+                    break
 
                 ret, frame = capture.read()
 
@@ -140,11 +166,12 @@ def main():
 
         # ffmpeg filter args
         crop_arg = f"crop={cropPos[0]}:{cropPos[1]}:{cropPos[2]}:{cropPos[3]},"
+        print(f"[CropInfo] (W : {cropPos[0]}), (H : {cropPos[1]}), (X : {cropPos[2]}), (Y : {cropPos[3]})")
 
-        # Debugging
+        target_ratio = target_width / target_height
-        print(f"[debug] crop={cropPos[0]}:{cropPos[1]}:{cropPos[2]}:{cropPos[3]}")
+        crop_ratio = cropPos[0] / cropPos[1]
 
-        if (target_height / target_height) >= (cropPos[0] / cropPos[1]):
+        if target_ratio < crop_ratio:
             tmp_height = target_width * (cropPos[1] / cropPos[0])
             scale_arg = f"scale=w={target_width}:h={tmp_height},"
         else:
@@ -154,17 +181,24 @@ def main():
         pad_arg = f"pad={target_width}:{target_height}:(ow-iw)/2:(ih-oh)/2,"
 
         print("Encoding \"%s\" ..." % vid_name)
-        if False:
+
-            p = subprocess.Popen(["ffmpeg", "-i", fpath, "-vf", crop_arg + scale_arg + pad_arg + "setsar=sar=1", "-y", "-pix_fmt",
+        if not os.path.exists(args.output_dir) :
-                                 "yuv420p", "-c:v", "libx264", "-crf", "15", "-preset", "medium", "-loglevel", "error", f"_{vid_name}_cropped.mp4"])
+            os.mkdir(args.output_dir)
-        else:
+            print(f'{args.output_dir} created')
+            
+        output_path = os.path.join(args.output_dir, f"{vid_name}_cropped.mp4")
+        if args.debug:
             # for Debug
             p = subprocess.Popen(["ffmpeg", "-t", "15", "-ss", "30", "-i", fpath, "-vf", crop_arg + scale_arg + pad_arg + "setsar=sar=1", "-y",
-                                 "-pix_fmt", "yuv420p", "-c:v", "libx264", "-crf", "15", "-preset", "medium", "-loglevel", "error", f"_{vid_name}_cropped.mp4"])
+                                 "-pix_fmt", "yuv420p", "-c:v", "libx264", "-crf", "15", "-preset", "medium", "-loglevel", "error", f"{output_path}"])
+        else:
+            p = subprocess.Popen(["ffmpeg", "-i", fpath, "-vf", crop_arg + scale_arg + pad_arg + "setsar=sar=1", "-y", "-pix_fmt",
+                                 "yuv420p", "-c:v", "libx264", "-crf", "15", "-preset", "medium", "-loglevel", "error", f"{output_path}"])
+
         print("Encoding Done.")
         
         # debugging
-        return
+        # return
 
 if __name__ == "__main__":
     main()