processors/body_swap.py

"""
Body swap processor.

Pipeline
--------
1. Segment both images with *rembg* (U²-Net) to isolate person masks.
2. Compute bounding boxes from the masks.
3. Scale the source person to match the target bounding-box dimensions.
4. Color-correct the source region to match target lighting/tone.
5. Feather-blend using the segmentation mask.
6. Apply Poisson seamless-clone for photorealistic edge merging.

MediaPipe is intentionally NOT used — its API changed in 0.10.14
(removed `solutions`) which breaks on Python 3.13.  Bounding-box
alignment alone is sufficient for clean body swaps.
"""

import cv2
import numpy as np
from PIL import Image

from utils.image_utils import (
    apply_color_correction,
    feather_mask,
    alpha_blend,
)


class BodySwapper:
    """
    Replaces the body in *target_bgr* with the body from *source_bgr*.
    """

    _rembg_session = None   # shared across instances; loaded once

    @classmethod
    def _get_session(cls):
        """Lazy-load the u2net_human_seg rembg session (better than general u2net)."""
        if cls._rembg_session is None:
            from rembg import new_session
            print("[BodySwapper] Loading u2net_human_seg segmentation model …")
            cls._rembg_session = new_session("u2net_human_seg")
        return cls._rembg_session

    # ── Private helpers ─────────────────────────────────────────────────

    @classmethod
    def _segment(cls, bgr: np.ndarray) -> np.ndarray:
        """Return uint8 single-channel person mask via rembg u2net_human_seg."""
        from rembg import remove

        pil    = Image.fromarray(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))
        result = remove(pil, only_mask=True, session=cls._get_session())
        mask   = np.array(result)
        if mask.ndim == 3:
            mask = mask[:, :, 0]
        return mask

    @staticmethod
    def _bbox(mask: np.ndarray):
        """Bounding box (x1, y1, x2, y2) of the non-zero region, or None."""
        ys, xs = np.where(mask > 128)
        if len(ys) == 0:
            return None
        return int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())

    @staticmethod
    def _vertical_center_of_mass(mask: np.ndarray) -> float:
        """Y coordinate of the mask centre of mass (for vertical alignment)."""
        ys, _ = np.where(mask > 128)
        return float(ys.mean()) if len(ys) > 0 else mask.shape[0] / 2.0

    # ── Public API ────────────────────────────────────────────────────────────

    def swap(self, source_bgr, target_bgr, blend_strength=0.85):
        """
        Swap the source person's body into the target scene.

        Returns:
            (result_bgr, status_message)
        """
        try:
            # ── 1. Segment ────────────────────────────────────────────────────
            src_mask = self._segment(source_bgr)
            tgt_mask = self._segment(target_bgr)

            src_bbox = self._bbox(src_mask)
            tgt_bbox = self._bbox(tgt_mask)

            if src_bbox is None:
                return None, "No person detected in source image."
            if tgt_bbox is None:
                return None, "No person detected in target image."

            sx1, sy1, sx2, sy2 = src_bbox
            tx1, ty1, tx2, ty2 = tgt_bbox
            tgt_w, tgt_h = tx2 - tx1, ty2 - ty1

            # ── 2. Crop + resize source to target dimensions ──────────────────
            src_person   = source_bgr[sy1:sy2, sx1:sx2]
            src_mask_roi = src_mask[sy1:sy2, sx1:sx2]

            src_resized  = cv2.resize(src_person,   (tgt_w, tgt_h), interpolation=cv2.INTER_LANCZOS4)
            mask_resized = cv2.resize(src_mask_roi, (tgt_w, tgt_h), interpolation=cv2.INTER_LINEAR)

            # ── 3. Vertical CoM alignment ─────────────────────────────────────
            src_com_y = self._vertical_center_of_mass(src_mask_roi)
            tgt_com_y = self._vertical_center_of_mass(tgt_mask[ty1:ty2, tx1:tx2])
            scale_y   = tgt_h / max(sy2 - sy1, 1)
            offset_y  = int(tgt_com_y - src_com_y * scale_y)

            # ── 4. Composite onto full canvas ─────────────────────────────────
            h_t, w_t  = target_bgr.shape[:2]
            canvas_fg   = np.zeros_like(target_bgr)
            canvas_mask = np.zeros((h_t, w_t), dtype=np.uint8)

            dst_x1 = int(np.clip(tx1,          0, w_t))
            dst_y1 = int(np.clip(ty1 + offset_y, 0, h_t))
            dst_x2 = int(np.clip(tx1 + tgt_w,  0, w_t))
            dst_y2 = int(np.clip(ty1 + offset_y + tgt_h, 0, h_t))

            src_x1 = dst_x1 - tx1
            src_y1 = dst_y1 - (ty1 + offset_y)
            src_x2 = src_x1 + (dst_x2 - dst_x1)
            src_y2 = src_y1 + (dst_y2 - dst_y1)

            if dst_x2 <= dst_x1 or dst_y2 <= dst_y1:
                return None, "Alignment offset moved body out of frame."

            canvas_fg  [dst_y1:dst_y2, dst_x1:dst_x2] = src_resized [src_y1:src_y2, src_x1:src_x2]
            canvas_mask[dst_y1:dst_y2, dst_x1:dst_x2] = mask_resized[src_y1:src_y2, src_x1:src_x2]

            # ── 5. Color correction ───────────────────────────────────────────
            canvas_fg = apply_color_correction(canvas_fg, target_bgr, canvas_mask)

            # ── 6. Feathered alpha blend ──────────────────────────────────────
            soft_mask = feather_mask(canvas_mask, blur_radius=25)
            soft_mask = (soft_mask.astype(float) * blend_strength).clip(0, 255).astype(np.uint8)
            result    = alpha_blend(canvas_fg, target_bgr, soft_mask)

            # ── 7. Seamless clone (best-effort) ───────────────────────────────
            try:
                cx = int((dst_x1 + dst_x2) / 2)
                cy = int((dst_y1 + dst_y2) / 2)
                sc_mask = (canvas_mask > 10).astype(np.uint8) * 255
                result  = cv2.seamlessClone(
                    canvas_fg, target_bgr, sc_mask,
                    (cx, cy), cv2.NORMAL_CLONE,
                )
            except Exception as e:
                print(f"[BodySwapper] seamlessClone skipped: {e}")

            return result, "Body swap completed successfully."

        except Exception as exc:
            return None, f"Body swap error: {exc}"

    # ── Private helpers ───────────────────────────────────────────────────────

    def _segment(self, bgr: np.ndarray) -> np.ndarray:
        """Return a uint8 single-channel person mask via rembg u2net_human_seg."""
        from rembg import remove
        pil    = Image.fromarray(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))
        result = remove(pil, only_mask=True, session=self._get_session())
        mask   = np.array(result)
        if mask.ndim == 3:
            mask = mask[:, :, 0]
        return mask