From 8086250d73a872b6e6dd01741a1ead87ccea7a78 Mon Sep 17 00:00:00 2001 From: gadget-man Date: Sat, 10 Feb 2024 16:05:52 +0000 Subject: [PATCH] Watched plates (#32) * clean up save image * fix frigate_url error * fix builds * Temp reduce debugging. * Enhancements to save_image Save full snapshot Get box co-ordinates from Events API call Move save_image routine to end of process to avoid delays updating mqtt whilst image processed. * Updated for comments * Update test.py * Test for status code on requests.get * watched_plates * watched_plates * watched_plates * watched_plates * watched_plates * check if first event before skipping due to top_score * take fuzzy_score into account when reviewing min_score * take fuzzy_score into account when reviewing min_score * watched plates debug best match and score * update frigate+ best plate score check * Additional duplication debugging * update frigate+ best plate score check * update frigate+ best plate score check * update frigate+ best plate score check * update frigate+ best plate score check * update frigate+ best plate score check * update frigate+ best plate score check * update frigate+ best plate score check * update frigate+ best plate score check * watched plates debug best match and score * Update index.py * Update index.py * Update index.py * only save image with not plate at end of event * only save image with not plate at end of event * Update index.py * add max_attempts to limit AI calls when Frigate_plus is true * add max_attempts to limit AI calls when Frigate_plus is true * starting to tidy up * Update DATETIME_FORMAT to avoid saving multiple versions of snapshot * update to match plates based on lower() in CP.AI * Tidy up watched_plates for PR * Test updates * Test updates * fix tests --------- Co-authored-by: Leonardo Merza --- .gitignore | 1 + README.md | 20 +++- index.py | 312 ++++++++++++++++++++++++++++++++++++++--------------- test.py | 120 ++++++++++++--------- 4 files changed, 311 insertions(+), 142 deletions(-) diff --git a/.gitignore b/.gitignore index c5999a8..9778724 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ __pycache__ venv +myenv config/* plates/* diff --git a/README.md b/README.md index b21c9d7..154e803 100644 --- a/README.md +++ b/README.md @@ -61,6 +61,7 @@ frigate: # ... frigate_plus: true license_plate_min_score: 0 # default is show all but can speficify a min score from 0 - 1 for example 0.8 + max_attempts: 20 # Optional: if set, will limit the number of snapshots sent for recognition for any particular event. ``` If you're using CodeProject.AI, you'll need to comment out plate_recognizer in your config. Then add and update "api_url" with your CodeProject.AI Service API URL. Your config should look like: @@ -106,10 +107,6 @@ logger_level: DEBUG Logs will be in `/config/frigate_plate_recognizer.log` -### Synology Diskstation - -Anyone trying this on Synology Diskstation, you need to set the volumes to point to `/usr/src/app/config` not just `/config` - ### Save Snapshot Images to Path If you want frigate-plate-recognizer to automatically save snapshots of recognized plates, add the following to your config.yml: @@ -136,3 +133,18 @@ services: - TZ=America/New_York ``` +### Monitor Watched Plates + +If you want frigate-plate-recognizer to check recognized plates against a list of watched plates for close matches (including fuzzy recognition), add the following to your config.yml: + +```yml +frigate: + watched_plates: #list of plates to watch. + - ABC123 + - DEF456 + fuzzy_match: 0.8 # default is test against plate-recognizer / CP.AI 'candidates' only, but can specify a min score for fuzzy matching if no candidates match watched plates from 0 - 1 for example 0.8 +``` + +If a watched plate is found in the list of candidates plates returned by plate-recognizer / CP.AI, the response will be updated to use that plate and it's score. The original plate will be added to the MQTT response as an additional `original_plate` field. + +If no candidates match and fuzzy_match is enabled with a value, the recognized plate is compared against each of the watched_plates using fuzzy matching. If a plate is found with a score > fuzzy_match, the response will be updated with that plate. The original plate and the associated fuzzy_score will be added to the MQTT response as additional fields `original_plate` and `fuzzy_score`. \ No newline at end of file diff --git a/index.py b/index.py index 442031f..fd5ec4c 100644 --- a/index.py +++ b/index.py @@ -15,23 +15,25 @@ import io from PIL import Image, ImageDraw, UnidentifiedImageError, ImageFont +import difflib mqtt_client = None config = None first_message = True _LOGGER = None -VERSION = '1.8.8' +VERSION = '1.8.10' -CONFIG_PATH = './config/config.yml' -DB_PATH = './config/frigate_plate_recogizer.db' -LOG_FILE = './config/frigate_plate_recogizer.log' +CONFIG_PATH = '/config/config.yml' +DB_PATH = '/config/frigate_plate_recogizer.db' +LOG_FILE = '/config/frigate_plate_recogizer.log' SNAPSHOT_PATH = '/plates' -DATETIME_FORMAT = "%Y-%m-%d_%H-%M-%S" +DATETIME_FORMAT = "%Y-%m-%d_%H-%M" PLATE_RECOGIZER_BASE_URL = 'https://api.platerecognizer.com/v1/plate-reader' DEFAULT_OBJECTS = ['car', 'motorcycle', 'bus'] +CURRENT_EVENTS = {} def on_connect(mqtt_client, userdata, flags, rc): @@ -60,17 +62,18 @@ def set_sublabel(frigate_url, frigate_event_id, sublabel, score): if len(sublabel) > 20: sublabel = sublabel[:20] + sublabel = str(sublabel).upper() # plates are always upper cased + # Submit the POST request with the JSON payload payload = { "subLabel": sublabel } headers = { "Content-Type": "application/json" } response = requests.post(post_url, data=json.dumps(payload), headers=headers) - - percentscore = "{:.1%}".format(score) + percent_score = "{:.1%}".format(score) # Check for a successful response if response.status_code == 200: - _LOGGER.info(f"Sublabel set successfully to: {sublabel} with {percentscore} confidence") + _LOGGER.info(f"Sublabel set successfully to: {sublabel} with {percent_score} confidence") else: _LOGGER.error(f"Failed to set sublabel. Status code: {response.status_code}") @@ -86,16 +89,22 @@ def code_project(image): if response.get('predictions') is None: _LOGGER.error(f"Failed to get plate number. Response: {response}") - return None, None + return None, None, None, None if len(response['predictions']) == 0: _LOGGER.debug(f"No plates found") - return None, None + return None, None, None, None plate_number = response['predictions'][0].get('plate') score = response['predictions'][0].get('confidence') - - return plate_number, score + + watched_plate, watched_score, fuzzy_score = check_watched_plates(plate_number, response['predictions']) + if fuzzy_score: + return plate_number, score, watched_plate, fuzzy_score + elif watched_plate: + return plate_number, watched_score, watched_plate, None + else: + return plate_number, score, None, None def plate_recognizer(image): api_url = config['plate_recognizer'].get('api_url') or PLATE_RECOGIZER_BASE_URL @@ -113,28 +122,101 @@ def plate_recognizer(image): if response.get('results') is None: _LOGGER.error(f"Failed to get plate number. Response: {response}") - return None, None + return None, None, None, None if len(response['results']) == 0: _LOGGER.debug(f"No plates found") - return None, None + return None, None, None, None plate_number = response['results'][0].get('plate') score = response['results'][0].get('score') + + watched_plate, watched_score, fuzzy_score = check_watched_plates(plate_number, response['results'][0].get('candidates')) + if fuzzy_score: + return plate_number, score, watched_plate, fuzzy_score + elif watched_plate: + return plate_number, watched_score, watched_plate, None + else: + return plate_number, score, None, None - return plate_number, score - -def send_mqtt_message(plate_number, plate_score, frigate_event_id, after_data, formatted_start_time): +def check_watched_plates(plate_number, response): + config_watched_plates = config['frigate'].get('watched_plates', []) + if not config_watched_plates: + _LOGGER.debug("Skipping checking Watched Plates because watched_plates is not set") + return None, None + + config_watched_plates = [str(x).lower() for x in config_watched_plates] #make sure watched_plates are all lower case + + #Step 1 - test if top plate is a watched plate + matching_plate = str(plate_number).lower() in config_watched_plates + if matching_plate: + _LOGGER.info(f"Recognised plate is a Watched Plate: {plate_number}") + return None, None, None + + #Step 2 - test against AI candidates: + for i, plate in enumerate(response): + matching_plate = plate.get('plate') in config_watched_plates + if matching_plate: + if config.get('plate_recognizer'): + score = plate.get('score') + else: + if i == 0: continue #skip first response for CodeProjet.AI as index 0 = original plate. + score = plate.get('confidence') + _LOGGER.info(f"Watched plate found from AI candidates: {plate.get('plate')} with score {score}") + return plate.get('plate'), score, None + + _LOGGER.debug("No Watched Plates found from AI candidates") + + #Step 3 - test against fuzzy match: + fuzzy_match = config['frigate'].get('fuzzy_match', 0) + + if fuzzy_match == 0: + _LOGGER.debug(f"Skipping fuzzy matching because fuzzy_match value not set in config") + return None, None, None + + max_score = 0 + best_match = None + for candidate in config_watched_plates: + seq = difflib.SequenceMatcher(a=str(plate_number).lower(), b=str(candidate).lower()) + if seq.ratio() > max_score: + max_score = seq.ratio() + best_match = candidate + + _LOGGER.debug(f"Best fuzzy_match: {best_match} ({max_score})") + + if max_score >= fuzzy_match: + _LOGGER.info(f"Watched plate found from fuzzy matching: {best_match} with score {max_score}") + return best_match, None, max_score + + + _LOGGER.debug("No matching Watched Plates found.") + #No watched_plate matches found + return None, None, None + + + +def send_mqtt_message(plate_number, plate_score, frigate_event_id, after_data, formatted_start_time, watched_plate, fuzzy_score): if not config['frigate'].get('return_topic'): return - message = { - 'plate_number': plate_number, - 'score': plate_score, - 'frigate_event_id': frigate_event_id, - 'camera_name': after_data['camera'], - 'start_time': formatted_start_time - } + if watched_plate: + message = { + 'plate_number': str(watched_plate).upper(), + 'score': plate_score, + 'frigate_event_id': frigate_event_id, + 'camera_name': after_data['camera'], + 'start_time': formatted_start_time, + 'fuzzy_score': fuzzy_score, + 'original_plate': str(plate_number).upper() + } + else: + message = { + 'plate_number': str(plate_number).upper(), + 'score': plate_score, + 'frigate_event_id': frigate_event_id, + 'camera_name': after_data['camera'], + 'start_time': formatted_start_time + } _LOGGER.debug(f"Sending MQTT message: {message}") @@ -147,37 +229,59 @@ def send_mqtt_message(plate_number, plate_score, frigate_event_id, after_data, f def has_common_value(array1, array2): return any(value in array2 for value in array1) -def save_image(config, after_data, image_content, license_plate_attribute, plate_number): +def save_image(config, after_data, frigate_url, frigate_event_id, plate_number): if not config['frigate'].get('save_snapshots', False): _LOGGER.debug(f"Skipping saving snapshot because save_snapshots is set to false") return + + # get latest Event Data from Frigate API + event_url = f"{frigate_url}/api/events/{frigate_event_id}" + + final_attribute = get_final_data(event_url) + + # get latest snapshot + snapshot = get_snapshot(frigate_event_id, frigate_url, False) + if not snapshot: + return - image = Image.open(io.BytesIO(bytearray(image_content))) + image = Image.open(io.BytesIO(bytearray(snapshot))) draw = ImageDraw.Draw(image) font = ImageFont.truetype("./Arial.ttf", size=14) - - # if given a plate number then draw it on the image along with the box around it - if license_plate_attribute and config['frigate'].get('draw_box', False): - vehicle = ( - license_plate_attribute[0]['box'][0], - license_plate_attribute[0]['box'][1], - license_plate_attribute[0]['box'][2], - license_plate_attribute[0]['box'][3] + + if final_attribute: + image_width, image_height = image.size + dimension_1 = int(final_attribute[0]['box'][0]) + dimension_2 = int(final_attribute[0]['box'][1]) + dimension_3 = int(final_attribute[0]['box'][2]) + dimension_4 = int(final_attribute[0]['box'][3]) + + plate = ( + dimension_1 * image_width, + dimension_2 * image_height, + (dimension_1 + dimension_3) * image_width, + (dimension_2 + dimension_4) * image_height ) - _LOGGER.debug(f"Drawing box: {vehicle}") - draw.rectangle(vehicle, outline="red", width=2) - + draw.rectangle(plate, outline="red", width=2) + _LOGGER.debug(f"Drawing Plate Box: {plate}") + if plate_number: - draw.text((license_plate_attribute[0]['box'][0]+5,license_plate_attribute[0]['box'][3]+5), plate_number, font=font) + draw.text( + ( + (dimension_1 * image_width)+ 5, + ((dimension_2 + dimension_4) * image_height) + 5 + ), + str(plate_number).upper(), + font=font + ) # save image timestamp = datetime.now().strftime(DATETIME_FORMAT) image_name = f"{after_data['camera']}_{timestamp}.png" if plate_number: - image_name = f"{plate_number}_{image_name}" + image_name = f"{str(plate_number).upper()}_{image_name}" image_path = f"{SNAPSHOT_PATH}/{image_name}" - _LOGGER.debug(f"Saving image with path: {image_path}") + _LOGGER.info(f"Saving image with path: {image_path}") image.save(image_path) def check_first_message(): @@ -208,18 +312,18 @@ def check_invalid_event(before_data, after_data): return True # limit api calls to plate checker api by only checking the best score for an event - if(before_data['top_score'] == after_data['top_score']): - _LOGGER.debug(f"duplicated snapshot from Frigate as top_score from before and after are the same: {after_data['top_score']}") + if(before_data['top_score'] == after_data['top_score'] and after_data['id'] in CURRENT_EVENTS) and not config['frigate'].get('frigate_plus', False): + _LOGGER.debug(f"duplicated snapshot from Frigate as top_score from before and after are the same: {after_data['top_score']} {after_data['id']}") return True return False -def get_snapshot(frigate_event_id, frigate_url): - _LOGGER.debug(f"Getting snapshot for event: {frigate_event_id}") +def get_snapshot(frigate_event_id, frigate_url, cropped): + _LOGGER.debug(f"Getting snapshot for event: {frigate_event_id}, Crop: {cropped}") snapshot_url = f"{frigate_url}/api/events/{frigate_event_id}/snapshot.jpg" _LOGGER.debug(f"event URL: {snapshot_url}") # get snapshot - response = requests.get(snapshot_url, params={ "crop": 1, "quality": 95 }) + response = requests.get(snapshot_url, params={ "crop": cropped, "quality": 95 }) # Check if the request was successful (HTTP status code 200) if response.status_code != 200: @@ -228,25 +332,44 @@ def get_snapshot(frigate_event_id, frigate_url): return response.content -def get_license_plate(after_data): +def get_license_plate_attribute(after_data): if config['frigate'].get('frigate_plus', False): attributes = after_data.get('current_attributes', []) license_plate_attribute = [attribute for attribute in attributes if attribute['label'] == 'license_plate'] return license_plate_attribute else: return None + +def get_final_data(event_url): + if config['frigate'].get('frigate_plus', False): + response = requests.get(event_url) + if response.status_code != 200: + _LOGGER.error(f"Error getting final data: {response.status_code}") + return + event_json = response.json() + event_data = event_json.get('data', {}) + + if event_data: + attributes = event_data.get('attributes', []) + final_attribute = [attribute for attribute in attributes if attribute['label'] == 'license_plate'] + return final_attribute + else: + return None + else: + return None + def is_valid_license_plate(after_data): # if user has frigate plus then check license plate attribute - license_plate_attribute = get_license_plate(after_data) - if not any(license_plate_attribute): + after_license_plate_attribute = get_license_plate_attribute(after_data) + if not any(after_license_plate_attribute): _LOGGER.debug(f"no license_plate attribute found in event attributes") return False # check min score of license plate attribute license_plate_min_score = config['frigate'].get('license_plate_min_score', 0) - if license_plate_attribute[0]['score'] < license_plate_min_score: - _LOGGER.debug(f"license_plate attribute score is below minimum: {license_plate_attribute[0]['score']}") + if after_license_plate_attribute[0]['score'] < license_plate_min_score: + _LOGGER.debug(f"license_plate attribute score is below minimum: {after_license_plate_attribute[0]['score']}") return False return True @@ -265,37 +388,28 @@ def is_duplicate_event(frigate_event_id): return False -def get_plate(snapshot, after_data, license_plate_attribute): +def get_plate(snapshot): # try to get plate number plate_number = None plate_score = None if config.get('plate_recognizer'): - plate_number, plate_score = plate_recognizer(snapshot) + plate_number, plate_score , watched_plate, fuzzy_score = plate_recognizer(snapshot) elif config.get('code_project'): - plate_number, plate_score = code_project(snapshot) + plate_number, plate_score, watched_plate, fuzzy_score = code_project(snapshot) else: _LOGGER.error("Plate Recognizer is not configured") - return None, None + return None, None, None, None # check Plate Recognizer score min_score = config['frigate'].get('min_score') score_too_low = min_score and plate_score and plate_score < min_score - if not score_too_low or config['frigate'].get('always_save_snapshot', False): - save_image( - config=config, - after_data=after_data, - image_content=snapshot, - license_plate_attribute=license_plate_attribute, - plate_number=plate_number - ) - - if score_too_low: - _LOGGER.info(f"Score is below minimum: {plate_score}") - return None, None + if not fuzzy_score and score_too_low: + _LOGGER.info(f"Score is below minimum: {plate_score} ({plate_number})") + return None, None, None, None - return plate_number, plate_score + return plate_number, plate_score, watched_plate, fuzzy_score def store_plate_in_db(plate_number, plate_score, frigate_event_id, after_data, formatted_start_time): conn = sqlite3.connect(DB_PATH) @@ -320,38 +434,62 @@ def on_message(client, userdata, message): before_data = payload_dict.get('before', {}) after_data = payload_dict.get('after', {}) - - if check_invalid_event(before_data, after_data): - return - + type = payload_dict.get('type','') + frigate_url = config['frigate']['frigate_url'] frigate_event_id = after_data['id'] - - if is_duplicate_event(frigate_event_id): + + if type == 'end' and after_data['id'] in CURRENT_EVENTS: + _LOGGER.debug(f"CLEARING EVENT: {frigate_event_id} after {CURRENT_EVENTS[frigate_event_id]} calls to AI engine") + del CURRENT_EVENTS[frigate_event_id] + + if check_invalid_event(before_data, after_data): return - snapshot = get_snapshot(frigate_event_id, frigate_url) - if not snapshot: + if is_duplicate_event(frigate_event_id): return frigate_plus = config['frigate'].get('frigate_plus', False) if frigate_plus and not is_valid_license_plate(after_data): return - - license_plate_attribute = get_license_plate(after_data) - - plate_number, plate_score = get_plate(snapshot, after_data, license_plate_attribute) - if not plate_number: + + if not type == 'end' and not after_data['id'] in CURRENT_EVENTS: + CURRENT_EVENTS[frigate_event_id] = 0 + + + snapshot = get_snapshot(frigate_event_id, frigate_url, True) + if not snapshot: + del CURRENT_EVENTS[frigate_event_id] # remove existing id from current events due to snapshot failure - will try again next frame return - start_time = datetime.fromtimestamp(after_data['start_time']) - formatted_start_time = start_time.strftime("%Y-%m-%d %H:%M:%S") - - store_plate_in_db(plate_number, plate_score, frigate_event_id, after_data, formatted_start_time) - set_sublabel(frigate_url, frigate_event_id, plate_number, plate_score) - - send_mqtt_message(plate_number, plate_score, frigate_event_id, after_data, formatted_start_time) + _LOGGER.debug(f"Getting plate for event: {frigate_event_id}") + if frigate_event_id in CURRENT_EVENTS: + if config['frigate'].get('max_attempts', 0) > 0 and CURRENT_EVENTS[frigate_event_id] > config['frigate'].get('max_attempts', 0): + _LOGGER.debug(f"Maximum number of AI attempts reached for event {frigate_event_id}: {CURRENT_EVENTS[frigate_event_id]}") + return + CURRENT_EVENTS[frigate_event_id] += 1 + plate_number, plate_score, watched_plate, fuzzy_score = get_plate(snapshot) + if plate_number: + start_time = datetime.fromtimestamp(after_data['start_time']) + formatted_start_time = start_time.strftime("%Y-%m-%d %H:%M:%S") + + if watched_plate: + store_plate_in_db(watched_plate, plate_score, frigate_event_id, after_data, formatted_start_time) + else: + store_plate_in_db(plate_number, plate_score, frigate_event_id, after_data, formatted_start_time) + set_sublabel(frigate_url, frigate_event_id, watched_plate if watched_plate else plate_number, plate_score) + + send_mqtt_message(plate_number, plate_score, frigate_event_id, after_data, formatted_start_time, watched_plate, fuzzy_score) + + if plate_number or config['frigate'].get('always_save_snapshot', False): + save_image( + config=config, + after_data=after_data, + frigate_url=frigate_url, + frigate_event_id=frigate_event_id, + plate_number=watched_plate if watched_plate else plate_number + ) def setup_db(): conn = sqlite3.connect(DB_PATH) diff --git a/test.py b/test.py index 53a0a44..622b136 100644 --- a/test.py +++ b/test.py @@ -3,11 +3,12 @@ import logging from pathlib import Path import os -import yaml -import logging import unittest from unittest.mock import patch, MagicMock, mock_open +from PIL import Image, ImageDraw +import yaml + import index class BaseTestCase(unittest.TestCase): @@ -33,12 +34,14 @@ class TestSaveImage(BaseTestCase): def setUp(self): index._LOGGER = logging.getLogger(__name__) + @patch('index.get_snapshot') + @patch('index.get_final_data') @patch('index.Image.open') @patch('index.ImageDraw.Draw') @patch('index.ImageFont.truetype') @patch('index.datetime') @patch('index.open', new_callable=mock_open) - def test_save_image(self, mock_file, mock_datetime, mock_truetype, mock_draw, mock_open): + def test_save_image_with_box(self, mock_file, mock_datetime, mock_truetype, mock_draw, mock_open, mock_get_final_data, mock_get_snapshot): # Mock current time mock_now = mock_datetime.now.return_value mock_now.strftime.return_value = '20210101_120000' @@ -46,23 +49,31 @@ def test_save_image(self, mock_file, mock_datetime, mock_truetype, mock_draw, mo # Setup configuration and input data index.config = {'frigate': {'save_snapshots': True, 'draw_box': True}} after_data = {'camera': 'test_camera'} - image_content = b'test_image_content' - license_plate_attribute = [{'box': [0, 0, 100, 100]}] + frigate_url = 'http://example.com' + frigate_event_id = 'test_event_id' plate_number = 'ABC123' # Mock PIL dependencies - mock_image = MagicMock() + mock_image = MagicMock(spec=Image.Image) + mock_image.size = (640, 480) # Example size + mock_image_draw = mock_draw.return_value mock_open.return_value = mock_image - mock_draw.return_value = MagicMock() mock_truetype.return_value = MagicMock() + mock_get_final_data.return_value = [{'box': [0, 0, 100, 100]}] + mock_get_snapshot.return_value = b'ImageBytes' + # Call the function - index.save_image(index.config, after_data, image_content, license_plate_attribute, plate_number) + index.save_image(index.config, after_data, frigate_url, frigate_event_id, plate_number) # Assert image operations - mock_image.save.assert_called_with(f'/plates/{plate_number}_test_camera_20210101_120000.png') - mock_draw.return_value.rectangle.assert_called_with((0, 0, 100, 100), outline='red', width=2) - mock_draw.return_value.text.assert_called_with((5, 105), 'ABC123', font=mock_truetype.return_value) + # Assert image operations + expected_path = '/plates/ABC123_test_camera_20210101_120000.png' + mock_image.save.assert_called_once_with(expected_path) + mock_image_draw.rectangle.assert_called_once_with( + (0, 0, 640 * 100, 480 * 100), outline="red", width=2 # Ensure this matches what's being called + ) + mock_image_draw.text.assert_called_once_with((5, 48005), 'ABC123', font=mock_truetype.return_value) @patch('index.Image.open') @patch('index.ImageDraw.Draw') @@ -108,7 +119,7 @@ def test_set_sublabel(self, mock_post): mock_post.assert_called_with( "http://example.com/api/events/123/sub_label", - data='{"subLabel": "test_label"}', + data='{"subLabel": "TEST_LABEL"}', headers={"Content-Type": "application/json"} ) @@ -121,7 +132,7 @@ def test_set_sublabel_shorten(self, mock_post): mock_post.assert_called_with( "http://example.com/api/events/123/sub_label", - data='{"subLabel": "test_label_too_long_"}', + data='{"subLabel": "TEST_LABEL_TOO_LONG_"}', headers={"Content-Type": "application/json"} ) @@ -189,25 +200,25 @@ def test_get_license_plate_with_frigate_plus_enabled(self): {'label': 'other_attribute', 'score': 0.8} ] } - result = index.get_license_plate(after_data) + result = index.get_license_plate_attribute(after_data) self.assertEqual(result, [{'label': 'license_plate', 'score': 0.9}]) def test_get_license_plate_with_frigate_plus_disabled(self): index.config = {'frigate': {'frigate_plus': False}} after_data = {'current_attributes': [{'label': 'license_plate', 'score': 0.9}]} - result = index.get_license_plate(after_data) + result = index.get_license_plate_attribute(after_data) self.assertIsNone(result) def test_get_license_plate_with_no_license_plate_attribute(self): index.config = {'frigate': {'frigate_plus': True}} after_data = {'current_attributes': [{'label': 'other_attribute', 'score': 0.8}]} - result = index.get_license_plate(after_data) + result = index.get_license_plate_attribute(after_data) self.assertEqual(result, []) def test_get_license_plate_with_empty_attributes(self): index.config = {'frigate': {'frigate_plus': True}} after_data = {'current_attributes': []} - result = index.get_license_plate(after_data) + result = index.get_license_plate_attribute(after_data) self.assertEqual(result, []) class TestCheckFirstMessage(BaseTestCase): @@ -279,7 +290,7 @@ class TestIsValidLicensePlate(BaseTestCase): def setUp(self): super().setUp() - @patch('index.get_license_plate') + @patch('index.get_license_plate_attribute') def test_no_license_plate_attribute(self, mock_get_license_plate): # Setup: No license plate attribute found mock_get_license_plate.return_value = [] @@ -292,7 +303,7 @@ def test_no_license_plate_attribute(self, mock_get_license_plate): self.assertFalse(result) self.mock_logger.debug.assert_called_with("no license_plate attribute found in event attributes") - @patch('index.get_license_plate') + @patch('index.get_license_plate_attribute') def test_license_plate_below_min_score(self, mock_get_license_plate): # Setup: License plate attribute found but below minimum score index.config = {'frigate': {'frigate_plus': True, 'license_plate_min_score': 0.5}} @@ -305,7 +316,7 @@ def test_license_plate_below_min_score(self, mock_get_license_plate): self.mock_logger.debug.assert_called_with("license_plate attribute score is below minimum: 0.4") - @patch('index.get_license_plate') + @patch('index.get_license_plate_attribute') def test_valid_license_plate(self, mock_get_license_plate): # Setup: Valid license plate attribute index.config = {'frigate': {'license_plate_min_score': 0.5}} @@ -330,12 +341,12 @@ def test_get_snapshot_successful(self, mock_requests_get): frigate_event_id = 'event123' frigate_url = 'http://example.com' - result = index.get_snapshot(frigate_event_id, frigate_url) + result = index.get_snapshot(frigate_event_id, frigate_url, True) self.assertEqual(result, b'image_data') mock_requests_get.assert_called_with(f"{frigate_url}/api/events/{frigate_event_id}/snapshot.jpg", params={"crop": 1, "quality": 95}) - self.mock_logger.debug.assert_any_call(f"Getting snapshot for event: {frigate_event_id}") + self.mock_logger.debug.assert_any_call(f"Getting snapshot for event: {frigate_event_id}, Crop: True") self.mock_logger.debug.assert_any_call(f"event URL: {frigate_url}/api/events/{frigate_event_id}/snapshot.jpg") @patch('index.requests.get') @@ -348,7 +359,7 @@ def test_get_snapshot_failure(self, mock_requests_get): frigate_event_id = 'event123' frigate_url = 'http://example.com' - result = index.get_snapshot(frigate_event_id, frigate_url) + result = index.get_snapshot(frigate_event_id, frigate_url, True) self.assertIsNone(result) mock_requests_get.assert_called_with(f"{frigate_url}/api/events/{frigate_event_id}/snapshot.jpg", @@ -393,19 +404,6 @@ def test_event_invalid_object(self): self.assertTrue(result) self.mock_logger.debug.assert_called_with("is not a correct label: tree") - def test_event_duplicate_top_score(self): - before_data = {'top_score': 0.8} - after_data = { - 'current_zones': ['zone1'], - 'camera': 'camera1', - 'label': 'car', - 'id': 'event123', - 'top_score': 0.8 - } - result = index.check_invalid_event(before_data, after_data) - self.assertTrue(result) - self.mock_logger.debug.assert_called_with("duplicated snapshot from Frigate as top_score from before and after are the same: 0.8") - def test_event_valid(self): before_data = {'top_score': 0.7} after_data = { @@ -425,37 +423,53 @@ def setUp(self): @patch('index.plate_recognizer') @patch('index.save_image') - def test_plate_recognizer_configured(self, mock_save_image, mock_plate_recognizer): + def test_plate_score_okay(self, mock_save_image, mock_plate_recognizer): # Set up configuration to use plate_recognizer index.config = {'plate_recognizer': True, 'frigate': {'min_score': 0.5, 'always_save_snapshot': False}} snapshot = b'image_data' - after_data = {'test_data': 'value'} - license_plate_attribute = [{'test_attr': 'value'}] # Mock the plate_recognizer to return a specific plate number and score - mock_plate_recognizer.return_value = ('ABC123', 0.6) - result, score = index.get_plate(snapshot, after_data, license_plate_attribute) + mock_plate_recognizer.return_value = ('ABC123', 0.6, None, None) + plate_number, plate_score, watched_plate, fuzzy_score = index.get_plate(snapshot) # Assert that the correct plate number is returned - self.assertEqual(result, 'ABC123') + self.assertEqual(plate_number, 'ABC123') + self.assertEqual(plate_score, 0.6) mock_plate_recognizer.assert_called_once_with(snapshot) - mock_save_image.assert_called_once() + mock_save_image.assert_not_called() # Assert that save_image is not called when plate_recognizer is used @patch('index.plate_recognizer') @patch('index.save_image') def test_plate_score_too_low(self, mock_save_image, mock_plate_recognizer): index.config = {'plate_recognizer': True, 'frigate': {'min_score': 0.7, 'always_save_snapshot': False}} snapshot = b'image_data' - after_data = {'test_data': 'value'} - license_plate_attribute = [{'test_attr': 'value'}] # Mock the plate_recognizer to return a plate number with a low score - mock_plate_recognizer.return_value = ('ABC123', 0.6) - result, score = index.get_plate(snapshot, after_data, license_plate_attribute) + mock_plate_recognizer.return_value = ('ABC123', 0.6, None, None) + plate_number, plate_score, watched_plate, fuzzy_score = index.get_plate(snapshot) # Assert that no plate number is returned due to low score - self.assertIsNone(result) - self.mock_logger.info.assert_called_with("Score is below minimum: 0.6") + self.assertIsNone(plate_number) + self.assertIsNone(plate_score) + self.mock_logger.info.assert_called_with("Score is below minimum: 0.6 (ABC123)") + mock_save_image.assert_not_called() + + @patch('index.plate_recognizer') + @patch('index.save_image') + def test_fuzzy_response(self, mock_save_image, mock_plate_recognizer): + index.config = {'plate_recognizer': True, 'frigate': {'min_score': 0.7, 'always_save_snapshot': False}} + snapshot = b'image_data' + + # Mock the plate_recognizer to return a plate number with a fuzzy score + mock_plate_recognizer.return_value = ('DEF456', 0.8, None, 0.9) + plate_number, plate_score, watched_plate, fuzzy_score = index.get_plate(snapshot) + + # Assert that plate number and score are returned despite the fuzzy score + self.assertEqual(plate_number, 'DEF456') + self.assertEqual(plate_score, 0.8) + self.assertIsNone(watched_plate) + self.assertEqual(fuzzy_score, 0.9) + self.mock_logger.error.assert_not_called() mock_save_image.assert_not_called() class TestSendMqttMessage(BaseTestCase): @@ -476,9 +490,11 @@ def test_send_mqtt_message(self, mock_mqtt_client): frigate_event_id = 'event123' after_data = {'camera': 'camera1'} formatted_start_time = '2021-01-01 12:00:00' + watched_plate = 'ABC123' + fuzzy_score = 0.8 # Call the function - index.send_mqtt_message(plate_number, plate_score, frigate_event_id, after_data, formatted_start_time) + index.send_mqtt_message(plate_number, plate_score, frigate_event_id, after_data, formatted_start_time, watched_plate, fuzzy_score) # Construct expected message expected_message = { @@ -486,7 +502,9 @@ def test_send_mqtt_message(self, mock_mqtt_client): 'score': plate_score, 'frigate_event_id': frigate_event_id, 'camera_name': after_data['camera'], - 'start_time': formatted_start_time + 'start_time': formatted_start_time, + 'fuzzy_score': fuzzy_score, + 'original_plate': watched_plate } # Assert that the MQTT client publish method is called correctly