resize_image_processor.py

import cv2
import numpy as np
from PIL import Image, ImageOps
from typing import Tuple, Optional


class ResizeImageProcessor:
    """
    Processeur d'images combinant DensePose cropping et ajustement de ratio.

    Le traitement se fait en deux étapes:
    1. Crop intelligent autour du masque DensePose avec ratio cible
    2. Ajustement final du ratio et redimensionnement
    """

    def __init__(
        self,
        densepose_model,
        # Paramètres par défaut pour process_image_with_densepose
        default_margin_ratio: float = 0.1,
        default_target_ratio: Tuple[int, int] = (4, 3),
        default_densepose_height: int = 1024,
        default_densepose_width: int = 768,
        # Paramètres par défaut pour adjust_aspect_ratio_then_resize
        default_final_height: int = 1024,
        default_final_width: int = 768,
        default_is_mask: bool = False
    ):
        """
        Initialise le processeur avec un modèle DensePose et des paramètres par défaut.

        Args:
            densepose_model: Modèle DensePose à utiliser
            default_margin_ratio: Marge autour du masque (0.1 = 10%)
            default_target_ratio: Ratio cible (h, w) pour le premier crop
            default_densepose_height: Hauteur de sortie du crop DensePose
            default_densepose_width: Largeur de sortie du crop DensePose
            default_final_height: Hauteur finale après ajustement
            default_final_width: Largeur finale après ajustement
            default_is_mask: Mode masque pour l'ajustement final
        """
        self.densepose_model = densepose_model

        # Stockage des paramètres par défaut
        self.default_margin_ratio = default_margin_ratio
        self.default_target_ratio = default_target_ratio
        self.default_densepose_height = default_densepose_height
        self.default_densepose_width = default_densepose_width
        self.default_final_height = default_final_height
        self.default_final_width = default_final_width
        self.default_is_mask = default_is_mask

    def __call__(
        self,
        image: Image.Image,
        # Paramètres optionnels pour process_image_with_densepose
        margin_ratio: Optional[float] = None,
        target_ratio: Optional[Tuple[int, int]] = None,
        densepose_height: Optional[int] = None,
        densepose_width: Optional[int] = None,
        # Paramètres optionnels pour adjust_aspect_ratio_then_resize
        final_height: Optional[int] = None,
        final_width: Optional[int] = None,
        is_mask: Optional[bool] = None
    ) -> Image.Image:
        """
        Traite une image avec DensePose puis ajuste le ratio et redimensionne.

        Args:
            image: Image PIL à traiter
            margin_ratio: Marge autour du masque (si None, utilise la valeur par défaut)
            target_ratio: Ratio cible (h, w) (si None, utilise la valeur par défaut)
            densepose_height: Hauteur du crop DensePose (si None, utilise la valeur par défaut)
            densepose_width: Largeur du crop DensePose (si None, utilise la valeur par défaut)
            final_height: Hauteur finale (si None, utilise la valeur par défaut)
            final_width: Largeur finale (si None, utilise la valeur par défaut)
            is_mask: Mode masque (si None, utilise la valeur par défaut)

        Returns:
            Image PIL traitée
        """
        # Utiliser les valeurs par défaut si non spécifiées
        margin_ratio = margin_ratio if margin_ratio is not None else self.default_margin_ratio
        target_ratio = target_ratio if target_ratio is not None else self.default_target_ratio
        densepose_height = densepose_height if densepose_height is not None else self.default_densepose_height
        densepose_width = densepose_width if densepose_width is not None else self.default_densepose_width
        final_height = final_height if final_height is not None else self.default_final_height
        final_width = final_width if final_width is not None else self.default_final_width
        is_mask = is_mask if is_mask is not None else self.default_is_mask

        # Étape 1: Process avec DensePose
        processed_image = self._process_image_with_densepose(
            image,
            margin_ratio=margin_ratio,
            target_ratio=target_ratio,
            final_height=densepose_height,
            final_width=densepose_width
        )

        # Étape 2: Ajustement final du ratio et redimensionnement
        final_image = self._adjust_aspect_ratio_then_resize(
            processed_image,
            target_height=final_height,
            target_width=final_width,
            is_mask=is_mask
        )

        return final_image

    def _process_image_with_densepose(
        self,
        img1: Image.Image,
        margin_ratio: float,
        target_ratio: Tuple[int, int],
        final_height: int,
        final_width: int
    ) -> Image.Image:
        """
        Traite une image PIL avec DensePose et retourne l'image croppée.
        """
        # Redimensionner l'image
        img1_resized = cv2.resize(
            np.array(img1),
            (int((1024 / img1.height) * img1.width), 1024)
        )

        # Créer une version PIL pour DensePose
        img_for_densepose = Image.fromarray(img1_resized)

        # Générer le masque DensePose
        densepose_output = self.densepose_model(img_for_densepose)

        # Convertir en numpy array et créer le masque binaire
        person_mask = (np.array(densepose_output) > 0).astype(np.uint8) * 255

        # Si le masque est en couleur (H, W, 3), prendre un seul canal
        if len(person_mask.shape) == 3:
            person_mask = person_mask[:, :, 0]

        # Cropper autour du masque
        result, bbox = self._crop_around_mask_bbox(
            img1_resized,
            person_mask,
            margin_ratio=margin_ratio,
            target_ratio=target_ratio,
        )

        # Convertir en PIL et retourner
        result = Image.fromarray(result)
        return ImageOps.contain(result, (final_width, final_height))

    @staticmethod
    def _crop_around_mask_bbox(
        img: np.ndarray,
        mask: np.ndarray,
        margin_ratio: float = 0.10,
        target_ratio: Optional[Tuple[int, int]] = None,
        output_size: Optional[Tuple[int, int]] = None
    ) -> Tuple[np.ndarray, Tuple[int, int, int, int]]:
        """
        Crop une image autour du masque avec gestion du ratio.
        """
        h, w = mask.shape[:2]
        m = (mask > 0).astype(np.uint8)

        num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(m, connectivity=8)
        if num_labels <= 1:
            x1, y1, x2, y2 = 0, 0, w-1, h-1
        else:
            areas = stats[1:, cv2.CC_STAT_AREA]
            biggest = 1 + np.argmax(areas)
            x, y, bw, bh, area = stats[biggest]
            x1, y1, x2, y2 = x, y, x + bw - 1, y + bh - 1

        # Ajouter une marge
        bw, bh = (x2 - x1 + 1), (y2 - y1 + 1)
        mx = int(bw * margin_ratio)
        my = int(bh * margin_ratio)
        x1 = max(0, x1 - mx)
        y1 = max(0, y1 - my)
        x2 = min(w-1, x2 + mx)
        y2 = min(h-1, y2 + my)

        # Ajuster au ratio cible si demandé
        if target_ratio is not None:
            rh, rw = target_ratio
            target_ar = rw / float(rh)
            cur_w, cur_h = (x2 - x1 + 1), (y2 - y1 + 1)
            cur_ar = cur_w / float(cur_h)
            cx = (x1 + x2) // 2
            cy = (y1 + y2) // 2

            if cur_ar < target_ar:
                new_w = int(round(cur_h * target_ar))
                half = new_w // 2
                x1 = max(0, cx - half)
                x2 = min(w-1, cx + (new_w - half - 1))
                cur_w = x2 - x1 + 1
                if cur_w < new_w:
                    shift = new_w - cur_w
                    x1 = max(0, x1 - shift//2)
                    x2 = min(w-1, x1 + new_w - 1)
            else:
                new_h = int(round(cur_w / target_ar))
                half = new_h // 2
                y1 = max(0, cy - half)
                y2 = min(h-1, cy + (new_h - half - 1))
                cur_h = y2 - y1 + 1
                if cur_h < new_h:
                    shift = new_h - cur_h
                    y1 = max(0, y1 - shift//2)
                    y2 = min(h-1, y1 + new_h - 1)

        # Crop final
        crop = img[y1:y2+1, x1:x2+1]
        if output_size is not None:
            H, W = output_size
            crop = cv2.resize(crop, (W, H), interpolation=cv2.INTER_AREA)

        return crop, (x1, y1, x2, y2)

    @staticmethod
    def _adjust_aspect_ratio_then_resize(
        image: Image.Image,
        target_height: int = 512,
        target_width: int = 384,
        is_mask: bool = False
    ) -> Image.Image:
        """
        Ajuste le ratio d'aspect puis redimensionne à la taille exacte.
        """
        original_width, original_height = image.size
        target_ratio = target_height / target_width
        original_ratio = original_height / original_width

        if is_mask:
            resampling = Image.Resampling.NEAREST
            fill_color = 'black'
        else:
            resampling = Image.Resampling.LANCZOS
            fill_color = 'white'

        if abs(original_ratio - target_ratio) < 0.001:
            return image.resize((target_width, target_height), resampling)

        if original_ratio < target_ratio:
            new_width = original_width
            new_height = int(original_width * target_ratio)

            if new_height > original_height:
                padding_needed = new_height - original_height
                top_padding = padding_needed // 2
                bottom_padding = padding_needed - top_padding
                image = ImageOps.expand(image, (0, top_padding, 0, bottom_padding), fill=fill_color)
            else:
                crop_amount = original_height - new_height
                top_crop = crop_amount // 2
                bottom_crop = crop_amount - top_crop
                image = image.crop((0, top_crop, original_width, original_height - bottom_crop))
        else:
            new_height = original_height
            new_width = int(original_height / target_ratio)

            if new_width > original_width:
                padding_needed = new_width - original_width
                left_padding = padding_needed // 2
                right_padding = padding_needed - left_padding
                image = ImageOps.expand(image, (left_padding, 0, right_padding, 0), fill=fill_color)
            else:
                crop_amount = original_width - new_width
                left_crop = crop_amount // 2
                right_crop = crop_amount - left_crop
                image = image.crop((left_crop, 0, original_width - right_crop, original_height))

        return image.resize((target_width, target_height), resampling)