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)