Merge pull request #184 from fastlabel/feature/codec-restriction-h264

SDKで動画のタスク登録するときに、動画のコーデックがH.264のものではない場合はエラーを返す

Author: takahiro-tamenishi

README.md

@@ -767,6 +767,8 @@ task_id = client.create_video_task(
 ##### Limitation
 
 - You can upload up to a size of 250 MB.
+- You can upload only videos with H.264 encoding.
+- You can upload only MP4 file format.
 
 #### Find Task
 

fastlabel/__init__.py

@@ -1252,6 +1252,10 @@ def create_video_task(
             raise FastLabelInvalidException(
                 "Supported video size is under 250 MB.", 422
             )
+        if not utils.is_video_supported_codec(file_path):
+            raise FastLabelInvalidException(
+                "Supported video encoding for registration through the SDK is only H.264.", 422
+            )
 
         file = utils.base64_encode(file_path)
         payload = {"project": project, "name": name, "file": file}
@@ -1316,6 +1320,10 @@ def create_video_classification_task(
             raise FastLabelInvalidException(
                 "Supported video size is under 250 MB.", 422
             )
+        if not utils.is_video_supported_codec(file_path):
+            raise FastLabelInvalidException(
+                "Supported video encoding for registration through the SDK is only H.264.", 422
+            )
 
         file = utils.base64_encode(file_path)
         payload = {"project": project, "name": name, "file": file}

fastlabel/const.py

@@ -222,6 +222,9 @@
 # API can accept under 250 MB
 SUPPORTED_OBJECT_SIZE = 250 * math.pow(1024, 2)
 
+# Only 'avc1' and 'H264' are supported for video task creation.
+SUPPORTED_FOURCC = ["avc1"]
+
 
 SUPPORTED_INFERENCE_IMAGE_SIZE = 6 * math.pow(1024, 2)
 

fastlabel/utils.py

@@ -5,6 +5,7 @@
 
 import geojson
 import numpy as np
+import cv2
 
 from fastlabel import const
 
@@ -18,6 +19,10 @@ def is_image_supported_ext(file_path: str) -> bool:
     return file_path.lower().endswith((".png", ".jpg", ".jpeg"))
 
 
+def is_video_supported_codec(file_path: str) -> bool:
+    return get_video_fourcc(file_path) in const.SUPPORTED_FOURCC
+
+
 def is_video_supported_ext(file_path: str) -> bool:
     return file_path.lower().endswith(".mp4")
 
@@ -164,3 +169,11 @@ def is_clockwise(points: list) -> bool:
 def get_json_length(value) -> int:
     json_str = json.dumps(value)
     return len(json_str)
+
+
+def get_video_fourcc(video_path: str) -> str:
+    cap = cv2.VideoCapture(video_path)
+    fourcc_code = int(cap.get(cv2.CAP_PROP_FOURCC))
+    fourcc_str = "".join([chr((fourcc_code >> 8 * i) & 0xFF) for i in range(4)])
+    cap.release()
+    return fourcc_str