""" ShortSmith v2 - Body Recognizer Module Full-body person recognition using OSNet for: - Identifying people when face is not visible - Back views, profile shots, masks, helmets - Clothing and appearance-based matching Complements face recognition for comprehensive person tracking. """ from pathlib import Path from typing import List, Optional, Tuple, Union from dataclasses import dataclass import numpy as np from utils.logger import get_logger, LogTimer from utils.helpers import ModelLoadError, InferenceError from config import get_config, ModelConfig logger = get_logger("models.body_recognizer") @dataclass class BodyDetection: """Represents a detected person body in an image.""" bbox: Tuple[int, int, int, int] # (x1, y1, x2, y2) confidence: float # Detection confidence embedding: Optional[np.ndarray] # Body appearance embedding track_id: Optional[int] = None # Tracking ID if available @property def center(self) -> Tuple[int, int]: """Center point of body bounding box.""" x1, y1, x2, y2 = self.bbox return ((x1 + x2) // 2, (y1 + y2) // 2) @property def area(self) -> int: """Area of bounding box.""" x1, y1, x2, y2 = self.bbox return (x2 - x1) * (y2 - y1) @property def width(self) -> int: return self.bbox[2] - self.bbox[0] @property def height(self) -> int: return self.bbox[3] - self.bbox[1] @property def aspect_ratio(self) -> float: """Height/width ratio (typical person is ~2.5-3.0).""" if self.width == 0: return 0 return self.height / self.width @dataclass class BodyMatch: """Result of body matching.""" detection: BodyDetection similarity: float is_match: bool reference_id: Optional[str] = None class BodyRecognizer: """ Body recognition using person re-identification models. Uses: - YOLO or similar for person detection - OSNet for body appearance embeddings Designed to work alongside FaceRecognizer for complete person identification across all viewing angles. """ def __init__( self, config: Optional[ModelConfig] = None, load_model: bool = True, ): """ Initialize body recognizer. Args: config: Model configuration load_model: Whether to load models immediately """ self.config = config or get_config().model self.detector = None self.reid_model = None self._reference_embeddings: dict = {} if load_model: self._load_models() logger.info(f"BodyRecognizer initialized (threshold={self.config.body_similarity_threshold})") def _load_models(self) -> None: """Load person detection and re-identification models.""" with LogTimer(logger, "Loading body recognition models"): self._load_detector() self._load_reid_model() def _load_detector(self) -> None: """Load person detector (YOLO).""" try: from ultralytics import YOLO # Use YOLOv8 for person detection self.detector = YOLO("yolov8n.pt") # Nano model for speed logger.info("YOLO detector loaded") except ImportError: logger.warning("ultralytics not installed, using fallback detection") self.detector = None except Exception as e: logger.warning(f"Failed to load YOLO detector: {e}") self.detector = None def _load_reid_model(self) -> None: """Load OSNet re-identification model.""" try: import torch import torchvision.transforms as T from torchvision.models import mobilenet_v2 # For simplicity, use MobileNetV2 as a feature extractor # In production, would use actual OSNet from torchreid self.reid_model = mobilenet_v2(pretrained=True) self.reid_model.classifier = torch.nn.Identity() # Remove classifier if self.config.device == "cuda" and torch.cuda.is_available(): self.reid_model = self.reid_model.cuda() self.reid_model.eval() # Transform for body crops self._transform = T.Compose([ T.ToPILImage(), T.Resize((256, 128)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) logger.info("Re-ID model loaded (MobileNetV2 backbone)") except Exception as e: logger.warning(f"Failed to load re-ID model: {e}") self.reid_model = None def detect_persons( self, image: Union[str, Path, np.ndarray], min_confidence: float = 0.5, min_area: int = 2000, ) -> List[BodyDetection]: """ Detect persons in an image. Args: image: Image path or numpy array (BGR format) min_confidence: Minimum detection confidence min_area: Minimum bounding box area Returns: List of BodyDetection objects """ import cv2 # Load image if path if isinstance(image, (str, Path)): img = cv2.imread(str(image)) if img is None: raise InferenceError(f"Could not load image: {image}") else: img = image detections = [] if self.detector is not None: try: # YOLO detection results = self.detector(img, classes=[0], verbose=False) # class 0 = person for result in results: for box in result.boxes: conf = float(box.conf[0]) if conf < min_confidence: continue bbox = tuple(map(int, box.xyxy[0].tolist())) area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]) if area < min_area: continue # Extract embedding embedding = self._extract_embedding(img, bbox) detections.append(BodyDetection( bbox=bbox, confidence=conf, embedding=embedding, )) except Exception as e: logger.warning(f"YOLO detection failed: {e}") else: # Fallback: assume full image is a person crop h, w = img.shape[:2] bbox = (0, 0, w, h) embedding = self._extract_embedding(img, bbox) detections.append(BodyDetection( bbox=bbox, confidence=1.0, embedding=embedding, )) logger.debug(f"Detected {len(detections)} persons") return detections def _extract_embedding( self, image: np.ndarray, bbox: Tuple[int, int, int, int], ) -> Optional[np.ndarray]: """Extract body appearance embedding.""" if self.reid_model is None: return None try: import torch x1, y1, x2, y2 = bbox crop = image[y1:y2, x1:x2] if crop.size == 0: return None # Convert BGR to RGB crop_rgb = crop[:, :, ::-1] # Transform tensor = self._transform(crop_rgb).unsqueeze(0) if self.config.device == "cuda" and torch.cuda.is_available(): tensor = tensor.cuda() # Extract features with torch.no_grad(): embedding = self.reid_model(tensor) embedding = embedding.cpu().numpy()[0] # Normalize embedding = embedding / (np.linalg.norm(embedding) + 1e-8) return embedding except Exception as e: logger.debug(f"Embedding extraction failed: {e}") return None def register_reference( self, reference_image: Union[str, Path, np.ndarray], reference_id: str = "target", bbox: Optional[Tuple[int, int, int, int]] = None, ) -> bool: """ Register a reference body appearance for matching. Args: reference_image: Image containing the reference person reference_id: Identifier for this reference bbox: Bounding box of person (auto-detected if None) Returns: True if registration successful """ with LogTimer(logger, f"Registering body reference '{reference_id}'"): import cv2 # Load image if isinstance(reference_image, (str, Path)): img = cv2.imread(str(reference_image)) else: img = reference_image if bbox is None: # Detect person detections = self.detect_persons(img, min_confidence=0.5) if not detections: raise InferenceError("No person detected in reference image") # Use largest detection detections.sort(key=lambda d: d.area, reverse=True) bbox = detections[0].bbox # Extract embedding embedding = self._extract_embedding(img, bbox) if embedding is None: raise InferenceError("Could not extract body embedding") self._reference_embeddings[reference_id] = embedding logger.info(f"Registered body reference: {reference_id}") return True def match_bodies( self, image: Union[str, Path, np.ndarray], reference_id: str = "target", threshold: Optional[float] = None, ) -> List[BodyMatch]: """ Find body matches for a registered reference. Args: image: Image to search reference_id: Reference to match against threshold: Similarity threshold Returns: List of BodyMatch objects """ threshold = threshold or self.config.body_similarity_threshold if reference_id not in self._reference_embeddings: logger.warning(f"Body reference '{reference_id}' not registered") return [] reference = self._reference_embeddings[reference_id] detections = self.detect_persons(image) matches = [] for detection in detections: if detection.embedding is None: continue similarity = self._cosine_similarity(reference, detection.embedding) matches.append(BodyMatch( detection=detection, similarity=similarity, is_match=similarity >= threshold, reference_id=reference_id, )) matches.sort(key=lambda m: m.similarity, reverse=True) return matches def find_target_in_frame( self, image: Union[str, Path, np.ndarray], reference_id: str = "target", threshold: Optional[float] = None, ) -> Optional[BodyMatch]: """ Find the best matching body in a frame. Args: image: Frame to search reference_id: Reference to match against threshold: Similarity threshold Returns: Best BodyMatch if found, None otherwise """ matches = self.match_bodies(image, reference_id, threshold) matching = [m for m in matches if m.is_match] if matching: return matching[0] return None def _cosine_similarity( self, embedding1: np.ndarray, embedding2: np.ndarray, ) -> float: """Compute cosine similarity.""" return float(np.dot(embedding1, embedding2)) def clear_references(self) -> None: """Clear all registered references.""" self._reference_embeddings.clear() logger.info("Cleared all body references") def get_registered_references(self) -> List[str]: """Get list of registered reference IDs.""" return list(self._reference_embeddings.keys()) # Export public interface __all__ = ["BodyRecognizer", "BodyDetection", "BodyMatch"]