mirror of
https://github.com/ashim-hq/ashim
synced 2026-04-21 13:37:52 +00:00
08a7ffe403
- Added model mismatch warnings in colorize, enhance-faces, and upscale routes. - Improved error handling in colorize, enhance_faces, remove_bg, restore, and upscale scripts with detailed logging. - Updated Dockerfile to align NCCL versions for compatibility. - Introduced a new full tool audit script to test all tools for functionality and GPU usage. - Created Playwright E2E tests for GPU-dependent tools to ensure proper functionality and performance.
654 lines
25 KiB
Python
654 lines
25 KiB
Python
"""AI photo restoration pipeline.
|
|
|
|
Multi-step pipeline for restoring old and damaged photos:
|
|
1. Scratch & damage detection (morphological analysis)
|
|
2. Damage inpainting (LaMa ONNX)
|
|
3. Face enhancement (CodeFormer ONNX)
|
|
4. Noise reduction (OpenCV NLMeans)
|
|
5. Optional B&W colorization (DDColor ONNX)
|
|
"""
|
|
import sys
|
|
import json
|
|
import os
|
|
import numpy as np
|
|
import cv2
|
|
from PIL import Image
|
|
|
|
|
|
def emit_progress(percent, stage):
|
|
"""Emit structured progress to stderr for bridge.ts to capture."""
|
|
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
|
|
|
|
|
|
# ── Model paths ───────────────────────────────────────────────────────
|
|
|
|
LAMA_MODEL_DIR = os.environ.get("LAMA_MODEL_DIR", "/opt/models/lama")
|
|
LAMA_MODEL_PATH = os.path.join(LAMA_MODEL_DIR, "lama_fp32.onnx")
|
|
LAMA_LOCAL_CACHE = os.path.join(os.path.expanduser("~"), ".cache", "ashim", "lama")
|
|
LAMA_LOCAL_PATH = os.path.join(LAMA_LOCAL_CACHE, "lama_fp32.onnx")
|
|
|
|
CODEFORMER_MODEL_DIR = os.environ.get("CODEFORMER_MODEL_DIR", "/opt/models/codeformer")
|
|
CODEFORMER_MODEL_PATH = os.path.join(CODEFORMER_MODEL_DIR, "codeformer.onnx")
|
|
CODEFORMER_LOCAL_CACHE = os.path.join(
|
|
os.path.expanduser("~"), ".cache", "ashim", "codeformer"
|
|
)
|
|
CODEFORMER_LOCAL_PATH = os.path.join(CODEFORMER_LOCAL_CACHE, "codeformer.onnx")
|
|
|
|
DDCOLOR_MODEL_PATH = os.environ.get(
|
|
"DDCOLOR_MODEL_PATH", "/opt/models/ddcolor/ddcolor.onnx"
|
|
)
|
|
|
|
LAMA_MODEL_SIZE = 512
|
|
CODEFORMER_SIZE = 512
|
|
|
|
|
|
# ── Scratch detection ─────────────────────────────────────────────────
|
|
|
|
def detect_scratches(img_bgr, sensitivity="medium"):
|
|
"""Detect scratches, tears, and spots using morphological analysis.
|
|
|
|
Uses top-hat and black-hat transforms with oriented line kernels
|
|
to find both bright and dark linear structures at multiple scales
|
|
and angles. Returns a binary mask (255 = damage, 0 = clean).
|
|
"""
|
|
gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
|
|
|
|
# CLAHE for local contrast enhancement to reveal faint scratches
|
|
clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
|
|
enhanced = clahe.apply(gray)
|
|
|
|
# Sensitivity controls the detection threshold
|
|
thresh_map = {"light": 170, "medium": 130, "heavy": 90}
|
|
thresh = thresh_map.get(sensitivity, 130)
|
|
|
|
h, w = gray.shape
|
|
base_dim = min(h, w)
|
|
|
|
# Scale kernel sizes to image resolution
|
|
kernel_sizes = [
|
|
max(9, base_dim // 80),
|
|
max(15, base_dim // 50),
|
|
max(25, base_dim // 30),
|
|
]
|
|
|
|
mask = np.zeros_like(gray)
|
|
|
|
for ksize in kernel_sizes:
|
|
ksize = ksize | 1 # ensure odd
|
|
for angle in [0, 45, 90, 135]:
|
|
kernel = _make_line_kernel(ksize, angle)
|
|
|
|
# Black-hat: detects dark structures (dark scratches on light areas)
|
|
blackhat = cv2.morphologyEx(enhanced, cv2.MORPH_BLACKHAT, kernel)
|
|
# Top-hat: detects bright structures (light scratches on dark areas)
|
|
tophat = cv2.morphologyEx(enhanced, cv2.MORPH_TOPHAT, kernel)
|
|
|
|
combined = cv2.add(blackhat, tophat)
|
|
_, binary = cv2.threshold(combined, thresh, 255, cv2.THRESH_BINARY)
|
|
mask = cv2.bitwise_or(mask, binary)
|
|
|
|
# Clean up: remove isolated noise pixels
|
|
kernel_open = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
|
|
mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel_open)
|
|
|
|
# Connect nearby scratch segments
|
|
kernel_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
|
|
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel_close)
|
|
|
|
# Dilate to include scratch edges for cleaner inpainting
|
|
kernel_dilate = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
|
|
mask = cv2.dilate(mask, kernel_dilate, iterations=1)
|
|
|
|
return mask
|
|
|
|
|
|
def _make_line_kernel(size, angle):
|
|
"""Create an oriented line structuring element."""
|
|
kernel = np.zeros((size, size), np.uint8)
|
|
mid = size // 2
|
|
if angle == 0:
|
|
kernel[mid, :] = 1
|
|
elif angle == 90:
|
|
kernel[:, mid] = 1
|
|
elif angle == 45:
|
|
for i in range(size):
|
|
kernel[i, i] = 1
|
|
elif angle == 135:
|
|
for i in range(size):
|
|
kernel[i, size - 1 - i] = 1
|
|
return kernel
|
|
|
|
|
|
# ── LaMa inpainting ──────────────────────────────────────────────────
|
|
|
|
def _get_lama_path():
|
|
"""Resolve LaMa model path, downloading if needed."""
|
|
if os.path.exists(LAMA_MODEL_PATH):
|
|
return LAMA_MODEL_PATH
|
|
if os.path.exists(LAMA_LOCAL_PATH):
|
|
return LAMA_LOCAL_PATH
|
|
# Auto-download for local dev
|
|
os.makedirs(LAMA_LOCAL_CACHE, exist_ok=True)
|
|
import urllib.request
|
|
url = "https://huggingface.co/Carve/LaMa-ONNX/resolve/main/lama_fp32.onnx"
|
|
urllib.request.urlretrieve(url, LAMA_LOCAL_PATH)
|
|
return LAMA_LOCAL_PATH
|
|
|
|
|
|
def inpaint_damage(img_bgr, mask):
|
|
"""Inpaint damaged areas using LaMa ONNX model.
|
|
|
|
Args:
|
|
img_bgr: Input BGR image as numpy array.
|
|
mask: Binary mask (255 = damage to repair, 0 = keep).
|
|
|
|
Returns:
|
|
Restored BGR image with damage inpainted.
|
|
"""
|
|
import onnxruntime as ort
|
|
|
|
model_path = _get_lama_path()
|
|
providers = ["CPUExecutionProvider"]
|
|
if "CUDAExecutionProvider" in ort.get_available_providers():
|
|
providers.insert(0, "CUDAExecutionProvider")
|
|
|
|
session = ort.InferenceSession(model_path, providers=providers)
|
|
|
|
orig_h, orig_w = img_bgr.shape[:2]
|
|
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
|
|
|
|
# Preprocess image: resize to 512x512, normalize to [0,1], NCHW
|
|
img_resized = cv2.resize(img_rgb, (LAMA_MODEL_SIZE, LAMA_MODEL_SIZE))
|
|
img_input = img_resized.astype(np.float32) / 255.0
|
|
img_input = np.transpose(img_input, (2, 0, 1))[np.newaxis, ...] # (1,3,512,512)
|
|
|
|
# Preprocess mask: resize to 512x512, binary, NCHW
|
|
mask_resized = cv2.resize(mask, (LAMA_MODEL_SIZE, LAMA_MODEL_SIZE),
|
|
interpolation=cv2.INTER_NEAREST)
|
|
mask_binary = (mask_resized > 127).astype(np.float32)
|
|
mask_input = mask_binary[np.newaxis, np.newaxis, ...] # (1,1,512,512)
|
|
|
|
# Run inference
|
|
outputs = session.run(None, {"image": img_input, "mask": mask_input})
|
|
result = outputs[0][0] # (3, 512, 512)
|
|
result = np.transpose(result, (1, 2, 0)) # (512, 512, 3)
|
|
result = np.clip(result, 0, 255).astype(np.uint8)
|
|
|
|
# Resize inpainted result back to original dimensions
|
|
inpainted = cv2.resize(result, (orig_w, orig_h), interpolation=cv2.INTER_LANCZOS4)
|
|
|
|
# Feathered composite: preserve quality outside mask, blend at edges
|
|
mask_full = mask.astype(np.float32) / 255.0
|
|
feather_r = max(3, min(orig_w, orig_h) // 200)
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (feather_r, feather_r))
|
|
dilated = cv2.dilate(mask_full, kernel, iterations=1)
|
|
blur_size = feather_r * 2 + 1
|
|
alpha = cv2.GaussianBlur(dilated, (blur_size, blur_size), 0)
|
|
alpha = np.clip(alpha, 0.0, 1.0)[:, :, np.newaxis]
|
|
|
|
# Composite in RGB space, then convert back to BGR
|
|
inpainted_rgb = inpainted
|
|
original_rgb = img_rgb
|
|
composited = (original_rgb.astype(np.float32) * (1.0 - alpha) +
|
|
inpainted_rgb.astype(np.float32) * alpha)
|
|
composited = np.clip(composited, 0, 255).astype(np.uint8)
|
|
|
|
return cv2.cvtColor(composited, cv2.COLOR_RGB2BGR)
|
|
|
|
|
|
# ── CodeFormer face enhancement ──────────────────────────────────────
|
|
|
|
def _get_codeformer_path():
|
|
"""Resolve CodeFormer ONNX model path, downloading if needed."""
|
|
if os.path.exists(CODEFORMER_MODEL_PATH):
|
|
return CODEFORMER_MODEL_PATH
|
|
if os.path.exists(CODEFORMER_LOCAL_PATH):
|
|
return CODEFORMER_LOCAL_PATH
|
|
|
|
# Auto-download for local dev
|
|
os.makedirs(CODEFORMER_LOCAL_CACHE, exist_ok=True)
|
|
emit_progress(35, "Downloading CodeFormer model")
|
|
from huggingface_hub import hf_hub_download
|
|
hf_hub_download(
|
|
repo_id="facefusion/models-3.0.0",
|
|
filename="codeformer.onnx",
|
|
local_dir=CODEFORMER_LOCAL_CACHE,
|
|
)
|
|
return CODEFORMER_LOCAL_PATH
|
|
|
|
|
|
# ── Model path for new mp.tasks API ─────────────────────────────────
|
|
|
|
_FACE_DETECT_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/latest/blaze_face_short_range.tflite"
|
|
_FACE_DETECT_DOCKER_PATH = "/opt/models/mediapipe/blaze_face_short_range.tflite"
|
|
_FACE_DETECT_LOCAL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models")
|
|
_FACE_DETECT_LOCAL_PATH = os.path.join(_FACE_DETECT_LOCAL_DIR, "blaze_face_short_range.tflite")
|
|
|
|
|
|
def _ensure_face_detect_model():
|
|
"""Resolve face detector model. Docker path first, then local dev."""
|
|
if os.path.exists(_FACE_DETECT_DOCKER_PATH):
|
|
return _FACE_DETECT_DOCKER_PATH
|
|
if os.path.exists(_FACE_DETECT_LOCAL_PATH):
|
|
return _FACE_DETECT_LOCAL_PATH
|
|
os.makedirs(_FACE_DETECT_LOCAL_DIR, exist_ok=True)
|
|
import urllib.request
|
|
emit_progress(15, "Downloading face detection model")
|
|
urllib.request.urlretrieve(_FACE_DETECT_MODEL_URL, _FACE_DETECT_LOCAL_PATH)
|
|
return _FACE_DETECT_LOCAL_PATH
|
|
|
|
|
|
def enhance_faces(img_bgr, fidelity=0.7):
|
|
"""Enhance faces in the image using CodeFormer ONNX.
|
|
|
|
1. Detect faces with MediaPipe
|
|
2. Crop each face with generous padding
|
|
3. Run CodeFormer ONNX inference
|
|
4. Paste enhanced face back with feathered blending
|
|
|
|
Args:
|
|
img_bgr: Input BGR image.
|
|
fidelity: 0.0 = aggressive enhancement, 1.0 = faithful to original.
|
|
|
|
Returns:
|
|
Tuple of (enhanced BGR image, number of faces found).
|
|
"""
|
|
import mediapipe as mp
|
|
import onnxruntime as ort
|
|
|
|
# Detect faces
|
|
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
|
|
ih, iw = img_bgr.shape[:2]
|
|
|
|
try:
|
|
mp_face = mp.solutions.face_detection
|
|
detections = []
|
|
for model_sel in [0, 1]:
|
|
detector = mp_face.FaceDetection(
|
|
model_selection=model_sel, min_detection_confidence=0.4
|
|
)
|
|
results = detector.process(img_rgb)
|
|
detector.close()
|
|
if results.detections:
|
|
detections = results.detections
|
|
break
|
|
|
|
if not detections:
|
|
return img_bgr, 0
|
|
|
|
face_boxes = []
|
|
for detection in detections:
|
|
bbox = detection.location_data.relative_bounding_box
|
|
face_boxes.append({
|
|
"x": int(bbox.xmin * iw),
|
|
"y": int(bbox.ymin * ih),
|
|
"w": int(bbox.width * iw),
|
|
"h": int(bbox.height * ih),
|
|
})
|
|
|
|
except AttributeError:
|
|
# mediapipe >= 0.10.30 removed mp.solutions, use tasks API
|
|
model_path = _ensure_face_detect_model()
|
|
options = mp.tasks.vision.FaceDetectorOptions(
|
|
base_options=mp.tasks.BaseOptions(model_asset_path=model_path),
|
|
running_mode=mp.tasks.vision.RunningMode.IMAGE,
|
|
min_detection_confidence=0.4,
|
|
)
|
|
fd = mp.tasks.vision.FaceDetector.create_from_options(options)
|
|
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=img_rgb)
|
|
result = fd.detect(mp_image)
|
|
fd.close()
|
|
|
|
if not result.detections:
|
|
return img_bgr, 0
|
|
|
|
face_boxes = []
|
|
for detection in result.detections:
|
|
bbox = detection.bounding_box
|
|
face_boxes.append({
|
|
"x": bbox.origin_x,
|
|
"y": bbox.origin_y,
|
|
"w": bbox.width,
|
|
"h": bbox.height,
|
|
})
|
|
|
|
# Load CodeFormer model
|
|
model_path = _get_codeformer_path()
|
|
providers = ["CPUExecutionProvider"]
|
|
if "CUDAExecutionProvider" in ort.get_available_providers():
|
|
providers.insert(0, "CUDAExecutionProvider")
|
|
|
|
session = ort.InferenceSession(model_path, providers=providers)
|
|
input_names = [inp.name for inp in session.get_inputs()]
|
|
|
|
result = img_bgr.copy()
|
|
faces_enhanced = 0
|
|
|
|
for face_box in face_boxes:
|
|
x = face_box["x"]
|
|
y = face_box["y"]
|
|
w = face_box["w"]
|
|
h = face_box["h"]
|
|
|
|
# Skip very small faces (under 48px) - enhancement won't help
|
|
if w < 48 or h < 48:
|
|
continue
|
|
|
|
# Expand bounding box by ~80% for hair, forehead, chin
|
|
pad_x = int(w * 0.8)
|
|
pad_y = int(h * 0.8)
|
|
x1 = max(0, x - pad_x)
|
|
y1 = max(0, y - pad_y)
|
|
x2 = min(iw, x + w + pad_x)
|
|
y2 = min(ih, y + h + pad_y)
|
|
|
|
# Crop face region
|
|
face_crop = img_bgr[y1:y2, x1:x2].copy()
|
|
crop_h, crop_w = face_crop.shape[:2]
|
|
|
|
# Resize to 512x512 for CodeFormer
|
|
face_resized = cv2.resize(face_crop, (CODEFORMER_SIZE, CODEFORMER_SIZE),
|
|
interpolation=cv2.INTER_LANCZOS4)
|
|
|
|
# Preprocess: BGR -> RGB, normalize to [-1, 1]
|
|
face_rgb = cv2.cvtColor(face_resized, cv2.COLOR_BGR2RGB)
|
|
face_input = face_rgb.astype(np.float32) / 255.0
|
|
face_input = (face_input - 0.5) / 0.5
|
|
face_input = np.transpose(face_input, (2, 0, 1))
|
|
face_input = np.expand_dims(face_input, 0) # (1, 3, 512, 512)
|
|
|
|
# Build model inputs
|
|
model_inputs = {}
|
|
for name in input_names:
|
|
if name == "input":
|
|
model_inputs[name] = face_input.astype(np.float32)
|
|
elif name == "weight":
|
|
model_inputs[name] = np.array([fidelity]).astype(np.float64)
|
|
|
|
# Run inference
|
|
try:
|
|
output = session.run(None, model_inputs)[0][0] # (3, 512, 512)
|
|
except Exception as e:
|
|
print(f"[restore] CodeFormer inference failed for face {i}: {e}", file=sys.stderr, flush=True)
|
|
continue
|
|
|
|
# Postprocess: [-1, 1] -> [0, 255], RGB -> BGR
|
|
output = np.clip(output, -1, 1)
|
|
output = (output + 1) / 2
|
|
output = np.transpose(output, (1, 2, 0)) # (512, 512, 3)
|
|
output = (output * 255.0).clip(0, 255).astype(np.uint8)
|
|
output_bgr = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
|
|
|
|
# Resize back to original crop size
|
|
enhanced_crop = cv2.resize(output_bgr, (crop_w, crop_h),
|
|
interpolation=cv2.INTER_LANCZOS4)
|
|
|
|
# Create feathered elliptical mask for smooth blending
|
|
blend_mask = np.zeros((crop_h, crop_w), dtype=np.float32)
|
|
center = (crop_w // 2, crop_h // 2)
|
|
axes = (int(crop_w * 0.42), int(crop_h * 0.45))
|
|
cv2.ellipse(blend_mask, center, axes, 0, 0, 360, 1.0, -1)
|
|
|
|
# Feather the mask edges
|
|
blur_r = max(5, min(crop_w, crop_h) // 8) | 1
|
|
blend_mask = cv2.GaussianBlur(blend_mask, (blur_r, blur_r), 0)
|
|
blend_mask = blend_mask[:, :, np.newaxis]
|
|
|
|
# Blend enhanced face into result
|
|
face_region = result[y1:y2, x1:x2].astype(np.float32)
|
|
blended = face_region * (1.0 - blend_mask) + enhanced_crop.astype(np.float32) * blend_mask
|
|
result[y1:y2, x1:x2] = np.clip(blended, 0, 255).astype(np.uint8)
|
|
faces_enhanced += 1
|
|
|
|
return result, faces_enhanced
|
|
|
|
|
|
# ── Denoising ─────────────────────────────────────────────────────────
|
|
|
|
def denoise_image(img_bgr, strength=40):
|
|
"""Apply noise reduction using Non-Local Means in LAB color space.
|
|
|
|
Processes luminance and chrominance channels independently
|
|
for better color preservation.
|
|
|
|
Args:
|
|
img_bgr: Input BGR image.
|
|
strength: 0-100, higher = more aggressive denoising.
|
|
|
|
Returns:
|
|
Denoised BGR image.
|
|
"""
|
|
if strength <= 0:
|
|
return img_bgr
|
|
|
|
# Map 0-100 to NLMeans filter strength
|
|
h = 3 + (strength / 100) * 12 # 3-15
|
|
|
|
lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)
|
|
l_ch, a_ch, b_ch = cv2.split(lab)
|
|
|
|
# Denoise luminance channel
|
|
l_ch = cv2.fastNlMeansDenoising(l_ch, None, h, 7, 21)
|
|
|
|
# Lightly denoise color channels to remove chroma noise
|
|
color_h = h * 0.5
|
|
if color_h > 1:
|
|
a_ch = cv2.fastNlMeansDenoising(a_ch, None, color_h, 7, 21)
|
|
b_ch = cv2.fastNlMeansDenoising(b_ch, None, color_h, 7, 21)
|
|
|
|
result = cv2.merge([l_ch, a_ch, b_ch])
|
|
return cv2.cvtColor(result, cv2.COLOR_LAB2BGR)
|
|
|
|
|
|
# ── B&W detection ────────────────────────────────────────────────────
|
|
|
|
def is_grayscale(img_bgr):
|
|
"""Detect if an image is grayscale/B&W.
|
|
|
|
Checks if color channels are nearly identical by measuring
|
|
the standard deviation of channel differences.
|
|
"""
|
|
if len(img_bgr.shape) == 2:
|
|
return True
|
|
if img_bgr.shape[2] == 1:
|
|
return True
|
|
|
|
b, g, r = cv2.split(img_bgr)
|
|
diff_rg = np.abs(r.astype(np.float32) - g.astype(np.float32)).mean()
|
|
diff_rb = np.abs(r.astype(np.float32) - b.astype(np.float32)).mean()
|
|
diff_gb = np.abs(g.astype(np.float32) - b.astype(np.float32)).mean()
|
|
avg_diff = (diff_rg + diff_rb + diff_gb) / 3
|
|
|
|
return bool(avg_diff < 5.0)
|
|
|
|
|
|
# ── DDColor colorization ─────────────────────────────────────────────
|
|
|
|
def colorize_bw(img_bgr, intensity=0.85):
|
|
"""Colorize a B&W image using DDColor ONNX.
|
|
|
|
Reuses the DDColor model that the colorize tool already downloads.
|
|
"""
|
|
import onnxruntime as ort
|
|
|
|
if not os.path.exists(DDCOLOR_MODEL_PATH):
|
|
return img_bgr, False
|
|
|
|
providers = ["CPUExecutionProvider"]
|
|
try:
|
|
from gpu import gpu_available
|
|
if gpu_available():
|
|
providers.insert(0, "CUDAExecutionProvider")
|
|
except ImportError as e:
|
|
print(f"[restore] GPU detection unavailable: {e}", file=sys.stderr, flush=True)
|
|
|
|
session = ort.InferenceSession(DDCOLOR_MODEL_PATH, providers=providers)
|
|
input_name = session.get_inputs()[0].name
|
|
input_shape = session.get_inputs()[0].shape
|
|
model_size = (
|
|
input_shape[2]
|
|
if len(input_shape) == 4 and isinstance(input_shape[2], int)
|
|
else 512
|
|
)
|
|
|
|
orig_h, orig_w = img_bgr.shape[:2]
|
|
img_lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)
|
|
orig_l = img_lab[:, :, 0].astype(np.float32)
|
|
|
|
# Prepare input
|
|
img_resized = cv2.resize(img_bgr, (model_size, model_size))
|
|
img_float = img_resized.astype(np.float32) / 255.0
|
|
img_nchw = np.transpose(img_float, (2, 0, 1))[np.newaxis, ...]
|
|
|
|
output = session.run(None, {input_name: img_nchw})[0]
|
|
ab_pred = output[0] # (2, model_size, model_size)
|
|
|
|
# Resize ab channels back to original
|
|
ab_resized = np.zeros((2, orig_h, orig_w), dtype=np.float32)
|
|
for i in range(2):
|
|
ab_resized[i] = cv2.resize(ab_pred[i], (orig_w, orig_h))
|
|
|
|
ab_a = np.clip(ab_resized[0], -128, 127)
|
|
ab_b = np.clip(ab_resized[1], -128, 127)
|
|
|
|
# Apply intensity blending
|
|
if intensity < 1.0:
|
|
orig_a = img_lab[:, :, 1].astype(np.float32) - 128.0
|
|
orig_b = img_lab[:, :, 2].astype(np.float32) - 128.0
|
|
ab_a = orig_a * (1 - intensity) + ab_a * intensity
|
|
ab_b = orig_b * (1 - intensity) + ab_b * intensity
|
|
|
|
result_lab = np.zeros((orig_h, orig_w, 3), dtype=np.uint8)
|
|
result_lab[:, :, 0] = np.clip(orig_l, 0, 255).astype(np.uint8)
|
|
result_lab[:, :, 1] = np.clip(ab_a + 128.0, 0, 255).astype(np.uint8)
|
|
result_lab[:, :, 2] = np.clip(ab_b + 128.0, 0, 255).astype(np.uint8)
|
|
|
|
return cv2.cvtColor(result_lab, cv2.COLOR_LAB2BGR), True
|
|
|
|
|
|
# ── Main pipeline ─────────────────────────────────────────────────────
|
|
|
|
def main():
|
|
input_path = sys.argv[1]
|
|
output_path = sys.argv[2]
|
|
settings = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {}
|
|
|
|
mode = settings.get("mode", "auto")
|
|
scratch_removal = settings.get("scratchRemoval", True)
|
|
face_enhancement = settings.get("faceEnhancement", True)
|
|
fidelity = float(settings.get("fidelity", 0.7))
|
|
do_denoise = settings.get("denoise", True)
|
|
denoise_strength = float(settings.get("denoiseStrength", 40))
|
|
do_colorize = settings.get("colorize", False)
|
|
|
|
# Mode presets override individual settings
|
|
if mode == "light":
|
|
scratch_sensitivity = "light"
|
|
if denoise_strength > 30:
|
|
denoise_strength = 30
|
|
elif mode == "heavy":
|
|
scratch_sensitivity = "heavy"
|
|
if denoise_strength < 60:
|
|
denoise_strength = 60
|
|
else:
|
|
scratch_sensitivity = "medium"
|
|
|
|
try:
|
|
emit_progress(5, "Opening image")
|
|
img_bgr = cv2.imread(input_path, cv2.IMREAD_COLOR)
|
|
if img_bgr is None:
|
|
pil_img = Image.open(input_path).convert("RGB")
|
|
img_bgr = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
|
|
|
|
orig_h, orig_w = img_bgr.shape[:2]
|
|
result = img_bgr.copy()
|
|
steps_applied = []
|
|
|
|
# ── Step 1: Analyze photo ────────────────────────────────
|
|
emit_progress(8, "Analyzing photo")
|
|
bw_detected = is_grayscale(img_bgr)
|
|
scratch_mask = None
|
|
scratch_coverage = 0.0
|
|
|
|
# ── Step 2: Scratch detection & inpainting ───────────────
|
|
if scratch_removal:
|
|
emit_progress(10, "Detecting damage")
|
|
scratch_mask = detect_scratches(result, scratch_sensitivity)
|
|
scratch_pixels = np.count_nonzero(scratch_mask)
|
|
total_pixels = scratch_mask.shape[0] * scratch_mask.shape[1]
|
|
scratch_coverage = float(scratch_pixels / total_pixels)
|
|
|
|
if scratch_coverage > 0.001: # At least 0.1% coverage
|
|
emit_progress(15, f"Repairing damage ({scratch_coverage:.1%} affected)")
|
|
result = inpaint_damage(result, scratch_mask)
|
|
steps_applied.append("scratch_removal")
|
|
emit_progress(30, "Damage repaired")
|
|
else:
|
|
emit_progress(15, "No significant damage detected")
|
|
else:
|
|
emit_progress(15, "Scratch removal disabled")
|
|
|
|
# ── Step 3: Face enhancement ─────────────────────────────
|
|
faces_found = 0
|
|
if face_enhancement:
|
|
emit_progress(35, "Detecting faces")
|
|
try:
|
|
result, faces_found = enhance_faces(result, fidelity)
|
|
if faces_found > 0:
|
|
steps_applied.append("face_enhancement")
|
|
emit_progress(65, f"Enhanced {faces_found} face{'s' if faces_found != 1 else ''}")
|
|
else:
|
|
emit_progress(65, "No faces detected")
|
|
except Exception as e:
|
|
emit_progress(65, f"Face enhancement skipped: {str(e)[:40]}")
|
|
else:
|
|
emit_progress(65, "Face enhancement disabled")
|
|
|
|
# ── Step 4: Noise reduction ──────────────────────────────
|
|
if do_denoise and denoise_strength > 0:
|
|
emit_progress(70, "Reducing noise")
|
|
result = denoise_image(result, denoise_strength)
|
|
steps_applied.append("denoise")
|
|
emit_progress(80, "Noise reduced")
|
|
else:
|
|
emit_progress(80, "Denoising disabled")
|
|
|
|
# ── Step 5: Colorization ─────────────────────────────────
|
|
colorized = False
|
|
if do_colorize and bw_detected:
|
|
emit_progress(82, "Colorizing B&W photo")
|
|
try:
|
|
result, colorized = colorize_bw(result, intensity=0.85)
|
|
if colorized:
|
|
steps_applied.append("colorize")
|
|
emit_progress(92, "Colorization complete")
|
|
else:
|
|
emit_progress(92, "Colorization model not available")
|
|
except Exception as e:
|
|
emit_progress(92, f"Colorization skipped: {str(e)[:40]}")
|
|
else:
|
|
emit_progress(92, "Colorization skipped")
|
|
|
|
# ── Save result ──────────────────────────────────────────
|
|
emit_progress(95, "Saving result")
|
|
cv2.imwrite(output_path, result)
|
|
|
|
print(json.dumps({
|
|
"success": True,
|
|
"width": orig_w,
|
|
"height": orig_h,
|
|
"steps": steps_applied,
|
|
"scratchCoverage": round(scratch_coverage * 100, 2),
|
|
"facesEnhanced": faces_found,
|
|
"isGrayscale": bw_detected,
|
|
"colorized": colorized,
|
|
"output_path": output_path,
|
|
}))
|
|
|
|
except Exception as e:
|
|
print(json.dumps({"success": False, "error": str(e)}))
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|