factor out end_frame calculations to utils

cvat-ai · Dec 19, 2024 · d9f0d5f · d9f0d5f
1 parent e41da1a
commit d9f0d5f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/tests/python/rest_api/test_tasks.py b/tests/python/rest_api/test_tasks.py
@@ -73,6 +73,7 @@
     export_task_dataset,
     parse_frame_step,
     wait_until_task_is_created,
+    calc_end_frame,
 )
 
 
@@ -3111,7 +3112,7 @@ def _compute_annotation_segment_params(self, task_spec: _TaskSpec) -> list[tuple
         stop_frame = getattr(task_spec, "stop_frame", None) or (
             start_frame + (task_spec.size - 1) * frame_step
         )
-        end_frame = stop_frame - ((stop_frame - start_frame) % frame_step) + frame_step
+        end_frame = calc_end_frame(start_frame, stop_frame, frame_step)
 
         validation_params = getattr(task_spec, "validation_params", None)
         if validation_params and validation_params.mode.value == "gt_pool":
@@ -6410,10 +6411,7 @@ def get_img_index(zinfo: zipfile.ZipInfo) -> int:
             task_meta["stop_frame"],
             spec.frame_step,
         )
-        src_end_frame = (
-            src_stop_frame - ((src_stop_frame - src_start_frame) % src_frame_step) + src_frame_step
-        )  # taken from TestTaskData._compute_annotation_segment_params
-
+        src_end_frame = calc_end_frame(src_start_frame, src_stop_frame, src_frame_step)
         assert len(frames) == spec.size == task_meta["size"], "Some frames were lost"
         assert np.all(
             frames == np.arange(src_start_frame, src_end_frame, src_frame_step)

diff --git a/tests/python/rest_api/utils.py b/tests/python/rest_api/utils.py
@@ -601,3 +601,7 @@ def _exclude_cb(obj, path):
 
 def parse_frame_step(frame_filter: str) -> int:
     return int((frame_filter or "step=1").split("=")[1])
+
+
+def calc_end_frame(start_frame: int, stop_frame: int, frame_step: int) -> int:
+    return stop_frame - ((stop_frame - start_frame) % frame_step) + frame_step