Server tests for context images for 2D tasks

cvat/apps/dataset_manager/tests/test_rest_api_formats.py

@@ -4,11 +4,14 @@
 
 import copy
 import json
+import os
 import os.path as osp
 import random
 import xml.etree.ElementTree as ET
 import zipfile
 from io import BytesIO
+import shutil
+import tempfile
 
 from datumaro.components.dataset import Dataset
 from datumaro.util.test_utils import TestDir, compare_datasets
@@ -103,6 +106,11 @@ def _create_task(self, data, image_data):
 
         return task
 
+    def _get_request(self, path, user):
+        with ForceLogin(user, self.client):
+            response = self.client.get(path)
+        return response
+
     def _get_request_with_data(self, path, data, user):
         with ForceLogin(user, self.client):
             response = self.client.get(path, data)
@@ -263,3 +271,57 @@ def test_api_v1_check_widerface_with_all_attributes(self):
                     extractor = CvatTaskDataExtractor(task_data, include_images=include_images)
                     data_from_task_after_upload = Dataset.from_extractors(extractor)
                     compare_datasets(self, data_from_task_before_upload, data_from_task_after_upload)
+
+
+class TaskAnnotation2DContext(_DbTestBase):
+    def create_zip_archive_with_related_images(self, file_name, test_dir, context_images_info):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            for img in context_images_info:
+                image = Image.new('RGB', size=(100, 50))
+                image.save(osp.join(tmp_dir, img), 'png')
+                if context_images_info[img]:
+                    related_path = osp.join(tmp_dir, "related_images", img.replace(".", "_"))
+                    os.makedirs(related_path)
+                    image.save(osp.join(related_path, f"related_{img}"), 'png')
+
+            zip_file_path = osp.join(test_dir, file_name)
+            shutil.make_archive(zip_file_path, 'zip', tmp_dir)
+        return f"{zip_file_path}.zip"
+
+    def test_check_flag_has_related_context(self):
+        with TestDir() as test_dir:
+            test_cases = {
+                "All images with context": {"image_1.png": True, "image_2.png": True},
+                "One image with context": {"image_1.png": True, "image_2.png": False}
+            }
+            for test_case, context_img_data in test_cases.items():
+                filename = self.create_zip_archive_with_related_images(test_case, test_dir, context_img_data)
+                img_data = {
+                    "client_files[0]": open(filename, 'rb'),
+                    "image_quality": 75,
+                }
+                task = self._create_task(tasks["main"], img_data)
+                task_id = task["id"]
+
+                response = self._get_request("/api/v1/tasks/%s/data/meta" % task_id, self.admin)
+                for frame in response.data["frames"]:
+                    self.assertEqual(context_img_data[frame["name"]], frame["has_related_context"])
+
+    def test_fetch_related_image_from_server(self):
+        test_name = self._testMethodName
+        context_img_data ={"image_1.png": True}
+        with TestDir() as test_dir:
+            filename = self.create_zip_archive_with_related_images(test_name, test_dir, context_img_data)
+            img_data = {
+                "client_files[0]": open(filename, 'rb'),
+                "image_quality": 75,
+            }
+            task = self._create_task(tasks["main"], img_data)
+            task_id = task["id"]
+            data = {
+                "quality": "original",
+                "type": "context_image",
+                "number": 0
+            }
+            response = self._get_request_with_data("/api/v1/tasks/%s/data" % task_id, data, self.admin)
+            self.assertEqual(response.status_code, status.HTTP_200_OK)