Ai-Swaper/core/image_engine.py

import shutil
from pathlib import Path
from core.config import cv2_imread_unicode, cv2_imwrite_unicode
from core.video_engine import is_output_path, is_within, safe_output_path, expected_output_path

IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tiff"}


class ImageSwapEngine:
    """Batch and single image face swapping engine."""

    def __init__(self, swapper, log_fn=print):
        self.swapper = swapper
        self.log = log_fn

    def find_images(self, input_dir, output_dir=None):
        """
        Findet alle unterstuetzten Bilddateien im Eingabeordner.

        Ausgabeordner werden ausgeschlossen - auch dann, wenn sie im
        Eingabeordner liegen. Der alte Filter pruefte nur auf die exakten
        Parts 'output'/'failed' und liess den echten Config-Ordner
        'Ausgabe' durch (Rekursion: Ergebnisse werden zu Eingaben).
        """
        indir = Path(input_dir)
        if not indir.is_dir():
            return []
        found = []
        for p in indir.rglob("*"):
            if not p.is_file() or p.suffix.lower() not in IMAGE_EXTENSIONS:
                continue
            if is_output_path(p, output_dir):
                continue
            if output_dir is not None and is_within(p, output_dir):
                continue
            found.append(p)
        return sorted(found)

    def process_image(self, img_path, src_face, out_path):
        """Swap face in a single image and save result to out_path."""
        original = cv2_imread_unicode(img_path)
        if original is None:
            return False, None, "Bild konnte nicht geladen werden"

        result, count = self.swapper.swap_single_frame(original, src_face)
        if count == 0:
            return False, original, "Kein Gesicht erkannt oder nicht modifiziert"

        Path(out_path).parent.mkdir(parents=True, exist_ok=True)
        ok = cv2_imwrite_unicode(out_path, result)
        if not ok:
            return False, result, "Konnte Ausgabedatei nicht schreiben"
        return True, result, None

    def process_batch(self, source_path, input_dir, output_dir, progress_cb=None, cancel_check=None):
        """
        Execute batch face swap across all images in input_dir.
        Yields progress and returns summary statistics.
        """
        src_p = Path(source_path)
        if not src_p.is_file():
            raise RuntimeError(f"Quellbild nicht gefunden: {source_path}")

        src_img = cv2_imread_unicode(src_p)
        if src_img is None:
            raise RuntimeError(f"Quellbild konnte nicht geoeffnet werden: {source_path}")

        self.swapper.init_models()
        src_face = self.swapper.get_first_face(src_img)
        if src_face is None:
            raise RuntimeError("Kein Gesicht im Quellbild gefunden!")

        if self.swapper.use_restoration and self.swapper.restorer:
            try:
                self.swapper.restorer.load()
            except Exception as e:
                self.log(f"  Hinweis: GFPGAN konnte nicht vorab geladen werden ({e}).")

        self.log(f"OK Quellgesicht erkannt in: {src_p.name}")

        images = self.find_images(input_dir, output_dir)
        if not images:
            self.log("Keine Bilder im Eingabe-Ordner gefunden.")
            return {"total": 0, "swapped": 0, "skipped": 0, "no_face": 0, "failed": 0}

        outdir = Path(output_dir)
        outdir.mkdir(parents=True, exist_ok=True)
        failed_dir = outdir / "failed"

        total = len(images)
        swapped = 0
        skipped = 0
        no_face = 0
        failed = 0
        used = set()

        self.log(f"Starte Batch-Swap fuer {total} Bilder ...")

        for idx, img_path in enumerate(images, 1):
            if cancel_check and cancel_check():
                self.log("Batch-Verarbeitung durch Benutzer abgebrochen.")
                break

            # Erwarteter Zielpfad (Unterordner-Struktur erhalten). Existiert er
            # schon, ist das eine Wiederaufnahme -> ueberspringen, NICHT umbenennen.
            out_file = expected_output_path(outdir, img_path, src_root=Path(input_dir))
            if out_file.exists():
                skipped += 1
                self.log(f"[{idx}/{total}] {img_path.name} -> UEBERSPRUNGEN (bereits vorhanden)")
                if progress_cb:
                    progress_cb(idx, total, None, img_path.name)
                continue

            # Namenskollision innerhalb DIESES Laufs aufloesen (zwei Eingaben mit
            # gleichem Namen). Ein bereits existierender Zielname wird nie
            # ueberschrieben.
            out_file = safe_output_path(outdir, img_path, used, src_root=Path(input_dir))

            try:
                success, res_img, err_msg = self.process_image(img_path, src_face, out_file)
            except Exception as e:
                # Eine kaputte Datei darf den Batch nicht abbrechen.
                success, res_img, err_msg = False, None, f"{type(e).__name__}: {e}"

            if success:
                swapped += 1
                self.log(f"[{idx}/{total}] {img_path.name} -> OK")
            else:
                failed += 1
                try:
                    failed_dir.mkdir(parents=True, exist_ok=True)
                    shutil.copy2(img_path, failed_dir / img_path.name)
                except Exception:
                    pass
                if "Kein Gesicht" in (err_msg or "") or "nicht modifiziert" in (err_msg or ""):
                    no_face += 1
                    self.log(f"[{idx}/{total}] {img_path.name} -> WARNUNG: kein Gesicht erkannt -> failed/{img_path.name}")
                else:
                    self.log(f"[{idx}/{total}] {img_path.name} -> FEHLER: {err_msg} -> failed/{img_path.name}")

            if progress_cb:
                progress_cb(idx, total, res_img, img_path.name)

        stats = {
            "total": total,
            "swapped": swapped,
            "skipped": skipped,
            "no_face": no_face,
            "failed": failed,
        }
        self.log(f"Fertig: {swapped} geswappt, {skipped} uebersprungen, {no_face} ohne Gesicht, {failed} fehlgeschlagen.")
        return stats