[GSoC2024] feature: allow downloading video frames in custom extension

changelog.d/20240410_200046_harsh_k_feature_custom_extension.md

@@ -0,0 +1,4 @@
+### Changed
+
+- Job and task `download_frames` now accepts custom extension for images
+  (<https://github.com/cvat-ai/cvat/pull/7697>)

cvat-sdk/cvat_sdk/core/proxies/jobs.py

@@ -119,6 +119,7 @@ def download_frames(
         self,
         frame_ids: Sequence[int],
         *,
+        image_extension: Optional[str] = None,
         outdir: StrPath = ".",
         quality: str = "original",
         filename_pattern: str = "frame_{frame_id:06d}{frame_ext}",
@@ -135,14 +136,17 @@ def download_frames(
             frame_bytes = self.get_frame(frame_id, quality=quality)
 
             im = Image.open(frame_bytes)
-            mime_type = im.get_format_mimetype() or "image/jpg"
-            im_ext = mimetypes.guess_extension(mime_type)
-
-            # FIXME It is better to use meta information from the server
-            # to determine the extension
-            # replace '.jpe' or '.jpeg' with a more used '.jpg'
-            if im_ext in (".jpe", ".jpeg", None):
-                im_ext = ".jpg"
+            if image_extension is None:
+                mime_type = im.get_format_mimetype() or "image/jpg"
+                im_ext = mimetypes.guess_extension(mime_type)
+
+                # FIXME It is better to use meta information from the server
+                # to determine the extension
+                # replace '.jpe' or '.jpeg' with a more used '.jpg'
+                if im_ext in (".jpe", ".jpeg", None):
+                    im_ext = ".jpg"
+            else:
+                im_ext = f".{image_extension.strip('.')}"
 
             outfile = filename_pattern.format(frame_id=frame_id, frame_ext=im_ext)
             im.save(outdir / outfile)

cvat-sdk/cvat_sdk/core/proxies/tasks.py

@@ -220,30 +220,33 @@ def download_frames(
         self,
         frame_ids: Sequence[int],
         *,
+        image_extension: Optional[str] = None,
         outdir: StrPath = ".",
         quality: str = "original",
         filename_pattern: str = "frame_{frame_id:06d}{frame_ext}",
     ) -> Optional[List[Image.Image]]:
         """
         Download the requested frame numbers for a task and save images as outdir/filename_pattern
         """
-        # TODO: add arg descriptions in schema
 
         outdir = Path(outdir)
-        outdir.mkdir(exist_ok=True)
+        outdir.mkdir(parents=True, exist_ok=True)
 
         for frame_id in frame_ids:
             frame_bytes = self.get_frame(frame_id, quality=quality)
 
             im = Image.open(frame_bytes)
-            mime_type = im.get_format_mimetype() or "image/jpg"
-            im_ext = mimetypes.guess_extension(mime_type)
-
-            # FIXME It is better to use meta information from the server
-            # to determine the extension
-            # replace '.jpe' or '.jpeg' with a more used '.jpg'
-            if im_ext in (".jpe", ".jpeg", None):
-                im_ext = ".jpg"
+            if image_extension is None:
+                mime_type = im.get_format_mimetype() or "image/jpg"
+                im_ext = mimetypes.guess_extension(mime_type)
+
+                # FIXME It is better to use meta information from the server
+                # to determine the extension
+                # replace '.jpe' or '.jpeg' with a more used '.jpg'
+                if im_ext in (".jpe", ".jpeg", None):
+                    im_ext = ".jpg"
+            else:
+                im_ext = f".{image_extension.strip('.')}"
 
             outfile = filename_pattern.format(frame_id=frame_id, frame_ext=im_ext)
             im.save(outdir / outfile)

tests/python/rest_api/test_jobs.py

@@ -576,7 +576,6 @@ def test_get_gt_job_in_org_task(
     "restore_db_per_class"
 )
 class TestGetGtJobData:
-
     def _delete_gt_job(self, user, gt_job_id):
         with make_api_client(user) as api_client:
             api_client.jobs_api.destroy(gt_job_id)

tests/python/sdk/test_jobs.py

@@ -141,18 +141,23 @@ def test_can_download_frame(self, fxt_new_task: Task, quality: str):
         assert self.stdout.getvalue() == ""
 
     @pytest.mark.parametrize("quality", ("compressed", "original"))
-    def test_can_download_frames(self, fxt_new_task: Task, quality: str):
+    @pytest.mark.parametrize("image_extension", (None, "bmp"))
+    def test_can_download_frames(self, fxt_new_task: Task, quality: str, image_extension: str):
         fxt_new_task.get_jobs()[0].download_frames(
             [0],
+            image_extension=image_extension,
             quality=quality,
             outdir=self.tmp_path,
             filename_pattern="frame-{frame_id}{frame_ext}",
         )
 
-        if quality == "original":
-            expected_frame_ext = "png"
+        if image_extension is not None:
+            expected_frame_ext = image_extension
         else:
-            expected_frame_ext = "jpg"
+            if quality == "original":
+                expected_frame_ext = "png"
+            else:
+                expected_frame_ext = "jpg"
 
         assert (self.tmp_path / f"frame-0.{expected_frame_ext}").is_file()
         assert self.stdout.getvalue() == ""

tests/python/sdk/test_tasks.py

@@ -322,18 +322,23 @@ def test_can_download_frame(self, fxt_new_task: Task, quality: str):
         assert self.stdout.getvalue() == ""
 
     @pytest.mark.parametrize("quality", ("compressed", "original"))
-    def test_can_download_frames(self, fxt_new_task: Task, quality: str):
+    @pytest.mark.parametrize("image_extension", (None, "bmp"))
+    def test_can_download_frames(self, fxt_new_task: Task, quality: str, image_extension: str):
         fxt_new_task.download_frames(
             [0],
+            image_extension=image_extension,
             quality=quality,
             outdir=self.tmp_path,
             filename_pattern="frame-{frame_id}{frame_ext}",
         )
 
-        if quality == "original":
-            expected_frame_ext = "png"
+        if image_extension is not None:
+            expected_frame_ext = image_extension
         else:
-            expected_frame_ext = "jpg"
+            if quality == "original":
+                expected_frame_ext = "png"
+            else:
+                expected_frame_ext = "jpg"
 
         assert (self.tmp_path / f"frame-0.{expected_frame_ext}").is_file()
         assert self.stdout.getvalue() == ""