66import base64
77import datetime
88import io
9- import random
9+ import math
1010import statistics
1111import uuid
1212from dataclasses import dataclass , field
1313from typing import TypedDict
1414from zoneinfo import ZoneInfo
1515
16+ import brisque
17+ import numpy as np
1618from PIL import Image , ImageStat
1719
1820from mock_vws ._constants import TargetStatuses
@@ -28,7 +30,6 @@ class TargetDict(TypedDict):
2830 image_base64 : str
2931 active_flag : bool
3032 processing_time_seconds : int | float
31- processed_tracking_rating : int
3233 application_metadata : str | None
3334 target_id : str
3435 last_modified_date : str
@@ -51,11 +52,24 @@ def _time_now() -> datetime.datetime:
5152 return datetime .datetime .now (tz = gmt )
5253
5354
54- def _random_tracking_rating ( ) -> int :
55+ def _quality ( image_content : bytes ) -> int :
5556 """
56- Return a random tracking rating.
57+ Args:
58+ image_content: The image content.
59+
60+ Returns:
61+ The quality of the image.
5762 """
58- return random .randint (0 , 5 )
63+ image_file = io .BytesIO (initial_bytes = image_content )
64+ image = Image .open (fp = image_file )
65+ image_array = np .asarray (a = image )
66+ obj = brisque .BRISQUE (url = False )
67+ # We avoid a barrage of warnings from the BRISQUE library.
68+ with np .errstate (divide = "ignore" , invalid = "ignore" ):
69+ score = obj .score (img = image_array )
70+ if math .isnan (score ):
71+ return 0
72+ return int (score / 20 )
5973
6074
6175@dataclass (frozen = True , eq = True )
@@ -75,9 +89,6 @@ class Target:
7589 delete_date : datetime .datetime | None = None
7690 last_modified_date : datetime .datetime = field (default_factory = _time_now )
7791 previous_month_recos : int = 0
78- processed_tracking_rating : int = field (
79- default_factory = _random_tracking_rating ,
80- )
8192 reco_rating : str = ""
8293 target_id : str = field (default_factory = _random_hex )
8394 total_recos : int = 0
@@ -158,7 +169,7 @@ def tracking_rating(self) -> int:
158169 return - 1
159170
160171 if self ._post_processing_status == TargetStatuses .SUCCESS :
161- return self .processed_tracking_rating
172+ return _quality ( image_content = self .image_value )
162173
163174 return 0
164175
@@ -219,7 +230,6 @@ def to_dict(self) -> TargetDict:
219230 "image_base64" : image_base64 ,
220231 "active_flag" : self .active_flag ,
221232 "processing_time_seconds" : self .processing_time_seconds ,
222- "processed_tracking_rating" : self .processed_tracking_rating ,
223233 "application_metadata" : self .application_metadata ,
224234 "target_id" : self .target_id ,
225235 "last_modified_date" : self .last_modified_date .isoformat (),
0 commit comments