backend/bg_remover/image_processor.py

import cv2
import numpy as np
from PIL import Image
from backend.utilities import pil_to_cv, cv_to_pil, resize_for_processing
from backend.bg_remover.segmentation import detect_automatic_bounding_box, run_grabcut
from backend.bg_remover.edge_detection import refine_mask
from backend.bg_remover.shadow_generator import generate_drop_shadow

try:
    from rembg import remove as rembg_remove
    REMBG_AVAILABLE = True
except ImportError:
    REMBG_AVAILABLE = False

_session_cache = {}

def get_rembg_session(model_name: str):
    """Retrieves or initializes a cached rembg model session to prevent reloading weights."""
    global _session_cache
    if model_name not in _session_cache:
        try:
            from rembg import new_session
            _session_cache[model_name] = new_session(model_name)
        except Exception as e:
            _session_cache[model_name] = None
    return _session_cache[model_name]

class ImageProcessor:
    """
    High-level orchestrator class to execute the background removal image processing pipeline.
    """
    @staticmethod
    def process_image(
        pil_image: Image.Image,
        rect: tuple = None,
        margin_percentage: float = 5.0,
        iter_count: int = 5,
        bg_seed_sensitivity: float = 35.0,
        closing_size: int = 5,
        keep_largest_only: bool = True,
        feather_radius: int = 3,
        matting_enabled: bool = True,
        matting_radius: int = 10,
        matting_eps: float = 1e-3,
        shadow_enabled: bool = False,
        shadow_opacity: float = 0.5,
        shadow_blur: int = 15,
        shadow_distance: int = 20,
        shadow_angle: float = 45.0,
        max_preview_dim: int = None,
        subject_mode: str = "AI Neural Network (U²-Net)"
    ) -> dict:
        """
        Processes the input PIL image and returns a dictionary of output PIL images.
        """
        # 1. Automatically normalize EXIF camera orientation tags
        from PIL import ImageOps
        pil_image = ImageOps.exif_transpose(pil_image)
        
        # Convert PIL to OpenCV (BGR)
        cv_raw = pil_to_cv(pil_image)
        
        # 2. Downscale for interactive preview speed if requested
        if max_preview_dim is not None:
            cv_img = resize_for_processing(cv_raw, max_preview_dim)
        else:
            cv_img = cv_raw.copy()
            
        h, w = cv_img.shape[:2]
        
        # 3. Bounding Box & Segmentation Determination
        is_neural = (subject_mode in ["AI BiRefNet (SOTA General)", "AI U²-Net (Legacy Neural)", "AI Neural Network (U²-Net)"] and REMBG_AVAILABLE)
        
        if is_neural:
            try:
                # Resolve correct model session
                if "BiRefNet" in subject_mode:
                    session = get_rembg_session("birefnet-general")
                else:
                    session = get_rembg_session("u2net")
                    
                if max_preview_dim is not None:
                    w_p, h_p = cv_img.shape[1], cv_img.shape[0]
                    pil_preview = pil_image.resize((w_p, h_p), Image.Resampling.LANCZOS)
                    cutout_pil = rembg_remove(pil_preview, session=session)
                else:
                    cutout_pil = rembg_remove(pil_image, session=session)
                    
                cv_cutout = pil_to_cv(cutout_pil)
                refined_mask = cv_cutout[:, :, 3].copy()
                
                # Resilient shape normalization to handle EXIF or transpose mismatches from rembg
                if refined_mask.shape[:2] != (h, w):
                    if refined_mask.shape[0] == w and refined_mask.shape[1] == h:
                        refined_mask = refined_mask.T
                    else:
                        refined_mask = cv2.resize(refined_mask, (w, h), interpolation=cv2.INTER_LINEAR)
                        
                actual_rect = (0, 0, w, h)
            except Exception as e:
                is_neural = False
                
        if not is_neural:
            if subject_mode == "Signature & Text (Ink)":
                pw = max(2, min(20, w // 20))
                ph = max(2, min(20, h // 20))
                c_tl = np.mean(cv_img[0:ph, 0:pw, :3], axis=(0, 1))
                c_tr = np.mean(cv_img[0:ph, w-pw:w, :3], axis=(0, 1))
                c_bg = (c_tl + c_tr) / 2.0
                
                dist = np.sqrt(np.sum((cv_img[:, :, :3] - c_bg) ** 2, axis=2))
                
                low_t = 15.0
                high_t = 45.0
                alpha = np.clip((dist - low_t) / (high_t - low_t) * 255.0, 0, 255).astype(np.uint8)
                
                refined_mask = alpha
                if closing_size > 0:
                    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (closing_size, closing_size))
                    refined_mask = cv2.morphologyEx(refined_mask, cv2.MORPH_CLOSE, kernel)
                    
                actual_rect = (0, 0, w, h)
            else:
                if rect is None:
                    actual_rect = detect_automatic_bounding_box(cv_img, margin_percentage)
                else:
                    if max_preview_dim is not None:
                        orig_h, orig_w = cv_raw.shape[:2]
                        scale_x = w / orig_w
                        scale_y = h / orig_h
                        rx, ry, rw, rh = rect
                        actual_rect = (
                            int(rx * scale_x),
                            int(ry * scale_y),
                            int(rw * scale_x),
                            int(rh * scale_y)
                        )
                    else:
                        actual_rect = rect
                        
                raw_mask = run_grabcut(cv_img, actual_rect, iter_count, bg_seed_sensitivity=bg_seed_sensitivity)
                
                refined_mask = refine_mask(
                    mask=raw_mask, 
                    img=cv_img,
                    closing_size=closing_size, 
                    keep_largest_only=keep_largest_only, 
                    feather_radius=feather_radius,
                    matting_enabled=matting_enabled,
                    matting_radius=matting_radius,
                    matting_eps=matting_eps
                )
        
        # 6. Generate Transparent PNG Cutout
        cutout = np.zeros((h, w, 4), dtype=np.uint8)
        cutout[:, :, :3] = cv_img[:, :, :3]
        cutout[:, :, 3] = refined_mask
        
        # 7. Generate Drop Shadow Composite
        if shadow_enabled:
            shadow_composite = generate_drop_shadow(
                cv_img,
                refined_mask,
                opacity=shadow_opacity,
                blur_radius=shadow_blur,
                distance=shadow_distance,
                angle_degrees=shadow_angle
            )
        else:
            shadow_composite = cutout
            
        # 8. Scale Bounding Box back to original coords if resized (for UI display overlay)
        if max_preview_dim is not None:
            orig_h, orig_w = cv_raw.shape[:2]
            scale_x = orig_w / w
            scale_y = orig_h / h
            ax, ay, aw, ah = actual_rect
            rect_out = (
                int(ax * scale_x),
                int(ay * scale_y),
                int(aw * scale_x),
                int(ah * scale_y)
            )
        else:
            rect_out = actual_rect
            
        # 9. Convert outputs back to PIL
        return {
            "original": cv_to_pil(cv_img),
            "mask": Image.fromarray(refined_mask).convert("L"),
            "transparent": cv_to_pil(cutout),
            "shadow": cv_to_pil(shadow_composite),
            "rect": rect_out
        }