diff --git a/models.py b/models.py index ffcc5e9..5b04b73 100644 --- a/models.py +++ b/models.py @@ -59,7 +59,7 @@ def load_models(models_dir: str) -> ModelBundle: raise ImportError("insightface is not installed. Run: pip install insightface") from e app = FaceAnalysis(name="buffalo_l", providers=["CPUExecutionProvider"]) - app.prepare(ctx_id=0, det_size=(1024, 1024)) + app.prepare(ctx_id=0, det_size=(640, 640)) logger.info("FaceAnalysis ready.") logger.info("Loading inswapper_128…") diff --git a/pipeline.py b/pipeline.py index 745f9a0..fef7a25 100644 --- a/pipeline.py +++ b/pipeline.py @@ -31,6 +31,7 @@ print(stats) """ +import gc import logging import os import shutil @@ -47,6 +48,10 @@ SUPPORTED_EXTENSIONS = (".jpg", ".jpeg", ".png", ".tiff", ".webp") +# Maximum pixel dimension (longest edge) before the image is downscaled. +# Override with the REFACER_MAX_DIMENSION environment variable. +MAX_DIMENSION = int(os.environ.get("REFACER_MAX_DIMENSION", "4096")) + # --------------------------------------------------------------------------- # Result types @@ -216,6 +221,17 @@ def _process_image( logger.error("Could not read %s, skipping.", filename) return result + # --- Downscale if needed --- + h, w = img.shape[:2] + if max(h, w) > MAX_DIMENSION: + scale = MAX_DIMENSION / max(h, w) + new_w, new_h = int(w * scale), int(h * scale) + logger.warning( + "%s — %dx%d exceeds MAX_DIMENSION=%d, downscaling to %dx%d", + filename, w, h, MAX_DIMENSION, new_w, new_h, + ) + img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA) + # --- Detect faces --- try: faces = models.app.get(img) @@ -333,6 +349,7 @@ def run( for filename in filenames: logger.info("── Processing: %s", filename) image_result = _process_image(filename, input_dir, output_dir, models) + gc.collect() stats.image_results.append(image_result) stats.total_faces += image_result.faces_detected