wip

.gitignore

@@ -162,3 +162,5 @@ cython_debug/
 .venv
 assets/*
 *.pt
+
+big-lama

README.md

@@ -13,7 +13,7 @@ I first realized that a normal mosaic algorithm isn't safe AT ALL seeing this pr
 ```bash
 # Step 1: Create and activate virtual environment
 python3 -m venv .venv
-source venv/bin/activate
+source .venv/bin/activate
 
 # Step 2: Install the local Python program add the -e flag for development
 pip install .
@@ -21,3 +21,21 @@ pip install .
 # Step 3: Run the secure-pixelation command
 secure-pixelation
 ```
+
+## Setup LaMa
+
+This is the generative ai model to impaint the blacked out areas.
+
+```
+# get the pretrained models
+mkdir -p ./big-lama
+wget https://huggingface.co/smartywu/big-lama/resolve/main/big-lama.zip
+unzip big-lama.zip -d ./big-lama
+rm big-lama.zip
+
+# get the code to run the models
+cd big-lama
+git clone https://github.com/advimman/lama.git
+cd lama
+pip install -r requirements.txt
+```

secure_pixelation/__main__.py

@@ -5,4 +5,6 @@ from .pixelation_process import pixelate
 def cli():
     print(f"Running secure_pixelation")
 
-    pixelate("assets/human_detection/test.png")
+    pixelate("assets/human_detection/test.png", generative_impaint=True)
+    # pixelate("assets/human_detection/humans.png", generative_impaint=False)
+    # pixelate("assets/human_detection/rev1.png", generative_impaint=False)

secure_pixelation/pixelation_process.py

@@ -2,6 +2,9 @@ from __future__ import annotations
 
 from typing import Optional
 from pathlib import Path
+import subprocess
+import sys
+import os
 
 import cv2
 import numpy as np
@@ -18,30 +21,57 @@ def blackout(raw_image: RawImage) -> np.ndarray:
     return image
 
 
-def 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)
-
+def get_mask(raw_image: RawImage) -> np.ndarray:
+    mask = np.zeros(raw_image.image.shape[:2], dtype=np.uint8)
     for (x, y, w, h) in raw_image.bounding_boxes:
         mask[y:y+h, x:x+w] = 255
 
+    return mask
+
+
+def quick_impaint(raw_image: RawImage, image: Optional[np.ndarray] = None) -> np.ndarray:
+    image = image if image is not None else raw_image.get_image()
+    
+    mask = get_mask(raw_image)
+
     # Apply inpainting using the Telea method
     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()
+    lama_dict = raw_image.get_dir("steps") / "lama"
+    lama_dict.mkdir(exist_ok=True)
+    lama_dict_in = lama_dict / "in"
+    lama_dict_in.mkdir(exist_ok=True)
+    lama_dict_out = lama_dict / "out"
+    lama_dict_out.mkdir(exist_ok=True)
     
-    # Create a mask where blacked-out areas are marked as 255 (white)
-    mask = np.zeros(image.shape[:2], dtype=np.uint8)
+    cv2.imwrite(str(lama_dict_in / "image.png"), raw_image.image)
+    mask = get_mask(raw_image)
+    cv2.imwrite(str(lama_dict_in / "mask.png"), mask)
 
-    for (x, y, w, h) in raw_image.bounding_boxes:
-        mask[y:y+h, x:x+w] = 255
+    # Run LaMa inference (adjust path if needed)
+    try:
+        pwd = os.getcwd()
+        subprocess.run([
+            sys.executable, "lama/bin/predict.py",
+            f"model.path={pwd}/lama/models/big-lama",
+            f"indir={pwd}/{str(lama_dict_in)}",
+            f"outdir={pwd}/{str(lama_dict_out)}"
+        ], check=True)
+    except subprocess.CalledProcessError as e:
+        print(f"Error running LaMa: {e}")
+        print("falling back to non generative inpaint")
+        return quick_impaint(raw_image=raw_image, image=image)
 
-    # Apply inpainting using the Telea method
-    return cv2.inpaint(image, mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)
+    # Load inpainted result
+    result_path = lama_dict_out / "image.png"
+    if result_path.exists():
+        return cv2.imread(str(result_path))
+    else:
+        print("Inpainted result not found, falling back to non generative inpaint")
+        return quick_impaint(raw_image=raw_image, image=image)
 
 
 
@@ -83,7 +113,7 @@ def pixelate(to_detect: str, generative_impaint: bool = True, debug_drawings: bo
     if generative_impaint:
         step_2 = do_generative_impaint(raw_image, image=step_1)
     else:
-        step_2 = impaint(raw_image, image=step_1)
+        step_2 = quick_impaint(raw_image, image=step_1)
     write_image(step_2, "step_2")
 
     step_3 = pixelate_regions(raw_image, image=step_2)