82 lines
2.7 KiB
Python
82 lines
2.7 KiB
Python
# Copyright 2026 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Perceptual-hash stability and distance behavior (M4 similarity hint)."""
|
|
|
|
import io
|
|
|
|
from PIL import Image, ImageDraw
|
|
|
|
from artemis.utils.image_hash import HASH_SIZE, dhash_hex, hamming_distance_hex
|
|
|
|
|
|
def _screen(draw_fn, size=(270, 600), quality=85) -> bytes:
|
|
img = Image.new("RGB", size, "white")
|
|
draw_fn(ImageDraw.Draw(img))
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="JPEG", quality=quality)
|
|
return buf.getvalue()
|
|
|
|
|
|
def _login_screen(d):
|
|
d.rectangle([20, 40, 250, 90], fill="#3355ff")
|
|
d.rectangle([30, 200, 240, 250], outline="#333333", width=4)
|
|
d.rectangle([30, 300, 240, 350], outline="#333333", width=4)
|
|
d.rectangle([60, 450, 210, 510], fill="#22aa66")
|
|
|
|
|
|
def _settings_screen(d):
|
|
for y in range(60, 560, 80):
|
|
d.rectangle([10, y, 260, y + 50], fill="#dddddd")
|
|
d.ellipse([200, 20, 250, 45], fill="#aa2222")
|
|
|
|
|
|
def test_dhash_is_stable_for_identical_bytes():
|
|
img = _screen(_login_screen)
|
|
assert dhash_hex(img) == dhash_hex(img)
|
|
assert len(dhash_hex(img)) == HASH_SIZE * HASH_SIZE // 4
|
|
|
|
|
|
def test_dhash_survives_jpeg_reencoding_with_small_distance():
|
|
original = _screen(_login_screen, quality=95)
|
|
with Image.open(io.BytesIO(original)) as img:
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="JPEG", quality=60)
|
|
reencoded = buf.getvalue()
|
|
|
|
distance = hamming_distance_hex(dhash_hex(original), dhash_hex(reencoded))
|
|
assert distance is not None and distance <= 4
|
|
|
|
|
|
def test_dhash_separates_different_screens():
|
|
distance = hamming_distance_hex(
|
|
dhash_hex(_screen(_login_screen)), dhash_hex(_screen(_settings_screen))
|
|
)
|
|
assert distance is not None and distance > 8
|
|
|
|
|
|
def test_dhash_best_effort_on_garbage():
|
|
assert dhash_hex(None) is None
|
|
assert dhash_hex(b"") is None
|
|
assert dhash_hex(b"not an image at all") is None
|
|
|
|
|
|
def test_hamming_distance_edge_cases():
|
|
assert hamming_distance_hex(None, "ff") is None
|
|
assert hamming_distance_hex("ff", None) is None
|
|
assert hamming_distance_hex("ff", "ffff") is None # length mismatch
|
|
assert hamming_distance_hex("zz", "aa") is None # not hex
|
|
assert hamming_distance_hex("00ff", "00ff") == 0
|
|
assert hamming_distance_hex("0000", "000f") == 4
|