From 11d87839257d3d31f76cce5bad3186277834f23a Mon Sep 17 00:00:00 2001 From: itisacloud Date: Wed, 14 Aug 2024 14:36:18 +0200 Subject: [PATCH] feat: add warning when no markings are beeing detected --- .../upload_processing/detect_markings.py | 4 ++++ .../upload_processing/test_detect_markings.py | 18 ++++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/sketch_map_tool/upload_processing/detect_markings.py b/sketch_map_tool/upload_processing/detect_markings.py index 39fb0f4f..2aa9bdc6 100644 --- a/sketch_map_tool/upload_processing/detect_markings.py +++ b/sketch_map_tool/upload_processing/detect_markings.py @@ -1,3 +1,5 @@ +import logging + import cv2 import numpy as np from numpy.typing import NDArray @@ -41,6 +43,8 @@ def detect_markings( yolo_cls, sam_predictor, ) + if len(colors) == 0: + logging.warning("No markings detected.") colors = [int(c) + 1 for c in colors] # +1 because 0 is background processed_markings = post_process(masks, bboxes, colors) return processed_markings diff --git a/tests/integration/upload_processing/test_detect_markings.py b/tests/integration/upload_processing/test_detect_markings.py index aba0a5bd..181fc743 100644 --- a/tests/integration/upload_processing/test_detect_markings.py +++ b/tests/integration/upload_processing/test_detect_markings.py @@ -1,3 +1,5 @@ +import logging + import numpy as np import pytest from PIL import Image, ImageEnhance @@ -77,3 +79,19 @@ def test_detect_markings( for m in markings: img = Image.fromarray(m) ImageEnhance.Contrast(img).enhance(10).show() + + +def test_detectec_markings_failure(yolo_osm_cls, yolo_osm_obj, sam_predictor, caplog): + # Empty map and template should not contain any markings + empty_map = np.zeros((1024, 1024, 3), dtype=np.uint8) + empty_template = np.zeros((1024, 1024, 3), dtype=np.uint8) + + with caplog.at_level(logging.WARNING): + detect_markings( + empty_map, + empty_template, + yolo_osm_obj, + yolo_osm_cls, + sam_predictor, + ) + assert "No markings detected." in caplog.text