Spaces:
Running on Zero
Running on Zero
| """Offline tests for app.scribble (scribble-guided exposure split). | |
| No network anywhere; asserts compare independently computed quantities. | |
| """ | |
| import numpy as np | |
| import pytest | |
| from app.scribble import ( | |
| PALETTE, | |
| parse_editor_scribbles, | |
| propagate_w, | |
| split_by_scribbles, | |
| ) | |
| def _two_region_image(h=96, w=96): | |
| """Left half bright / right half dark with a crisp vertical edge at w//2.""" | |
| img = np.full((h, w), 0.25, np.float32) | |
| img[:, : w // 2] = 0.85 | |
| return np.stack([img] * 3, axis=-1) | |
| def test_propagate_w_saturates_per_region(): | |
| """Seeds in each half -> w saturates toward 1 (A side) / 0 (B side), edge-aware.""" | |
| rgb = _two_region_image() | |
| h, w = rgb.shape[:2] | |
| seeds_a = np.zeros((h, w), bool) | |
| seeds_b = np.zeros((h, w), bool) | |
| seeds_a[h // 2 - 2 : h // 2 + 2, 8:20] = True # stroke in left (bright) half | |
| seeds_b[h // 2 - 2 : h // 2 + 2, w - 20 : w - 8] = True # stroke in right half | |
| field = propagate_w(rgb, seeds_a, seeds_b) | |
| # Independent references: mean of the field over each region interior | |
| left = field[:, : w // 2 - 4].mean() | |
| right = field[:, w // 2 + 4 :].mean() | |
| assert left > 0.8, f"left (A-seeded) region mean w = {left:.3f}, want > 0.8" | |
| assert right < 0.2, f"right (B-seeded) region mean w = {right:.3f}, want < 0.2" | |
| # Teeth: swapping the seeds must flip the field, not reproduce it | |
| flipped = propagate_w(rgb, seeds_b, seeds_a) | |
| assert flipped[:, : w // 2 - 4].mean() < 0.2 | |
| def test_split_by_scribbles_sum_exact_and_shapes(): | |
| """H_A + H_B == H_total everywhere (by construction — verify independently).""" | |
| rgb = _two_region_image() | |
| h, w = rgb.shape[:2] | |
| h_total = np.linspace(0.1, 2.0, h * w, dtype=np.float32).reshape(h, w) | |
| cm = np.ones((h, w), np.uint8) | |
| seeds_a = np.zeros((h, w), bool) | |
| seeds_b = np.zeros((h, w), bool) | |
| seeds_a[10:14, 10:30] = True | |
| seeds_b[10:14, 60:80] = True | |
| lay_a, lay_b, field = split_by_scribbles(rgb, h_total, cm, seeds_a, seeds_b) | |
| assert lay_a.shape == rgb.shape and lay_b.shape == rgb.shape | |
| # Reconstruct the split from the returned w and compare to h_total independently | |
| recon = field * h_total + (1.0 - field) * h_total | |
| assert np.abs(recon - h_total).max() < 1e-5 | |
| # The two layers must differ (teeth: not the same render twice) | |
| assert float(np.abs(lay_a - lay_b).mean()) > 0.01 | |
| def test_parse_editor_scribbles_colors_and_empty(): | |
| """Red strokes -> seeds_a, blue -> seeds_b; empty/None values -> all-False.""" | |
| h, w = 32, 40 | |
| layer = np.zeros((h, w, 4), np.uint8) | |
| layer[4:8, 4:12, :3] = PALETTE["red"] | |
| layer[4:8, 4:12, 3] = 255 | |
| layer[20:24, 20:30, :3] = PALETTE["blue"] | |
| layer[20:24, 20:30, 3] = 255 | |
| value = {"background": None, "layers": [layer], "composite": None} | |
| sa, sb = parse_editor_scribbles(value, (h, w)) | |
| # Independent references: the exact painted boxes | |
| ref_a = np.zeros((h, w), bool); ref_a[4:8, 4:12] = True | |
| ref_b = np.zeros((h, w), bool); ref_b[20:24, 20:30] = True | |
| assert np.array_equal(sa, ref_a) | |
| assert np.array_equal(sb, ref_b) | |
| for empty in (None, {}, {"layers": []}, {"layers": [None]}): | |
| sa, sb = parse_editor_scribbles(empty, (h, w)) | |
| assert not sa.any() and not sb.any() | |
| def test_parse_editor_scribbles_resizes_layer(): | |
| """A layer at a different resolution is mapped onto the target grid.""" | |
| h, w = 40, 40 | |
| layer = np.zeros((20, 20, 4), np.uint8) # half-res canvas | |
| layer[2:6, 2:6, :3] = PALETTE["red"] | |
| layer[2:6, 2:6, 3] = 255 | |
| sa, sb = parse_editor_scribbles({"layers": [layer]}, (h, w)) | |
| assert sa.any() and not sb.any() | |
| ys, xs = np.where(sa) | |
| # The painted box (rows/cols 2..6 of 20) must land in the upper-left quadrant | |
| assert ys.max() < h // 2 and xs.max() < w // 2 | |
| def test_recover_scribble_path_uses_bundle_for_both_layers(): | |
| """WP-19: with seeds, both calls are evidence-bundle prompts on the OBSERVED frame. | |
| Each prompt recovers its own scene and names the OTHER as removable | |
| contamination; the primary input stays the observed frame (references carry | |
| the physics layers on the live path — the offline hook sees only the primary). | |
| """ | |
| from app.restore import recover | |
| rgb = _two_region_image() | |
| h, w = rgb.shape[:2] | |
| h_total = np.full((h, w), 1.0, np.float32) | |
| cm = np.ones((h, w), np.uint8) | |
| seeds_h = np.zeros((h, w), bool); seeds_h[10:14, 5:25] = True | |
| seeds_o = np.zeros((h, w), bool); seeds_o[10:14, 60:85] = True | |
| prompts: list[str] = [] | |
| inputs: list[np.ndarray] = [] | |
| def hook(img, prompt): | |
| prompts.append(prompt) | |
| inputs.append(np.asarray(img)) | |
| return img | |
| res = recover(rgb, h_total, cm, restore_fn=hook, | |
| scene_headline="scene one", scene_other="scene two", | |
| seeds_headline=seeds_h, seeds_other=seeds_o) | |
| assert res.dominant is not None and res.second is not None | |
| assert len(prompts) == 2 | |
| assert "Recover this photo: scene one" in prompts[0] | |
| assert "scene two" in prompts[0] and "contamination" in prompts[0] | |
| assert "Recover this photo: scene two" in prompts[1] | |
| assert "scene one" in prompts[1] | |
| # the primary edit input is the observed frame, not the muddy layer | |
| assert all(np.array_equal(i, rgb) for i in inputs) | |
| # dreamed_frac reports the contested-w fraction: recompute independently | |
| from app.scribble import propagate_w | |
| field = propagate_w(rgb, seeds_h, seeds_o) | |
| expected = 100.0 * float(np.mean((field > 0.3) & (field < 0.7))) | |
| assert res.dreamed_frac == pytest.approx(expected, abs=1.0) | |
| def test_recover_empty_seeds_falls_back_to_default_path(monkeypatch): | |
| """All-False seed masks must not trigger the scribble path (separate mode leads). | |
| Discriminated via a kontext-family model override: kontext keeps the legacy | |
| single-image prompts, where only the DEFAULT path opens with _SEPARATE (the | |
| scribble path would use _RESTORE for both layers). | |
| """ | |
| from app.restore import recover, _SEPARATE | |
| monkeypatch.setenv("REPLICATE_RESTORE_MODEL", "black-forest-labs/flux-kontext-pro") | |
| rgb = _two_region_image() | |
| h, w = rgb.shape[:2] | |
| prompts: list[str] = [] | |
| res = recover(rgb, np.ones((h, w), np.float32), np.ones((h, w), np.uint8), | |
| restore_fn=lambda img, p: (prompts.append(p) or img), | |
| scene_headline="x", scene_other="y", | |
| seeds_headline=np.zeros((h, w), bool), | |
| seeds_other=np.zeros((h, w), bool)) | |
| assert res.dominant is not None | |
| assert prompts[0].startswith(_SEPARATE.split("{")[0]) | |
| def test_recover_context_folded_into_prompts(): | |
| """The whole-photo context string reaches both layer prompts.""" | |
| from app.restore import recover | |
| rgb = _two_region_image() | |
| h, w = rgb.shape[:2] | |
| prompts: list[str] = [] | |
| recover(rgb, np.ones((h, w), np.float32), np.ones((h, w), np.uint8), | |
| restore_fn=lambda img, p: (prompts.append(p) or img), | |
| scene_headline="a pool", scene_other="a frame", | |
| context="honeymoon roll, backyard over France") | |
| assert len(prompts) == 2 | |
| assert all("honeymoon roll, backyard over France" in p for p in prompts) | |
| def test_parse_tagged_scribbles_assignments_and_hints(): | |
| """Tagged colors route to their assigned scenes and yield located hints; an | |
| override (blue -> Scene 1) is honored over the default.""" | |
| from app.scribble import parse_tagged_scribbles, PALETTE | |
| h, w = 60, 90 | |
| layer = np.zeros((h, w, 4), np.uint8) | |
| layer[5:12, 5:20, :3] = PALETTE["red"] # top-left red stroke | |
| layer[5:12, 5:20, 3] = 255 | |
| layer[45:52, 60:80, :3] = PALETTE["blue"] # bottom-right blue stroke | |
| layer[45:52, 60:80, 3] = 255 | |
| value = {"background": None, "layers": [layer], "composite": None} | |
| assignments = { | |
| "red": {"scene": "Scene 1", "tag": "pool"}, | |
| "blue": {"scene": "Scene 1", "tag": "frame"}, # override: blue joins scene 1 | |
| } | |
| s1, s2, h1, h2 = parse_tagged_scribbles(value, (h, w), assignments) | |
| assert s1[8, 10] and s1[48, 70] # both strokes landed in scene 1 | |
| assert not s2.any() | |
| assert "pool" in h1 and "frame" in h1 and h2 == "" | |
| assert "top left" in h1 and "bottom right" in h1 # located hints | |
| # Defaults (no assignments): red->1, blue->2, no hints without tags | |
| s1d, s2d, h1d, h2d = parse_tagged_scribbles(value, (h, w), None) | |
| assert s1d[8, 10] and not s1d[48, 70] | |
| assert s2d[48, 70] | |
| assert h1d == "" and h2d == "" | |
| def test_region_phrase_thirds(): | |
| from app.scribble import region_phrase | |
| m = np.zeros((30, 30), bool) | |
| m[2:5, 2:5] = True | |
| assert region_phrase(m) == "top left" | |
| m2 = np.zeros((30, 30), bool) | |
| m2[13:17, 13:17] = True | |
| assert region_phrase(m2) == "center" | |
| assert region_phrase(np.zeros((30, 30), bool)) == "" | |
| def test_recover_hints_reach_prompts(): | |
| """Tagged-stroke hints are folded into the layer prompts on the scribble path.""" | |
| from app.restore import recover | |
| from tests.test_restore import _toy_scan | |
| pre = _toy_scan() | |
| h, w = pre.rgb.shape[:2] | |
| sa = np.zeros((h, w), bool); sa[10:20, 10:30] = True | |
| sb = np.zeros((h, w), bool); sb[60:70, 60:80] = True | |
| prompts: list[str] = [] | |
| recover(pre.rgb, pre.h_total, pre.confidence_mask, | |
| restore_fn=lambda rgb, p: prompts.append(p) or rgb, | |
| seeds_headline=sa, seeds_other=sb, | |
| hints_headline="the pool (top left)", hints_other="the frame (bottom right)") | |
| assert len(prompts) == 2 | |
| assert "the pool (top left)" in prompts[0] | |
| assert "the frame (bottom right)" in prompts[1] | |
| def test_recover_best_pair_used_as_anchor(): | |
| """A supplied best separation pair becomes the physics anchor (not the p45 split).""" | |
| from app.restore import recover | |
| from tests.test_restore import _toy_scan | |
| pre = _toy_scan() | |
| # Distinctive dominant layer: bright constant; other: dark constant | |
| bright = np.full_like(pre.rgb, 0.9) | |
| dark = np.full_like(pre.rgb, 0.1) | |
| res = recover(pre.rgb, pre.h_total, pre.confidence_mask, | |
| restore_fn=lambda rgb, p: rgb, best_pair=(bright, dark)) | |
| assert np.allclose(res.anchor_rgb, bright) # picked the larger-share member | |
| res2 = recover(pre.rgb, pre.h_total, pre.confidence_mask, | |
| restore_fn=lambda rgb, p: rgb, best_pair=(dark, bright)) | |
| assert np.allclose(res2.anchor_rgb, bright) # order-independent | |
| def test_hard_attribution_of_painted_pixels(): | |
| """Seed pixels are attributed EXACTLY (w=1/0) in the rendered split, while | |
| unpainted pixels keep the soft clamp (no full erasure).""" | |
| from app.scribble import split_by_scribbles | |
| from tests.test_restore import _toy_scan | |
| pre = _toy_scan() | |
| h, w = pre.rgb.shape[:2] | |
| sa = np.zeros((h, w), bool); sa[10:14, 10:14] = True | |
| sb = np.zeros((h, w), bool); sb[60:64, 60:64] = True | |
| lay_a, lay_b, _wf = split_by_scribbles(pre.rgb, pre.h_total, pre.confidence_mask, sa, sb) | |
| # Recover the implied w_r from the H split identity: h_a = w_r * h_total. | |
| # Independent check at seed pixels: layer A owns ALL of H at sa, none at sb. | |
| # Render is monotone in h, so compare via the pre-render arrays' proxy: | |
| # rebuild h arrays through the same identity the function guarantees. | |
| from densitometry import phi_display | |
| # At sa, layer B's render must carry ~zero luminance relative to layer A's; | |
| # at sb the reverse. Compare within the same location across the two layers. | |
| la, lb = phi_display(lay_a), phi_display(lay_b) | |
| assert float(lb[sa].mean()) < 0.05 * max(float(la[sa].mean()), 1e-6) or float(lb[sa].mean()) < 1e-3 | |
| assert float(la[sb].mean()) < 0.05 * max(float(lb[sb].mean()), 1e-6) or float(la[sb].mean()) < 1e-3 | |
| # --------------------------------------------------------------------------- | |
| # WP-18 D1 — scribble correctness | |
| # --------------------------------------------------------------------------- | |
| def test_both_scene_overlap_pixels_stay_contested(): | |
| """D1a: a pixel painted with BOTH scenes' colors keeps the soft value in both | |
| renders and counts toward the contested fraction; exclusive pixels still pin.""" | |
| from app.scribble import split_by_scribbles, propagate_w | |
| from tests.test_restore import _toy_scan | |
| pre = _toy_scan() | |
| h, w = pre.rgb.shape[:2] | |
| sa = np.zeros((h, w), bool); sa[10:20, 10:20] = True # A-only | |
| sb = np.zeros((h, w), bool); sb[60:70, 60:70] = True # B-only | |
| sa[40:50, 40:50] = True; sb[40:50, 40:50] = True # painted BOTH | |
| lay_a, lay_b, wf = split_by_scribbles(pre.rgb, pre.h_total, pre.confidence_mask, sa, sb) | |
| # Recompute the render weights independently to check the pin rule | |
| contested = sa & sb | |
| w_r = np.clip(wf, 0.12, 0.88) | |
| w_r[sa & ~contested] = 1.0 | |
| w_r[sb & ~contested] = 0.0 | |
| h_a = w_r * pre.h_total | |
| # A-only pixels: layer A owns the full exposure; B-only: none of it | |
| assert np.allclose(h_a[15, 15], pre.h_total[15, 15]) | |
| assert np.allclose(h_a[65, 65], 0.0) | |
| # Contested pixels: NEITHER side owns them fully in the render weights | |
| assert 0.12 - 1e-6 <= float(w_r[45, 45]) <= 0.88 + 1e-6 | |
| assert not np.isclose(float(w_r[45, 45]), 1.0) and not np.isclose(float(w_r[45, 45]), 0.0) | |
| # propagate_w does not hard-pin contested pixels to 0/1 either | |
| assert 0.0 < float(wf[45, 45]) < 1.0 | |
| def test_marks_unreadable_distinguishes_blended_from_empty(): | |
| """D1b: blended off-axis paint => unreadable warning; clean stroke => readable; | |
| nothing painted => not flagged.""" | |
| from app.scribble import marks_unreadable, parse_tagged_scribbles, PALETTE | |
| h, w = 40, 40 | |
| # NOTE: an equal red+blue blend (128,0,128) is exactly magenta's hue and is | |
| # legitimately read as a magenta stroke — the palette's known residual risk. | |
| # A muddy multi-color blend (gray-ish) is off EVERY palette axis: | |
| blended = np.zeros((h, w, 4), np.uint8) | |
| blended[5:30, 5:30, :3] = (120, 120, 120) | |
| blended[5:30, 5:30, 3] = 255 | |
| v_blend = {"background": None, "layers": [blended], "composite": None} | |
| s1, s2, _h1, _h2 = parse_tagged_scribbles(v_blend, (h, w)) | |
| assert not s1.any() and not s2.any() # gate rejected everything | |
| assert marks_unreadable(v_blend, (h, w)) # ...and we can SAY so | |
| clean = np.zeros((h, w, 4), np.uint8) | |
| clean[5:15, 5:15, :3] = PALETTE["red"] | |
| clean[5:15, 5:15, 3] = 255 | |
| v_clean = {"background": None, "layers": [clean], "composite": None} | |
| assert not marks_unreadable(v_clean, (h, w)) | |
| assert not marks_unreadable({"background": None, "layers": [], "composite": None}, (h, w)) | |
| assert not marks_unreadable(None, (h, w)) | |
| def test_trim_bbox_crops_stroke_layers_registered(): | |
| """D1c: with a trim bbox, a stroke at a known untrimmed landmark lands on the | |
| same landmark in trimmed coordinates (compared against a hand-computed crop).""" | |
| from app.scribble import parse_tagged_scribbles, PALETTE | |
| # Untrimmed canvas 100x100; working image = central crop [10:90, 20:80] -> 80x60 | |
| trim = (0.10, 0.90, 0.20, 0.80) | |
| th, tw = 80, 60 | |
| layer = np.zeros((100, 100, 4), np.uint8) | |
| layer[50:54, 50:54, :3] = PALETTE["red"] # landmark at untrimmed (50..54)^2 | |
| layer[50:54, 50:54, 3] = 255 | |
| v = {"background": None, "layers": [layer], "composite": None} | |
| s1, _s2, _h1, _h2 = parse_tagged_scribbles(v, (th, tw), None, trim_bbox_frac=trim) | |
| # Hand-computed: crop rows 10:90 cols 20:80 puts the stroke at rows 40:44, cols 30:34 | |
| assert s1[42, 32], "stroke missing at the hand-computed trimmed location" | |
| assert not s1[42, 50], "stroke leaked to an untrimmed-coordinate location" | |
| # Teeth: WITHOUT the bbox the same stroke lands misregistered (squashed resize) | |
| s1_no, _s2n, _h1n, _h2n = parse_tagged_scribbles(v, (th, tw), None) | |
| assert not s1_no[42, 32] or s1_no.sum() != s1.sum() | |
| # --------------------------------------------------------------------------- | |
| # WP-19 — annotated-copy markup + per-scene legends | |
| # --------------------------------------------------------------------------- | |
| def test_render_markup_touches_only_masked_pixels(): | |
| from app.scribble import render_markup, PALETTE | |
| rgb = np.full((20, 20, 3), 0.5, np.float32) | |
| m = np.zeros((20, 20), bool); m[5:8, 5:8] = True | |
| out = render_markup(rgb, {"red": m}, alpha=0.5) | |
| assert np.array_equal(out[~m], rgb[~m]) | |
| expect = 0.5 * 0.5 + 0.5 * np.asarray(PALETTE["red"], np.float32) / 255.0 | |
| assert np.allclose(out[m], expect, atol=1e-5) | |
| def test_markup_and_legends_perspective_flip(): | |
| """The same strokes read 'this photo' from their scene and 'the other photo' | |
| from the opposite scene; nothing painted -> no markup image.""" | |
| from app.scribble import markup_and_legends | |
| h, w = 24, 24 | |
| rgb = np.full((h, w, 3), 0.4, np.float32) | |
| layer = np.zeros((h, w, 4), np.uint8) | |
| layer[3:7, 3:12] = (255, 0, 0, 255) # red -> scene 1 | |
| layer[15:19, 12:20] = (0, 0, 255, 255) # blue -> scene 2 | |
| val = {"layers": [layer]} | |
| asg = {"red": {"scene": "1", "tag": "paintings"}, | |
| "blue": {"scene": "2", "tag": "women"}} | |
| mk, l1, l2 = markup_and_legends(val, (h, w), asg, rgb=rgb) | |
| assert mk is not None and mk.shape == rgb.shape | |
| assert "red strokes mark 'paintings' — belongs to this photo" in l1 | |
| assert "blue strokes mark 'women' — belongs to the other photo" in l1 | |
| assert "red strokes mark 'paintings' — belongs to the other photo" in l2 | |
| assert "blue strokes mark 'women' — belongs to this photo" in l2 | |
| mk_none, l1e, l2e = markup_and_legends({"layers": []}, (h, w), asg, rgb=rgb) | |
| assert mk_none is None and l1e == "" and l2e == "" | |
| # --------------------------------------------------------------------------- | |
| # WP-22 — tapped-object guidance (objects_guidance) and the tap UI handlers | |
| # --------------------------------------------------------------------------- | |
| def _toy_objects(h=40, w=60): | |
| m1 = np.zeros((h, w), bool); m1[5:12, 5:15] = True | |
| m2 = np.zeros((h, w), bool); m2[5:12, 20:30] = True | |
| m3 = np.zeros((h, w), bool); m3[25:35, 40:55] = True | |
| return [ | |
| {"mask": m1, "tag": "painting", "scene": "2"}, | |
| {"mask": m2, "tag": "painting", "scene": "2"}, | |
| {"mask": m3, "tag": "pool", "scene": "1"}, | |
| ] | |
| def test_objects_guidance_seeds_counts_and_legends(): | |
| from app.scribble import objects_guidance | |
| objs = _toy_objects() | |
| s1, s2, h1, h2, mk, l1, l2 = objects_guidance( | |
| objs, (40, 60), base_rgb=np.full((40, 60, 3), 0.5, np.float32) | |
| ) | |
| assert s1.sum() == objs[2]["mask"].sum() | |
| assert s2.sum() == (objs[0]["mask"] | objs[1]["mask"]).sum() | |
| assert "2× painting" in h2 and "the pool" in h1 | |
| # Legends: same objects, opposite perspectives, count included | |
| assert "this photo: 2× painting" in l2 and "other photo: 2× painting" in l1 | |
| assert "green shapes" in l1 and "cyan shapes" in l2 | |
| # Markup: fills only where masks are; elsewhere untouched | |
| assert mk is not None | |
| untouched = ~(s1 | s2) | |
| assert np.allclose(mk[untouched], 0.5, atol=1e-5) | |
| assert not np.allclose(mk[s1], 0.5, atol=1e-2) | |
| def test_objects_guidance_empty_and_geometry(): | |
| from app.scribble import objects_guidance | |
| s1, s2, h1, h2, mk, l1, l2 = objects_guidance( | |
| None, (40, 60), base_rgb=np.zeros((40, 60, 3), np.float32) | |
| ) | |
| assert not s1.any() and not s2.any() and mk is None and l1 == "" and h2 == "" | |
| # A mask in click geometry (80x120) lands registered in target (40x60) | |
| big = np.zeros((80, 120), bool); big[10:24, 10:30] = True | |
| s1, _s2, *_ = objects_guidance( | |
| [{"mask": big, "tag": "t", "scene": "1"}], (40, 60) | |
| ) | |
| assert s1.any() and abs(s1.mean() - big.mean()) < 0.02 | |
| def test_sam_ui_handlers_flow(monkeypatch): | |
| """Tap -> refine -> add -> summary, with a stubbed segmenter.""" | |
| import app.main as m | |
| import app.segment as seg | |
| calls = {"n": 0} | |
| def fake_point_mask(rgb, points): | |
| calls["n"] += 1 | |
| h, w = rgb.shape[:2] | |
| mask = np.zeros((h, w), bool) | |
| x, y = int(points[-1][0]), int(points[-1][1]) | |
| mask[max(0, y - 3):y + 3, max(0, x - 3):x + 3] = True | |
| return mask, 0.9 | |
| monkeypatch.setattr(seg, "point_mask", fake_point_mask) | |
| frame = np.full((50, 70, 3), 128, np.uint8) | |
| class Evt: # gr.SelectData stand-in | |
| index = [30, 20] | |
| disp, pts, pending, note = m.sam_click(frame, None, None, "Magic select (auto edges)", Evt()) | |
| assert len(pts) == 1 and pending.any() and "confidence 0.90" in note | |
| # Commit with a tag | |
| disp, pts, pending, objects, summary, note, tag_out = m.sam_add( | |
| frame, pts, pending, None, "painting", "Scene 2" | |
| ) | |
| assert len(objects) == 1 and objects[0]["scene"] == "2" and tag_out == "" | |
| assert "painting" in summary | |
| # Undo with no pending points clears cleanly | |
| disp, pts, pending, note = m.sam_undo(frame, [], objects, "Magic select (auto edges)") | |
| assert pts == [] and pending is None | |
| # Clear drops everything | |
| disp, pts, pending, objects, summary, note = m.sam_clear(frame) | |
| assert objects == [] and summary == "" | |
| def test_restore_handler_merges_tapped_objects(monkeypatch): | |
| """sam_objects reach recover as seeds + legend even with no brush strokes.""" | |
| import app.main as m | |
| import app.restore as r | |
| got = {} | |
| def fake_recover(observed_rgb, h_total=None, confidence_mask=None, **kw): | |
| got.update(kw) | |
| return r.RecoverResult(notes=kw.get("notes", [])) | |
| monkeypatch.setattr(r, "recover", fake_recover) | |
| from PIL import Image as PILImage | |
| rng = np.random.default_rng(3) | |
| upload = PILImage.fromarray((rng.random((80, 120, 3)) * 255).astype(np.uint8)) | |
| mask = np.zeros((80, 120), bool); mask[10:30, 10:40] = True | |
| objs = [{"mask": mask, "tag": "painting", "scene": "1"}] | |
| m.restore_best_scene( | |
| None, "Scene 1", "a", "b", "", None, | |
| "", "Scene 1", "", "Scene 1", "", "Scene 2", "", "Scene 2", | |
| upload, "Generic", "Auto", "auto-exposed", False, | |
| best_of_3=False, ref_photo_1=None, ref_photo_2=None, sam_objects=objs, | |
| ) | |
| assert got["seeds_headline"] is not None and got["seeds_headline"].any() | |
| assert "tapped objects" in (got["legend_headline"] or "") | |
| assert got["markup_rgb"] is not None | |
| # --------------------------------------------------------------------------- | |
| # WP-23 — dots-to-fill shapes and the targeted repair loop | |
| # --------------------------------------------------------------------------- | |
| def test_fill_mode_dots_close_a_shape(): | |
| import app.main as m | |
| frame = np.full((60, 90, 3), 100, np.uint8) | |
| class Evt: | |
| def __init__(self, xy): self.index = xy | |
| mode = "Fill shape from dots" | |
| disp, pts, mask, note = m.sam_click(frame, None, None, mode, Evt([10, 10])) | |
| assert mask is None and "add 2 more" in note | |
| disp, pts, mask, note = m.sam_click(frame, pts, None, mode, Evt([50, 10])) | |
| assert mask is None | |
| disp, pts, mask, note = m.sam_click(frame, pts, None, mode, Evt([30, 40])) | |
| assert mask is not None and mask.any() and "Shape filled" in note | |
| # The filled triangle centroid is inside; far corner is out | |
| assert mask[20, 30] and not mask[55, 85] | |
| # Undo reopens the shape | |
| disp, pts, mask, note = m.sam_undo(frame, pts, None, mode) | |
| assert len(pts) == 2 and mask is None | |
| def test_repair_region_outside_pixels_untouched(): | |
| from app.restore import repair_region | |
| base = np.full((80, 80, 3), 0.30, np.float32) | |
| mask = np.zeros((80, 80), bool) | |
| mask[20:40, 20:40] = True | |
| prompts = {} | |
| def hook(rgb, prompt): | |
| prompts["p"] = prompt | |
| return np.full_like(rgb, 0.90) | |
| notes: list[str] = [] | |
| out, meta = repair_region( | |
| base, mask, "a second woman in the chair", | |
| observed_rgb=np.zeros((80, 80, 3), np.float32), | |
| restore_fn=hook, notes=notes, | |
| ) | |
| assert out is not None | |
| assert "a second woman in the chair" in prompts["p"] | |
| assert "Image 3 is the original damaged" in prompts["p"] | |
| # Far from the region: EXACTLY the base (hard composite, feather decayed) | |
| assert np.allclose(out[:5, :5], 0.30, atol=1e-4) | |
| assert np.allclose(out[70:, 70:], 0.30, atol=1e-4) | |
| # Region core took the new content | |
| assert np.allclose(out[29:31, 29:31], 0.90, atol=0.02) | |
| assert any("pixel-identical" in n for n in notes) | |
| def test_repair_region_refuses_without_shape(): | |
| from app.restore import repair_region | |
| base = np.zeros((40, 40, 3), np.float32) | |
| notes: list[str] = [] | |
| out, _ = repair_region(base, np.zeros((40, 40), bool), "x", | |
| restore_fn=lambda r, p: r, notes=notes) | |
| assert out is None and any("3+ dots" in n for n in notes) | |
| def test_repair_apply_handler(monkeypatch): | |
| import app.main as m | |
| import app.restore as r | |
| from PIL import Image as PILImage | |
| def fake_repair(base, mask, instruction, observed=None, notes=None, **kw): | |
| (notes or []).append("Repaired the marked region (…pixel-identical…).") | |
| out = base.copy(); out[np.asarray(mask, bool)] = 0.9 | |
| return out, {"api_contacted": True} | |
| monkeypatch.setattr(r, "repair_region", fake_repair) | |
| main = PILImage.fromarray(np.full((50, 60, 3), 80, np.uint8)) | |
| second = PILImage.fromarray(np.full((50, 60, 3), 40, np.uint8)) | |
| pts = [[5, 5], [40, 5], [20, 30]] | |
| pil, sec_u, status, note = m.repair_apply( | |
| "Main scene", main, second, pts, "fix it", None, "old status") | |
| assert pil is not None and "🩹 Repaired the main scene" in status and "Done" in note | |
| # WP-24: the second scene is repairable too — only ITS image updates | |
| main_u, sec_pil, status2, note2 = m.repair_apply( | |
| "Second scene", main, second, pts, "fix it", None, "") | |
| assert sec_pil is not None and "🩹 Repaired the second scene" in status2 | |
| # Guards | |
| _p, _s, _st, note3 = m.repair_apply("Main scene", main, second, [[1, 1]], "x", None, "") | |
| assert "3+ dots" in note3 | |