feat: added impaint

secure_pixelation/pixelation_process.py

@@ -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)
 
 
+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:
     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
 
 
-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)
 
     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)
     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")
 
     step_3 = pixelate_regions(raw_image, image=step_2)