feat: added impaint

This commit is contained in:
Hazel 2025-04-24 11:52:38 +02:00
parent ad38eef03b
commit 529e1af517

View File

@ -31,6 +31,19 @@ def impaint(raw_image: RawImage, image: Optional[np.ndarray] = None) -> np.ndarr
return cv2.inpaint(image, mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA) return cv2.inpaint(image, mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)
def do_generative_impaint(raw_image: RawImage, image: Optional[np.ndarray] = None) -> np.ndarray:
image = image if image is not None else raw_image.get_image()
# Create a mask where blacked-out areas are marked as 255 (white)
mask = np.zeros(image.shape[:2], dtype=np.uint8)
for (x, y, w, h) in raw_image.bounding_boxes:
mask[y:y+h, x:x+w] = 255
# Apply inpainting using the Telea method
return cv2.inpaint(image, mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)
def pixelate_regions(raw_image: RawImage, image: Optional[np.ndarray] = None, pixel_size: int = 10) -> np.ndarray: def pixelate_regions(raw_image: RawImage, image: Optional[np.ndarray] = None, pixel_size: int = 10) -> np.ndarray:
image = image.copy() if image is not None else raw_image.get_image().copy() image = image.copy() if image is not None else raw_image.get_image().copy()
@ -47,7 +60,7 @@ def pixelate_regions(raw_image: RawImage, image: Optional[np.ndarray] = None, pi
return image return image
def pixelate(to_detect: str, debug_drawings: bool = True): def pixelate(to_detect: str, generative_impaint: bool = True, debug_drawings: bool = True):
raw_image = RawImage(to_detect) raw_image = RawImage(to_detect)
step_dir = raw_image.get_dir("steps") step_dir = raw_image.get_dir("steps")
@ -67,7 +80,10 @@ def pixelate(to_detect: str, debug_drawings: bool = True):
step_1 = blackout(raw_image) step_1 = blackout(raw_image)
write_image(step_1, "step_1") write_image(step_1, "step_1")
step_2 = impaint(raw_image, image=step_1) if generative_impaint:
step_2 = do_generative_impaint(raw_image, image=step_1)
else:
step_2 = impaint(raw_image, image=step_1)
write_image(step_2, "step_2") write_image(step_2, "step_2")
step_3 = pixelate_regions(raw_image, image=step_2) step_3 = pixelate_regions(raw_image, image=step_2)