mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
fix: resolve 5 bugs found during comprehensive tool testing
1. split batch 404: register split tool in batch registry via registerToolProcessFn() so /api/v1/tools/split/batch works 2. CodeFormer crash: inference_app() expects a file path, not a numpy array. Save to temp file before calling, read result back. 3. OCR fallback chain: fix case-sensitive "Segmentation fault" match that prevented PaddleOCR crash from triggering Tesseract fallback. Also add "process crashed" check. Upgrade ARM paddlepaddle to >=3.2.1. 4. blur-faces large images: downscale to 1920px max before MediaPipe detection, scale coordinates back. Also add rotation retry for portrait-oriented images where BlazeFace misses faces. Applied to detect_faces.py, enhance_faces.py, and restore.py. 5. color-adjustments tool ID: fix mismatch in index.ts registration array (was "color-adjustments", should be "adjust-colors").
This commit is contained in:
@@ -92,7 +92,7 @@ export async function registerToolRoutes(app: FastifyInstance): Promise<void> {
|
|||||||
{ id: "compress", register: registerCompress },
|
{ id: "compress", register: registerCompress },
|
||||||
{ id: "strip-metadata", register: registerStripMetadata },
|
{ id: "strip-metadata", register: registerStripMetadata },
|
||||||
{ id: "edit-metadata", register: registerEditMetadata },
|
{ id: "edit-metadata", register: registerEditMetadata },
|
||||||
{ id: "color-adjustments", register: registerColorAdjustments },
|
{ id: "adjust-colors", register: registerColorAdjustments },
|
||||||
{ id: "sharpening", register: registerSharpening },
|
{ id: "sharpening", register: registerSharpening },
|
||||||
|
|
||||||
// Watermark & Overlay
|
// Watermark & Overlay
|
||||||
|
|||||||
@@ -165,12 +165,13 @@ export function registerOcr(app: FastifyInstance) {
|
|||||||
});
|
});
|
||||||
} catch (err) {
|
} catch (err) {
|
||||||
lastError = err;
|
lastError = err;
|
||||||
const msg = err instanceof Error ? err.message : String(err);
|
const msg = (err instanceof Error ? err.message : String(err)).toLowerCase();
|
||||||
// If the Python process crashed (segfault, dispatcher exit), try next tier
|
// If the Python process crashed (segfault, dispatcher exit), try next tier
|
||||||
if (
|
if (
|
||||||
msg.includes("exited unexpectedly") ||
|
msg.includes("exited unexpectedly") ||
|
||||||
msg.includes("exited with code") ||
|
msg.includes("exited with code") ||
|
||||||
msg.includes("Segmentation fault")
|
msg.includes("segmentation fault") ||
|
||||||
|
msg.includes("process crashed")
|
||||||
) {
|
) {
|
||||||
request.log.warn(
|
request.log.warn(
|
||||||
{ toolId: "ocr", quality: tier, err },
|
{ toolId: "ocr", quality: tier, err },
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ import { z } from "zod";
|
|||||||
import { autoOrient } from "../../lib/auto-orient.js";
|
import { autoOrient } from "../../lib/auto-orient.js";
|
||||||
import { formatZodErrors } from "../../lib/errors.js";
|
import { formatZodErrors } from "../../lib/errors.js";
|
||||||
import { ensureSharpCompat } from "../../lib/heic-converter.js";
|
import { ensureSharpCompat } from "../../lib/heic-converter.js";
|
||||||
|
import { registerToolProcessFn } from "../tool-factory.js";
|
||||||
|
|
||||||
const settingsSchema = z.object({
|
const settingsSchema = z.object({
|
||||||
columns: z.number().min(1).max(100).default(3),
|
columns: z.number().min(1).max(100).default(3),
|
||||||
@@ -159,4 +160,87 @@ export function registerSplit(app: FastifyInstance) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
|
|
||||||
|
registerToolProcessFn({
|
||||||
|
toolId: "split",
|
||||||
|
settingsSchema,
|
||||||
|
process: async (inputBuffer, _settings, filename) => {
|
||||||
|
const settings = _settings as z.infer<typeof settingsSchema>;
|
||||||
|
const metadata = await sharp(inputBuffer).metadata();
|
||||||
|
const fullW = metadata.width ?? 0;
|
||||||
|
const fullH = metadata.height ?? 0;
|
||||||
|
|
||||||
|
let cols = settings.columns;
|
||||||
|
let rows = settings.rows;
|
||||||
|
if (settings.tileWidth && settings.tileHeight) {
|
||||||
|
cols = Math.max(1, Math.ceil(fullW / settings.tileWidth));
|
||||||
|
rows = Math.max(1, Math.ceil(fullH / settings.tileHeight));
|
||||||
|
}
|
||||||
|
cols = Math.min(cols, 100);
|
||||||
|
rows = Math.min(rows, 100);
|
||||||
|
|
||||||
|
const cellW = Math.floor(fullW / cols);
|
||||||
|
const cellH = Math.floor(fullH / rows);
|
||||||
|
const originalExt = extname(filename) || ".png";
|
||||||
|
const baseName = filename.replace(/\.[^.]+$/, "");
|
||||||
|
const { sharpFormat, ext: outputExt } = resolveOutputFormat(
|
||||||
|
settings.outputFormat,
|
||||||
|
originalExt,
|
||||||
|
);
|
||||||
|
|
||||||
|
const archive = archiver("zip", { zlib: { level: 5 } });
|
||||||
|
const chunks: Buffer[] = [];
|
||||||
|
archive.on("data", (chunk: Buffer) => chunks.push(chunk));
|
||||||
|
const done = new Promise<void>((resolve, reject) => {
|
||||||
|
archive.on("end", resolve);
|
||||||
|
archive.on("error", reject);
|
||||||
|
});
|
||||||
|
|
||||||
|
for (let row = 0; row < rows; row++) {
|
||||||
|
for (let col = 0; col < cols; col++) {
|
||||||
|
let left: number;
|
||||||
|
let top: number;
|
||||||
|
let w: number;
|
||||||
|
let h: number;
|
||||||
|
|
||||||
|
if (settings.tileWidth && settings.tileHeight) {
|
||||||
|
left = col * settings.tileWidth;
|
||||||
|
top = row * settings.tileHeight;
|
||||||
|
w = col === cols - 1 ? fullW - left : Math.min(settings.tileWidth, fullW - left);
|
||||||
|
h = row === rows - 1 ? fullH - top : Math.min(settings.tileHeight, fullH - top);
|
||||||
|
} else {
|
||||||
|
left = col * cellW;
|
||||||
|
top = row * cellH;
|
||||||
|
w = col === cols - 1 ? fullW - left : cellW;
|
||||||
|
h = row === rows - 1 ? fullH - top : cellH;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (left >= fullW || top >= fullH || w <= 0 || h <= 0) continue;
|
||||||
|
|
||||||
|
let pipeline = sharp(inputBuffer).extract({ left, top, width: w, height: h });
|
||||||
|
if (sharpFormat) {
|
||||||
|
const formatOpts: Record<string, unknown> = {};
|
||||||
|
if (sharpFormat === "jpeg" || sharpFormat === "webp") {
|
||||||
|
formatOpts.quality = settings.quality;
|
||||||
|
}
|
||||||
|
pipeline = pipeline.toFormat(sharpFormat, formatOpts);
|
||||||
|
}
|
||||||
|
|
||||||
|
const partBuffer = await pipeline.toBuffer();
|
||||||
|
archive.append(partBuffer, {
|
||||||
|
name: `${baseName}_r${row + 1}_c${col + 1}${outputExt}`,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
await archive.finalize();
|
||||||
|
await done;
|
||||||
|
|
||||||
|
return {
|
||||||
|
buffer: Buffer.concat(chunks),
|
||||||
|
filename: `${baseName}_split.zip`,
|
||||||
|
contentType: "application/zip",
|
||||||
|
};
|
||||||
|
},
|
||||||
|
});
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -283,7 +283,7 @@
|
|||||||
"paddlepaddle-gpu>=3.2.1 --extra-index-url https://www.paddlepaddle.org.cn/packages/stable/cu126/",
|
"paddlepaddle-gpu>=3.2.1 --extra-index-url https://www.paddlepaddle.org.cn/packages/stable/cu126/",
|
||||||
"paddleocr[doc-parser]>=3.4.0,<3.5.0"
|
"paddleocr[doc-parser]>=3.4.0,<3.5.0"
|
||||||
],
|
],
|
||||||
"arm64": ["paddlepaddle==3.0.0", "paddleocr[doc-parser]>=3.4.0,<3.5.0"]
|
"arm64": ["paddlepaddle>=3.2.1", "paddleocr[doc-parser]>=3.4.0,<3.5.0"]
|
||||||
},
|
},
|
||||||
"pipFlags": {},
|
"pipFlags": {},
|
||||||
"postInstall": [],
|
"postInstall": [],
|
||||||
|
|||||||
@@ -120,14 +120,106 @@ def _detect_with_tasks(img_array, min_confidence):
|
|||||||
return faces
|
return faces
|
||||||
|
|
||||||
|
|
||||||
def _detect_faces(img_array, min_confidence):
|
_MAX_DETECT_DIM = 1920
|
||||||
"""Detect faces, trying legacy API first then falling back to tasks API."""
|
|
||||||
|
|
||||||
|
def _downscale_for_detection(img_array):
|
||||||
|
"""Downscale image if needed so MediaPipe can detect faces reliably.
|
||||||
|
|
||||||
|
Returns (scaled_array, scale_factor). Coordinates from detection on
|
||||||
|
the scaled image must be multiplied by scale_factor to map back to
|
||||||
|
the original resolution.
|
||||||
|
"""
|
||||||
|
h, w = img_array.shape[:2]
|
||||||
|
longest = max(h, w)
|
||||||
|
if longest <= _MAX_DETECT_DIM:
|
||||||
|
return img_array, 1.0
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
scale = _MAX_DETECT_DIM / longest
|
||||||
|
new_w = int(w * scale)
|
||||||
|
new_h = int(h * scale)
|
||||||
|
resized = cv2.resize(img_array, (new_w, new_h), interpolation=cv2.INTER_AREA)
|
||||||
|
return resized, 1.0 / scale
|
||||||
|
|
||||||
|
|
||||||
|
def _detect_single_orientation(img_array, min_confidence):
|
||||||
|
"""Run face detection on a single image orientation."""
|
||||||
try:
|
try:
|
||||||
return _detect_with_solutions(img_array, min_confidence)
|
return _detect_with_solutions(img_array, min_confidence)
|
||||||
except AttributeError:
|
except AttributeError:
|
||||||
return _detect_with_tasks(img_array, min_confidence)
|
return _detect_with_tasks(img_array, min_confidence)
|
||||||
|
|
||||||
|
|
||||||
|
def _remap_faces_from_rotation(faces, rotation, orig_h, orig_w):
|
||||||
|
"""Map face coordinates from a rotated image back to the original.
|
||||||
|
|
||||||
|
rotation: 90, 180, or 270 (clockwise degrees applied to the original).
|
||||||
|
orig_h, orig_w: dimensions of the original (un-rotated) image.
|
||||||
|
|
||||||
|
When the original is (orig_h, orig_w):
|
||||||
|
90 CW -> rotated is (orig_w, orig_h). (rx, ry) -> (ry, orig_h - rx - rw_box)
|
||||||
|
180 -> rotated is (orig_h, orig_w). (rx, ry) -> (orig_w - rx - rw_box, orig_h - ry - rh_box)
|
||||||
|
270 CW -> rotated is (orig_w, orig_h). (rx, ry) -> (orig_w - ry - rh_box, rx)
|
||||||
|
"""
|
||||||
|
remapped = []
|
||||||
|
for f in faces:
|
||||||
|
x, y, w, h = f["x"], f["y"], f["w"], f["h"]
|
||||||
|
if rotation == 90:
|
||||||
|
remapped.append({"x": y, "y": orig_h - x - w, "w": h, "h": w})
|
||||||
|
elif rotation == 180:
|
||||||
|
remapped.append({"x": orig_w - x - w, "y": orig_h - y - h, "w": w, "h": h})
|
||||||
|
elif rotation == 270:
|
||||||
|
remapped.append({"x": orig_w - y - h, "y": x, "w": h, "h": w})
|
||||||
|
else:
|
||||||
|
remapped.append(f)
|
||||||
|
return remapped
|
||||||
|
|
||||||
|
|
||||||
|
def _detect_faces(img_array, min_confidence):
|
||||||
|
"""Detect faces, trying multiple orientations if needed.
|
||||||
|
|
||||||
|
MediaPipe BlazeFace can miss faces in portrait-oriented or rotated
|
||||||
|
images. We first try the image as-is; if no faces are found we
|
||||||
|
retry at 90, 180, and 270-degree rotations and map the coordinates
|
||||||
|
back. Large images are downscaled before detection.
|
||||||
|
"""
|
||||||
|
import cv2
|
||||||
|
|
||||||
|
orig_h, orig_w = img_array.shape[:2]
|
||||||
|
scaled, inv_scale = _downscale_for_detection(img_array)
|
||||||
|
|
||||||
|
faces = _detect_single_orientation(scaled, min_confidence)
|
||||||
|
|
||||||
|
# If nothing found, try rotated copies and merge all detections.
|
||||||
|
# Different rotations can catch different faces, so we union them
|
||||||
|
# and de-duplicate with NMS.
|
||||||
|
if not faces:
|
||||||
|
rotations = [
|
||||||
|
(cv2.ROTATE_90_CLOCKWISE, 90),
|
||||||
|
(cv2.ROTATE_180, 180),
|
||||||
|
(cv2.ROTATE_90_COUNTERCLOCKWISE, 270),
|
||||||
|
]
|
||||||
|
all_rotated_faces = []
|
||||||
|
sh, sw = scaled.shape[:2]
|
||||||
|
for cv2_flag, degrees in rotations:
|
||||||
|
rotated = cv2.rotate(scaled, cv2_flag)
|
||||||
|
found = _detect_single_orientation(rotated, min_confidence)
|
||||||
|
if found:
|
||||||
|
remapped = _remap_faces_from_rotation(found, degrees, sh, sw)
|
||||||
|
all_rotated_faces.extend(remapped)
|
||||||
|
faces = _nms_faces(all_rotated_faces)
|
||||||
|
|
||||||
|
if inv_scale != 1.0:
|
||||||
|
for f in faces:
|
||||||
|
f["x"] = int(f["x"] * inv_scale)
|
||||||
|
f["y"] = int(f["y"] * inv_scale)
|
||||||
|
f["w"] = int(f["w"] * inv_scale)
|
||||||
|
f["h"] = int(f["h"] * inv_scale)
|
||||||
|
|
||||||
|
return faces
|
||||||
|
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
input_path = sys.argv[1]
|
input_path = sys.argv[1]
|
||||||
output_path = sys.argv[2]
|
output_path = sys.argv[2]
|
||||||
|
|||||||
@@ -62,51 +62,66 @@ def _ensure_face_detect_model():
|
|||||||
return _LOCAL_MODEL_PATH
|
return _LOCAL_MODEL_PATH
|
||||||
|
|
||||||
|
|
||||||
|
_MAX_DETECT_DIM = 1920
|
||||||
|
|
||||||
|
|
||||||
|
def _downscale_for_detection(img_array):
|
||||||
|
"""Downscale image if needed so MediaPipe can detect faces reliably."""
|
||||||
|
h, w = img_array.shape[:2]
|
||||||
|
longest = max(h, w)
|
||||||
|
if longest <= _MAX_DETECT_DIM:
|
||||||
|
return img_array, 1.0
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
scale = _MAX_DETECT_DIM / longest
|
||||||
|
new_w = int(w * scale)
|
||||||
|
new_h = int(h * scale)
|
||||||
|
resized = cv2.resize(img_array, (new_w, new_h), interpolation=cv2.INTER_AREA)
|
||||||
|
return resized, 1.0 / scale
|
||||||
|
|
||||||
|
|
||||||
def detect_faces_mediapipe(img_array, sensitivity):
|
def detect_faces_mediapipe(img_array, sensitivity):
|
||||||
"""Detect faces using MediaPipe with dual-model approach.
|
"""Detect faces using MediaPipe with dual-model approach.
|
||||||
|
|
||||||
Returns a list of {x, y, w, h} dicts for each detected face.
|
Returns a list of {x, y, w, h} dicts for each detected face.
|
||||||
Tries legacy mp.solutions API first, falls back to mp.tasks.
|
Tries legacy mp.solutions API first, falls back to mp.tasks.
|
||||||
|
Large images are downscaled before detection for reliability.
|
||||||
"""
|
"""
|
||||||
import mediapipe as mp
|
import mediapipe as mp
|
||||||
|
|
||||||
min_confidence = max(0.1, 1.0 - sensitivity)
|
min_confidence = max(0.1, 1.0 - sensitivity)
|
||||||
|
scaled, inv_scale = _downscale_for_detection(img_array)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
mp_face = mp.solutions.face_detection
|
mp_face = mp.solutions.face_detection
|
||||||
|
|
||||||
# Try short-range model first (model_selection=0, best for faces
|
all_detections = []
|
||||||
# within ~2m which covers most photos), then fall back to
|
|
||||||
# full-range model (model_selection=1) for distant/group shots.
|
|
||||||
detections = []
|
|
||||||
for model_sel in [0, 1]:
|
for model_sel in [0, 1]:
|
||||||
detector = mp_face.FaceDetection(
|
detector = mp_face.FaceDetection(
|
||||||
model_selection=model_sel,
|
model_selection=model_sel,
|
||||||
min_detection_confidence=min_confidence,
|
min_detection_confidence=min_confidence,
|
||||||
)
|
)
|
||||||
results = detector.process(img_array)
|
results = detector.process(scaled)
|
||||||
detector.close()
|
detector.close()
|
||||||
if results.detections:
|
if results.detections:
|
||||||
detections = results.detections
|
all_detections.extend(results.detections)
|
||||||
break
|
|
||||||
|
|
||||||
if not detections:
|
if not all_detections:
|
||||||
return []
|
return []
|
||||||
|
|
||||||
ih, iw = img_array.shape[:2]
|
ih, iw = scaled.shape[:2]
|
||||||
faces = []
|
faces = []
|
||||||
for detection in detections:
|
for detection in all_detections:
|
||||||
bbox = detection.location_data.relative_bounding_box
|
bbox = detection.location_data.relative_bounding_box
|
||||||
faces.append({
|
faces.append({
|
||||||
"x": int(bbox.xmin * iw),
|
"x": int(bbox.xmin * iw * inv_scale),
|
||||||
"y": int(bbox.ymin * ih),
|
"y": int(bbox.ymin * ih * inv_scale),
|
||||||
"w": int(bbox.width * iw),
|
"w": int(bbox.width * iw * inv_scale),
|
||||||
"h": int(bbox.height * ih),
|
"h": int(bbox.height * ih * inv_scale),
|
||||||
})
|
})
|
||||||
return faces
|
return faces
|
||||||
|
|
||||||
except AttributeError:
|
except AttributeError:
|
||||||
# mediapipe >= 0.10.30 removed mp.solutions, use tasks API
|
|
||||||
model_path = _ensure_face_detect_model()
|
model_path = _ensure_face_detect_model()
|
||||||
options = mp.tasks.vision.FaceDetectorOptions(
|
options = mp.tasks.vision.FaceDetectorOptions(
|
||||||
base_options=mp.tasks.BaseOptions(model_asset_path=model_path),
|
base_options=mp.tasks.BaseOptions(model_asset_path=model_path),
|
||||||
@@ -114,7 +129,7 @@ def detect_faces_mediapipe(img_array, sensitivity):
|
|||||||
min_detection_confidence=min_confidence,
|
min_detection_confidence=min_confidence,
|
||||||
)
|
)
|
||||||
detector = mp.tasks.vision.FaceDetector.create_from_options(options)
|
detector = mp.tasks.vision.FaceDetector.create_from_options(options)
|
||||||
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=img_array)
|
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=scaled)
|
||||||
result = detector.detect(mp_image)
|
result = detector.detect(mp_image)
|
||||||
detector.close()
|
detector.close()
|
||||||
|
|
||||||
@@ -122,10 +137,10 @@ def detect_faces_mediapipe(img_array, sensitivity):
|
|||||||
for detection in result.detections:
|
for detection in result.detections:
|
||||||
bbox = detection.bounding_box
|
bbox = detection.bounding_box
|
||||||
faces.append({
|
faces.append({
|
||||||
"x": bbox.origin_x,
|
"x": int(bbox.origin_x * inv_scale),
|
||||||
"y": bbox.origin_y,
|
"y": int(bbox.origin_y * inv_scale),
|
||||||
"w": bbox.width,
|
"w": int(bbox.width * inv_scale),
|
||||||
"h": bbox.height,
|
"h": int(bbox.height * inv_scale),
|
||||||
})
|
})
|
||||||
return faces
|
return faces
|
||||||
|
|
||||||
@@ -166,23 +181,18 @@ def enhance_with_codeformer(img_array, fidelity_weight):
|
|||||||
face detection, alignment, restoration, and paste-back internally.
|
face detection, alignment, restoration, and paste-back internally.
|
||||||
fidelity_weight controls quality vs fidelity (0 = quality, 1 = fidelity).
|
fidelity_weight controls quality vs fidelity (0 = quality, 1 = fidelity).
|
||||||
|
|
||||||
NOTE: codeformer-pip's app.py runs heavy module-level initialization
|
NOTE: inference_app() expects a file path, not a numpy array. We save
|
||||||
(model downloads, GPU setup) on import. The Docker image must place
|
to a temp file and pass the path. The function returns a file path to
|
||||||
model weights where the package expects them, or set environment
|
the result which we read back.
|
||||||
variables so the download step succeeds. If the import or inference
|
|
||||||
fails, the auto model selection will fall back to GFPGAN.
|
|
||||||
"""
|
"""
|
||||||
|
import tempfile
|
||||||
|
import cv2
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import torch
|
import torch
|
||||||
from gpu import gpu_available
|
from gpu import gpu_available
|
||||||
|
|
||||||
use_gpu = gpu_available()
|
use_gpu = gpu_available()
|
||||||
|
|
||||||
# CodeFormer selects its device during module-level init and inside
|
|
||||||
# inference_app(). It has no device= parameter, so to respect
|
|
||||||
# ASHIM_GPU=false we temporarily override torch.cuda.is_available
|
|
||||||
# so all internal device checks see False. When use_gpu is True
|
|
||||||
# (the common path) no override happens.
|
|
||||||
_orig_cuda_check = torch.cuda.is_available
|
_orig_cuda_check = torch.cuda.is_available
|
||||||
if not use_gpu:
|
if not use_gpu:
|
||||||
torch.cuda.is_available = lambda: False
|
torch.cuda.is_available = lambda: False
|
||||||
@@ -190,19 +200,31 @@ def enhance_with_codeformer(img_array, fidelity_weight):
|
|||||||
from codeformer.app import inference_app
|
from codeformer.app import inference_app
|
||||||
|
|
||||||
img_bgr = img_array[:, :, ::-1].copy()
|
img_bgr = img_array[:, :, ::-1].copy()
|
||||||
restored_bgr = inference_app(
|
|
||||||
image=img_bgr,
|
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_in:
|
||||||
background_enhance=False,
|
cv2.imwrite(tmp_in.name, img_bgr)
|
||||||
face_upsample=False,
|
tmp_in_path = tmp_in.name
|
||||||
upscale=1,
|
|
||||||
codeformer_fidelity=fidelity_weight,
|
try:
|
||||||
)
|
result_path = inference_app(
|
||||||
|
image=tmp_in_path,
|
||||||
|
background_enhance=False,
|
||||||
|
face_upsample=False,
|
||||||
|
upscale=1,
|
||||||
|
codeformer_fidelity=fidelity_weight,
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
os.unlink(tmp_in_path)
|
||||||
|
|
||||||
|
if result_path is None:
|
||||||
|
raise RuntimeError("CodeFormer returned no result (face detection may have failed)")
|
||||||
|
|
||||||
|
restored_bgr = cv2.imread(str(result_path), cv2.IMREAD_COLOR)
|
||||||
|
if restored_bgr is None:
|
||||||
|
raise RuntimeError("CodeFormer output file could not be read")
|
||||||
finally:
|
finally:
|
||||||
torch.cuda.is_available = _orig_cuda_check
|
torch.cuda.is_available = _orig_cuda_check
|
||||||
|
|
||||||
if restored_bgr is None:
|
|
||||||
raise RuntimeError("CodeFormer returned no result (face detection may have failed)")
|
|
||||||
|
|
||||||
restored_rgb = restored_bgr[:, :, ::-1].copy()
|
restored_rgb = restored_bgr[:, :, ::-1].copy()
|
||||||
return restored_rgb
|
return restored_rgb
|
||||||
|
|
||||||
|
|||||||
@@ -259,38 +259,51 @@ def enhance_faces(img_bgr, fidelity=0.7):
|
|||||||
import mediapipe as mp
|
import mediapipe as mp
|
||||||
from gpu import safe_onnx_session
|
from gpu import safe_onnx_session
|
||||||
|
|
||||||
# Detect faces
|
# Detect faces — downscale large images for reliable MediaPipe detection
|
||||||
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
|
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
|
||||||
ih, iw = img_bgr.shape[:2]
|
ih, iw = img_bgr.shape[:2]
|
||||||
|
|
||||||
|
max_dim = 1920
|
||||||
|
longest = max(ih, iw)
|
||||||
|
if longest > max_dim:
|
||||||
|
det_scale = max_dim / longest
|
||||||
|
det_rgb = cv2.resize(
|
||||||
|
img_rgb,
|
||||||
|
(int(iw * det_scale), int(ih * det_scale)),
|
||||||
|
interpolation=cv2.INTER_AREA,
|
||||||
|
)
|
||||||
|
inv_scale = 1.0 / det_scale
|
||||||
|
else:
|
||||||
|
det_rgb = img_rgb
|
||||||
|
inv_scale = 1.0
|
||||||
|
|
||||||
try:
|
try:
|
||||||
mp_face = mp.solutions.face_detection
|
mp_face = mp.solutions.face_detection
|
||||||
detections = []
|
all_detections = []
|
||||||
for model_sel in [0, 1]:
|
for model_sel in [0, 1]:
|
||||||
detector = mp_face.FaceDetection(
|
detector = mp_face.FaceDetection(
|
||||||
model_selection=model_sel, min_detection_confidence=0.4
|
model_selection=model_sel, min_detection_confidence=0.4
|
||||||
)
|
)
|
||||||
results = detector.process(img_rgb)
|
results = detector.process(det_rgb)
|
||||||
detector.close()
|
detector.close()
|
||||||
if results.detections:
|
for detection in (results.detections or []):
|
||||||
detections = results.detections
|
all_detections.append(detection)
|
||||||
break
|
|
||||||
|
|
||||||
if not detections:
|
if not all_detections:
|
||||||
return img_bgr, 0
|
return img_bgr, 0
|
||||||
|
|
||||||
|
dh, dw = det_rgb.shape[:2]
|
||||||
face_boxes = []
|
face_boxes = []
|
||||||
for detection in detections:
|
for detection in all_detections:
|
||||||
bbox = detection.location_data.relative_bounding_box
|
bbox = detection.location_data.relative_bounding_box
|
||||||
face_boxes.append({
|
face_boxes.append({
|
||||||
"x": int(bbox.xmin * iw),
|
"x": int(bbox.xmin * dw * inv_scale),
|
||||||
"y": int(bbox.ymin * ih),
|
"y": int(bbox.ymin * dh * inv_scale),
|
||||||
"w": int(bbox.width * iw),
|
"w": int(bbox.width * dw * inv_scale),
|
||||||
"h": int(bbox.height * ih),
|
"h": int(bbox.height * dh * inv_scale),
|
||||||
})
|
})
|
||||||
|
|
||||||
except AttributeError:
|
except AttributeError:
|
||||||
# mediapipe >= 0.10.30 removed mp.solutions, use tasks API
|
|
||||||
model_path = _ensure_face_detect_model()
|
model_path = _ensure_face_detect_model()
|
||||||
options = mp.tasks.vision.FaceDetectorOptions(
|
options = mp.tasks.vision.FaceDetectorOptions(
|
||||||
base_options=mp.tasks.BaseOptions(model_asset_path=model_path),
|
base_options=mp.tasks.BaseOptions(model_asset_path=model_path),
|
||||||
@@ -298,7 +311,7 @@ def enhance_faces(img_bgr, fidelity=0.7):
|
|||||||
min_detection_confidence=0.4,
|
min_detection_confidence=0.4,
|
||||||
)
|
)
|
||||||
fd = mp.tasks.vision.FaceDetector.create_from_options(options)
|
fd = mp.tasks.vision.FaceDetector.create_from_options(options)
|
||||||
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=img_rgb)
|
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=det_rgb)
|
||||||
result = fd.detect(mp_image)
|
result = fd.detect(mp_image)
|
||||||
fd.close()
|
fd.close()
|
||||||
|
|
||||||
@@ -309,10 +322,10 @@ def enhance_faces(img_bgr, fidelity=0.7):
|
|||||||
for detection in result.detections:
|
for detection in result.detections:
|
||||||
bbox = detection.bounding_box
|
bbox = detection.bounding_box
|
||||||
face_boxes.append({
|
face_boxes.append({
|
||||||
"x": bbox.origin_x,
|
"x": int(bbox.origin_x * inv_scale),
|
||||||
"y": bbox.origin_y,
|
"y": int(bbox.origin_y * inv_scale),
|
||||||
"w": bbox.width,
|
"w": int(bbox.width * inv_scale),
|
||||||
"h": bbox.height,
|
"h": int(bbox.height * inv_scale),
|
||||||
})
|
})
|
||||||
|
|
||||||
# Load CodeFormer model
|
# Load CodeFormer model
|
||||||
|
|||||||
Reference in New Issue
Block a user