Compare commits

..

No commits in common. "94b641cbd6d23788aeb506699f2703e1243f003e" and "529e1af517e05712875eb1dd1c9d580fb341d7bc" have entirely different histories.

4 changed files with 16 additions and 68 deletions

2
.gitignore vendored
View File

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

View File

@ -13,7 +13,7 @@ I first realized that a normal mosaic algorithm isn't safe AT ALL seeing this pr
```bash ```bash
# Step 1: Create and activate virtual environment # Step 1: Create and activate virtual environment
python3 -m venv .venv 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 # Step 2: Install the local Python program add the -e flag for development
pip install . pip install .
@ -21,21 +21,3 @@ pip install .
# Step 3: Run the secure-pixelation command # Step 3: Run the secure-pixelation command
secure-pixelation 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
```

View File

@ -5,6 +5,4 @@ from .pixelation_process import pixelate
def cli(): def cli():
print(f"Running secure_pixelation") 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)

View File

@ -2,9 +2,6 @@ from __future__ import annotations
from typing import Optional from typing import Optional
from pathlib import Path from pathlib import Path
import subprocess
import sys
import os
import cv2 import cv2
import numpy as np import numpy as np
@ -21,18 +18,14 @@ def blackout(raw_image: RawImage) -> np.ndarray:
return image 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() 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 # Apply inpainting using the Telea method
return cv2.inpaint(image, mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA) 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: 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() 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: if generative_impaint:
step_2 = do_generative_impaint(raw_image, image=step_1) step_2 = do_generative_impaint(raw_image, image=step_1)
else: else:
step_2 = quick_impaint(raw_image, image=step_1) 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)