Merge pull request #1512 from VWS-Python/test-image-quality

Test image quality

Author: adamtheturtle

pyproject.toml

@@ -141,6 +141,7 @@ warn_untyped_fields = true
 
 module = [
     "brisque",
+    "cv2",
     "docker",
     "docker.errors",
     "docker.models.networks",
@@ -252,6 +253,7 @@ requires-python = ">=3.10"
 dependencies = [
     "Pillow",
     "brisque",
+    "opencv-python",
     "flask",
     "multipart",
     "numpy",

src/mock_vws/target.py

@@ -14,6 +14,7 @@
 from zoneinfo import ZoneInfo
 
 import brisque
+import cv2
 import numpy as np
 from PIL import Image, ImageStat
 
@@ -54,6 +55,12 @@ def _time_now() -> datetime.datetime:
 
 def _quality(image_content: bytes) -> int:
     """
+    Get a quality score for an image.
+
+    This is a rough approximation of the quality score used by Vuforia, but
+    is not accurate. For example, our "corrupted_image" fixture is rated as -2
+    by Vuforia, but is rated as 0 by this function.
+
     Args:
         image_content: The image content.
 
@@ -66,7 +73,10 @@ def _quality(image_content: bytes) -> int:
     obj = brisque.BRISQUE(url=False)
     # We avoid a barrage of warnings from the BRISQUE library.
     with np.errstate(divide="ignore", invalid="ignore"):
-        score = obj.score(img=image_array)
+        try:
+            score = obj.score(img=image_array)
+        except cv2.error:  # pylint: disable=no-member
+            return 0
     if math.isnan(score):
         return 0
     return int(score / 20)
@@ -168,10 +178,7 @@ def tracking_rating(self) -> int:
         if time_since_upload <= pre_rating_time:
             return -1
 
-        if self._post_processing_status == TargetStatuses.SUCCESS:
-            return _quality(image_content=self.image_value)
-
-        return 0
+        return _quality(image_content=self.image_value)
 
     @classmethod
     def from_dict(cls, target_dict: TargetDict) -> Target:

tests/mock_vws/test_get_target.py

@@ -119,6 +119,62 @@ def test_success_status(
         assert new_tracking_rating == tracking_rating
 
 
+def _get_target_tracking_rating(
+    vws_client: VWS,
+    image_file: io.BytesIO,
+) -> int:
+    """
+    Get the tracking rating of a target.
+    """
+    target_id = vws_client.add_target(
+        name=f"example_{uuid.uuid4().hex}",
+        width=1,
+        image=image_file,
+        active_flag=True,
+        application_metadata=None,
+    )
+
+    vws_client.wait_for_target_processed(target_id=target_id)
+    target_details = vws_client.get_target_record(target_id=target_id)
+    return target_details.target_record.tracking_rating
+
+
+@pytest.mark.usefixtures("verify_mock_vuforia")
+class TestTargetTrackingRating:
+    """
+    Tests which exercise the target tracking_rating, and check the image
+    fixtures we use.
+    """
+
+    @staticmethod
+    def test_target_quality(
+        vws_client: VWS,
+        high_quality_image: io.BytesIO,
+        image_file_success_state_low_rating: io.BytesIO,
+        corrupted_image_file: io.BytesIO,
+    ) -> None:
+        """
+        The target tracking rating is as expected.
+        """
+        high_quality_image_tracking_rating = _get_target_tracking_rating(
+            vws_client=vws_client,
+            image_file=high_quality_image,
+        )
+        low_quality_image_tracking_rating = _get_target_tracking_rating(
+            vws_client=vws_client,
+            image_file=image_file_success_state_low_rating,
+        )
+        corrupted_image_file_tracking_rating = _get_target_tracking_rating(
+            vws_client=vws_client,
+            image_file=corrupted_image_file,
+        )
+        assert (
+            high_quality_image_tracking_rating
+            > low_quality_image_tracking_rating
+            >= corrupted_image_file_tracking_rating
+        )
+
+
 @pytest.mark.usefixtures("verify_mock_vuforia")
 class TestInactiveProject:
     """