.gitignore

@@ -161,6 +161,4 @@ cython_debug/
 #.idea/
 .venv
 assets/*
-*.pt
+*.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,21 +21,3 @@ 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,6 +5,4 @@ from .pixelation_process import pixelate
 def cli():
     print(f"Running secure_pixelation")
 
-    pixelate("assets/human_detection/test.png", generative_impaint=True)
+    pixelate("assets/human_detection/test.png")
-    # pixelate("assets/human_detection/humans.png", generative_impaint=False)
-    # pixelate("assets/human_detection/rev1.png", generative_impaint=False)

secure_pixelation/pixelation_process.py

@@ -2,9 +2,6 @@ from __future__ import annotations
 
 from typing import Optional
 from pathlib import Path
-import subprocess
-import sys
-import os
 
 import cv2
 import numpy as np
@@ -21,18 +18,14 @@ def blackout(raw_image: RawImage) -> np.ndarray:
     return image
 
 
-def get_mask(raw_image: RawImage) -> np.ndarray:
+def impaint(raw_image: RawImage, image: Optional[np.ndarray] = None) -> 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)
+    # 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)
@@ -40,38 +33,15 @@ def quick_impaint(raw_image: RawImage, image: Optional[np.ndarray] = None) -> np
 
 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)
     
-    cv2.imwrite(str(lama_dict_in / "image.png"), raw_image.image)
+    # Create a mask where blacked-out areas are marked as 255 (white)
-    mask = get_mask(raw_image)
+    mask = np.zeros(image.shape[:2], dtype=np.uint8)
-    cv2.imwrite(str(lama_dict_in / "mask.png"), mask)
 
-    # Run LaMa inference (adjust path if needed)
+    for (x, y, w, h) in raw_image.bounding_boxes:
-    try:
+        mask[y:y+h, x:x+w] = 255
-        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)
 
-    # Load inpainted result
+    # Apply inpainting using the Telea method
-    result_path = lama_dict_out / "image.png"
+    return cv2.inpaint(image, mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)
-    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)
 
 
 
@@ -113,7 +83,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 = quick_impaint(raw_image, image=step_1)
+        step_2 = impaint(raw_image, image=step_1)
     write_image(step_2, "step_2")
 
     step_3 = pixelate_regions(raw_image, image=step_2)