ICRA/tests/test_image_processor.py

175 lines
5.8 KiB
Python

import numpy as np
import pytest
from PIL import Image
from app.qt.image_processor import Stats, _rgb_to_hsv_numpy, QtImageProcessor
def test_stats_summary():
s = Stats(
matches_all=50, total_all=100,
matches_keep=40, total_keep=80,
matches_excl=10, total_excl=20
)
def mock_t(key, **kwargs):
if key == "stats.placeholder":
return "Placeholder"
if not kwargs:
return key
return f"{kwargs['with_pct']:.1f} {kwargs['without_pct']:.1f} {kwargs['excluded_pct']:.1f}"
weights = {"match_all": 30, "match_keep": 50, "brightness": 10, "grouping": 10}
res = s.summary(mock_t, weights)
# with_pct: 40/80 = 50.0
# without_pct: 50/100 = 50.0
# excluded_pct: 20/100 = 20.0
assert res == "50.0 50.0 20.0"
def test_stats_empty():
s = Stats()
weights = {"match_all": 30, "match_keep": 50, "brightness": 10, "grouping": 10}
assert s.summary(lambda k, **kw: "Empty", weights) == "Empty"
def test_rgb_to_hsv_numpy():
# Test red
arr = np.array([[[1.0, 0.0, 0.0]]], dtype=np.float32)
hsv = _rgb_to_hsv_numpy(arr)
assert np.allclose(hsv[0, 0], [0.0, 100.0, 100.0])
# Test green
arr = np.array([[[0.0, 1.0, 0.0]]], dtype=np.float32)
hsv = _rgb_to_hsv_numpy(arr)
assert np.allclose(hsv[0, 0], [120.0, 100.0, 100.0])
# Test blue
arr = np.array([[[0.0, 0.0, 1.0]]], dtype=np.float32)
hsv = _rgb_to_hsv_numpy(arr)
assert np.allclose(hsv[0, 0], [240.0, 100.0, 100.0])
# Test white
arr = np.array([[[1.0, 1.0, 1.0]]], dtype=np.float32)
hsv = _rgb_to_hsv_numpy(arr)
assert np.allclose(hsv[0, 0], [0.0, 0.0, 100.0])
# Test black
arr = np.array([[[0.0, 0.0, 0.0]]], dtype=np.float32)
hsv = _rgb_to_hsv_numpy(arr)
assert np.allclose(hsv[0, 0], [0.0, 0.0, 0.0])
def test_qt_processor_matches_legacy():
proc = QtImageProcessor()
proc.hue_min = 350
proc.hue_max = 10
proc.sat_min = 50
proc.val_min = 50
proc.val_max = 100
# Red wraps around 360, so H=0 -> ok
assert proc._matches(255, 0, 0) is True
# Green H=120 -> fail
assert proc._matches(0, 255, 0) is False
# Dark red S=100, V=25 -> fail because val_min=50
assert proc._matches(64, 0, 0) is False
def test_set_overlay_color():
proc = QtImageProcessor()
# default red
assert proc.overlay_r == 255
assert proc.overlay_g == 0
assert proc.overlay_b == 0
proc.set_overlay_color("#00ff00")
assert proc.overlay_r == 0
assert proc.overlay_g == 255
assert proc.overlay_b == 0
# invalid hex does nothing
proc.set_overlay_color("blue")
assert proc.overlay_r == 0
def test_coordinate_scaling():
proc = QtImageProcessor()
# Create a 200x200 image where everything is red
red_img_small = Image.new("RGBA", (200, 200), (255, 0, 0, 255))
proc.orig_img = red_img_small # satisfy preview logic
proc.preview_img = red_img_small
# All red. Thresholds cover all red.
proc.hue_min = 0
proc.hue_max = 360
proc.sat_min = 10
proc.val_min = 10
# Exclude the right half (100-200)
proc.set_exclusions([{"kind": "rect", "coords": (100, 0, 200, 200)}])
# Verify small stats
s_small = proc.get_stats_headless(red_img_small)
# total=40000, keep=20000, excl=20000
assert s_small.total_all == 40000
assert s_small.total_keep == 20000
assert s_small.total_excl == 20000
# Now check on a 1000x1000 image (5x scale)
red_img_large = Image.new("RGBA", (1000, 1000), (255, 0, 0, 255))
s_large = proc.get_stats_headless(red_img_large)
# total=1,000,000. If scaling works, keep=500,000, excl=500,000.
# If scaling FAILED, the mask is still 100x200 (20,000 px) -> excl=20,000.
assert s_large.total_all == 1000000
assert s_large.total_keep == 500000
assert s_large.total_excl == 500000
def test_calculate_grouping_score():
proc = QtImageProcessor()
# 1. Empty mask
mask_empty = np.zeros((20, 20), dtype=bool)
assert proc._calculate_grouping_score(mask_empty) == 0.0
# 2. Single mask pixel (0 neighbors)
mask_single = np.zeros((20, 20), dtype=bool)
mask_single[10, 10] = True
assert proc._calculate_grouping_score(mask_single) == 0.0
# 3. 2x2 block
# each pixel in 2x2 has 3 neighbors in 3x3, 3 neighbors in 5x5, 3 neighbors in 9x9.
# score = ((3/80)^2) * 100
expected_2x2 = ((3/80.0)**2) * 100.0
mask_block = np.zeros((20, 20), dtype=bool)
mask_block[10:12, 10:12] = True
assert pytest.approx(proc._calculate_grouping_score(mask_block)) == expected_2x2
# 4. 9x9 block
# center pixel has 80 neighbors (100% density).
# many pixels have high density.
mask_9x9 = np.zeros((20, 20), dtype=bool)
mask_9x9[5:14, 5:14] = True
res_9x9 = proc._calculate_grouping_score(mask_9x9)
assert res_9x9 > expected_2x2
# For a 9x9 block, the center pixel is 100%. Boundary pixels are less.
# 1 center pixel = 80/80 = 1.0.
# Overall it should be a healthy percentage.
assert res_9x9 > 10.0 # significant grouping
def test_export_worker_error():
from app.qt.image_processor import _export_worker
# 1. Provide a missing file to trigger an exception during Image.open()
res1 = _export_worker(("missing_file.png", {
"hue_min": 0, "hue_max": 360, "sat_min": 0, "sat_max": 100,
"val_min": 0, "val_max": 100, "exclude_bg": False,
"exclude_bg_rgb": (0, 0, 0), "exclude_bg_tolerance": 5,
"prefer_dark": False, "exclude_shapes": [], "exclude_ref_size": None,
"weights": {}
}))
assert res1 == ("missing_file.png", None, None, None, None, None)
# 2. Provide an empty params dict to trigger KeyError before opening image
res2 = _export_worker(("dummy.png", {}))
assert res2 == ("dummy.png", None, None, None, None, None)