AI-Talent-Force/dev_caio

1"""2ShortSmith v2 - Body Recognizer Module3 4Full-body person recognition using OSNet for:5- Identifying people when face is not visible6- Back views, profile shots, masks, helmets7- Clothing and appearance-based matching8 9Complements face recognition for comprehensive person tracking.10"""11 12from pathlib import Path13from typing import List, Optional, Tuple, Union14from dataclasses import dataclass15import numpy as np16 17from utils.logger import get_logger, LogTimer18from utils.helpers import ModelLoadError, InferenceError19from config import get_config, ModelConfig20 21logger = get_logger("models.body_recognizer")22 23 24@dataclass25class BodyDetection:26    """Represents a detected person body in an image."""27    bbox: Tuple[int, int, int, int]  # (x1, y1, x2, y2)28    confidence: float                 # Detection confidence29    embedding: Optional[np.ndarray]   # Body appearance embedding30    track_id: Optional[int] = None    # Tracking ID if available31 32    @property33    def center(self) -> Tuple[int, int]:34        """Center point of body bounding box."""35        x1, y1, x2, y2 = self.bbox36        return ((x1 + x2) // 2, (y1 + y2) // 2)37 38    @property39    def area(self) -> int:40        """Area of bounding box."""41        x1, y1, x2, y2 = self.bbox42        return (x2 - x1) * (y2 - y1)43 44    @property45    def width(self) -> int:46        return self.bbox[2] - self.bbox[0]47 48    @property49    def height(self) -> int:50        return self.bbox[3] - self.bbox[1]51 52    @property53    def aspect_ratio(self) -> float:54        """Height/width ratio (typical person is ~2.5-3.0)."""55        if self.width == 0:56            return 057        return self.height / self.width58 59 60@dataclass61class BodyMatch:62    """Result of body matching."""63    detection: BodyDetection64    similarity: float65    is_match: bool66    reference_id: Optional[str] = None67 68 69class BodyRecognizer:70    """71    Body recognition using person re-identification models.72 73    Uses:74    - YOLO or similar for person detection75    - OSNet for body appearance embeddings76 77    Designed to work alongside FaceRecognizer for complete78    person identification across all viewing angles.79    """80 81    def __init__(82        self,83        config: Optional[ModelConfig] = None,84        load_model: bool = True,85    ):86        """87        Initialize body recognizer.88 89        Args:90            config: Model configuration91            load_model: Whether to load models immediately92        """93        self.config = config or get_config().model94        self.detector = None95        self.reid_model = None96        self._reference_embeddings: dict = {}97 98        if load_model:99            self._load_models()100 101        logger.info(f"BodyRecognizer initialized (threshold={self.config.body_similarity_threshold})")102 103    def _load_models(self) -> None:104        """Load person detection and re-identification models."""105        with LogTimer(logger, "Loading body recognition models"):106            self._load_detector()107            self._load_reid_model()108 109    def _load_detector(self) -> None:110        """Load person detector (YOLO)."""111        try:112            from ultralytics import YOLO113 114            # Use YOLOv8 for person detection115            self.detector = YOLO("yolov8n.pt")  # Nano model for speed116            logger.info("YOLO detector loaded")117 118        except ImportError:119            logger.warning("ultralytics not installed, using fallback detection")120            self.detector = None121 122        except Exception as e:123            logger.warning(f"Failed to load YOLO detector: {e}")124            self.detector = None125 126    def _load_reid_model(self) -> None:127        """Load OSNet re-identification model."""128        try:129            import torch130            import torchvision.transforms as T131            from torchvision.models import mobilenet_v2132 133            # For simplicity, use MobileNetV2 as a feature extractor134            # In production, would use actual OSNet from torchreid135            self.reid_model = mobilenet_v2(pretrained=True)136            self.reid_model.classifier = torch.nn.Identity()  # Remove classifier137 138            if self.config.device == "cuda" and torch.cuda.is_available():139                self.reid_model = self.reid_model.cuda()140 141            self.reid_model.eval()142 143            # Transform for body crops144            self._transform = T.Compose([145                T.ToPILImage(),146                T.Resize((256, 128)),147                T.ToTensor(),148                T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),149            ])150 151            logger.info("Re-ID model loaded (MobileNetV2 backbone)")152 153        except Exception as e:154            logger.warning(f"Failed to load re-ID model: {e}")155            self.reid_model = None156 157    def detect_persons(158        self,159        image: Union[str, Path, np.ndarray],160        min_confidence: float = 0.5,161        min_area: int = 2000,162    ) -> List[BodyDetection]:163        """164        Detect persons in an image.165 166        Args:167            image: Image path or numpy array (BGR format)168            min_confidence: Minimum detection confidence169            min_area: Minimum bounding box area170 171        Returns:172            List of BodyDetection objects173        """174        import cv2175 176        # Load image if path177        if isinstance(image, (str, Path)):178            img = cv2.imread(str(image))179            if img is None:180                raise InferenceError(f"Could not load image: {image}")181        else:182            img = image183 184        detections = []185 186        if self.detector is not None:187            try:188                # YOLO detection189                results = self.detector(img, classes=[0], verbose=False)  # class 0 = person190 191                for result in results:192                    for box in result.boxes:193                        conf = float(box.conf[0])194                        if conf < min_confidence:195                            continue196 197                        bbox = tuple(map(int, box.xyxy[0].tolist()))198                        area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])199 200                        if area < min_area:201                            continue202 203                        # Extract embedding204                        embedding = self._extract_embedding(img, bbox)205 206                        detections.append(BodyDetection(207                            bbox=bbox,208                            confidence=conf,209                            embedding=embedding,210                        ))211 212            except Exception as e:213                logger.warning(f"YOLO detection failed: {e}")214        else:215            # Fallback: assume full image is a person crop216            h, w = img.shape[:2]217            bbox = (0, 0, w, h)218            embedding = self._extract_embedding(img, bbox)219 220            detections.append(BodyDetection(221                bbox=bbox,222                confidence=1.0,223                embedding=embedding,224            ))225 226        logger.debug(f"Detected {len(detections)} persons")227        return detections228 229    def _extract_embedding(230        self,231        image: np.ndarray,232        bbox: Tuple[int, int, int, int],233    ) -> Optional[np.ndarray]:234        """Extract body appearance embedding."""235        if self.reid_model is None:236            return None237 238        try:239            import torch240 241            x1, y1, x2, y2 = bbox242            crop = image[y1:y2, x1:x2]243 244            if crop.size == 0:245                return None246 247            # Convert BGR to RGB248            crop_rgb = crop[:, :, ::-1]249 250            # Transform251            tensor = self._transform(crop_rgb).unsqueeze(0)252 253            if self.config.device == "cuda" and torch.cuda.is_available():254                tensor = tensor.cuda()255 256            # Extract features257            with torch.no_grad():258                embedding = self.reid_model(tensor)259                embedding = embedding.cpu().numpy()[0]260 261            # Normalize262            embedding = embedding / (np.linalg.norm(embedding) + 1e-8)263 264            return embedding265 266        except Exception as e:267            logger.debug(f"Embedding extraction failed: {e}")268            return None269 270    def register_reference(271        self,272        reference_image: Union[str, Path, np.ndarray],273        reference_id: str = "target",274        bbox: Optional[Tuple[int, int, int, int]] = None,275    ) -> bool:276        """277        Register a reference body appearance for matching.278 279        Args:280            reference_image: Image containing the reference person281            reference_id: Identifier for this reference282            bbox: Bounding box of person (auto-detected if None)283 284        Returns:285            True if registration successful286        """287        with LogTimer(logger, f"Registering body reference '{reference_id}'"):288            import cv2289 290            # Load image291            if isinstance(reference_image, (str, Path)):292                img = cv2.imread(str(reference_image))293            else:294                img = reference_image295 296            if bbox is None:297                # Detect person298                detections = self.detect_persons(img, min_confidence=0.5)299                if not detections:300                    raise InferenceError("No person detected in reference image")301 302                # Use largest detection303                detections.sort(key=lambda d: d.area, reverse=True)304                bbox = detections[0].bbox305 306            # Extract embedding307            embedding = self._extract_embedding(img, bbox)308 309            if embedding is None:310                raise InferenceError("Could not extract body embedding")311 312            self._reference_embeddings[reference_id] = embedding313            logger.info(f"Registered body reference: {reference_id}")314            return True315 316    def match_bodies(317        self,318        image: Union[str, Path, np.ndarray],319        reference_id: str = "target",320        threshold: Optional[float] = None,321    ) -> List[BodyMatch]:322        """323        Find body matches for a registered reference.324 325        Args:326            image: Image to search327            reference_id: Reference to match against328            threshold: Similarity threshold329 330        Returns:331            List of BodyMatch objects332        """333        threshold = threshold or self.config.body_similarity_threshold334 335        if reference_id not in self._reference_embeddings:336            logger.warning(f"Body reference '{reference_id}' not registered")337            return []338 339        reference = self._reference_embeddings[reference_id]340        detections = self.detect_persons(image)341 342        matches = []343        for detection in detections:344            if detection.embedding is None:345                continue346 347            similarity = self._cosine_similarity(reference, detection.embedding)348 349            matches.append(BodyMatch(350                detection=detection,351                similarity=similarity,352                is_match=similarity >= threshold,353                reference_id=reference_id,354            ))355 356        matches.sort(key=lambda m: m.similarity, reverse=True)357        return matches358 359    def find_target_in_frame(360        self,361        image: Union[str, Path, np.ndarray],362        reference_id: str = "target",363        threshold: Optional[float] = None,364    ) -> Optional[BodyMatch]:365        """366        Find the best matching body in a frame.367 368        Args:369            image: Frame to search370            reference_id: Reference to match against371            threshold: Similarity threshold372 373        Returns:374            Best BodyMatch if found, None otherwise375        """376        matches = self.match_bodies(image, reference_id, threshold)377        matching = [m for m in matches if m.is_match]378 379        if matching:380            return matching[0]381        return None382 383    def _cosine_similarity(384        self,385        embedding1: np.ndarray,386        embedding2: np.ndarray,387    ) -> float:388        """Compute cosine similarity."""389        return float(np.dot(embedding1, embedding2))390 391    def clear_references(self) -> None:392        """Clear all registered references."""393        self._reference_embeddings.clear()394        logger.info("Cleared all body references")395 396    def get_registered_references(self) -> List[str]:397        """Get list of registered reference IDs."""398        return list(self._reference_embeddings.keys())399 400 401# Export public interface402__all__ = ["BodyRecognizer", "BodyDetection", "BodyMatch"]403