mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-10-14 13:30:55 +02:00
perf(visible): crop MI-GAN around the mask so peak RAM is bounded by mark size
erase_migan fed the whole frame to the ONNX model, so peak RSS scaled with the upload (~0.6 GB at 4 MP up to ~2.4 GB at 25 MP). Mirror erase_lama: crop a padded region around the mask (pad = max(256, 2*bbox)), feed only that crop (at native resolution -- MI-GAN accepts arbitrary dims, unlike LaMa's fixed 512 square), and paste only masked pixels back. Peak RSS is now bounded by the mark size (~0.6-0.9 GB), so a memory-tight host (a 1-2 GB web worker) can run MI-GAN on a 25 MP upload. Fill quality is unchanged: verified by eye on real Gemini/Doubao marks plus a ground-truth reconstruction sweep -- a tighter view if anything reduces the GAN's hallucination of large background structure. Extract the shared padded-crop-box math into _padded_crop_box (used by both erase_lama and erase_migan). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
1 parent
f69fedcf2d
commit
190dc89d23
4 files changed
+84
-26
No files matched your search
@@ -171,6 +171,7 @@ class TestMiganWrapper:
|
||||
|
||||
def run(self, _outputs, feeds):
|
||||
self.outer.captured["mask"] = feeds["mask"]
|
||||
self.outer.captured["image_shape"] = feeds["image"].shape
|
||||
img = feeds["image"] # (1,3,H,W) RGB
|
||||
red = np.zeros_like(img)
|
||||
red[:, 0] = 255 # pure red in RGB
|
||||
@@ -191,6 +192,21 @@ class TestMiganWrapper:
|
||||
assert m[50, 50] == 0
|
||||
assert m[10, 10] == 255
|
||||
|
||||
@pytest.mark.usefixtures("_fake_migan")
|
||||
def test_crops_around_mask_so_onnx_input_is_bounded(self):
|
||||
# Large frame, small corner mark: the tensor fed to MI-GAN is the padded
|
||||
# CROP (pad = max(256, 2*bbox)), not the full image -- this is what holds the
|
||||
# ONNX working set roughly constant on big uploads instead of scaling with
|
||||
# the image (the memory fix). Untouched pixels stay exact; the mark is filled.
|
||||
img = np.full((2000, 3000, 3), 120, np.uint8)
|
||||
out = erase(img, boxes=[(2900, 1900, 60, 60)], backend="migan", dilate=0)
|
||||
assert out.shape == img.shape
|
||||
_, _, fh, fw = self.captured["image_shape"]
|
||||
assert fh < 700 # crop height, not the 2000px frame
|
||||
assert fw < 700 # crop width, not the 3000px frame
|
||||
assert tuple(int(v) for v in out[1930, 2930]) == (0, 0, 255) # mark -> red fill
|
||||
assert np.array_equal(out[:100, :100], img[:100, :100]) # far corner untouched
|
||||
|
||||
@pytest.mark.usefixtures("_fake_migan")
|
||||
def test_grayscale_2d_does_not_raise(self):
|
||||
gray = np.full((100, 100), 120, np.uint8)
|
||||
|
||||
Reference in new issue
Block a user