image-cleanup / create_test_images.py
kuko6's picture
Add image cleanup app
0e6ebc7
Raw
History Blame Contribute Delete
1.26 kB
import cv2
import numpy as np
from skimage import io
def add_debris(img, n_debris=40, seed=None):
rng = np.random.default_rng(seed)
h, w = img.shape[:2]
out = img.copy()
mask = np.zeros((h, w), dtype=np.uint8)
for _ in range(n_debris):
cx = rng.integers(10, w - 10)
cy = rng.integers(10, h - 10)
radius = int(rng.choice(np.concatenate([
rng.integers(20, 40, size=6), # small specks
rng.integers(40, 60, size=3), # medium
rng.integers(100, 120, size=1), # occasional big chunk
])))
# irregular blob via random polygon
n_pts = rng.integers(6, 14)
angles = np.linspace(0, 2 * np.pi, n_pts, endpoint=False)
angles += rng.uniform(-0.2, 0.2, n_pts)
radii = radius * rng.uniform(0.5, 1.4, n_pts)
pts = np.stack([
cx + radii * np.cos(angles),
cy + radii * np.sin(angles)
], axis=1).astype(np.int32)
cv2.fillPoly(out, [pts], color=(10, 10, 10))
cv2.fillPoly(mask, [pts], color=255)
return out, mask
img = io.imread("data/original.tiff")[:, :, :3]
for i in range(10):
dirty, mask = add_debris(img, n_debris=60)
io.imsave(f"data/test/dirty{i}.tiff", dirty)