diff --git a/.gitignore b/.gitignore index 16062a91..b30e578a 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,24 @@ weight/ wandb/ *.pth *.onnx +*.mp4 +datasets/ +*.db +outputs/ +*.engine +venv/ + +# Editor +.vscode/ +.history/ +.codex + +# Large binary / data folders +Traning/ +segmentation_sivert/data/ + +# Cache +__pycache__/ +*.egg-info/ +.pytest_cache/ +DETRPose/ diff --git a/configs/dfine_hgnetv2_x_obj2coco.yml b/configs/dfine_hgnetv2_x_obj2coco.yml new file mode 100644 index 00000000..b1f66072 --- /dev/null +++ b/configs/dfine_hgnetv2_x_obj2coco.yml @@ -0,0 +1,61 @@ +__include__: [ + './dataset/coco_detection.yml', + './runtime.yml', + './dfine/include/dataloader.yml', + './dfine/include/optimizer.yml', + './dfine/include/dfine_hgnetv2.yml', +] + +output_dir: ./output/dfine_hgnetv2_x_obj2coco + + +DFINE: + backbone: HGNetv2 + +HGNetv2: + name: 'B5' + return_idx: [1, 2, 3] + freeze_stem_only: True + freeze_at: 0 + freeze_norm: True + +HybridEncoder: + # intra + hidden_dim: 384 + dim_feedforward: 2048 + +DFINETransformer: + feat_channels: [384, 384, 384] + reg_scale: 8 + +optimizer: + type: AdamW + params: + - + params: '^(?=.*backbone)(?!.*norm|bn).*$' + lr: 0.0000025 + - + params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$' + weight_decay: 0. + + lr: 0.00025 + betas: [0.9, 0.999] + weight_decay: 0.000125 + + +epochs: 36 # Early stop +train_dataloader: + dataset: + transforms: + policy: + epoch: 30 + collate_fn: + stop_epoch: 30 + ema_restart_decay: 0.9999 + base_size_repeat: 3 + +ema: + warmups: 0 + +lr_warmup_scheduler: + warmup_duration: 0 diff --git a/detection_training/configs/soldier_finetune.yml b/detection_training/configs/soldier_finetune.yml new file mode 100644 index 00000000..03ac9b0b --- /dev/null +++ b/detection_training/configs/soldier_finetune.yml @@ -0,0 +1,62 @@ +__include__: [ + '../../configs/dataset/coco_detection.yml', + '../../configs/runtime.yml', + '../../configs/dfine/include/dataloader.yml', + '../../configs/dfine/include/optimizer.yml', + '../../configs/dfine/include/dfine_hgnetv2.yml', +] + +output_dir: ./outputs/soldier_finetune + +num_classes: 2 +remap_mscoco_category: False + +HGNetv2: + name: 'B5' + return_idx: [1, 2, 3] + freeze_stem_only: True + freeze_at: 0 + freeze_norm: True + +HybridEncoder: + hidden_dim: 384 + dim_feedforward: 2048 + +DFINETransformer: + feat_channels: [384, 384, 384] + reg_scale: 8 + num_classes: 2 + +epochs: 60 +lr_warmup_scheduler: + warmup_duration: 100 + +optimizer: + type: AdamW + params: + - params: '^(?=.*backbone)(?!.*norm|bn).*$' + lr: 0.000005 + - params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$' + weight_decay: 0. + lr: 0.00005 + betas: [0.9, 0.999] + weight_decay: 0.0001 + +train_dataloader: + dataset: + img_folder: ./datasets/soldier_coco/images/train + ann_file: ./datasets/soldier_coco/annotations/instances_train.json + total_batch_size: 2 + collate_fn: + stop_epoch: 50 + shuffle: True + +val_dataloader: + dataset: + img_folder: ./datasets/soldier_coco/images/val + ann_file: ./datasets/soldier_coco/annotations/instances_val.json + total_batch_size: 2 + +evaluator: + type: CocoEvaluator + iou_types: ['bbox'] diff --git a/detection_training/merge_checkpoints.py b/detection_training/merge_checkpoints.py new file mode 100644 index 00000000..0f731625 --- /dev/null +++ b/detection_training/merge_checkpoints.py @@ -0,0 +1,95 @@ +#!/usr/bin/env python3 +""" +Model Surgery: merge new soldier-finetuned backbone into original combined checkpoint. + +Strategy: BACKBONE-ONLY transfer from the new soldier-finetuned model. + - backbone → from new fine-tune (learned to find camouflaged soldiers) + - encoder → from original (keeps encoder features compatible with pose decoder) + - det_decoder → from original (80-class COCO, detects 'person') + - pose_decoder → from original (unchanged, compatible with original encoder) + - seg_head → from original (unchanged) + +Why backbone-only: the pose decoder was trained on the original encoder's feature +representations. Swapping encoder+decoder breaks pose (encoder features change → +pose decoder hallucinates keypoints everywhere). Backbone-only surgery preserves +all downstream compatibility while still improving low-level feature extraction +for camouflaged persons. +""" + +import argparse +import torch +from pathlib import Path + + +def load_state(path: str): + ckpt = torch.load(path, map_location="cpu") + if isinstance(ckpt, dict): + if "ema" in ckpt and isinstance(ckpt["ema"], dict): + m = ckpt["ema"].get("module") + if m is not None: + return m + if "model" in ckpt and isinstance(ckpt["model"], dict): + return ckpt["model"] + if all(isinstance(v, torch.Tensor) for v in ckpt.values()): + return ckpt + raise RuntimeError(f"Cannot extract state dict from {path}") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--new-det", default="outputs/soldier_finetune/best_stg2.pth", + help="New soldier-finetuned detection checkpoint") + ap.add_argument("--orig", default="Traning/best_modelsurgery.pth", + help="Original combined checkpoint (all heads)") + ap.add_argument("--out", default="outputs/soldier_finetune/merged_surgery.pth", + help="Output merged checkpoint path") + args = ap.parse_args() + + print(f"Loading new detection checkpoint: {args.new_det}") + new_state = load_state(args.new_det) + + print(f"Loading original surgery checkpoint: {args.orig}") + orig_state = load_state(args.orig) + + merged = {} + counts = {"backbone_new": 0, "encoder_orig": 0, "det_decoder_orig": 0, + "pose_decoder_orig": 0, "seg_head_orig": 0} + + # 1. Backbone: from NEW model — learns camouflage-aware low-level features + for k, v in new_state.items(): + if k.startswith("backbone."): + merged[k] = v + counts["backbone_new"] += 1 + + # 2-5. Everything else: from ORIGINAL — preserves encoder↔pose_decoder compatibility + for k, v in orig_state.items(): + if k.startswith("encoder."): + merged[k] = v + counts["encoder_orig"] += 1 + elif k.startswith("det_decoder."): + merged[k] = v + counts["det_decoder_orig"] += 1 + elif k.startswith("pose_decoder."): + merged[k] = v + counts["pose_decoder_orig"] += 1 + elif k.startswith("seg_head."): + merged[k] = v + counts["seg_head_orig"] += 1 + + print("\n=== Merge summary ===") + for part, n in counts.items(): + print(f" {part}: {n} tensors") + print(f" TOTAL: {sum(counts.values())} tensors") + + score_key = "det_decoder.enc_score_head.weight" + if score_key in merged: + print(f"\nDetection class head shape: {merged[score_key].shape}") + + out_path = Path(args.out) + out_path.parent.mkdir(parents=True, exist_ok=True) + torch.save({"model": merged}, str(out_path)) + print(f"\nSaved merged checkpoint → {out_path}") + + +if __name__ == "__main__": + main() diff --git a/detection_training/train.py b/detection_training/train.py new file mode 100644 index 00000000..c9538475 --- /dev/null +++ b/detection_training/train.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python3 +""" +Detection finetuning for camouflaged soldier detection. +Finetunes the DFINE detection backbone on soldier/civilian dataset. + +Usage: + python detection_training/train.py \ + --config detection_training/configs/soldier_finetune.yml \ + --resume outputs/phase2_run/best_modelsurgery.pth \ + --finetune +""" + +import argparse +import os +import sys + +# Ensure repo root is on path +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, REPO_ROOT) + +import torch +from src.core import YAMLConfig +from src.solver import DetSolver + + +def get_args(): + parser = argparse.ArgumentParser('DFINE soldier detection finetuning') + parser.add_argument('--config', '-c', type=str, + default='detection_training/configs/soldier_finetune.yml') + parser.add_argument('--resume', '-r', type=str, + default='outputs/phase2_run/best_modelsurgery.pth', + help='Checkpoint to resume or finetune from') + parser.add_argument('--finetune', action='store_true', + help='Load weights only (ignore optimizer/scheduler state)') + parser.add_argument('--test-only', action='store_true') + parser.add_argument('--amp', action='store_true', default=True) + return parser.parse_args() + + +def main(): + args = get_args() + + cfg = YAMLConfig(args.config, resume=args.resume if not args.finetune else None) + + if args.finetune and args.resume: + print(f'Finetuning from: {args.resume}') + ckpt = torch.load(args.resume, map_location='cpu') + # The surgery checkpoint stores state under 'model' key + state = ckpt.get('model', ckpt) + # Strip 'module.' prefix if present + state = {k.replace('module.', ''): v for k, v in state.items()} + # Load only matching keys (detection head; ignore pose/seg heads) + missing, unexpected = cfg.model.load_state_dict(state, strict=False) + print(f' Loaded weights — missing: {len(missing)}, unexpected: {len(unexpected)}') + if missing: + print(f' Missing keys (first 5): {missing[:5]}') + + if args.test_only: + solver = DetSolver(cfg) + solver.val() + return + + solver = DetSolver(cfg) + solver.fit() + + +if __name__ == '__main__': + main() diff --git a/pose_estimation_berna/base_dfine/dfine_hgnetv2_x_obj2coco_detrpose_paper.yml b/pose_estimation_berna/base_dfine/dfine_hgnetv2_x_obj2coco_detrpose_paper.yml new file mode 100644 index 00000000..93becf8d --- /dev/null +++ b/pose_estimation_berna/base_dfine/dfine_hgnetv2_x_obj2coco_detrpose_paper.yml @@ -0,0 +1,55 @@ +__include__: [ + '../../configs/dataset/coco_detection.yml', + '../../configs/runtime.yml', + '../../configs/dfine/include/dataloader.yml', + '../../configs/dfine/include/optimizer.yml', + '../../configs/dfine/include/dfine_hgnetv2.yml', +] + +output_dir: ./output/dfine_hgnetv2_x_obj2coco_detrpose_paper + +# IMPORTANT: +# coco_detection.yml sets num_classes=80 globally. DETRPoseTransformer shares `num_classes`, +# so without overriding it here, the decoder will silently become 80-class. +num_classes: 2 + +DFINE: + backbone: HGNetv2 + encoder: HybridEncoder + decoder: DETRPoseTransformer + +HGNetv2: + name: 'B5' + return_idx: [1, 2, 3] + freeze_stem_only: True + freeze_at: 0 + freeze_norm: True + +# X-size encoder: hidden_dim=384 to match D-FINE-X detection config +HybridEncoder: + hidden_dim: 384 + dim_feedforward: 2048 + +# Paper-like DETRPose defaults — hidden_dim=384 to match encoder output +DETRPoseTransformer: + num_classes: 2 + num_queries: 60 + hidden_dim: 384 + num_decoder_layers: 6 + dim_feedforward: 2048 + dropout: 0.0 + activation: relu + num_feature_levels: 3 + dec_n_points: 4 + nhead: 8 + aux_loss: True + num_body_points: 17 + feat_strides: [8, 16, 32] + eval_spatial_size: [640, 640] + reg_max: 32 + reg_scale: 8.0 + dn_number: 20 + dn_label_noise_ratio: 0.5 + +DFINEPostProcessor: + num_top_queries: 60 \ No newline at end of file diff --git a/pose_estimation_berna/core/__init__.py b/pose_estimation_berna/core/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/pose_estimation_berna/core/calibrators.py b/pose_estimation_berna/core/calibrators.py new file mode 100644 index 00000000..48ce6648 --- /dev/null +++ b/pose_estimation_berna/core/calibrators.py @@ -0,0 +1,350 @@ +#!/usr/bin/env python3 +""" +Core TensorRT Calibrators for DFINE Segmentation +Reusable INT8 calibration classes for different datasets +""" + +import os +import numpy as np +import random +from typing import List, Optional +import tensorrt as trt +import pycuda.driver as cuda +import pycuda.autoinit +import cv2 +from PIL import Image +import torchvision.transforms as transforms +from abc import ABC, abstractmethod + + +class BaseCalibrator(trt.IInt8EntropyCalibrator2, ABC): + """Base class for INT8 calibrators""" + + def __init__(self, + batch_size: int = 1, + cache_file: str = "calibration.cache", + max_calibration_images: int = 500, + input_shape: tuple = (640, 640)): + super().__init__() + + self.batch_size = batch_size + self.cache_file = cache_file + self.max_calibration_images = max_calibration_images + self.input_shape = input_shape + self.current_index = 0 + + # Image preprocessing + self.transform = transforms.Compose([ + transforms.Resize(input_shape), + transforms.ToTensor(), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + ]) + + # Load image paths + self.image_files = self._load_image_paths() + print(f"📊 Loaded {len(self.image_files)} images for calibration") + + # Allocate GPU memory + self._allocate_memory() + + @abstractmethod + def _load_image_paths(self) -> List[str]: + """Load paths to calibration images""" + pass + + def _allocate_memory(self): + """Allocate GPU memory for calibration""" + # Input image memory + image_size = self.batch_size * 3 * self.input_shape[0] * self.input_shape[1] * 4 # float32 + self.device_input_image = cuda.mem_alloc(image_size) + + # Input sizes memory (for DFINE model) + sizes_size = self.batch_size * 2 * 8 # int64 + self.device_input_sizes = cuda.mem_alloc(sizes_size) + + # Host memory for sizes + self.host_sizes = np.array([[self.input_shape[1], self.input_shape[0]] for _ in range(self.batch_size)], dtype=np.int64) + + # CUDA stream + self.stream = cuda.Stream() + + def get_batch_size(self): + """Return batch size""" + return self.batch_size + + def get_batch(self, names): + """Get next batch for calibration""" + if self.current_index + self.batch_size > len(self.image_files): + print(f"📊 Calibration complete. Processed {self.current_index} images.") + return None + + batch_progress = self.current_index // self.batch_size + 1 + total_batches = len(self.image_files) // self.batch_size + 1 + print(f"📊 Calibration batch {batch_progress}/{total_batches}") + + # Prepare batch data + batch_data = np.zeros((self.batch_size, 3, *self.input_shape), dtype=np.float32) + + valid_images = 0 + for i in range(self.batch_size): + if self.current_index + i >= len(self.image_files): + break + + image_path = self.image_files[self.current_index + i] + + try: + pil_image = Image.open(image_path).convert('RGB') + processed_image = self.transform(pil_image) + batch_data[i] = processed_image.numpy() + valid_images += 1 + + if i == 0: + print(f" Processing: {os.path.basename(image_path)}") + + except Exception as e: + print(f" Error processing {image_path}: {e}") + continue + + if valid_images == 0: + self.current_index += self.batch_size + return self.get_batch(names) + + # Copy to GPU + cuda.memcpy_htod_async(self.device_input_image, batch_data.ravel(), self.stream) + cuda.memcpy_htod_async(self.device_input_sizes, self.host_sizes.ravel(), self.stream) + self.stream.synchronize() + + self.current_index += self.batch_size + + # Return bindings based on input names + bindings = [] + for name in names: + if name in ["images", "input"]: + bindings.append(int(self.device_input_image)) + elif name in ["orig_target_sizes", "sizes"]: + bindings.append(int(self.device_input_sizes)) + else: + print(f"⚠️ Unknown binding name: {name}") + bindings.append(0) + + return bindings + + def read_calibration_cache(self): + """Read calibration cache if exists""" + if os.path.exists(self.cache_file): + with open(self.cache_file, "rb") as f: + cache_data = f.read() + print(f"📁 Read {len(cache_data)} bytes from calibration cache: {self.cache_file}") + return cache_data + return None + + def write_calibration_cache(self, cache): + """Write calibration cache""" + with open(self.cache_file, "wb") as f: + f.write(cache) + print(f"💾 Wrote {len(cache)} bytes to calibration cache: {self.cache_file}") + + +class PascalPersonPartsCalibrator(BaseCalibrator): + """Pascal Person Parts calibrator for human segmentation optimization""" + + def __init__(self, + pascal_data_dir: str, + batch_size: int = 1, + cache_file: str = "pascal_calibration.cache", + max_calibration_images: int = 500, + input_shape: tuple = (640, 640)): + + self.pascal_data_dir = pascal_data_dir + super().__init__(batch_size, cache_file, max_calibration_images, input_shape) + + def _load_image_paths(self) -> List[str]: + """Load Pascal Person Parts image paths""" + image_files = [] + + # Search in multiple possible directories + search_dirs = [ + os.path.join(self.pascal_data_dir, 'images', 'val'), + os.path.join(self.pascal_data_dir, 'images', 'train'), + os.path.join(self.pascal_data_dir, 'images'), + self.pascal_data_dir + ] + + for img_dir in search_dirs: + if os.path.exists(img_dir): + print(f"🔍 Searching Pascal images in: {img_dir}") + for root, dirs, files in os.walk(img_dir): + for file in files: + if file.lower().endswith(('.jpg', '.jpeg', '.png')): + image_files.append(os.path.join(root, file)) + + if len(image_files) >= self.max_calibration_images: + break + + if len(image_files) == 0: + raise ValueError(f"No Pascal images found in {self.pascal_data_dir}") + + # Limit and shuffle + if len(image_files) > self.max_calibration_images: + random.shuffle(image_files) + image_files = image_files[:self.max_calibration_images] + + return image_files + + +class COCOCalibrator(BaseCalibrator): + """COCO calibrator for general detection optimization""" + + def __init__(self, + coco_data_dir: str, + batch_size: int = 1, + cache_file: str = "coco_calibration.cache", + max_calibration_images: int = 500, + input_shape: tuple = (640, 640)): + + self.coco_data_dir = coco_data_dir + super().__init__(batch_size, cache_file, max_calibration_images, input_shape) + + def _load_image_paths(self) -> List[str]: + """Load COCO image paths""" + image_files = [] + + if os.path.exists(self.coco_data_dir): + print(f"🔍 Searching COCO images in: {self.coco_data_dir}") + for root, dirs, files in os.walk(self.coco_data_dir): + for file in files: + if file.lower().endswith(('.jpg', '.jpeg', '.png')): + image_files.append(os.path.join(root, file)) + + if len(image_files) == 0: + raise ValueError(f"No COCO images found in {self.coco_data_dir}") + + # Limit and shuffle + if len(image_files) > self.max_calibration_images: + random.shuffle(image_files) + image_files = image_files[:self.max_calibration_images] + + return image_files + + +class MixedDatasetCalibrator(BaseCalibrator): + """Mixed Pascal + COCO calibrator for balanced optimization""" + + def __init__(self, + pascal_data_dir: str, + coco_data_dir: str, + pascal_ratio: float = 0.7, + batch_size: int = 1, + cache_file: str = "mixed_calibration.cache", + max_calibration_images: int = 500, + input_shape: tuple = (640, 640)): + + self.pascal_data_dir = pascal_data_dir + self.coco_data_dir = coco_data_dir + self.pascal_ratio = pascal_ratio + super().__init__(batch_size, cache_file, max_calibration_images, input_shape) + + def _load_image_paths(self) -> List[str]: + """Load mixed dataset image paths""" + # Collect Pascal images + pascal_images = [] + pascal_search_dirs = [ + os.path.join(self.pascal_data_dir, 'images', 'val'), + os.path.join(self.pascal_data_dir, 'images', 'train'), + os.path.join(self.pascal_data_dir, 'images'), + self.pascal_data_dir + ] + + for img_dir in pascal_search_dirs: + if os.path.exists(img_dir): + print(f"🔍 Searching Pascal images in: {img_dir}") + for root, dirs, files in os.walk(img_dir): + for file in files: + if file.lower().endswith(('.jpg', '.jpeg', '.png')): + pascal_images.append(os.path.join(root, file)) + break + + # Collect COCO images + coco_images = [] + if os.path.exists(self.coco_data_dir): + print(f"🔍 Searching COCO images in: {self.coco_data_dir}") + for root, dirs, files in os.walk(self.coco_data_dir): + for file in files: + if file.lower().endswith(('.jpg', '.jpeg', '.png')): + coco_images.append(os.path.join(root, file)) + + # Mix datasets according to ratio + num_pascal = min(len(pascal_images), int(self.max_calibration_images * self.pascal_ratio)) + num_coco = min(len(coco_images), self.max_calibration_images - num_pascal) if coco_images else 0 + + # Sample images + selected_pascal = [] + if pascal_images: + random.shuffle(pascal_images) + selected_pascal = pascal_images[:num_pascal] + + selected_coco = [] + if coco_images and num_coco > 0: + random.shuffle(coco_images) + selected_coco = coco_images[:num_coco] + + # Combine and shuffle + image_files = selected_pascal + selected_coco + random.shuffle(image_files) + + print(f"📊 Mixed calibration dataset:") + print(f" Pascal images: {len(selected_pascal)}") + print(f" COCO images: {len(selected_coco)}") + print(f" Total: {len(image_files)}") + + if len(image_files) == 0: + raise ValueError("No calibration images found in either dataset") + + return image_files + + +# Factory functions for creating calibrators +def create_pascal_calibrator(pascal_data_dir: str, **kwargs) -> PascalPersonPartsCalibrator: + """Create Pascal Person Parts calibrator""" + return PascalPersonPartsCalibrator(pascal_data_dir, **kwargs) + + +def create_coco_calibrator(coco_data_dir: str, **kwargs) -> COCOCalibrator: + """Create COCO calibrator""" + return COCOCalibrator(coco_data_dir, **kwargs) + + +def create_mixed_calibrator(pascal_data_dir: str, coco_data_dir: str, **kwargs) -> MixedDatasetCalibrator: + """Create mixed dataset calibrator""" + return MixedDatasetCalibrator(pascal_data_dir, coco_data_dir, **kwargs) + + +def create_calibrator(mode: str, + pascal_data_dir: Optional[str] = None, + coco_data_dir: Optional[str] = None, + **kwargs): + """Factory function to create calibrators based on mode""" + + if mode == 'pascal': + if not pascal_data_dir: + raise ValueError("pascal_data_dir required for Pascal mode") + return create_pascal_calibrator(pascal_data_dir, **kwargs) + + elif mode == 'coco': + if not coco_data_dir: + raise ValueError("coco_data_dir required for COCO mode") + return create_coco_calibrator(coco_data_dir, **kwargs) + + elif mode == 'mixed': + if not pascal_data_dir or not coco_data_dir: + raise ValueError("Both pascal_data_dir and coco_data_dir required for mixed mode") + return create_mixed_calibrator(pascal_data_dir, coco_data_dir, **kwargs) + + else: + raise ValueError(f"Unknown calibration mode: {mode}") + + +# Legacy aliases for backward compatibility +class COCOOnlyCalibrator(COCOCalibrator): + """Legacy alias""" + pass \ No newline at end of file diff --git a/pose_estimation_berna/core/datasets.py b/pose_estimation_berna/core/datasets.py new file mode 100644 index 00000000..045019b0 --- /dev/null +++ b/pose_estimation_berna/core/datasets.py @@ -0,0 +1,220 @@ +#!/usr/bin/env python3 +""" +COCO Keypoints Dataset for D-FINE Pose Estimation (COCO-17). + +This supports the **per-detection / per-query** approach: +- D-FINE predicts boxes per query +- We predict 17 keypoints per query, **bbox-relative**: [B, Q, 17, 3] + where 3 = (x_rel, y_rel, visibility_logit) + +Targets returned: +- boxes: [N, 4] (cxcywh normalized to [0,1]) +- labels: [N] (COCO contiguous label, person is 0) +- keypoints: [N, 17, 3] (x_px, y_px, v) in resized image pixels, v in {0,1,2} +- size: [2] (h, w) in resized image pixels +- orig_size: [2] (w, h) in original image pixels +""" + +import json +import os +import random +from dataclasses import dataclass +from typing import Dict, List + +import numpy as np +import torch +from torch.utils.data import Dataset +from PIL import Image + + +COCO_KEYPOINT_FLIP_INDEX = [ + 0, # nose + 2, 1, # left_eye <-> right_eye + 4, 3, # left_ear <-> right_ear + 6, 5, # left_shoulder <-> right_shoulder + 8, 7, # left_elbow <-> right_elbow + 10, 9, # left_wrist <-> right_wrist + 12, 11, # left_hip <-> right_hip + 14, 13, # left_knee <-> right_knee + 16, 15, # left_ankle <-> right_ankle +] + + +@dataclass +class CocoImageRecord: + id: int + file_name: str + width: int + height: int + + +class CocoKeypointsDataset(Dataset): + def __init__( + self, + root_dir: str, + split: str = "train", + image_size: int = 640, + num_keypoints: int = 17, + flip_prob: float = 0.5, + return_dict: bool = False, + ): + assert split in ("train", "val"), "split must be 'train' or 'val'" + self.root_dir = root_dir + self.split = split + self.image_size = int(image_size) + self.num_keypoints = int(num_keypoints) + self.flip_prob = float(flip_prob if split == "train" else 0.0) + self.return_dict = bool(return_dict) + + self.img_dir = os.path.join(root_dir, f"{split}2017") + self.ann_path = os.path.join( + root_dir, "annotations", f"person_keypoints_{split}2017.json" + ) + + with open(self.ann_path, "r") as f: + coco = json.load(f) + + self.images: Dict[int, CocoImageRecord] = { + img["id"]: CocoImageRecord( + id=img["id"], + file_name=img["file_name"], + width=img["width"], + height=img["height"], + ) + for img in coco["images"] + } + + anns_by_img: Dict[int, List[dict]] = {} + for ann in coco["annotations"]: + if ann.get("iscrowd", 0) == 1: + continue + if ann.get("category_id") != 1: # person + continue + if "keypoints" not in ann: + continue + if int(ann.get("num_keypoints", 0)) <= 0: + continue + anns_by_img.setdefault(ann["image_id"], []).append(ann) + + self.ids: List[int] = [img_id for img_id in anns_by_img.keys() if img_id in self.images] + self.anns_by_img = anns_by_img + + # ImageNet normalization + self.mean = torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32) + self.std = torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32) + + print(f"📊 Loaded COCO-{split} keypoints: {len(self.ids)} images from {self.ann_path}") + + def __len__(self): + return len(self.ids) + + @staticmethod + def _xywh_to_xyxy(box: np.ndarray) -> np.ndarray: + x, y, w, h = box.tolist() + return np.array([x, y, x + w, y + h], dtype=np.float32) + + @staticmethod + def _xyxy_to_cxcywh(box: np.ndarray) -> np.ndarray: + x1, y1, x2, y2 = box.tolist() + cx = (x1 + x2) * 0.5 + cy = (y1 + y2) * 0.5 + w = (x2 - x1) + h = (y2 - y1) + return np.array([cx, cy, w, h], dtype=np.float32) + + def __getitem__(self, idx: int): + image_id = self.ids[idx] + rec = self.images[image_id] + + img_path = os.path.join(self.img_dir, rec.file_name) + try: + with Image.open(img_path) as im: + im = im.convert("RGB") + orig_w, orig_h = im.size + scale_x = self.image_size / float(orig_w) + scale_y = self.image_size / float(orig_h) + im = im.resize((self.image_size, self.image_size), resample=Image.BILINEAR) + image = np.array(im, dtype=np.uint8) + except Exception as e: + raise FileNotFoundError(f"Could not read image: {img_path} ({e})") from e + + anns = self.anns_by_img.get(image_id, []) + boxes_xyxy = [] + kpts = [] + + for ann in anns: + bbox_xywh = np.array(ann["bbox"], dtype=np.float32) + bbox_xywh[[0, 2]] *= scale_x + bbox_xywh[[1, 3]] *= scale_y + box_xyxy = self._xywh_to_xyxy(bbox_xywh) + + box_xyxy[0::2] = np.clip(box_xyxy[0::2], 0, self.image_size - 1) + box_xyxy[1::2] = np.clip(box_xyxy[1::2], 0, self.image_size - 1) + + kp = np.array(ann["keypoints"], dtype=np.float32).reshape(-1, 3) + kp[:, 0] *= scale_x + kp[:, 1] *= scale_y + + boxes_xyxy.append(box_xyxy) + kpts.append(kp) + + boxes_xyxy = np.stack(boxes_xyxy, axis=0).astype(np.float32) # [N,4] + kpts = np.stack(kpts, axis=0).astype(np.float32) # [N,17,3] + + if self.flip_prob > 0 and random.random() < self.flip_prob: + image = np.ascontiguousarray(image[:, ::-1, :]) + + x1 = boxes_xyxy[:, 0].copy() + x2 = boxes_xyxy[:, 2].copy() + boxes_xyxy[:, 0] = (self.image_size - 1) - x2 + boxes_xyxy[:, 2] = (self.image_size - 1) - x1 + + kpts[:, :, 0] = (self.image_size - 1) - kpts[:, :, 0] + kpts = kpts[:, COCO_KEYPOINT_FLIP_INDEX, :] + + # boxes to normalized cxcywh + boxes_cxcywh = np.stack([self._xyxy_to_cxcywh(b) for b in boxes_xyxy], axis=0) + boxes_cxcywh[:, 0] /= self.image_size + boxes_cxcywh[:, 1] /= self.image_size + boxes_cxcywh[:, 2] /= self.image_size + boxes_cxcywh[:, 3] /= self.image_size + boxes_cxcywh = np.clip(boxes_cxcywh, 0.0, 1.0) + + image_t = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0 + image_t = (image_t - self.mean[:, None, None]) / self.std[:, None, None] + + target = { + "boxes": torch.from_numpy(boxes_cxcywh).float(), + "labels": torch.zeros((boxes_cxcywh.shape[0],), dtype=torch.int64), + "keypoints": torch.from_numpy(kpts).float(), + "image_id": torch.tensor([image_id], dtype=torch.int64), + "orig_size": torch.tensor([orig_w, orig_h], dtype=torch.int64), + "size": torch.tensor([self.image_size, self.image_size], dtype=torch.int64), + } + + if self.return_dict: + return {"image": image_t, **target} + + return image_t, target + + +def create_coco_pose_dataset( + root_dir: str, + split: str = "train", + image_size: int = 640, + num_keypoints: int = 17, + tier: str = "standard", + return_dict: bool = False, +): + # tier can drive aug strength later; for now only flip in train + flip_prob = 0.5 if split == "train" else 0.0 + return CocoKeypointsDataset( + root_dir=root_dir, + split=split, + image_size=image_size, + num_keypoints=num_keypoints, + flip_prob=flip_prob, + return_dict=return_dict, + ) + + diff --git a/pose_estimation_berna/core/depth.py b/pose_estimation_berna/core/depth.py new file mode 100644 index 00000000..7b2c173f --- /dev/null +++ b/pose_estimation_berna/core/depth.py @@ -0,0 +1,260 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal, Optional, Tuple + +import numpy as np + + +DepthModelName = Literal["midas_small", "dpt_hybrid", "dpt_large"] +DepthRegion = Literal["inner", "torso"] + + +def _clip_xyxy(box_xyxy: np.ndarray, w: int, h: int) -> Tuple[int, int, int, int]: + x1, y1, x2, y2 = [float(v) for v in box_xyxy.tolist()] + # allow partially out-of-frame, then clip + x1i = int(np.floor(np.clip(x1, 0.0, float(max(0, w - 1))))) + y1i = int(np.floor(np.clip(y1, 0.0, float(max(0, h - 1))))) + x2i = int(np.ceil(np.clip(x2, 0.0, float(w)))) + y2i = int(np.ceil(np.clip(y2, 0.0, float(h)))) + if x2i <= x1i: + x2i = min(w, x1i + 1) + if y2i <= y1i: + y2i = min(h, y1i + 1) + return x1i, y1i, x2i, y2i + + +def _shrink_region(x1: int, y1: int, x2: int, y2: int, region: DepthRegion) -> Tuple[int, int, int, int]: + bw = max(1, int(x2 - x1)) + bh = max(1, int(y2 - y1)) + # Conservative crop to reduce edge bleed from occluders/background. + if region == "inner": + fx1, fx2 = 0.20, 0.80 + fy1, fy2 = 0.20, 0.80 + else: # "torso" + fx1, fx2 = 0.20, 0.80 + fy1, fy2 = 0.25, 0.80 + rx1 = x1 + int(round(fx1 * bw)) + rx2 = x1 + int(round(fx2 * bw)) + ry1 = y1 + int(round(fy1 * bh)) + ry2 = y1 + int(round(fy2 * bh)) + # keep at least 2x2 + if rx2 <= rx1 + 1: + rx1 = x1 + rx2 = x2 + if ry2 <= ry1 + 1: + ry1 = y1 + ry2 = y2 + return rx1, ry1, rx2, ry2 + + +@dataclass +class DepthStats: + inv_depth_med: float + inv_depth_mad: float + inv_depth_center: float + valid: bool + + +def bbox_inv_depth_stats( + inv_depth: np.ndarray, + box_xyxy: np.ndarray, + *, + region: DepthRegion = "torso", +) -> DepthStats: + """ + Robust inverse-depth stats inside bbox. + + inv_depth: [H,W] float32, larger => closer (typical MiDaS convention). + Returns median + MAD (robust spread) + center sample. + """ + if inv_depth is None or getattr(inv_depth, "size", 0) == 0: + return DepthStats(inv_depth_med=float("nan"), inv_depth_mad=float("nan"), inv_depth_center=float("nan"), valid=False) + if inv_depth.ndim != 2: + return DepthStats(inv_depth_med=float("nan"), inv_depth_mad=float("nan"), inv_depth_center=float("nan"), valid=False) + + h, w = int(inv_depth.shape[0]), int(inv_depth.shape[1]) + x1, y1, x2, y2 = _clip_xyxy(box_xyxy, w=w, h=h) + rx1, ry1, rx2, ry2 = _shrink_region(x1, y1, x2, y2, region=region) + patch = inv_depth[ry1:ry2, rx1:rx2] + if patch.size < 4: + return DepthStats(inv_depth_med=float("nan"), inv_depth_mad=float("nan"), inv_depth_center=float("nan"), valid=False) + + # robust center sample (clipped) + cx = int(np.clip(round(0.5 * (x1 + x2)), 0, w - 1)) + cy = int(np.clip(round(0.5 * (y1 + y2)), 0, h - 1)) + zc = float(inv_depth[cy, cx]) + + # median + MAD + vals = patch.reshape(-1).astype(np.float32) + vals = vals[np.isfinite(vals)] + if vals.size < 8: + return DepthStats(inv_depth_med=float("nan"), inv_depth_mad=float("nan"), inv_depth_center=zc, valid=False) + med = float(np.median(vals)) + mad = float(np.median(np.abs(vals - med))) + 1e-6 + return DepthStats(inv_depth_med=med, inv_depth_mad=mad, inv_depth_center=zc, valid=True) + + +def occluder_front_fraction( + inv_depth: np.ndarray, + box_xyxy: np.ndarray, + *, + inv_depth_expected: float, + region: DepthRegion = "torso", + margin_abs: float = 0.02, + margin_rel: float = 0.08, +) -> float: + """ + Fraction of pixels inside bbox that are "in front of" expected depth. + + Using inverse depth, "in front" means inv_depth is larger than expected by a margin. + """ + if inv_depth is None or getattr(inv_depth, "size", 0) == 0 or not np.isfinite(float(inv_depth_expected)): + return 0.0 + if inv_depth.ndim != 2: + return 0.0 + + h, w = int(inv_depth.shape[0]), int(inv_depth.shape[1]) + x1, y1, x2, y2 = _clip_xyxy(box_xyxy, w=w, h=h) + rx1, ry1, rx2, ry2 = _shrink_region(x1, y1, x2, y2, region=region) + patch = inv_depth[ry1:ry2, rx1:rx2] + if patch.size < 16: + return 0.0 + vals = patch.reshape(-1).astype(np.float32) + vals = vals[np.isfinite(vals)] + if vals.size < 16: + return 0.0 + + zexp = float(inv_depth_expected) + thr = zexp + float(margin_abs) + float(margin_rel) * abs(zexp) + frac = float(np.mean(vals > thr)) + if not np.isfinite(frac): + return 0.0 + return float(np.clip(frac, 0.0, 1.0)) + + +class DepthEstimator: + """ + Minimal MiDaS/DPT wrapper (torch.hub) for inverse depth. + + - Inference-only; no retraining + - Runs at configurable inference resolution for speed + - Returns inverse depth in original frame size (float32) + """ + + def __init__( + self, + *, + model_name: DepthModelName = "midas_small", + device: str = "cuda", + input_short_side: int = 320, + hub_repo: str = "intel-isl/MiDaS", + ): + self.model_name = str(model_name) + self.device_str = str(device) + self.input_short_side = int(max(128, input_short_side)) + self.hub_repo = str(hub_repo) + + self._model = None + self._transform = None + self._torch = None + + def _lazy_init(self): + if self._model is not None: + return + + # MiDaS hub currently depends on timm (even for MiDaS_small via some backbones). + # Make the failure mode obvious. + try: + import timm # noqa: F401 + except Exception as e: + raise RuntimeError( + "Monocular depth requires the 'timm' package. Install it with:\n" + " pip install timm\n" + "or in this repo's venv:\n" + " /home/berna/D-FINE-NEWBRINGER/venv/bin/python3 -m pip install timm\n" + f"Original import error: {e}" + ) from e + + import torch # local import to keep tracker lightweight when depth is disabled + + self._torch = torch + device = torch.device(self.device_str if (self.device_str == "cpu" or torch.cuda.is_available()) else "cpu") + + # NOTE: torch.hub will download weights on first run. + if self.model_name == "midas_small": + hub_model = "MiDaS_small" + elif self.model_name == "dpt_hybrid": + hub_model = "DPT_Hybrid" + elif self.model_name == "dpt_large": + hub_model = "DPT_Large" + else: + raise ValueError(f"Unknown depth model_name: {self.model_name}") + + model = torch.hub.load(self.hub_repo, hub_model) + model = model.to(device).eval() + + transforms = torch.hub.load(self.hub_repo, "transforms") + if hub_model == "MiDaS_small": + tfm = transforms.small_transform + else: + tfm = transforms.dpt_transform + + self._model = model + self._transform = tfm + self._device = device + + def predict_inv_depth(self, frame_rgb_u8: np.ndarray) -> np.ndarray: + """ + frame_rgb_u8: [H,W,3] uint8 RGB + returns inv_depth: [H,W] float32 (larger => closer) + """ + self._lazy_init() + assert self._model is not None + assert self._transform is not None + assert self._torch is not None + + import torch.nn.functional as F # local import for same reason as torch + + if frame_rgb_u8 is None or frame_rgb_u8.ndim != 3 or frame_rgb_u8.shape[2] != 3: + raise ValueError("frame_rgb_u8 must be HxWx3 RGB uint8") + H, W = int(frame_rgb_u8.shape[0]), int(frame_rgb_u8.shape[1]) + + torch = self._torch + with torch.no_grad(): + # The hub transform expects RGB uint8 numpy. + inp = self._transform(frame_rgb_u8) + # MiDaS hub transforms vary by version: + # - sometimes return Tensor [3,h,w] + # - sometimes return Tensor [1,3,h,w] + # - sometimes return dict with key 'image' + if isinstance(inp, dict): + inp = inp.get("image", None) + if inp is None: + raise RuntimeError("MiDaS transform returned None (unexpected).") + if not hasattr(inp, "ndim"): + raise RuntimeError(f"MiDaS transform returned unexpected type: {type(inp)}") + if int(inp.ndim) == 3: + inp = inp.unsqueeze(0) + elif int(inp.ndim) == 4: + pass + else: + raise RuntimeError(f"MiDaS transform returned tensor with unsupported ndim={int(inp.ndim)}") + + inp = inp.to(self._device) # [1,3,h,w] + pred = self._model(inp) # [1,h',w'] or [1,1,h',w'] + if pred.ndim == 4: + pred = pred[:, 0, :, :] + + # Upsample to original resolution + pred_up = F.interpolate(pred.unsqueeze(1), size=(H, W), mode="bicubic", align_corners=False)[:, 0] + inv_depth = pred_up.squeeze(0).float().detach().cpu().numpy().astype(np.float32) + + # Normalize per-frame to stabilize thresholds (keeps ordering, not metric). + # This makes margin_abs usable across different scenes. + lo = float(np.percentile(inv_depth, 2.0)) + hi = float(np.percentile(inv_depth, 98.0)) + if np.isfinite(lo) and np.isfinite(hi) and hi > lo + 1e-6: + inv_depth = (inv_depth - lo) / (hi - lo) + return inv_depth.astype(np.float32) + diff --git a/pose_estimation_berna/core/engines.py b/pose_estimation_berna/core/engines.py new file mode 100644 index 00000000..2c413e61 --- /dev/null +++ b/pose_estimation_berna/core/engines.py @@ -0,0 +1,453 @@ +#!/usr/bin/env python3 +""" +Core TensorRT Engine Classes for DFINE Segmentation +Reusable TensorRT engine management and inference +""" + +import os +import time +import numpy as np +import cv2 +import tensorrt as trt +import pycuda.driver as cuda +import pycuda.autoinit +from typing import List, Dict, Tuple, Optional, Any +from abc import ABC, abstractmethod +import logging + + +class BaseTensorRTEngine(ABC): + """Base class for TensorRT engines with common functionality""" + + def __init__(self, engine_path: str, logger_level: int = trt.Logger.WARNING): + self.engine_path = engine_path + self.logger = trt.Logger(logger_level) + self.runtime = None + self.engine = None + self.context = None + self.stream = None + + # Memory management + self.host_inputs = [] + self.host_outputs = [] + self.device_inputs = [] + self.device_outputs = [] + self.input_shapes = {} + self.output_shapes = {} + self.input_names = [] + self.output_names = [] + + # Performance tracking + self.inference_times = [] + + # Load and setup engine + self._load_engine() + self._allocate_memory() + + def _load_engine(self): + """Load TensorRT engine from file""" + print(f"🔧 Loading TensorRT engine: {self.engine_path}") + + if not os.path.exists(self.engine_path): + raise FileNotFoundError(f"Engine file not found: {self.engine_path}") + + self.runtime = trt.Runtime(self.logger) + + with open(self.engine_path, 'rb') as f: + engine_bytes = f.read() + + self.engine = self.runtime.deserialize_cuda_engine(engine_bytes) + if not self.engine: + raise RuntimeError(f"Failed to load TensorRT engine from {self.engine_path}") + + self.context = self.engine.create_execution_context() + if not self.context: + raise RuntimeError("Failed to create TensorRT execution context") + + self.stream = cuda.Stream() + print(f"✅ Engine loaded successfully") + + def _allocate_memory(self): + """Pre-allocate host and device memory""" + print(f"🧠 Allocating memory buffers...") + + # Check TensorRT version for API compatibility + use_new_api = hasattr(self.engine, 'num_io_tensors') + + if use_new_api: + self._allocate_memory_new_api() + else: + self._allocate_memory_legacy_api() + + def _allocate_memory_new_api(self): + """Allocate memory using TensorRT 10.x+ API""" + for i in range(self.engine.num_io_tensors): + tensor_name = self.engine.get_tensor_name(i) + shape = self.engine.get_tensor_shape(tensor_name) + dtype = self.engine.get_tensor_dtype(tensor_name) + is_input = self.engine.get_tensor_mode(tensor_name) == trt.TensorIOMode.INPUT + + # Handle dynamic shapes + if -1 in shape: + actual_shape = tuple(1 if dim == -1 else dim for dim in shape) + self.context.set_input_shape(tensor_name, actual_shape) + shape = actual_shape + + # Convert dtype and calculate size + np_dtype, element_size = self._get_numpy_dtype(dtype) + size = trt.volume(shape) + + # Allocate memory + host_mem = cuda.pagelocked_empty(size, np_dtype) + device_mem = cuda.mem_alloc(size * element_size) + + if is_input: + self.host_inputs.append(host_mem) + self.device_inputs.append(device_mem) + self.input_shapes[tensor_name] = shape + self.input_names.append(tensor_name) + else: + self.host_outputs.append(host_mem) + self.device_outputs.append(device_mem) + self.output_shapes[tensor_name] = shape + self.output_names.append(tensor_name) + + # Set tensor address + self.context.set_tensor_address(tensor_name, int(device_mem)) + + print(f" {'📥' if is_input else '📤'} {tensor_name}: {shape} ({size * element_size / 1024 / 1024:.1f} MB)") + + def _allocate_memory_legacy_api(self): + """Allocate memory using legacy TensorRT API""" + self.bindings = [] + + for i in range(self.engine.num_bindings): + binding_name = self.engine.get_binding_name(i) + shape = self.engine.get_binding_shape(i) + dtype = self.engine.get_binding_dtype(i) + is_input = self.engine.binding_is_input(i) + + # Handle dynamic shapes + if -1 in shape: + actual_shape = tuple(1 if dim == -1 else dim for dim in shape) + self.context.set_binding_shape(i, actual_shape) + shape = actual_shape + + # Convert dtype and calculate size + np_dtype, element_size = self._get_numpy_dtype(dtype) + size = trt.volume(shape) + + # Allocate memory + host_mem = cuda.pagelocked_empty(size, np_dtype) + device_mem = cuda.mem_alloc(size * element_size) + + if is_input: + self.host_inputs.append(host_mem) + self.device_inputs.append(device_mem) + self.input_shapes[binding_name] = shape + self.input_names.append(binding_name) + else: + self.host_outputs.append(host_mem) + self.device_outputs.append(device_mem) + self.output_shapes[binding_name] = shape + self.output_names.append(binding_name) + + self.bindings.append(int(device_mem)) + + print(f" {'📥' if is_input else '📤'} {binding_name}: {shape} ({size * element_size / 1024 / 1024:.1f} MB)") + + def _get_numpy_dtype(self, trt_dtype) -> Tuple[np.dtype, int]: + """Convert TensorRT dtype to numpy dtype and element size""" + if trt_dtype == trt.DataType.FLOAT: + return np.float32, 4 + elif trt_dtype == trt.DataType.HALF: + return np.float16, 2 + elif trt_dtype == trt.DataType.INT32: + return np.int32, 4 + elif trt_dtype == trt.DataType.INT64: + return np.int64, 8 + elif trt_dtype == trt.DataType.INT8: + return np.int8, 1 + else: + return np.float32, 4 # Default fallback + + def infer(self, inputs: List[np.ndarray]) -> List[np.ndarray]: + """Perform inference with timing""" + start_time = time.perf_counter() + + # Copy inputs to device + for i, input_data in enumerate(inputs): + np.copyto(self.host_inputs[i][:input_data.size], input_data.ravel()) + cuda.memcpy_htod_async(self.device_inputs[i], self.host_inputs[i], self.stream) + + # Execute inference + if hasattr(self.context, 'execute_async_v3'): + success = self.context.execute_async_v3(stream_handle=self.stream.handle) + elif hasattr(self.context, 'execute_async_v2'): + success = self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle) + else: + success = self.context.execute_async(bindings=self.bindings, stream_handle=self.stream.handle) + + if not success: + raise RuntimeError("TensorRT inference failed") + + # Copy outputs from device + outputs = [] + for i, (host_output, device_output, shape) in enumerate( + zip(self.host_outputs, self.device_outputs, self.output_shapes.values()) + ): + cuda.memcpy_dtoh_async(host_output, device_output, self.stream) + self.stream.synchronize() + output = host_output[:np.prod(shape)].reshape(shape).copy() + outputs.append(output) + + # Track inference time + inference_time = time.perf_counter() - start_time + self.inference_times.append(inference_time) + + return outputs + + def get_engine_info(self) -> Dict[str, Any]: + """Get engine information""" + info = { + 'engine_path': self.engine_path, + 'num_inputs': len(self.input_names), + 'num_outputs': len(self.output_names), + 'input_shapes': self.input_shapes, + 'output_shapes': self.output_shapes, + 'input_names': self.input_names, + 'output_names': self.output_names + } + + if self.inference_times: + info.update({ + 'avg_inference_time_ms': np.mean(self.inference_times) * 1000, + 'min_inference_time_ms': np.min(self.inference_times) * 1000, + 'max_inference_time_ms': np.max(self.inference_times) * 1000, + 'avg_fps': 1.0 / np.mean(self.inference_times) + }) + + return info + + def print_engine_info(self): + """Print engine information""" + info = self.get_engine_info() + print(f"📊 TensorRT Engine Info:") + print(f" Engine: {os.path.basename(info['engine_path'])}") + print(f" Inputs: {info['num_inputs']}") + print(f" Outputs: {info['num_outputs']}") + + if 'avg_inference_time_ms' in info: + print(f" Avg inference time: {info['avg_inference_time_ms']:.2f}ms") + print(f" Avg FPS: {info['avg_fps']:.1f}") + + @abstractmethod + def preprocess(self, *args, **kwargs) -> List[np.ndarray]: + """Preprocess inputs for inference""" + pass + + @abstractmethod + def postprocess(self, outputs: List[np.ndarray], *args, **kwargs) -> Any: + """Postprocess inference outputs""" + pass + + +class SegmentationTensorRTEngine(BaseTensorRTEngine): + """TensorRT engine for segmentation models""" + + def __init__(self, engine_path: str, target_size: Tuple[int, int] = (640, 640)): + self.target_size = target_size + super().__init__(engine_path) + + # Validate expected outputs for segmentation + self._validate_segmentation_outputs() + + def _validate_segmentation_outputs(self): + """Validate that engine has expected segmentation outputs""" + expected_outputs = ['labels', 'boxes', 'scores', 'seg_probs', 'seg_preds'] + + if len(self.output_names) != 5: + print(f"⚠️ Warning: Expected 5 outputs for segmentation model, found {len(self.output_names)}") + print(f" Expected: {expected_outputs}") + print(f" Found: {self.output_names}") + else: + print(f"✅ Segmentation model outputs validated") + + def preprocess(self, frame: np.ndarray) -> List[np.ndarray]: + """Preprocess frame for segmentation inference""" + # Resize frame + resized = cv2.resize(frame, self.target_size, interpolation=cv2.INTER_LINEAR) + + # Convert to RGB and normalize + rgb_frame = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB) + normalized = rgb_frame.astype(np.float32) / 255.0 + + # Apply ImageNet normalization + mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) + std = np.array([0.229, 0.224, 0.225], dtype=np.float32) + normalized = (normalized - mean) / std + + # Convert to NCHW format + input_tensor = np.transpose(normalized, (2, 0, 1)) + input_tensor = np.expand_dims(input_tensor, axis=0) # Add batch dimension + + # Target size tensor + target_size_tensor = np.array([[self.target_size[0], self.target_size[1]]], dtype=np.int64) + + return [input_tensor, target_size_tensor] + + def postprocess(self, outputs: List[np.ndarray]) -> Dict[str, np.ndarray]: + """Postprocess segmentation inference outputs""" + if len(outputs) != 5: + raise ValueError(f"Expected 5 outputs, got {len(outputs)}") + + # Extract outputs + labels = outputs[0].flatten() # Detection labels + boxes = outputs[1][0] # Detection boxes (N, 4) + scores = outputs[2].flatten() # Detection scores + seg_probs = outputs[3][0] # Segmentation probabilities (7, 640, 640) + seg_preds = outputs[4][0] # Segmentation predictions (640, 640) + + return { + 'labels': labels, + 'boxes': boxes, + 'scores': scores, + 'seg_probs': seg_probs, + 'seg_preds': seg_preds + } + + def infer_frame(self, frame: np.ndarray) -> Dict[str, np.ndarray]: + """Complete inference pipeline for a single frame""" + # Preprocess + inputs = self.preprocess(frame) + + # Inference + outputs = self.infer(inputs) + + # Postprocess + results = self.postprocess(outputs) + + return results + + +class LightweightSegmentationEngine(SegmentationTensorRTEngine): + """Lightweight segmentation engine optimized for speed""" + + def __init__(self, engine_path: str, target_size: Tuple[int, int] = (512, 512)): + # Use smaller input size for lightweight model + super().__init__(engine_path, target_size) + print(f"🚀 Lightweight segmentation engine ready") + print(f" Target size: {target_size}") + print(f" Optimized for: Speed and low memory") + + +class StandardSegmentationEngine(SegmentationTensorRTEngine): + """Standard segmentation engine with balanced performance""" + + def __init__(self, engine_path: str, target_size: Tuple[int, int] = (640, 640)): + super().__init__(engine_path, target_size) + print(f"🚀 Standard segmentation engine ready") + print(f" Target size: {target_size}") + print(f" Optimized for: Balanced performance") + + +class AdvancedSegmentationEngine(SegmentationTensorRTEngine): + """Advanced segmentation engine with maximum accuracy""" + + def __init__(self, engine_path: str, target_size: Tuple[int, int] = (768, 768)): + # Use larger input size for advanced model + super().__init__(engine_path, target_size) + print(f"🚀 Advanced segmentation engine ready") + print(f" Target size: {target_size}") + print(f" Optimized for: Maximum accuracy") + + def infer_multiscale(self, frame: np.ndarray, scales: List[float] = [0.75, 1.0, 1.25]) -> Dict[str, np.ndarray]: + """Multi-scale inference for better accuracy""" + all_results = [] + + for scale in scales: + # Scale frame + h, w = frame.shape[:2] + scaled_h, scaled_w = int(h * scale), int(w * scale) + scaled_frame = cv2.resize(frame, (scaled_w, scaled_h)) + + # Infer on scaled frame + result = self.infer_frame(scaled_frame) + all_results.append(result) + + # Ensemble results (simple averaging for seg_probs) + ensemble_seg_probs = np.mean([r['seg_probs'] for r in all_results], axis=0) + ensemble_seg_preds = np.argmax(ensemble_seg_probs, axis=0) + + # Use results from scale 1.0 for detection + base_result = all_results[1] if len(all_results) > 1 else all_results[0] + + return { + 'labels': base_result['labels'], + 'boxes': base_result['boxes'], + 'scores': base_result['scores'], + 'seg_probs': ensemble_seg_probs, + 'seg_preds': ensemble_seg_preds + } + + +# Factory functions +def create_segmentation_engine(tier: str, engine_path: str, **kwargs) -> SegmentationTensorRTEngine: + """Factory function to create segmentation engines""" + + if tier == 'lightweight': + return LightweightSegmentationEngine(engine_path, **kwargs) + elif tier == 'standard': + return StandardSegmentationEngine(engine_path, **kwargs) + elif tier == 'advanced': + return AdvancedSegmentationEngine(engine_path, **kwargs) + else: + raise ValueError(f"Unknown engine tier: {tier}") + + +# Utility functions +def benchmark_engine(engine: BaseTensorRTEngine, num_iterations: int = 100) -> Dict[str, float]: + """Benchmark engine performance""" + print(f"🏃 Benchmarking engine for {num_iterations} iterations...") + + # Create dummy input + dummy_inputs = [] + for shape in engine.input_shapes.values(): + dummy_input = np.random.randn(*shape).astype(np.float32) + dummy_inputs.append(dummy_input) + + # Warmup + for _ in range(10): + _ = engine.infer(dummy_inputs) + + # Benchmark + start_time = time.perf_counter() + for _ in range(num_iterations): + _ = engine.infer(dummy_inputs) + + total_time = time.perf_counter() - start_time + avg_time = total_time / num_iterations + fps = 1.0 / avg_time + + print(f"✅ Benchmark completed:") + print(f" Average inference time: {avg_time*1000:.2f}ms") + print(f" Average FPS: {fps:.1f}") + + return { + 'avg_inference_time_ms': avg_time * 1000, + 'avg_fps': fps, + 'total_time_s': total_time, + 'iterations': num_iterations + } + + +# Export classes +__all__ = [ + 'BaseTensorRTEngine', + 'SegmentationTensorRTEngine', + 'LightweightSegmentationEngine', + 'StandardSegmentationEngine', + 'AdvancedSegmentationEngine', + 'create_segmentation_engine', + 'benchmark_engine' +] \ No newline at end of file diff --git a/pose_estimation_berna/core/hit_registration.py b/pose_estimation_berna/core/hit_registration.py new file mode 100644 index 00000000..08ba0440 --- /dev/null +++ b/pose_estimation_berna/core/hit_registration.py @@ -0,0 +1,141 @@ +""" +Hit registration logic for occluded players (geometry + tracker uncertainty only). + +This module intentionally does NOT depend on any CV models or training-time changes. +It consumes tracker state (predicted bbox + time_since_update + uncertainty) and +applies deterministic, explainable gameplay rules. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Optional, Tuple + +import math + + +@dataclass(frozen=True) +class HitDecision: + """ + Result of hit registration. + - kind: + - "person_direct": player was visible; normal hitbox logic should be used (not handled here) + - "person_transferred": hit on occluder counts as hit on person (occluded) + - "occluder_only": hit is absorbed by occluder + """ + + kind: str + p: float # transfer probability/plausibility in [0,1] (0 for non-transfers) + damage_multiplier: float # multiply base damage by this (0 for occluder_only) + + +def _clamp(x: float, lo: float, hi: float) -> float: + return float(min(max(float(x), float(lo)), float(hi))) + + +def _point_in_bbox_xyxy(pt_xy: Tuple[float, float], box_xyxy) -> bool: + x, y = float(pt_xy[0]), float(pt_xy[1]) + x1, y1, x2, y2 = [float(v) for v in box_xyxy] + return (x >= x1) and (x <= x2) and (y >= y1) and (y <= y2) + + +def _expand_bbox_xyxy(box_xyxy, dx: float, dy: float): + x1, y1, x2, y2 = [float(v) for v in box_xyxy] + return (x1 - float(dx), y1 - float(dy), x2 + float(dx), y2 + float(dy)) + + +def _bbox_center_xy(box_xyxy) -> Tuple[float, float]: + x1, y1, x2, y2 = [float(v) for v in box_xyxy] + return (0.5 * (x1 + x2), 0.5 * (y1 + y2)) + + +def register_occluder_hit_2d( + *, + hit_xy: Tuple[float, float], + occluder_id: Any, + track: Any, + ttl_frames: int, + p_min: float = 0.4, + k_sigma: float = 2.0, + transfer_damage_min: float = 0.6, + transfer_damage_max: float = 1.0, + require_occluder_match: bool = True, + track_occluder_attr: str = "current_occluder_id", +) -> HitDecision: + """ + 2D occluder-assisted hit registration. + + Inputs: + - hit_xy: screen-space hit point on the occluder (pixels). + - occluder_id: identifier of the occluder object hit by the bullet (game collision system). + - track: tracker track object. Expected fields: + - time_since_update: int + - box_xyxy: iterable (x1,y1,x2,y2) for predicted/last bbox in pixels + - kf: optional; if present and has get_center_std() -> float, used as sigma + - ttl_frames: max frames to allow transfer while occluded (<= ~fps for 1s). + + Deterministic logic: + - Only transfers when track is occluded (time_since_update > 0) and within TTL. + - Only transfers when hit is inside an expanded region around predicted bbox (k*sigma). + - Transfer plausibility p decays with distance from predicted center and occlusion age. + - Damage is scaled by p (bounded) for fairness. + """ + # Visible players should use your normal direct hitbox logic. + tsu = int(getattr(track, "time_since_update", 0)) + if tsu <= 0: + return HitDecision(kind="person_direct", p=0.0, damage_multiplier=1.0) + + ttl = max(1, int(ttl_frames)) + if tsu > ttl: + return HitDecision(kind="occluder_only", p=0.0, damage_multiplier=0.0) + + if require_occluder_match: + track_occ = getattr(track, track_occluder_attr, None) + if track_occ is not None and occluder_id != track_occ: + return HitDecision(kind="occluder_only", p=0.0, damage_multiplier=0.0) + + box = getattr(track, "box_xyxy", None) + if box is None: + return HitDecision(kind="occluder_only", p=0.0, damage_multiplier=0.0) + + # sigma from Kalman if available, else fall back to a fraction of bbox size. + sigma = None + kf = getattr(track, "kf", None) + if kf is not None and hasattr(kf, "get_center_std"): + try: + sigma = float(kf.get_center_std()) + except Exception: + sigma = None + if sigma is None or not math.isfinite(sigma) or sigma <= 1e-3: + x1, y1, x2, y2 = [float(v) for v in box] + bw = max(1.0, x2 - x1) + bh = max(1.0, y2 - y1) + sigma = 0.15 * min(bw, bh) + + sigma = max(1.0, float(sigma)) + k = float(max(0.0, k_sigma)) + region = _expand_bbox_xyxy(box, dx=k * sigma, dy=k * sigma) + if not _point_in_bbox_xyxy(hit_xy, region): + return HitDecision(kind="occluder_only", p=0.0, damage_multiplier=0.0) + + cx, cy = _bbox_center_xy(box) + dx = float(hit_xy[0]) - float(cx) + dy = float(hit_xy[1]) - float(cy) + d2 = dx * dx + dy * dy + + # Gaussian plausibility around predicted center + linear time decay to 0 at TTL. + w_sigma = math.exp(-d2 / (2.0 * sigma * sigma)) + w_t = max(0.0, 1.0 - (float(tsu) / float(ttl))) + p = _clamp(w_sigma * w_t, 0.0, 1.0) + + if p < float(p_min): + return HitDecision(kind="occluder_only", p=p, damage_multiplier=0.0) + + dmg_lo = float(transfer_damage_min) + dmg_hi = float(transfer_damage_max) + if dmg_hi < dmg_lo: + dmg_hi, dmg_lo = dmg_lo, dmg_hi + dmg = _clamp(dmg_lo + (dmg_hi - dmg_lo) * p, 0.0, 1.0) + return HitDecision(kind="person_transferred", p=p, damage_multiplier=dmg) + + diff --git a/pose_estimation_berna/core/losses.py b/pose_estimation_berna/core/losses.py new file mode 100644 index 00000000..d0c94f06 --- /dev/null +++ b/pose_estimation_berna/core/losses.py @@ -0,0 +1,274 @@ +#!/usr/bin/env python3 +""" +Loss/criterion factory for D-FINE pose estimation. + +We re-use D-FINE's built-in matcher + criterion, and add the "keypoints" loss +that we implemented in `src/zoo/dfine/dfine_criterion.py`. +""" + +from typing import Dict, List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +try: + import cv2 # used by segmentation boundary utilities (optional) + import numpy as np +except Exception: # pragma: no cover + cv2 = None + np = None + + +def create_pose_criterion( + num_classes: int = 80, + weight_dict: Dict[str, float] | None = None, + losses: List[str] | None = None, +): + """ + Create a DFINECriterion that includes keypoint losses. + + Notes: + - Keypoints are supervised using the SAME Hungarian matching as the boxes (no extra matching). + - Keypoints are bbox-relative: outputs["pred_keypoints"][..., :2] are in [0,1] within bbox. + """ + # Two pose training modes exist in this repo: + # 1) DFINE pose head: pred_boxes + bbox-relative pred_keypoints [B,Q,K,3] + # 2) DETRPose-style: pose-only decoder producing pred_keypoints [B,Q,2K] (image-normalized) + # + # We auto-select criterion based on the model outputs at runtime is hard here, + # so selection is done by configuration: + # - if cfg["dfine"]["config_path"] contains "detrpose" -> DETRPose-style criterion + # - else -> DFINECriterion + from src.zoo.dfine.matcher import HungarianMatcher + from src.zoo.dfine.dfine_criterion import DFINECriterion + + if weight_dict is None: + weight_dict = { + "loss_vfl": 1.0, + "loss_bbox": 5.0, + "loss_giou": 2.0, + "loss_fgl": 0.15, + "loss_ddf": 1.5, + "loss_keypoints": 10.0, + "loss_keypoints_vis": 1.0, + "loss_pose_lqe": 1.0, + } + + if losses is None: + losses = ["vfl", "boxes", "local", "keypoints", "pose_lqe"] + + matcher = HungarianMatcher( + weight_dict={"cost_class": 2, "cost_bbox": 5, "cost_giou": 2, "cost_oks": 2.0}, + alpha=0.25, + gamma=2.0, + ) + + return DFINECriterion( + matcher=matcher, + weight_dict=weight_dict, + losses=losses, + alpha=0.75, + gamma=2.0, + num_classes=num_classes, + reg_max=32, + boxes_weight_format=None, + keypoints_box_mode="pred_detached", + vfl_target="oks", + ) + + +def create_pose_criterion_detrpose( + num_classes: int = 2, + weight_dict: Dict[str, float] | None = None, +): + """ + DETRPose-style pose-only criterion (no boxes). + Outputs expected: + - pred_logits: [B,Q,C] + - pred_keypoints: [B,Q,2K] in [0,1] (resized image-normalized) + """ + from src.zoo.dfine.detrpose.matcher import HungarianMatcherDETRPose + from src.zoo.dfine.detrpose.criterion import DETRPoseCriterion + + if weight_dict is None: + # Match DETRPose default config (see DETRPose/configs/detrpose/include/detrpose_hgnetv2.py) + weight_dict = {"loss_vfl": 2.0, "loss_keypoints": 10.0, "loss_oks": 4.0} + + matcher = HungarianMatcherDETRPose(cost_class=2.0, cost_keypoints=10.0, cost_oks=4.0, num_keypoints=17) + return DETRPoseCriterion( + num_classes=num_classes, + matcher=matcher, + weight_dict=weight_dict, + focal_alpha=0.25, + gamma=2.0, + num_keypoints=17, + ) + + +class BaseLoss(nn.Module): + """Base loss class for different model tiers""" + + def __init__(self, + num_classes: int = 7, + ignore_index: int = 255, + tier: str = 'standard'): + super().__init__() + self.num_classes = num_classes + self.ignore_index = ignore_index + self.tier = tier + + # Configure loss based on tier + self._configure_loss() + + def _configure_loss(self): + """Configure loss components based on tier""" + if self.tier == 'lightweight': + self._configure_lightweight_loss() + elif self.tier == 'standard': + self._configure_standard_loss() + elif self.tier == 'advanced': + self._configure_advanced_loss() + else: + self._configure_standard_loss() + + def _configure_lightweight_loss(self): + """Configure lightweight loss (simple, fast)""" + self.focal_loss = FocalLoss(alpha=0.25, gamma=2.0, ignore_index=self.ignore_index) + self.dice_loss = DiceLoss(ignore_index=self.ignore_index) + + self.weights = { + 'focal': 1.0, + 'dice': 0.3, + 'ce': 0.1 + } + + def _configure_standard_loss(self): + """Configure standard loss (enhanced to match original)""" + self.focal_loss = FocalLoss(alpha=0.25, gamma=2.0, ignore_index=self.ignore_index) + self.dice_loss = DiceLoss(ignore_index=self.ignore_index) + + # Use the same weights as the original AdvancedSegmentationLoss + self.weights = { + 'focal': 1.0, + 'dice': 0.4, # Same as original + 'ce': 0.2, # Same as original + 'boundary': 0.3, # Same as original + 'size': 0.1 # Same as original + } + + def _configure_advanced_loss(self): + """Configure advanced loss (comprehensive)""" + self.focal_loss = FocalLoss(alpha=0.25, gamma=2.0, ignore_index=self.ignore_index) + self.dice_loss = DiceLoss(ignore_index=self.ignore_index) + + self.weights = { + 'focal': 1.0, + 'dice': 0.4, + 'ce': 0.2, + 'boundary': 0.3, + 'size': 0.1 + } + + def size_sensitive_loss(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """Loss that gives more weight to smaller objects (distant people)""" + # Calculate object sizes + unique_labels = torch.unique(target) + size_weights = torch.ones_like(target, dtype=torch.float32) + + for label in unique_labels: + if label == self.ignore_index: + continue + + mask = (target == label) + size = mask.sum().float() + + # Inverse size weighting - smaller objects get higher weight + if size > 0: + weight = 1.0 / (torch.sqrt(size) + 1e-6) + size_weights[mask] = weight + + # Apply size weights to cross entropy loss + ce_loss = F.cross_entropy(pred, target, ignore_index=self.ignore_index, reduction='none') + weighted_ce = (ce_loss * size_weights).mean() + + return weighted_ce + + def create_boundary_target(self, target: torch.Tensor) -> torch.Tensor: + """Create boundary target from segmentation mask""" + target_np = target.cpu().numpy() + boundary_targets = [] + + for i in range(target_np.shape[0]): + mask = target_np[i] + + # Detect edges using morphological gradient + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)) + boundary = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_GRADIENT, kernel) + boundary = (boundary > 0).astype(np.float32) + + boundary_targets.append(torch.from_numpy(boundary)) + + return torch.stack(boundary_targets).to(target.device) + + def forward(self, + pred: torch.Tensor, + target: torch.Tensor, + boundary_pred: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, Dict[str, float]]: + """Forward pass of loss computation""" + + loss_dict = {} + total_loss = 0.0 + + # Focal loss (main loss) + if 'focal' in self.weights: + focal = self.focal_loss(pred, target) + loss_dict['focal'] = focal.item() + total_loss += self.weights['focal'] * focal + + # Dice loss + if 'dice' in self.weights: + dice = self.dice_loss(pred, target) + loss_dict['dice'] = dice.item() + total_loss += self.weights['dice'] * dice + + # Standard CE loss + if 'ce' in self.weights: + ce = F.cross_entropy(pred, target, ignore_index=self.ignore_index) + loss_dict['ce'] = ce.item() + total_loss += self.weights['ce'] * ce + + # Size-sensitive loss + if 'size' in self.weights: + size_loss = self.size_sensitive_loss(pred, target) + loss_dict['size'] = size_loss.item() + total_loss += self.weights['size'] * size_loss + + # Boundary loss + if 'boundary' in self.weights: + if boundary_pred is not None: + # Use provided boundary prediction + boundary_target = self.create_boundary_target(target) + boundary_loss_val = F.binary_cross_entropy(boundary_pred.squeeze(1), boundary_target) + else: + # No boundary prediction, skip boundary loss + boundary_loss_val = torch.tensor(0.0, device=pred.device) + + loss_dict['boundary'] = boundary_loss_val.item() + total_loss += self.weights['boundary'] * boundary_loss_val + + loss_dict['total'] = total_loss.item() + + return total_loss, loss_dict + + +# Factory function for creating losses +def create_loss(tier: str, num_classes: int = 7, ignore_index: int = 255, **kwargs) -> BaseLoss: + """Factory function to create loss based on tier""" + return BaseLoss(num_classes=num_classes, ignore_index=ignore_index, tier=tier) + + +# Keep the original AdvancedSegmentationLoss for backward compatibility +def create_advanced_loss(num_classes=7, **kwargs): + """Create the original advanced loss""" + return AdvancedSegmentationLoss(num_classes=num_classes, **kwargs) \ No newline at end of file diff --git a/pose_estimation_berna/core/metrics.py b/pose_estimation_berna/core/metrics.py new file mode 100644 index 00000000..5b86028b --- /dev/null +++ b/pose_estimation_berna/core/metrics.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python3 +""" +Pose metrics for COCO keypoints. + +Lightweight OKS tracking on matched pairs (mainly for debugging). +Full COCO AP should be computed via COCO evaluator when integrated. +""" + +from typing import Dict, List + +import torch + + +# COCO sigmas for the 17 keypoints +COCO_SIGMAS = torch.tensor( + [ + 0.026, 0.025, 0.025, 0.035, 0.035, + 0.079, 0.079, 0.072, 0.072, 0.062, + 0.062, 0.107, 0.107, 0.087, 0.087, + 0.089, 0.089, + ], + dtype=torch.float32, +) + + +def _cxcywh_to_xyxy(boxes: torch.Tensor) -> torch.Tensor: + cx, cy, w, h = boxes.unbind(-1) + x1 = cx - 0.5 * w + y1 = cy - 0.5 * h + x2 = cx + 0.5 * w + y2 = cy + 0.5 * h + return torch.stack([x1, y1, x2, y2], dim=-1) + + +class PoseMetricsTracker: + def __init__(self, num_keypoints: int = 17): + self.num_keypoints = int(num_keypoints) + self.reset() + + def reset(self): + self._oks: List[float] = [] + + @torch.no_grad() + def update(self, outputs: Dict[str, torch.Tensor], targets: List[Dict[str, torch.Tensor]], indices): + """ + Args: + Mode A (DFINE pose): + - outputs["pred_keypoints"]: [B,Q,K,3] (x_rel,y_rel,vis_logit) bbox-relative + - outputs["pred_boxes"]: [B,Q,4] cxcywh normalized + + Mode B (DETRPose-style pose-only): + - outputs["pred_keypoints"]: [B,Q,2K] image-normalized in [0,1] (resized image space) + - targets[i]["keypoints"]: [N,K,3] (x_px,y_px,v) + - targets[i]["boxes"]: [N,4] cxcywh normalized + - indices: list of (src_idx, tgt_idx) from Hungarian matching + """ + if "pred_keypoints" not in outputs: + return + + pred_kpts = outputs["pred_keypoints"] + pred_boxes = outputs.get("pred_boxes", None) + + sigmas = COCO_SIGMAS.to(pred_kpts.device)[: self.num_keypoints] + vars_ = (sigmas * 2) ** 2 # [K] + + for b, (src, tgt) in enumerate(indices): + if src.numel() == 0: + continue + + if pred_boxes is not None and pred_kpts.ndim == 4: + # Mode A: bbox-relative + pb = pred_boxes[b, src] # [M,4] + pb_xyxy = _cxcywh_to_xyxy(pb) + pb_wh = (pb_xyxy[:, 2:] - pb_xyxy[:, :2]).clamp(min=1e-6) # [M,2] + pk = pred_kpts[b, src] # [M,K,3] + pk_xy = pb_xyxy[:, None, :2] + pk[..., :2] * pb_wh[:, None, :] # [M,K,2] normalized + else: + # Mode B: image-normalized keypoints [M,2K] + pk = pred_kpts[b, src] # [M,2K] + pk_xy = pk.view(pk.shape[0], self.num_keypoints, 2) + + tk = targets[b]["keypoints"][tgt].to(pk_xy.device) # [M,K,3] + if "size" in targets[b]: + h, w = targets[b]["size"].tolist() + else: + w, h = targets[b]["orig_size"].tolist() + wh = torch.tensor([w, h], device=pk_xy.device, dtype=pk_xy.dtype) + tk_xy = tk[..., :2] / wh[None, None, :] + v = (tk[..., 2] > 0).to(pk_xy.dtype) # [M,K] + + tb = targets[b]["boxes"][tgt].to(pk_xy.device) + tb_xyxy = _cxcywh_to_xyxy(tb) + area = ((tb_xyxy[:, 2] - tb_xyxy[:, 0]) * (tb_xyxy[:, 3] - tb_xyxy[:, 1])).clamp(min=1e-6) + + dx2 = (pk_xy[..., 0] - tk_xy[..., 0]) ** 2 + dy2 = (pk_xy[..., 1] - tk_xy[..., 1]) ** 2 + e = (dx2 + dy2) / (vars_[None, :] * area[:, None] * 2 + 1e-6) + + oks = torch.exp(-e) * v + denom = v.sum(dim=1).clamp(min=1.0) + oks = oks.sum(dim=1) / denom # [M] + self._oks.extend(oks.detach().cpu().tolist()) + + def compute(self) -> Dict[str, float]: + if not self._oks: + return {"OKS": 0.0} + return {"OKS": float(sum(self._oks) / len(self._oks))} + + +def create_metrics_tracker(num_keypoints: int = 17) -> PoseMetricsTracker: + """Factory function for compatibility with the segmentation_sivert template.""" + return PoseMetricsTracker(num_keypoints=num_keypoints) + + diff --git a/pose_estimation_berna/core/models.py b/pose_estimation_berna/core/models.py new file mode 100644 index 00000000..f7a16d1c --- /dev/null +++ b/pose_estimation_berna/core/models.py @@ -0,0 +1,184 @@ +#!/usr/bin/env python3 +""" +Core model helpers for D-FINE pose estimation. + +Pose is implemented as a per-query keypoint branch inside `DFINETransformer`: +- outputs["pred_boxes"]: [B,Q,4] (cxcywh normalized) +- outputs["pred_keypoints"]: [B,Q,17,3] (x_rel,y_rel,vis_logit), bbox-relative +""" + +import os +import sys +from typing import List, Tuple + +import torch +import torch.nn as nn + +# Ensure `src/` imports work when running from repo root +current_dir = os.path.dirname(os.path.abspath(__file__)) +project_root = os.path.dirname(os.path.dirname(current_dir)) +src_path = os.path.join(project_root, "src") +if os.path.exists(src_path) and src_path not in sys.path: + sys.path.insert(0, src_path) + sys.path.insert(0, project_root) + + +def load_pretrained_dfine(config_path: str, checkpoint_path: str) -> nn.Module: + """Load DFINE model from YAMLConfig + checkpoint.""" + from src.core import YAMLConfig + + cfg = YAMLConfig(config_path) + model = cfg.model + + checkpoint = torch.load(checkpoint_path, map_location="cpu") + if isinstance(checkpoint, dict) and "ema" in checkpoint and "module" in checkpoint.get("ema", {}): + state_dict = checkpoint["ema"]["module"] + elif isinstance(checkpoint, dict) and "model" in checkpoint: + state_dict = checkpoint["model"] + else: + state_dict = checkpoint + + model.load_state_dict(state_dict, strict=False) + return model + + +def freeze_except_keypoints(model: nn.Module) -> None: + """Freeze all params except pose-related decoder heads (keypoints + optional Pose-LQE).""" + for p in model.parameters(): + # Only floating/complex params can require gradients. + if p.is_floating_point() or p.is_complex(): + p.requires_grad = False + + dec = getattr(model, "decoder", None) + if dec is None: + return + kpt_head = getattr(dec, "dec_keypoint_head", None) + if kpt_head is None: + return + for p in kpt_head.parameters(): + if p.is_floating_point() or p.is_complex(): + p.requires_grad = True + + pose_q = getattr(dec, "dec_pose_lqe_head", None) + if pose_q is not None: + for p in pose_q.parameters(): + if p.is_floating_point() or p.is_complex(): + p.requires_grad = True + + dn_kpt_proj = getattr(dec, "denoising_kpt_proj", None) + if dn_kpt_proj is not None: + for p in dn_kpt_proj.parameters(): + if p.is_floating_point() or p.is_complex(): + p.requires_grad = True + + +def unfreeze_all(model: nn.Module) -> None: + """Unfreeze all parameters.""" + for p in model.parameters(): + # Only floating/complex params can require gradients (PyTorch constraint). + if p.is_floating_point() or p.is_complex(): + p.requires_grad = True + + +def set_trainable_by_name(model: nn.Module, name_substrings: List[str], trainable: bool) -> int: + """ + Set requires_grad for parameters whose name contains any of the provided substrings. + Returns number of parameters affected. + """ + affected = 0 + for n, p in model.named_parameters(): + if any(s in n for s in name_substrings): + if p.is_floating_point() or p.is_complex(): + p.requires_grad = trainable + affected += 1 + return affected + + +def freeze_backbone(model: nn.Module) -> None: + """Freeze backbone parameters (if present).""" + affected = set_trainable_by_name(model, ["backbone."], trainable=False) + _ = affected + + +def unfreeze_backbone_stages(model: nn.Module, stage_indices: List[int]) -> int: + """ + Unfreeze specific HGNetv2 backbone stages by index (e.g. [3] for stage4). + Works by matching parameter names containing 'backbone.stages.{idx}'. + """ + subs = [f"backbone.stages.{i}." for i in stage_indices] + return set_trainable_by_name(model, subs, trainable=True) + + +def build_param_groups( + model: nn.Module, + base_lr: float, + weight_decay: float, + backbone_lr_factor: float = 0.1, + keypoint_lr_factor: float = 1.0, +): + """ + Create AdamW param groups with different LRs: + - backbone.* uses base_lr * backbone_lr_factor + - decoder.dec_keypoint_head.* uses base_lr * keypoint_lr_factor + - everything else uses base_lr + Only includes params with requires_grad=True. + """ + backbone_params = [] + keypoint_params = [] + other_params = [] + + for name, p in model.named_parameters(): + if not p.requires_grad: + continue + if name.startswith("backbone."): + backbone_params.append(p) + elif ( + ("decoder.dec_keypoint_head" in name) + or ("decoder.dec_pose_lqe_head" in name) + or ("decoder.denoising_kpt_proj" in name) + ): + keypoint_params.append(p) + else: + other_params.append(p) + + groups = [] + if other_params: + groups.append({"params": other_params, "lr": base_lr, "weight_decay": weight_decay}) + if keypoint_params: + groups.append( + { + "params": keypoint_params, + "lr": base_lr * float(keypoint_lr_factor), + "weight_decay": weight_decay, + } + ) + if backbone_params: + groups.append( + { + "params": backbone_params, + "lr": base_lr * float(backbone_lr_factor), + "weight_decay": weight_decay, + } + ) + return groups + + +def get_trainable_params(model: nn.Module): + return [p for p in model.parameters() if p.requires_grad] + + +def get_device(device: str = "auto"): + if device == "auto": + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + return torch.device(device) + + +def get_actual_backbone_channels(model: nn.Module, input_size: Tuple[int, int] = (640, 640)) -> List[int]: + """Utility: inspect backbone output channel sizes.""" + model.eval() + dummy_input = torch.randn(1, 3, *input_size) + with torch.no_grad(): + feats = model.backbone(dummy_input) + return [f.shape[1] for f in feats] + + diff --git a/pose_estimation_berna/core/postprocess_box_fusion.py b/pose_estimation_berna/core/postprocess_box_fusion.py new file mode 100644 index 00000000..b9645a4d --- /dev/null +++ b/pose_estimation_berna/core/postprocess_box_fusion.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +from typing import Dict +import torch + + +@torch.no_grad() +def keypoints_to_box_xyxy(kpts_xy: torch.Tensor) -> torch.Tensor: + """ + kpts_xy: [B,Q,K,2] in pixel coords + returns: [B,Q,4] xyxy + """ + x = kpts_xy[..., 0] + y = kpts_xy[..., 1] + x1 = x.min(dim=-1).values + y1 = y.min(dim=-1).values + x2 = x.max(dim=-1).values + y2 = y.max(dim=-1).values + return torch.stack([x1, y1, x2, y2], dim=-1) + + +@torch.no_grad() +def expand_xyxy(boxes: torch.Tensor, ex: float = 0.15, ey_top: float = 0.20, ey_bot: float = 0.30) -> torch.Tensor: + """ + boxes: [B,Q,4] + """ + x1, y1, x2, y2 = boxes.unbind(dim=-1) + bw = (x2 - x1).clamp(min=1.0) + bh = (y2 - y1).clamp(min=1.0) + x1 = x1 - ex * bw + x2 = x2 + ex * bw + y1 = y1 - ey_top * bh + y2 = y2 + ey_bot * bh + return torch.stack([x1, y1, x2, y2], dim=-1) + + +@torch.no_grad() +def union_xyxy(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + x1 = torch.minimum(a[..., 0], b[..., 0]) + y1 = torch.minimum(a[..., 1], b[..., 1]) + x2 = torch.maximum(a[..., 2], b[..., 2]) + y2 = torch.maximum(a[..., 3], b[..., 3]) + return torch.stack([x1, y1, x2, y2], dim=-1) + + +@torch.no_grad() +def clamp_xyxy(boxes: torch.Tensor, orig_wh: torch.Tensor) -> torch.Tensor: + """ + boxes: [B,Q,4] + orig_wh: [B,2] -> (w,h) + """ + w = orig_wh[:, 0].view(-1, 1) + h = orig_wh[:, 1].view(-1, 1) + x1 = boxes[..., 0].clamp(min=0.0) + y1 = boxes[..., 1].clamp(min=0.0) + x2 = boxes[..., 2].clamp(min=0.0) + y2 = boxes[..., 3].clamp(min=0.0) + x1 = torch.minimum(x1, w) + x2 = torch.minimum(x2, w) + y1 = torch.minimum(y1, h) + y2 = torch.minimum(y2, h) + x_min = torch.minimum(x1, x2) + y_min = torch.minimum(y1, y2) + x_max = torch.maximum(x1, x2) + y_max = torch.maximum(y1, y2) + return torch.stack([x_min, y_min, x_max, y_max], dim=-1) + + +@torch.no_grad() +def fuse_det_and_pose_boxes( + det_boxes_xyxy: torch.Tensor, # [B,Q,4] + pose_keypoints_xy_norm: torch.Tensor, # [B,Q,K,2] normalized + orig_wh: torch.Tensor, # [B,2] (w,h) + ex: float = 0.15, + ey_top: float = 0.20, + ey_bot: float = 0.30, +) -> torch.Tensor: + """ + Return fused full-body box = union(det_box, expanded_pose_box). + """ + kpts_px = pose_keypoints_xy_norm * orig_wh[:, None, None, :] + pose_box = keypoints_to_box_xyxy(kpts_px) + pose_box = expand_xyxy(pose_box, ex=ex, ey_top=ey_top, ey_bot=ey_bot) + fused = union_xyxy(det_boxes_xyxy, pose_box) + fused = clamp_xyxy(fused, orig_wh) + return fused diff --git a/pose_estimation_berna/core/postprocess_detrpose.py b/pose_estimation_berna/core/postprocess_detrpose.py new file mode 100644 index 00000000..9b39d80f --- /dev/null +++ b/pose_estimation_berna/core/postprocess_detrpose.py @@ -0,0 +1,135 @@ +#!/usr/bin/env python3 +""" +Post-processing for DETRPose-style outputs (pose-only). + +Expected model outputs: + - pred_logits: [B, Q, C] (logits) + - pred_keypoints: [B, Q, 2K] (x,y) normalized in [0,1] in resized image space + +This produces a DFINEPostProcessor-like dict so existing COCOeval plumbing can run: + - boxes: [B, top, 4] xyxy in original pixel coords (tight box around keypoints) + - scores: [B, top] + - labels: [B, top] (MSCOCO category_id if remap enabled) + - keypoints: [B, top, K, 3] in original pixel coords, with v=1 +""" + +from __future__ import annotations + +import torch +import torch.nn as nn + + +def _remap_labels_to_coco_category(labels: torch.Tensor) -> torch.Tensor: + # D-FINE labels are contiguous 0..79; COCO evaluator expects category_id (person==1). + from src.data.dataset import mscoco_label2category + + flat = labels.flatten() + mapped = torch.tensor([mscoco_label2category[int(x.item())] for x in flat], device=labels.device) + return mapped.view_as(labels) + + +class DETRPosePostProcessor(nn.Module): + def __init__( + self, + num_classes: int = 80, + num_keypoints: int = 17, + num_top_queries: int = 300, + remap_mscoco_category: bool = True, + ): + super().__init__() + self.num_classes = int(num_classes) + self.num_keypoints = int(num_keypoints) + self.num_top_queries = int(num_top_queries) + self.remap_mscoco_category = bool(remap_mscoco_category) + + @torch.no_grad() + def forward(self, outputs: dict, orig_target_sizes: torch.Tensor) -> list[dict]: + logits = outputs["pred_logits"] # [B,Q,C] + kpts = outputs["pred_keypoints"] # [B,Q,2K] + if kpts.ndim != 3: + raise ValueError(f"DETRPosePostProcessor expects pred_keypoints [B,Q,2K], got shape={tuple(kpts.shape)}") + + B, Q, _ = kpts.shape + K = self.num_keypoints + if int(kpts.shape[-1]) != int(K * 2): + raise ValueError(f"pred_keypoints last dim must be 2K={K*2}, got {int(kpts.shape[-1])}") + + # DETRPose repo selects top-k over (Q*C). That is fine when labels are stable. + # For person-only training, it is more robust to select per-query best class first, + # then pick top-k queries. This avoids accidentally selecting the "wrong class slot" + # for an otherwise good pose query early in training. + scores_all = torch.sigmoid(logits) # [B,Q,C] + per_q_scores, per_q_labels = scores_all.max(dim=-1) # [B,Q] + topk = min(int(self.num_top_queries), int(Q)) + scores, q_index = torch.topk(per_q_scores, topk, dim=-1) # [B,top] + labels = per_q_labels.gather(dim=1, index=q_index) # [B,top] + + kpts_sel = kpts.gather(dim=1, index=q_index.unsqueeze(-1).repeat(1, 1, kpts.shape[-1])) # [B,top,2K] + kpts_xy = kpts_sel.view(B, -1, K, 2) # normalized in resized image space + + # Convert normalized coords to original pixel coords: + # resized pixel = norm * image_size ; original pixel = resized / scale = norm * orig_size + orig_wh = orig_target_sizes.to(dtype=kpts_xy.dtype) # [B,2] (w,h) + kpts_px = kpts_xy * orig_wh[:, None, None, :] # [B,top,K,2] + + # Build a robust box around keypoints (in original px). + # Using strict min/max is very sensitive to a single outlier keypoint. + x = kpts_px[..., 0] + y = kpts_px[..., 1] + try: + x1 = torch.quantile(x, 0.02, dim=-1) + y1 = torch.quantile(y, 0.02, dim=-1) + x2 = torch.quantile(x, 0.98, dim=-1) + y2 = torch.quantile(y, 0.98, dim=-1) + except Exception: + x1 = x.min(dim=-1).values + y1 = y.min(dim=-1).values + x2 = x.max(dim=-1).values + y2 = y.max(dim=-1).values + + # small padding for visualization stability + pad = 0.05 + bw = (x2 - x1).clamp(min=1.0) + bh = (y2 - y1).clamp(min=1.0) + x1 = x1 - pad * bw + x2 = x2 + pad * bw + y1 = y1 - pad * bh + y2 = y2 + pad * bh + + w_img = orig_wh[:, 0].view(B, 1) + h_img = orig_wh[:, 1].view(B, 1) + x1 = torch.clamp(x1, min=0.0) + y1 = torch.clamp(y1, min=0.0) + x2 = torch.clamp(x2, min=0.0) + y2 = torch.clamp(y2, min=0.0) + x1 = torch.minimum(x1, w_img) + x2 = torch.minimum(x2, w_img) + y1 = torch.minimum(y1, h_img) + y2 = torch.minimum(y2, h_img) + x_min = torch.minimum(x1, x2) + y_min = torch.minimum(y1, y2) + x_max = torch.maximum(x1, x2) + y_max = torch.maximum(y1, y2) + boxes = torch.stack([x_min, y_min, x_max, y_max], dim=-1) # [B,top,4] + + # Keypoints: [x,y,v] where v is a visibility flag; set to 1 for all predicted keypoints. + v = torch.ones_like(kpts_px[..., :1]) + keypoints = torch.cat([kpts_px, v], dim=-1) # [B,top,K,3] + + if self.remap_mscoco_category: + # For person-only pose training we typically use 1 or 2 "classes" just to keep indices valid. + # In that case, mapping through mscoco_label2category (0..79 -> category_id) is not meaningful + # for label==1 (it becomes bicycle==2). That can cause almost all predictions to be filtered out + # downstream (we keep only category_id==1 for keypoint eval). + # + # Make eval robust: force all predictions to COCO category_id=1 (person) when num_classes <= 2. + if int(self.num_classes) <= 2: + labels = torch.ones_like(labels) + else: + labels = _remap_labels_to_coco_category(labels) + + out: list[dict] = [] + for bi in range(B): + out.append({"labels": labels[bi], "boxes": boxes[bi], "scores": scores[bi], "keypoints": keypoints[bi]}) + return out + diff --git a/pose_estimation_berna/core/posture.py b/pose_estimation_berna/core/posture.py new file mode 100644 index 00000000..31f82c94 --- /dev/null +++ b/pose_estimation_berna/core/posture.py @@ -0,0 +1,267 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Dict, Optional, Tuple + +import numpy as np + + +COCO17 = { + "nose": 0, + "l_eye": 1, + "r_eye": 2, + "l_ear": 3, + "r_ear": 4, + "l_sho": 5, + "r_sho": 6, + "l_elb": 7, + "r_elb": 8, + "l_wri": 9, + "r_wri": 10, + "l_hip": 11, + "r_hip": 12, + "l_kne": 13, + "r_kne": 14, + "l_ank": 15, + "r_ank": 16, +} + + +def _softmax(x: np.ndarray) -> np.ndarray: + x = x - np.max(x) + e = np.exp(x) + return e / (np.sum(e) + 1e-9) + + +def _kpt_xy(kpts: np.ndarray, idx: int) -> Tuple[float, float, float]: + """Return (x,y,score) for keypoint idx.""" + x, y, s = kpts[idx] + return float(x), float(y), float(s) + + +@dataclass +class PostureResult: + probs: Dict[str, float] # stand/crouch/lie + state: str + conf: float + valid: bool + + +class PostureEstimator: + """ + Heuristic posture estimator for COCO-17 keypoints. + + Outputs probabilities over: + - stand + - crouch + - lie + + Key design: occlusion-robust via hold/decay: + - if too few keypoints are confident, keep previous probabilities and decay them slowly. + """ + + def __init__( + self, + kpt_thr: float = 0.2, + hold_frames: int = 15, + decay: float = 0.9, + switch_frames: int = 3, + signal_ema: float = 0.0, + ): + self.kpt_thr = float(kpt_thr) + self.hold_frames = int(hold_frames) + self.decay = float(decay) + self.switch_frames = max(1, int(switch_frames)) + # EMA smoothing for posture input signals (0 disables smoothing) + self.signal_ema = float(signal_ema) + if not np.isfinite(self.signal_ema): + self.signal_ema = 0.0 + self.signal_ema = float(np.clip(self.signal_ema, 0.0, 0.999)) + + # per track state + self._last_probs: Dict[int, np.ndarray] = {} + self._missing: Dict[int, int] = {} + self._state: Dict[int, str] = {} + self._pending_state: Dict[int, str] = {} + self._pending_count: Dict[int, int] = {} + self._feat_ema: Dict[int, np.ndarray] = {} # [vert_extent, hip_knee, hip_to_ank, knee_norm, aspect] + + def reset(self): + self._last_probs = {} + self._missing = {} + self._state = {} + self._pending_state = {} + self._pending_count = {} + self._feat_ema = {} + + def estimate(self, track_id: int, box_xyxy: np.ndarray, kpts: Optional[np.ndarray]) -> PostureResult: + if kpts is None or kpts.size == 0: + return self._fallback(track_id, valid=False) + + # collect core joints + nose = _kpt_xy(kpts, COCO17["nose"]) + lsho = _kpt_xy(kpts, COCO17["l_sho"]) + rsho = _kpt_xy(kpts, COCO17["r_sho"]) + lhip = _kpt_xy(kpts, COCO17["l_hip"]) + rhip = _kpt_xy(kpts, COCO17["r_hip"]) + lkne = _kpt_xy(kpts, COCO17["l_kne"]) + rkne = _kpt_xy(kpts, COCO17["r_kne"]) + lank = _kpt_xy(kpts, COCO17["l_ank"]) + rank = _kpt_xy(kpts, COCO17["r_ank"]) + + core = [nose, lsho, rsho, lhip, rhip, lkne, rkne, lank, rank] + n_good = sum(1 for (_, _, s) in core if s >= self.kpt_thr) + + # if too occluded, hold previous + if n_good < 4: + return self._fallback(track_id, valid=False) + + # bbox geometry + x1, y1, x2, y2 = [float(v) for v in box_xyxy.tolist()] + bw = max(1.0, x2 - x1) + bh = max(1.0, y2 - y1) + + # compute midpoints when both sides available + def mid(a, b): + ax, ay, as_ = a + bx, by, bs_ = b + if as_ >= self.kpt_thr and bs_ >= self.kpt_thr: + return (0.5 * (ax + bx), 0.5 * (ay + by), 0.5 * (as_ + bs_)) + # fall back to whichever is visible + if as_ >= self.kpt_thr: + return (ax, ay, as_) + if bs_ >= self.kpt_thr: + return (bx, by, bs_) + return (float("nan"), float("nan"), 0.0) + + sho = mid(lsho, rsho) + hip = mid(lhip, rhip) + kne = mid(lkne, rkne) + ank = mid(lank, rank) + + _, y_sho, s_sho = sho + _, y_hip, s_hip = hip + _, y_kne, s_kne = kne + _, y_ank, s_ank = ank + + # Knee angle feature (more robust for stand vs crouch than raw y-deltas). + def _angle_deg(a, b, c): + # angle at point b for triangle a-b-c + ax, ay, as_ = a + bx, by, bs_ = b + cx, cy, cs_ = c + if as_ < self.kpt_thr or bs_ < self.kpt_thr or cs_ < self.kpt_thr: + return None + v1 = np.array([ax - bx, ay - by], dtype=np.float32) + v2 = np.array([cx - bx, cy - by], dtype=np.float32) + n1 = float(np.linalg.norm(v1)) + n2 = float(np.linalg.norm(v2)) + if n1 < 1e-3 or n2 < 1e-3: + return None + cosang = float(np.dot(v1, v2) / (n1 * n2)) + cosang = float(np.clip(cosang, -1.0, 1.0)) + return float(np.degrees(np.arccos(cosang))) + + l_knee_ang = _angle_deg(lhip, lkne, lank) + r_knee_ang = _angle_deg(rhip, rkne, rank) + knee_angles = [a for a in [l_knee_ang, r_knee_ang] if a is not None] + if knee_angles: + knee_ang = float(sum(knee_angles) / len(knee_angles)) # deg + else: + knee_ang = 170.0 # default to "straight-ish" (avoid false crouch when noisy) + knee_norm = float(np.clip(knee_ang / 180.0, 0.0, 1.0)) + + # features (scale-invariant-ish) + # - vertical extent proxy: (ankle_y - shoulder_y) / bbox_h + # - crouch proxy: (knee_y - hip_y) / bbox_h (smaller when crouched) + knee angle + # - lie proxy: bbox aspect ratio + low vertical extent + vert_extent = (y_ank - y_sho) / bh if (s_ank >= self.kpt_thr and s_sho >= self.kpt_thr) else (bh / bh) + hip_knee = (y_kne - y_hip) / bh if (s_kne >= self.kpt_thr and s_hip >= self.kpt_thr) else 0.2 + hip_to_ank = (y_ank - y_hip) / bh if (s_ank >= self.kpt_thr and s_hip >= self.kpt_thr) else 0.4 + aspect = bw / bh + + # Optional EMA smoothing of input signals to reduce frame-to-frame flips. + if self.signal_ema > 0.0: + cur = np.array([vert_extent, hip_knee, hip_to_ank, knee_norm, aspect], dtype=np.float32) + prev = self._feat_ema.get(track_id, None) + if prev is None or prev.shape != cur.shape: + sm = cur + else: + a = float(self.signal_ema) + sm = a * prev + (1.0 - a) * cur + self._feat_ema[track_id] = sm.astype(np.float32) + vert_extent, hip_knee, hip_to_ank, knee_norm, aspect = [float(x) for x in sm.tolist()] + + # logits (hand-tuned, now with knee angle to reduce stand/crouch flips) + # stand: straight knees + large vertical extent + hips not too low. + stand_logit = ( + 3.0 * (vert_extent - 0.62) + + 2.0 * (knee_norm - 0.86) + + 1.0 * (hip_to_ank - 0.32) + + 0.8 * (0.35 - aspect) + ) + # crouch: bent knees + hips closer to knees/ankles, but still vertical-ish. + crouch_logit = ( + 3.0 * (0.24 - hip_knee) + + 2.0 * (0.83 - knee_norm) + + 1.0 * (0.34 - hip_to_ank) + + 1.2 * (vert_extent - 0.45) + ) + # lie: wide aspect and small vertical extent + lie_logit = 3.0 * (aspect - 0.55) + 2.0 * (0.45 - vert_extent) + + probs_arr = _softmax(np.array([stand_logit, crouch_logit, lie_logit], dtype=np.float32)) + + self._last_probs[track_id] = probs_arr + self._missing[track_id] = 0 + + probs = {"stand": float(probs_arr[0]), "crouch": float(probs_arr[1]), "lie": float(probs_arr[2])} + raw_state = max(probs, key=probs.get) + raw_conf = float(probs[raw_state]) + + # Hysteresis: require N consecutive frames before switching state (reduces flip-flop). + cur_state = self._state.get(track_id, raw_state) + if raw_state == cur_state: + self._pending_state.pop(track_id, None) + self._pending_count.pop(track_id, None) + final_state = cur_state + else: + pend = self._pending_state.get(track_id, None) + if pend != raw_state: + self._pending_state[track_id] = raw_state + self._pending_count[track_id] = 1 + else: + self._pending_count[track_id] = int(self._pending_count.get(track_id, 0)) + 1 + if int(self._pending_count.get(track_id, 0)) >= int(self.switch_frames): + final_state = raw_state + self._state[track_id] = final_state + self._pending_state.pop(track_id, None) + self._pending_count.pop(track_id, None) + else: + final_state = cur_state + + # ensure state exists + self._state.setdefault(track_id, final_state) + conf = float(probs.get(final_state, raw_conf)) + return PostureResult(probs=probs, state=final_state, conf=conf, valid=True) + + def _fallback(self, track_id: int, valid: bool) -> PostureResult: + prev = self._last_probs.get(track_id, None) + miss = int(self._missing.get(track_id, 0)) + 1 + self._missing[track_id] = miss + + if prev is None: + probs_arr = np.array([1 / 3, 1 / 3, 1 / 3], dtype=np.float32) + else: + # hold for a while, then decay towards uniform + if miss <= self.hold_frames: + probs_arr = prev + else: + probs_arr = self.decay * prev + (1.0 - self.decay) * np.array([1 / 3, 1 / 3, 1 / 3], dtype=np.float32) + probs_arr = probs_arr / (np.sum(probs_arr) + 1e-9) + self._last_probs[track_id] = probs_arr + + probs = {"stand": float(probs_arr[0]), "crouch": float(probs_arr[1]), "lie": float(probs_arr[2])} + state = max(probs, key=probs.get) + conf = float(probs[state]) + return PostureResult(probs=probs, state=state, conf=conf, valid=valid) \ No newline at end of file diff --git a/pose_estimation_berna/core/tracking.py b/pose_estimation_berna/core/tracking.py new file mode 100644 index 00000000..5b27671d --- /dev/null +++ b/pose_estimation_berna/core/tracking.py @@ -0,0 +1,1577 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Optional, Tuple + +import numpy as np + +from pose_estimation_berna.core.depth import DepthRegion, occluder_front_fraction + + +COCO17 = { + "nose": 0, + "l_eye": 1, + "r_eye": 2, + "l_ear": 3, + "r_ear": 4, + "l_sho": 5, + "r_sho": 6, + "l_elb": 7, + "r_elb": 8, + "l_wri": 9, + "r_wri": 10, + "l_hip": 11, + "r_hip": 12, + "l_kne": 13, + "r_kne": 14, + "l_ank": 15, + "r_ank": 16, +} + + +def _softmax(x: np.ndarray) -> np.ndarray: + x = x.astype(np.float32) + x = x - float(np.max(x)) + e = np.exp(x) + return (e / (float(np.sum(e)) + 1e-9)).astype(np.float32) + + +def _posture_probs_from_kpts( + box_xyxy: np.ndarray, + kpts: Optional[np.ndarray], + kpt_thr: float = 0.2, +) -> Optional[np.ndarray]: + """ + Heuristic stand/crouch/lie probabilities from COCO-17 keypoints. + Returns probs [3] or None if too occluded/invalid. + """ + if kpts is None or kpts.size == 0: + return None + if kpts.shape[0] < 17 or kpts.shape[1] < 3: + return None + + def _kpt(idx: int): + x, y, s = kpts[int(idx)] + return float(x), float(y), float(s) + + nose = _kpt(COCO17["nose"]) + lsho = _kpt(COCO17["l_sho"]) + rsho = _kpt(COCO17["r_sho"]) + lhip = _kpt(COCO17["l_hip"]) + rhip = _kpt(COCO17["r_hip"]) + lkne = _kpt(COCO17["l_kne"]) + rkne = _kpt(COCO17["r_kne"]) + lank = _kpt(COCO17["l_ank"]) + rank = _kpt(COCO17["r_ank"]) + core = [nose, lsho, rsho, lhip, rhip, lkne, rkne, lank, rank] + n_good = sum(1 for (_, _, s) in core if s >= float(kpt_thr)) + if n_good < 4: + return None + + x1, y1, x2, y2 = [float(v) for v in box_xyxy.tolist()] + bw = max(1.0, x2 - x1) + bh = max(1.0, y2 - y1) + + def mid(a, b): + ax, ay, as_ = a + bx, by, bs_ = b + if as_ >= float(kpt_thr) and bs_ >= float(kpt_thr): + return (0.5 * (ax + bx), 0.5 * (ay + by), 0.5 * (as_ + bs_)) + if as_ >= float(kpt_thr): + return (ax, ay, as_) + if bs_ >= float(kpt_thr): + return (bx, by, bs_) + return (float("nan"), float("nan"), 0.0) + + sho = mid(lsho, rsho) + hip = mid(lhip, rhip) + kne = mid(lkne, rkne) + ank = mid(lank, rank) + + _, y_sho, s_sho = sho + _, y_hip, s_hip = hip + _, y_kne, s_kne = kne + _, y_ank, s_ank = ank + + vert_extent = (y_ank - y_sho) / bh if (s_ank >= kpt_thr and s_sho >= kpt_thr) else 1.0 + hip_knee = (y_kne - y_hip) / bh if (s_kne >= kpt_thr and s_hip >= kpt_thr) else 0.2 + aspect = bw / bh + + stand_logit = 3.0 * (vert_extent - 0.65) + 1.0 * (0.25 - aspect) + crouch_logit = 3.0 * (0.20 - hip_knee) + 1.5 * (vert_extent - 0.45) + lie_logit = 3.0 * (aspect - 0.55) + 2.0 * (0.45 - vert_extent) + return _softmax(np.array([stand_logit, crouch_logit, lie_logit], dtype=np.float32)) + + +def _posture_decay_update( + probs: Optional[np.ndarray], + missing: int, + hold_frames: int, + decay: float, +) -> Tuple[np.ndarray, int]: + """ + Hold posture probs for a while when missing, then decay towards uniform. + Returns (new_probs, new_missing). + """ + miss = int(missing) + 1 + if probs is None or getattr(probs, "size", 0) != 3: + probs_arr = np.array([1 / 3, 1 / 3, 1 / 3], dtype=np.float32) + return probs_arr, miss + + probs_arr = probs.astype(np.float32) + if miss <= int(hold_frames): + return probs_arr, miss + d = float(decay) + uni = np.array([1 / 3, 1 / 3, 1 / 3], dtype=np.float32) + probs_arr = d * probs_arr + (1.0 - d) * uni + probs_arr = probs_arr / (float(np.sum(probs_arr)) + 1e-9) + return probs_arr.astype(np.float32), miss + + +COCO_SIGMAS = np.array( + [ + 0.026, 0.025, 0.025, 0.035, 0.035, + 0.079, 0.079, 0.072, 0.072, 0.062, + 0.062, 0.107, 0.107, 0.087, 0.087, + 0.089, 0.089, + ], + dtype=np.float32, +) + + +def _iou_xyxy(a: np.ndarray, b: np.ndarray) -> float: + """IoU for single boxes in xyxy.""" + x1 = max(float(a[0]), float(b[0])) + y1 = max(float(a[1]), float(b[1])) + x2 = min(float(a[2]), float(b[2])) + y2 = min(float(a[3]), float(b[3])) + iw = max(0.0, x2 - x1) + ih = max(0.0, y2 - y1) + inter = iw * ih + if inter <= 0: + return 0.0 + area_a = max(0.0, float(a[2] - a[0])) * max(0.0, float(a[3] - a[1])) + area_b = max(0.0, float(b[2] - b[0])) * max(0.0, float(b[3] - b[1])) + denom = area_a + area_b - inter + return float(inter / denom) if denom > 1e-9 else 0.0 + + +def _area_xyxy(b: np.ndarray) -> float: + return float(max(0.0, float(b[2] - b[0])) * max(0.0, float(b[3] - b[1]))) + + +def _oks_like( + kpts_a: Optional[np.ndarray], + kpts_b: Optional[np.ndarray], + area: float, + kpt_thr: float = 0.2, + sigmas: np.ndarray = COCO_SIGMAS, +) -> float: + """ + Lightweight OKS-like similarity in [0,1] between two keypoint sets in pixel coords. + + kpts_*: [K,3] where last dim is (x_px, y_px, score) + area: object area in pixel^2 (used to normalize distances) + """ + if kpts_a is None or kpts_b is None: + return 0.0 + if kpts_a.size == 0 or kpts_b.size == 0: + return 0.0 + if kpts_a.shape[-1] < 2 or kpts_b.shape[-1] < 2: + return 0.0 + + K = int(min(kpts_a.shape[0], kpts_b.shape[0], sigmas.shape[0])) + if K <= 0: + return 0.0 + + xa = kpts_a[:K, 0].astype(np.float32) + ya = kpts_a[:K, 1].astype(np.float32) + xb = kpts_b[:K, 0].astype(np.float32) + yb = kpts_b[:K, 1].astype(np.float32) + + if kpts_a.shape[1] >= 3 and kpts_b.shape[1] >= 3: + sa = kpts_a[:K, 2].astype(np.float32) + sb = kpts_b[:K, 2].astype(np.float32) + vis = (sa >= float(kpt_thr)) & (sb >= float(kpt_thr)) + else: + vis = np.ones((K,), dtype=bool) + + vis_count = int(np.sum(vis)) + if vis_count <= 0: + return 0.0 + + dx = (xa - xb).astype(np.float32) + dy = (ya - yb).astype(np.float32) + d2 = dx * dx + dy * dy # [K] + + area = float(max(1.0, area)) + vars_ = (sigmas[:K] * 2.0) ** 2 # [K] + denom = (2.0 * vars_ * area) + denom = np.maximum(1e-6, denom).astype(np.float32) + + e = d2 / denom # [K] + oks_per_kpt = np.exp(-e) # [K] + oks = float(np.sum(oks_per_kpt[vis]) / float(vis_count)) + if not np.isfinite(oks): + return 0.0 + return float(np.clip(oks, 0.0, 1.0)) + + +@dataclass +class Track: + track_id: int + box_xyxy: np.ndarray # [4] + score: float + keypoints: Optional[np.ndarray] # [K,3] in pixels + pose_quality: float = 0.0 + posture_probs: Optional[np.ndarray] = None # [3] stand/crouch/lie + posture_missing: int = 0 + age: int = 0 + time_since_update: int = 0 + + # Smoothed state (optional) + box_smooth: Optional[np.ndarray] = None + kpt_smooth: Optional[np.ndarray] = None + # Occlusion bookkeeping (Kalman tracker primarily, but safe to keep here) + last_meas_cxcywh: Optional[np.ndarray] = None # [4] last measurement (cx,cy,w,h) in pixels + y_foot_ref: Optional[float] = None # last measured footline (y2) in pixels; update only when visible + oof_count: int = 0 # consecutive "out of frame" predictions + last_visible_box_xyxy: Optional[np.ndarray] = None # [4] xyxy from last visible frame + freeze_active: bool = False + freeze_box_xyxy: Optional[np.ndarray] = None # [4] xyxy frozen at first occlusion frame + # Anti-FP / stability: require a few matches before "confirming" a track. + hits: int = 0 + confirmed: bool = False + + # Depth / occlusion reasoning (optional; inference-only) + inv_depth: Optional[float] = None # inverse depth proxy (larger => closer), per-frame estimate + inv_depth_vel: float = 0.0 + inv_depth_conf: float = 0.0 # 0..1 confidence/consistency proxy + occluded_conf: float = 0.0 # 0..1 "likely behind occluder" confidence + occluder_front_frac: float = 0.0 # 0..1 fraction of nearer pixels in predicted region + occlusion_dir: str = "none" # "none" | "bottom_up" (extendable) + + # Full-body reference (helps prevent bbox collapsing during partial occlusion). + # Stored in cxcywh pixel space; updated only when we believe the full body is visible. + full_ref_cxcywh: Optional[np.ndarray] = None # [4] + + # Side-occluder boundary (e.g. wall edge) in image x-coordinates. + # Used to "snap" the occluded side of the bbox into the occluder for a more realistic full-body region. + wall_edge_x: Optional[float] = None + + # Occlusion latch: once we detect partial occlusion (bbox collapse / kpt loss), + # keep drawing an occluded "full-body" box for a while even if evidence disappears next frame. + occ_latch: int = 0 # frames remaining + occ_latch_box_xyxy: Optional[np.ndarray] = None # [4] stable occluded box to draw while latched + + +_UPPER_KPT_IDXS = np.array( + [ + COCO17["nose"], + COCO17["l_eye"], + COCO17["r_eye"], + COCO17["l_ear"], + COCO17["r_ear"], + COCO17["l_sho"], + COCO17["r_sho"], + COCO17["l_elb"], + COCO17["r_elb"], + COCO17["l_wri"], + COCO17["r_wri"], + ], + dtype=np.int64, +) +_LOWER_KPT_IDXS = np.array( + [ + COCO17["l_hip"], + COCO17["r_hip"], + COCO17["l_kne"], + COCO17["r_kne"], + COCO17["l_ank"], + COCO17["r_ank"], + ], + dtype=np.int64, +) + +_LEFT_SIDE_KPT_IDXS = np.array( + [ + COCO17["l_sho"], + COCO17["l_elb"], + COCO17["l_wri"], + COCO17["l_hip"], + COCO17["l_kne"], + COCO17["l_ank"], + ], + dtype=np.int64, +) +_RIGHT_SIDE_KPT_IDXS = np.array( + [ + COCO17["r_sho"], + COCO17["r_elb"], + COCO17["r_wri"], + COCO17["r_hip"], + COCO17["r_kne"], + COCO17["r_ank"], + ], + dtype=np.int64, +) + + +def _count_visible_kpts(kpts: Optional[np.ndarray], idxs: np.ndarray, kpt_thr: float) -> int: + if kpts is None or getattr(kpts, "size", 0) == 0: + return 0 + if kpts.ndim != 2 or kpts.shape[1] < 3: + return 0 + K = int(kpts.shape[0]) + idxs = idxs[(idxs >= 0) & (idxs < K)] + if idxs.size == 0: + return 0 + s = kpts[idxs, 2].astype(np.float32) + return int(np.sum(s >= float(kpt_thr))) + + +def _estimate_vertical_edge_x( + gray_u8: Optional[np.ndarray], + *, + x_start: int, + x_end: int, + y_start: int, + y_end: int, +) -> Optional[float]: + """ + Estimate a strong vertical edge position (x) by summing |dI/dx| over a y-range. + Returns x in image coordinates or None if not enough data. + """ + if gray_u8 is None or getattr(gray_u8, "size", 0) == 0: + return None + if gray_u8.ndim != 2: + return None + H, W = int(gray_u8.shape[0]), int(gray_u8.shape[1]) + xs = int(np.clip(int(x_start), 0, max(0, W - 2))) + xe = int(np.clip(int(x_end), 1, max(1, W - 1))) + if xe <= xs + 1: + return None + ys = int(np.clip(int(y_start), 0, max(0, H - 2))) + ye = int(np.clip(int(y_end), 1, max(1, H - 1))) + if ye <= ys + 2: + return None + + patch = gray_u8[ys:ye, xs:xe].astype(np.int16) # [h,w] + # gradient along x between neighboring columns => shape [h, w-1] + gx = np.abs(patch[:, 1:] - patch[:, :-1]).astype(np.int32) + col_score = np.sum(gx, axis=0) # [w-1] + if col_score.size < 2: + return None + j = int(np.argmax(col_score)) + # edge lies between columns (xs+j) and (xs+j+1); return as float x + return float(xs + j + 0.5) + + +class _KalmanFilter2D: + """ + Minimal constant-velocity Kalman filter for 1D scalar (e.g. inverse depth): + x = [z, vz] + z_meas = [z] + """ + + def __init__(self, init_z: float, process_noise: float = 1e-3, measurement_noise: float = 2e-2): + self.x = np.zeros((2, 1), dtype=np.float32) + self.x[0, 0] = np.float32(init_z) + self.x[1, 0] = np.float32(0.0) + + self.P = np.eye(2, dtype=np.float32) * 1.0 + self.F = np.array([[1.0, 1.0], [0.0, 1.0]], dtype=np.float32) + self.H = np.array([[1.0, 0.0]], dtype=np.float32) + self.Q0 = np.eye(2, dtype=np.float32) * float(process_noise) + self.Q = self.Q0.copy() + self.R0 = np.eye(1, dtype=np.float32) * float(measurement_noise) + self.R = self.R0.copy() + + def predict(self, q_scale: float = 1.0) -> float: + self.x = self.F @ self.x + self.Q = self.Q0 * float(max(0.0, q_scale)) + self.P = (self.F @ self.P @ self.F.T) + self.Q + return float(self.x[0, 0]) + + def update(self, z_meas: float, r_scale: float = 1.0): + z = np.array([[float(z_meas)]], dtype=np.float32) + self.R = self.R0 * float(max(1e-6, r_scale)) + y = z - (self.H @ self.x) + S = self.H @ self.P @ self.H.T + self.R + K = self.P @ self.H.T @ np.linalg.inv(S) + self.x = self.x + (K @ y) + I = np.eye(2, dtype=np.float32) + self.P = (I - (K @ self.H)) @ self.P + + def get(self) -> float: + return float(self.x[0, 0]) + + def get_vel(self) -> float: + return float(self.x[1, 0]) + + +class IoUTracker: + """ + Simple online tracker: + - matches detections to existing tracks by IoU (greedy) + - keeps tracks alive for max_age frames without update + - optional EMA smoothing for boxes/keypoints (reduces jitter) + + This is intentionally lightweight (no Kalman) to keep dependencies minimal. + """ + + def __init__( + self, + iou_threshold: float = 0.3, + max_age: int = 30, + smooth_alpha: float = 0.8, + smooth_boxes: bool = True, + smooth_keypoints: bool = True, + min_area_ratio: float = 0.6, + posture_kpt_thr: float = 0.2, + posture_poseq_thr: float = 0.0, + posture_hold_frames: int = 15, + posture_decay: float = 0.9, + score_ema: float = 0.8, + new_track_score_thr: float = 0.25, + confirm_hits: int = 3, + unconfirmed_max_age: int = 2, + ): + self.iou_threshold = float(iou_threshold) + self.max_age = int(max_age) + self.smooth_alpha = float(smooth_alpha) + self.smooth_boxes = bool(smooth_boxes) + self.smooth_keypoints = bool(smooth_keypoints) + self.min_area_ratio = float(min_area_ratio) + self.posture_kpt_thr = float(posture_kpt_thr) + self.posture_poseq_thr = float(posture_poseq_thr) + self.posture_hold_frames = int(posture_hold_frames) + self.posture_decay = float(posture_decay) + self.score_ema = float(np.clip(float(score_ema), 0.0, 0.999)) + self.new_track_score_thr = float(new_track_score_thr) + self.confirm_hits = max(1, int(confirm_hits)) + self.unconfirmed_max_age = max(0, int(unconfirmed_max_age)) + self._next_id = 1 + self._tracks: List[Track] = [] + + @property + def tracks(self) -> List[Track]: + return list(self._tracks) + + def reset(self): + self._tracks = [] + self._next_id = 1 + + def _ema(self, prev: np.ndarray, cur: np.ndarray) -> np.ndarray: + a = self.smooth_alpha + return (a * prev + (1.0 - a) * cur).astype(np.float32) + + def update( + self, + det_boxes_xyxy: np.ndarray, # [N,4] + det_scores: np.ndarray, # [N] + det_keypoints: Optional[np.ndarray] = None, # [N,K,3] + det_pose_quality: Optional[np.ndarray] = None, # [N] + ) -> List[Track]: + # age all tracks + for t in self._tracks: + t.age += 1 + t.time_since_update += 1 + + n = int(det_boxes_xyxy.shape[0]) if det_boxes_xyxy is not None else 0 + if n == 0: + # prune dead tracks + self._tracks = [t for t in self._tracks if t.time_since_update <= self.max_age] + return self.tracks + + # greedy match by IoU + unmatched_dets = set(range(n)) + unmatched_tracks = set(range(len(self._tracks))) + matches: List[Tuple[int, int]] = [] # (track_idx, det_idx) + + if self._tracks: + # Build all candidate pairs and greedily pick highest IoU first (global 1-1 matching). + pairs: List[Tuple[float, int, int]] = [] + for ti, t in enumerate(self._tracks): + for di in range(n): + iou = _iou_xyxy(t.box_xyxy, det_boxes_xyxy[di]) + if iou >= self.iou_threshold: + pairs.append((iou, ti, di)) + pairs.sort(key=lambda x: x[0], reverse=True) + + used_tracks, used_dets = set(), set() + for iou, ti, di in pairs: + if ti in used_tracks or di in used_dets: + continue + used_tracks.add(ti) + used_dets.add(di) + matches.append((ti, di)) + + for ti, di in matches: + unmatched_tracks.discard(ti) + unmatched_dets.discard(di) + + # update matched tracks + for ti, di in matches: + t = self._tracks[ti] + cur_box = det_boxes_xyxy[di].astype(np.float32) + + # Occlusion robustness: don't let the box collapse too much vs previous. + prev_box = t.box_smooth if t.box_smooth is not None else t.box_xyxy + prev_area = _area_xyxy(prev_box) + cur_area = _area_xyxy(cur_box) + if prev_area > 1.0 and cur_area / prev_area < self.min_area_ratio: + # Keep previous box as the measurement to avoid shrink under occlusion. + cur_box = prev_box.astype(np.float32) + + t.box_xyxy = cur_box + det_s = float(det_scores[di]) + a = float(self.score_ema) + t.score = float(a * float(t.score) + (1.0 - a) * det_s) + t.hits = int(t.hits) + 1 + if int(t.hits) >= int(self.confirm_hits): + t.confirmed = True + if det_pose_quality is not None and det_pose_quality.size > di: + t.pose_quality = float(det_pose_quality[di]) + t.time_since_update = 0 + + cur_kpt = None + if det_keypoints is not None: + cur_kpt = det_keypoints[di].astype(np.float32) + t.keypoints = cur_kpt + + # posture update gated by pose_quality (optional) + pq_ok = (self.posture_poseq_thr <= 0.0) or (float(t.pose_quality) >= float(self.posture_poseq_thr)) + probs = _posture_probs_from_kpts(cur_box, cur_kpt, kpt_thr=self.posture_kpt_thr) if pq_ok else None + if probs is not None and probs.size == 3: + t.posture_probs = probs + t.posture_missing = 0 + else: + t.posture_probs, t.posture_missing = _posture_decay_update( + t.posture_probs, t.posture_missing, self.posture_hold_frames, self.posture_decay + ) + + if self.smooth_boxes: + if t.box_smooth is None: + t.box_smooth = cur_box.copy() + else: + t.box_smooth = self._ema(t.box_smooth, cur_box) + else: + t.box_smooth = None + + if self.smooth_keypoints and cur_kpt is not None: + if t.kpt_smooth is None: + t.kpt_smooth = cur_kpt.copy() + else: + # smooth xy only; keep score channel as-is + prev = t.kpt_smooth + sm_xy = self._ema(prev[..., :2], cur_kpt[..., :2]) + t.kpt_smooth = np.concatenate([sm_xy, cur_kpt[..., 2:3]], axis=-1) + else: + t.kpt_smooth = None + + # create new tracks for unmatched detections + for di in sorted(unmatched_dets): + box = det_boxes_xyxy[di].astype(np.float32) + score = float(det_scores[di]) + if float(score) < float(self.new_track_score_thr): + continue + kpt = det_keypoints[di].astype(np.float32) if det_keypoints is not None else None + pq = float(det_pose_quality[di]) if det_pose_quality is not None and det_pose_quality.size > di else 0.0 + t = Track( + track_id=self._next_id, + box_xyxy=box, + score=score, + keypoints=kpt, + pose_quality=pq, + ) + t.hits = 1 + t.confirmed = int(t.hits) >= int(self.confirm_hits) + pq_ok = (self.posture_poseq_thr <= 0.0) or (float(t.pose_quality) >= float(self.posture_poseq_thr)) + t.posture_probs = _posture_probs_from_kpts(box, kpt, kpt_thr=self.posture_kpt_thr) if pq_ok else None + t.posture_missing = 0 if (t.posture_probs is not None) else 1 + if self.smooth_boxes: + t.box_smooth = box.copy() + if self.smooth_keypoints and kpt is not None: + t.kpt_smooth = kpt.copy() + self._tracks.append(t) + self._next_id += 1 + + # occlusion handling for unmatched tracks: hold/decay posture probs + for ti in sorted(unmatched_tracks): + t = self._tracks[ti] + t.posture_probs, t.posture_missing = _posture_decay_update( + t.posture_probs, t.posture_missing, self.posture_hold_frames, self.posture_decay + ) + + # prune dead tracks + kept = [] + for t in self._tracks: + if int(t.time_since_update) > int(self.max_age): + continue + # Unconfirmed tracks are allowed only briefly if they miss updates (suppresses false positives). + if (not bool(t.confirmed)) and int(t.time_since_update) > int(self.unconfirmed_max_age): + continue + kept.append(t) + self._tracks = kept + return self.tracks + + +def _xyxy_to_cxcywh(b: np.ndarray) -> np.ndarray: + x1, y1, x2, y2 = [float(x) for x in b.tolist()] + w = max(1.0, x2 - x1) + h = max(1.0, y2 - y1) + cx = x1 + 0.5 * w + cy = y1 + 0.5 * h + return np.array([cx, cy, w, h], dtype=np.float32) + + +def _cxcywh_to_xyxy(s: np.ndarray) -> np.ndarray: + cx, cy, w, h = [float(x) for x in s.tolist()] + w = max(1.0, w) + h = max(1.0, h) + x1 = cx - 0.5 * w + y1 = cy - 0.5 * h + x2 = cx + 0.5 * w + y2 = cy + 0.5 * h + return np.array([x1, y1, x2, y2], dtype=np.float32) + + +class _KalmanFilter8D: + """ + Minimal constant-velocity Kalman filter for bbox in (cx,cy,w,h) with velocities: + x = [cx, cy, w, h, vx, vy, vw, vh] + z = [cx, cy, w, h] + """ + + def __init__( + self, + init_cxcywh: np.ndarray, + process_noise: float = 1e-2, + measurement_noise: float = 1e-1, + ): + x0 = init_cxcywh.astype(np.float32).reshape(4) + self.x = np.zeros((8, 1), dtype=np.float32) + self.x[0:4, 0] = x0 + + self.P = np.eye(8, dtype=np.float32) * 10.0 + self.P[4:, 4:] *= 100.0 # velocity uncertainty higher + + self.F = np.eye(8, dtype=np.float32) + for i in range(4): + self.F[i, i + 4] = 1.0 # dt=1 + + self.H = np.zeros((4, 8), dtype=np.float32) + self.H[0, 0] = 1.0 + self.H[1, 1] = 1.0 + self.H[2, 2] = 1.0 + self.H[3, 3] = 1.0 + + q = float(process_noise) + r = float(measurement_noise) + self.Q0 = np.eye(8, dtype=np.float32) * q + self.Q = self.Q0.copy() + self.R = np.eye(4, dtype=np.float32) * r + + def predict(self, q_scale: float = 1.0, p_inflate: float = 1.0) -> np.ndarray: + self.x = self.F @ self.x + qs = float(max(0.0, q_scale)) + self.Q = self.Q0 * qs + pi = float(max(1.0, p_inflate)) + self.P = (pi * (self.F @ self.P @ self.F.T)) + self.Q + return self.x[0:4, 0].copy() + + def update(self, z_cxcywh: np.ndarray): + z = z_cxcywh.astype(np.float32).reshape(4, 1) + y = z - (self.H @ self.x) + S = self.H @ self.P @ self.H.T + self.R + K = self.P @ self.H.T @ np.linalg.inv(S) + self.x = self.x + (K @ y) + I = np.eye(8, dtype=np.float32) + self.P = (I - (K @ self.H)) @ self.P + + def get_state(self) -> np.ndarray: + return self.x[0:4, 0].copy() + + def get_center_std(self) -> float: + # sqrt(var_cx + var_cy) + var = float(max(0.0, self.P[0, 0])) + float(max(0.0, self.P[1, 1])) + return float(np.sqrt(max(0.0, var))) + + def clamp_size(self, w_min: float, w_max: float, h_min: float, h_max: float): + self.x[2, 0] = np.float32(np.clip(float(self.x[2, 0]), float(w_min), float(w_max))) + self.x[3, 0] = np.float32(np.clip(float(self.x[3, 0]), float(h_min), float(h_max))) + + +@dataclass +class KalmanTrack(Track): + kf: Optional[_KalmanFilter8D] = None + zf: Optional[_KalmanFilter2D] = None + + +class KalmanTracker: + """ + Occlusion-robust tracker: + - Kalman predict keeps a bbox even when detector returns nothing (short occlusions) + - matching is still IoU-based, but on the *predicted* box + - optional EMA smoothing is preserved for drawing stability + """ + + def __init__( + self, + iou_threshold: float = 0.2, + oks_threshold: float = 1.0, + max_age: int = 60, + smooth_alpha: float = 0.8, + smooth_boxes: bool = True, + smooth_keypoints: bool = True, + min_area_ratio: float = 0.6, + process_noise: float = 1e-2, + measurement_noise: float = 1e-1, + score_decay: float = 0.98, + score_decay_grace: int = 0, + score_ema: float = 0.8, + iou_weight: float = 1.0, + oks_weight: float = 0.0, + kpt_thr: float = 0.2, + kpt_age_decay: float = 0.97, + posture_kpt_thr: float = 0.2, + posture_poseq_thr: float = 0.0, + posture_hold_frames: int = 15, + posture_decay: float = 0.9, + # Step 2: occlusion-aware update (deterministic, inference-only) + fps: float = 30.0, + occlusion_ttl_sec: float = 1.0, + q_inflate_alpha: float = 0.02, + p_inflate_per_frame: float = 1.02, + size_clamp_min_scale: float = 0.7, + size_clamp_max_scale: float = 1.3, + uncertainty_rel: float = 0.35, + uncertainty_abs_px: float = 80.0, + out_of_frame_max: int = 5, + new_track_score_thr: float = 0.25, + freeze_bbox: bool = True, + confirm_hits: int = 3, + unconfirmed_max_age: int = 2, + # Depth-based occlusion reasoning (optional; inference-only) + depth_enabled: bool = False, + depth_region: DepthRegion = "torso", + depth_occ_margin_abs: float = 0.02, + depth_occ_margin_rel: float = 0.08, + depth_occ_frac_thr: float = 0.35, + depth_occ_conf_ema: float = 0.85, + depth_ttl_mult: float = 3.0, + depth_score_decay_occluded: float = 0.995, + depth_process_noise: float = 1e-3, + depth_measurement_noise: float = 2e-2, + # Keypoint-based "full-body bbox" under partial occlusion (no retraining). + full_body_bbox: bool = False, + full_body_upper_min: int = 4, + full_body_lower_max: int = 1, + full_body_lower_update_min: int = 3, + full_body_min_width_ratio: float = 0.70, + full_body_wall_snap: bool = False, + wall_snap_search_px: int = 80, + wall_snap_pad_px: int = 6, + wall_snap_ema: float = 0.85, + occ_latch_expand_x: float = 0.25, + occ_latch_expand_y: float = 0.25, + ): + self.iou_threshold = float(iou_threshold) + self.oks_threshold = float(oks_threshold) + self.max_age = int(max_age) + self.smooth_alpha = float(smooth_alpha) + self.smooth_boxes = bool(smooth_boxes) + self.smooth_keypoints = bool(smooth_keypoints) + self.min_area_ratio = float(min_area_ratio) + self.process_noise = float(process_noise) + self.measurement_noise = float(measurement_noise) + self.score_decay = float(score_decay) + self.score_decay_grace = int(score_decay_grace) + self.score_ema = float(np.clip(float(score_ema), 0.0, 0.999)) + self.iou_weight = float(iou_weight) + self.oks_weight = float(oks_weight) + self.kpt_thr = float(kpt_thr) + self.kpt_age_decay = float(kpt_age_decay) + self.posture_kpt_thr = float(posture_kpt_thr) + self.posture_poseq_thr = float(posture_poseq_thr) + self.posture_hold_frames = int(posture_hold_frames) + self.posture_decay = float(posture_decay) + + self.fps = float(fps) if float(fps) > 1e-3 else 30.0 + self.occlusion_ttl_sec = float(max(0.0, occlusion_ttl_sec)) + self.q_inflate_alpha = float(max(0.0, q_inflate_alpha)) + self.p_inflate_per_frame = float(max(1.0, p_inflate_per_frame)) + self.size_clamp_min_scale = float(max(0.1, size_clamp_min_scale)) + self.size_clamp_max_scale = float(max(self.size_clamp_min_scale, size_clamp_max_scale)) + self.uncertainty_rel = float(max(0.0, uncertainty_rel)) + self.uncertainty_abs_px = float(max(0.0, uncertainty_abs_px)) + self.out_of_frame_max = int(max(0, out_of_frame_max)) + self.new_track_score_thr = float(new_track_score_thr) + self.freeze_bbox = bool(freeze_bbox) + self.confirm_hits = max(1, int(confirm_hits)) + self.unconfirmed_max_age = max(0, int(unconfirmed_max_age)) + + self.depth_enabled = bool(depth_enabled) + self.depth_region = depth_region + self.depth_occ_margin_abs = float(max(0.0, depth_occ_margin_abs)) + self.depth_occ_margin_rel = float(max(0.0, depth_occ_margin_rel)) + self.depth_occ_frac_thr = float(np.clip(float(depth_occ_frac_thr), 0.0, 1.0)) + self.depth_occ_conf_ema = float(np.clip(float(depth_occ_conf_ema), 0.0, 0.999)) + self.depth_ttl_mult = float(max(1.0, depth_ttl_mult)) + self.depth_score_decay_occluded = float(np.clip(float(depth_score_decay_occluded), 0.90, 0.9999)) + self.depth_process_noise = float(max(1e-8, depth_process_noise)) + self.depth_measurement_noise = float(max(1e-8, depth_measurement_noise)) + + self.full_body_bbox = bool(full_body_bbox) + self.full_body_upper_min = max(0, int(full_body_upper_min)) + self.full_body_lower_max = max(0, int(full_body_lower_max)) + self.full_body_lower_update_min = max(0, int(full_body_lower_update_min)) + self.full_body_min_width_ratio = float(np.clip(float(full_body_min_width_ratio), 0.1, 1.0)) + self.full_body_wall_snap = bool(full_body_wall_snap) + self.wall_snap_search_px = int(max(10, wall_snap_search_px)) + self.wall_snap_pad_px = int(max(0, wall_snap_pad_px)) + self.wall_snap_ema = float(np.clip(float(wall_snap_ema), 0.0, 0.999)) + # When occlusion latch triggers, optionally expand the latched box in the occlusion direction. + # Values are fractions of the stored full-body width/height (fallback to current box size). + self.occ_latch_expand_x = float(np.clip(float(occ_latch_expand_x), 0.0, 2.0)) + self.occ_latch_expand_y = float(np.clip(float(occ_latch_expand_y), 0.0, 2.0)) + + self._next_id = 1 + self._tracks: List[KalmanTrack] = [] + + @property + def tracks(self) -> List[KalmanTrack]: + return list(self._tracks) + + def reset(self): + self._tracks = [] + self._next_id = 1 + + def set_fps(self, fps: float): + if fps is None: + return + f = float(fps) + if f > 1e-3 and np.isfinite(f): + self.fps = f + + def _posture_label(self, t: KalmanTrack) -> str: + # posture_probs order: [stand, crouch, lie(prone)] + if t.posture_probs is None or getattr(t.posture_probs, "size", 0) != 3: + return "stand" + idx = int(np.argmax(t.posture_probs)) + return "stand" if idx == 0 else ("crouch" if idx == 1 else "prone") + + def _clamp_center_to_image(self, cx: float, cy: float, w: float, h: float, img_wh: Tuple[int, int]) -> Tuple[float, float]: + iw, ih = int(img_wh[0]), int(img_wh[1]) + # clamp so bbox stays inside frame (reduces jitter at edges) + cx = float(np.clip(cx, 0.5 * w, max(0.5 * w, float(iw) - 0.5 * w))) + cy = float(np.clip(cy, 0.5 * h, max(0.5 * h, float(ih) - 0.5 * h))) + return cx, cy + + def _ema(self, prev: np.ndarray, cur: np.ndarray) -> np.ndarray: + a = self.smooth_alpha + return (a * prev + (1.0 - a) * cur).astype(np.float32) + + def update( + self, + det_boxes_xyxy: np.ndarray, # [N,4] + det_scores: np.ndarray, # [N] + det_keypoints: Optional[np.ndarray] = None, # [N,K,3] + det_pose_quality: Optional[np.ndarray] = None, # [N] + img_wh: Optional[Tuple[int, int]] = None, # (w,h) for deterministic out-of-frame termination + det_inv_depth: Optional[np.ndarray] = None, # [N] inverse depth proxy per detection (aligned with det_boxes) + inv_depth_map: Optional[np.ndarray] = None, # [H,W] inverse depth proxy for current frame + frame_gray: Optional[np.ndarray] = None, # [H,W] uint8 grayscale (for wall-edge snapping) + ) -> List[KalmanTrack]: + # 1) predict all tracks forward + age (occlusion-aware) + ttl_frames = int(min(float(self.max_age), round(self.occlusion_ttl_sec * float(self.fps)))) + ttl_frames = max(0, ttl_frames) + for t in self._tracks: + t.age += 1 + t.time_since_update += 1 + # decay occlusion latch timer + if hasattr(t, "occ_latch"): + t.occ_latch = int(max(0, int(getattr(t, "occ_latch", 0)) - 1)) + if t.kf is not None: + # Inflate uncertainty during occlusion (time_since_update > 0) + miss = int(max(0, t.time_since_update)) + q_scale = 1.0 + (self.q_inflate_alpha * float(miss * miss)) + p_inflate = float(self.p_inflate_per_frame) if miss > 0 else 1.0 + pred = t.kf.predict(q_scale=q_scale, p_inflate=p_inflate) + + # Clamp predicted size to last measured size to avoid explosion/collapse during short hallucination. + if t.last_meas_cxcywh is not None: + mw = float(max(1.0, float(t.last_meas_cxcywh[2]))) + mh = float(max(1.0, float(t.last_meas_cxcywh[3]))) + t.kf.clamp_size( + w_min=self.size_clamp_min_scale * mw, + w_max=self.size_clamp_max_scale * mw, + h_min=self.size_clamp_min_scale * mh, + h_max=self.size_clamp_max_scale * mh, + ) + pred = t.kf.get_state() + + # Step 3/4: Freeze-BBox policy (critical for demo): + # - On first missing-detection frame: freeze width and y1/y2 from last visible frame. + # - During occlusion: allow ONLY horizontal motion (cx) from Kalman prediction. + if bool(self.freeze_bbox) and miss > 0 and t.last_visible_box_xyxy is not None: + # Activate freeze on first occluded frame only. + if miss == 1 and not bool(t.freeze_active): + t.freeze_active = True + t.freeze_box_xyxy = t.last_visible_box_xyxy.astype(np.float32).copy() + + if bool(t.freeze_active) and t.freeze_box_xyxy is not None: + fx1, fy1, fx2, fy2 = [float(v) for v in t.freeze_box_xyxy.tolist()] + w = float(max(1.0, fx2 - fx1)) + h = float(max(1.0, fy2 - fy1)) + cx = float(pred[0]) # allow only horizontal motion + cy = 0.5 * (fy1 + fy2) # freeze vertical extent + + if img_wh is not None: + cx, _ = self._clamp_center_to_image(cx, cy, w, h, img_wh) + cy = float(np.clip(cy, 0.5 * h, max(0.5 * h, float(img_wh[1]) - 0.5 * h))) + + # Write back into KF state for consistency (keep velocities as predicted) + t.kf.x[0, 0] = np.float32(cx) + t.kf.x[1, 0] = np.float32(cy) + t.kf.x[2, 0] = np.float32(w) + t.kf.x[3, 0] = np.float32(h) + pred = t.kf.get_state() + + # Optional: snap the occluded side of the bbox to a stored wall edge (sideways occlusion), + # while keeping the full-body width from the reference. + if ( + self.full_body_wall_snap + and miss > 0 + and t.wall_edge_x is not None + and t.full_ref_cxcywh is not None + and str(getattr(t, "occlusion_dir", "none")).startswith("sideways_") + ): + try: + full_w = float(max(1.0, float(t.full_ref_cxcywh[2]))) + except Exception: + full_w = 0.0 + if full_w > 1.0: + x_wall = float(t.wall_edge_x) + pad = float(self.wall_snap_pad_px) + # Use current predicted y1/y2 from KF, but adjust x1/x2. + x1p, y1p, x2p, y2p = [float(v) for v in _cxcywh_to_xyxy(pred).tolist()] + if str(t.occlusion_dir) == "sideways_left": + # occluder on left: snap left edge near wall edge, extend right by full_w + x1p = x_wall - pad + x2p = x1p + full_w + else: + # sideways_right: occluder on right + x2p = x_wall + pad + x1p = x2p - full_w + if img_wh is not None: + iw, ih = int(img_wh[0]), int(img_wh[1]) + x1p = float(np.clip(x1p, 0.0, float(iw - 1))) + x2p = float(np.clip(x2p, 0.0, float(iw))) + y1p = float(np.clip(y1p, 0.0, float(ih - 1))) + y2p = float(np.clip(y2p, 0.0, float(ih))) + t.box_xyxy = np.array([x1p, y1p, x2p, y2p], dtype=np.float32) + # Also reflect into KF state so following predictions are consistent. + pred2 = _xyxy_to_cxcywh(t.box_xyxy) + t.kf.x[0, 0] = np.float32(pred2[0]) + t.kf.x[1, 0] = np.float32(pred2[1]) + t.kf.x[2, 0] = np.float32(pred2[2]) + t.kf.x[3, 0] = np.float32(pred2[3]) + pred = t.kf.get_state() + + t.box_xyxy = _cxcywh_to_xyxy(pred) + + # Out-of-frame termination bookkeeping (deterministic). + if img_wh is not None: + w_img, h_img = int(img_wh[0]), int(img_wh[1]) + x1, y1, x2, y2 = [float(v) for v in t.box_xyxy.tolist()] + oof = (x2 < 0.0) or (y2 < 0.0) or (x1 > float(w_img)) or (y1 > float(h_img)) + t.oof_count = int(t.oof_count + 1) if oof else 0 + + # predict depth if enabled + if self.depth_enabled and t.zf is not None: + miss = int(max(0, t.time_since_update)) + q_scale_z = 1.0 + 0.02 * float(miss * miss) + try: + t.inv_depth = float(t.zf.predict(q_scale=q_scale_z)) + t.inv_depth_vel = float(t.zf.get_vel()) + except Exception: + pass + + # update occluder evidence while occluded (per-track) + if self.depth_enabled and inv_depth_map is not None and int(t.time_since_update) > 0 and t.inv_depth is not None: + frac = occluder_front_fraction( + inv_depth_map, + t.box_xyxy, + inv_depth_expected=float(t.inv_depth), + region=self.depth_region, + margin_abs=self.depth_occ_margin_abs, + margin_rel=self.depth_occ_margin_rel, + ) + t.occluder_front_frac = float(frac) + occ_like = 1.0 if float(frac) >= float(self.depth_occ_frac_thr) else 0.0 + a = float(self.depth_occ_conf_ema) + t.occluded_conf = float(a * float(t.occluded_conf) + (1.0 - a) * float(occ_like)) + + # decay score when not updated (makes "ghost" tracks fade) + if t.time_since_update > int(self.score_decay_grace): + decay = float(self.score_decay) + if self.depth_enabled and float(getattr(t, "occluded_conf", 0.0)) > 0.5: + decay = float(self.depth_score_decay_occluded) + t.score = float(t.score * decay) + + n = int(det_boxes_xyxy.shape[0]) if det_boxes_xyxy is not None else 0 + if n == 0: + # no detections => occlusion: decay posture towards unknown + for t in self._tracks: + t.posture_probs, t.posture_missing = _posture_decay_update( + t.posture_probs, t.posture_missing, self.posture_hold_frames, self.posture_decay + ) + # Step 2 termination rules while occluded: + # - hard TTL (<= 1s by default) + # - uncertainty too large + # - predicted bbox out-of-frame for too long + kept = [] + for t in self._tracks: + # Depth-aware TTL extension: keep alive longer when we have strong occluder evidence. + ttl_eff = int(ttl_frames) + if self.depth_enabled and ttl_frames > 0: + ttl_eff = int(round(float(ttl_frames) * (1.0 + float(t.occluded_conf) * (float(self.depth_ttl_mult) - 1.0)))) + if ttl_eff > 0 and int(t.time_since_update) > int(ttl_eff): + continue + if t.kf is not None and t.last_meas_cxcywh is not None: + min_wh = float(max(1.0, min(float(t.last_meas_cxcywh[2]), float(t.last_meas_cxcywh[3])))) + thr = max(self.uncertainty_abs_px, self.uncertainty_rel * min_wh) + if float(t.kf.get_center_std()) > float(thr): + continue + if self.out_of_frame_max > 0 and int(t.oof_count) > int(self.out_of_frame_max): + continue + if int(t.time_since_update) <= int(self.max_age): + kept.append(t) + self._tracks = kept + return self.tracks + + unmatched_dets = set(range(n)) + unmatched_tracks = set(range(len(self._tracks))) + matches: List[Tuple[int, int]] = [] + + # 2) greedy global 1-1 matching by combined score: + # IoU(pred_box, det_box) + OKS-like(track_kpts, det_kpts) + if self._tracks: + pairs: List[Tuple[float, int, int]] = [] + for ti, t in enumerate(self._tracks): + for di in range(n): + iou = _iou_xyxy(t.box_xyxy, det_boxes_xyxy[di]) + oks = 0.0 + use_oks = (self.oks_weight > 0.0) and (det_keypoints is not None) and (t.keypoints is not None) + if use_oks: + # Use a conservative area proxy to normalize distances. + area = min(_area_xyxy(t.box_xyxy), _area_xyxy(det_boxes_xyxy[di])) + oks = _oks_like( + t.keypoints, + det_keypoints[di], + area=area, + kpt_thr=self.kpt_thr, + ) + # If the track is "lost" for a while, its stored keypoints get stale: + # decay OKS contribution with time_since_update to avoid wrong long-range re-association. + oks = float(oks * (self.kpt_age_decay ** int(max(0, t.time_since_update)))) + + # Default behavior (stable baseline): IoU-only matching. + # If OKS association is enabled (oks_weight>0), allow OKS to also create a candidate match. + if (iou >= self.iou_threshold) or (use_oks and (oks >= self.oks_threshold)): + comb = (self.iou_weight * float(iou)) + (self.oks_weight * float(oks)) + pairs.append((comb, ti, di)) + + pairs.sort(key=lambda x: x[0], reverse=True) + + used_tracks, used_dets = set(), set() + for comb, ti, di in pairs: + if ti in used_tracks or di in used_dets: + continue + used_tracks.add(ti) + used_dets.add(di) + matches.append((ti, di)) + + for ti, di in matches: + unmatched_tracks.discard(ti) + unmatched_dets.discard(di) + + # 3) update matched tracks with measurements + for ti, di in matches: + t = self._tracks[ti] + cur_box = det_boxes_xyxy[di].astype(np.float32) + raw_box = cur_box.copy() # keep detector measurement before any occlusion expansion + cur_kpt = None + if det_keypoints is not None: + cur_kpt = det_keypoints[di].astype(np.float32) + + # Partial-occlusion handling (bottom-up): if upper body is visible but legs disappear, + # keep drawing / tracking a full-body bbox "through" the occluder. + # This prevents footline + full-body reference from collapsing upwards when someone walks behind a chair. + t.occlusion_dir = "none" + if self.full_body_bbox and cur_kpt is not None: + upper_vis = _count_visible_kpts(cur_kpt, _UPPER_KPT_IDXS, kpt_thr=self.kpt_thr) + lower_vis = _count_visible_kpts(cur_kpt, _LOWER_KPT_IDXS, kpt_thr=self.kpt_thr) + left_vis = _count_visible_kpts(cur_kpt, _LEFT_SIDE_KPT_IDXS, kpt_thr=self.kpt_thr) + right_vis = _count_visible_kpts(cur_kpt, _RIGHT_SIDE_KPT_IDXS, kpt_thr=self.kpt_thr) + + # Update "full-body reference" only when both sides + lower body are sufficiently visible. + # This is the key fix for sideways occlusion: don't let the reference shrink to a sliver. + full_ref_ok = ( + (upper_vis >= int(self.full_body_upper_min)) + and (lower_vis >= int(self.full_body_lower_update_min)) + and (left_vis >= 2) + and (right_vis >= 2) + ) + if full_ref_ok: + meas_raw = _xyxy_to_cxcywh(raw_box) + t.full_ref_cxcywh = meas_raw.astype(np.float32) + # Also refresh stable footline + last_visible full-body box for freeze policy. + t.y_foot_ref = float(raw_box[3]) + t.last_visible_box_xyxy = _cxcywh_to_xyxy(t.full_ref_cxcywh).astype(np.float32) + + bottom_up = (upper_vis >= int(self.full_body_upper_min)) and (lower_vis <= int(self.full_body_lower_max)) + # Sideways occlusion: one side visible, the other missing (e.g. behind a wall edge). + sideways_left_hidden = (right_vis >= 3) and (left_vis <= 1) + sideways_right_hidden = (left_vis >= 3) and (right_vis <= 1) + + if bottom_up: + t.occlusion_dir = "bottom_up" + # Prefer last measured full-body size if available. + ref = t.full_ref_cxcywh if t.full_ref_cxcywh is not None else t.last_meas_cxcywh + if ref is not None: + last_w = float(max(1.0, float(ref[2]))) + last_h = float(max(1.0, float(ref[3]))) + cx = 0.5 * float(cur_box[0] + cur_box[2]) + y1 = float(cur_box[1]) # keep current top (head/upper torso) + y2 = y1 + last_h + if t.y_foot_ref is not None and np.isfinite(float(t.y_foot_ref)): + y2 = max(y2, float(t.y_foot_ref)) + x1 = cx - 0.5 * last_w + x2 = cx + 0.5 * last_w + cur_box = np.array([x1, y1, x2, y2], dtype=np.float32) + elif t.y_foot_ref is not None and np.isfinite(float(t.y_foot_ref)): + # At least keep a stable footline if we have it. + cur_box = cur_box.copy() + cur_box[3] = max(float(cur_box[3]), float(t.y_foot_ref)) + + # Clamp to image if known + if img_wh is not None: + iw, ih = int(img_wh[0]), int(img_wh[1]) + cur_box[0] = np.float32(np.clip(float(cur_box[0]), 0.0, float(iw - 1))) + cur_box[2] = np.float32(np.clip(float(cur_box[2]), 0.0, float(iw))) + cur_box[1] = np.float32(np.clip(float(cur_box[1]), 0.0, float(ih - 1))) + cur_box[3] = np.float32(np.clip(float(cur_box[3]), 0.0, float(ih))) + elif sideways_left_hidden or sideways_right_hidden: + # Expand horizontally back to last known full-body width, anchored on the visible side. + ref = t.full_ref_cxcywh if t.full_ref_cxcywh is not None else t.last_meas_cxcywh + if ref is not None: + last_w = float(max(1.0, float(ref[2]))) + x1, y1, x2, y2 = [float(v) for v in cur_box.tolist()] + if sideways_left_hidden: + t.occlusion_dir = "sideways_left" + # right side visible => keep x2, expand x1 into the occluder (left) + x1 = x2 - last_w + else: + t.occlusion_dir = "sideways_right" + # left side visible => keep x1, expand x2 into the occluder (right) + x2 = x1 + last_w + cur_box = np.array([x1, y1, x2, y2], dtype=np.float32) + + if img_wh is not None: + iw, ih = int(img_wh[0]), int(img_wh[1]) + cur_box[0] = np.float32(np.clip(float(cur_box[0]), 0.0, float(iw - 1))) + cur_box[2] = np.float32(np.clip(float(cur_box[2]), 0.0, float(iw))) + cur_box[1] = np.float32(np.clip(float(cur_box[1]), 0.0, float(ih - 1))) + cur_box[3] = np.float32(np.clip(float(cur_box[3]), 0.0, float(ih))) + + # Estimate occluder boundary (wall edge) from grayscale edge. + if self.full_body_wall_snap and frame_gray is not None: + bh = float(max(1.0, float(raw_box[3] - raw_box[1]))) + ys = int(round(float(raw_box[1]) + 0.20 * bh)) + ye = int(round(float(raw_box[3]) - 0.10 * bh)) + spx = int(self.wall_snap_search_px) + if str(t.occlusion_dir) == "sideways_right": + xs = int(round(float(raw_box[2]))) + xe = xs + spx + else: + xe = int(round(float(raw_box[0]))) + xs = xe - spx + x_edge = _estimate_vertical_edge_x(frame_gray, x_start=xs, x_end=xe, y_start=ys, y_end=ye) + if x_edge is not None and np.isfinite(float(x_edge)): + if t.wall_edge_x is None or (not np.isfinite(float(t.wall_edge_x))): + t.wall_edge_x = float(x_edge) + else: + a = float(self.wall_snap_ema) + t.wall_edge_x = float(a * float(t.wall_edge_x) + (1.0 - a) * float(x_edge)) + + # Generic sideways collapse guard (works even if left/right kpts are unreliable): + # If current measured width shrinks a lot vs stored full-body reference, + # keep the reference for freeze and (if possible) infer which side is occluded. + if t.full_ref_cxcywh is not None: + ref_x1, _, ref_x2, _ = _cxcywh_to_xyxy(t.full_ref_cxcywh).tolist() + ref_w = float(max(1.0, float(ref_x2 - ref_x1))) + cur_w = float(max(1.0, float(raw_box[2] - raw_box[0]))) + if cur_w < (self.full_body_min_width_ratio * ref_w): + # Decide which side is occluded. + # Prefer a detected occluder boundary (wall_edge_x) if available; otherwise fall back to reference-edge heuristic. + wx = getattr(t, "wall_edge_x", None) + if wx is not None and np.isfinite(float(wx)): + wl = abs(float(wx) - float(raw_box[0])) + wr = abs(float(wx) - float(raw_box[2])) + t.occlusion_dir = "sideways_left" if wl < wr else "sideways_right" + else: + d_left = abs(float(raw_box[0]) - float(ref_x1)) + d_right = abs(float(raw_box[2]) - float(ref_x2)) + t.occlusion_dir = "sideways_left" if d_right < d_left else "sideways_right" + + # Update last_visible box to the full reference box (do NOT overwrite with the sliver). + t.last_visible_box_xyxy = _cxcywh_to_xyxy(t.full_ref_cxcywh).astype(np.float32) + t.y_foot_ref = float(t.last_visible_box_xyxy[3]) + + # Also try to estimate occluder boundary near the sliver edge. + if self.full_body_wall_snap and frame_gray is not None: + bh = float(max(1.0, float(raw_box[3] - raw_box[1]))) + ys = int(round(float(raw_box[1]) + 0.20 * bh)) + ye = int(round(float(raw_box[3]) - 0.10 * bh)) + spx = int(self.wall_snap_search_px) + if str(t.occlusion_dir) == "sideways_right": + xs = int(round(float(raw_box[2]))) + xe = xs + spx + else: + xe = int(round(float(raw_box[0]))) + xs = xe - spx + x_edge = _estimate_vertical_edge_x(frame_gray, x_start=xs, x_end=xe, y_start=ys, y_end=ye) + if x_edge is not None and np.isfinite(float(x_edge)): + if t.wall_edge_x is None or (not np.isfinite(float(t.wall_edge_x))): + t.wall_edge_x = float(x_edge) + else: + a = float(self.wall_snap_ema) + t.wall_edge_x = float(a * float(t.wall_edge_x) + (1.0 - a) * float(x_edge)) + + # Fallback: sideways occlusion without reliable keypoints. + # If the detector bbox collapses to a narrow sliver, expand to full_ref width and mark occlusion_dir, + # so visualization (orange box) and wall snapping can kick in. + if self.full_body_bbox and t.full_ref_cxcywh is not None: + ref_xyxy = _cxcywh_to_xyxy(t.full_ref_cxcywh).astype(np.float32) + ref_x1, ref_y1, ref_x2, ref_y2 = [float(v) for v in ref_xyxy.tolist()] + ref_w = float(max(1.0, ref_x2 - ref_x1)) + cur_w_raw = float(max(1.0, float(raw_box[2] - raw_box[0]))) + if cur_w_raw < (self.full_body_min_width_ratio * ref_w): + # Decide which side is occluded. + # Prefer wall_edge_x if available; otherwise fall back to reference-edge heuristic. + wx = getattr(t, "wall_edge_x", None) + use_wall = wx is not None and np.isfinite(float(wx)) + if use_wall: + wl = abs(float(wx) - float(raw_box[0])) + wr = abs(float(wx) - float(raw_box[2])) + side_left = wl < wr + else: + d_left = abs(float(raw_box[0]) - ref_x1) + d_right = abs(float(raw_box[2]) - ref_x2) + side_left = d_right < d_left + x1, y1, x2, y2 = [float(v) for v in cur_box.tolist()] + if side_left: + t.occlusion_dir = "sideways_left" + x1 = x2 - ref_w + else: + t.occlusion_dir = "sideways_right" + x2 = x1 + ref_w + cur_box = np.array([x1, y1, x2, y2], dtype=np.float32) + + # Keep last visible as full-body reference (avoid overwriting with sliver) + t.last_visible_box_xyxy = ref_xyxy.astype(np.float32) + t.y_foot_ref = float(ref_y2) + + # Estimate boundary near sliver edge. + if self.full_body_wall_snap and frame_gray is not None: + bh = float(max(1.0, float(raw_box[3] - raw_box[1]))) + ys = int(round(float(raw_box[1]) + 0.20 * bh)) + ye = int(round(float(raw_box[3]) - 0.10 * bh)) + spx = int(self.wall_snap_search_px) + if str(t.occlusion_dir) == "sideways_right": + xs = int(round(float(raw_box[2]))) + xe = xs + spx + else: + xe = int(round(float(raw_box[0]))) + xs = xe - spx + x_edge = _estimate_vertical_edge_x(frame_gray, x_start=xs, x_end=xe, y_start=ys, y_end=ye) + if x_edge is not None and np.isfinite(float(x_edge)): + if t.wall_edge_x is None or (not np.isfinite(float(t.wall_edge_x))): + t.wall_edge_x = float(x_edge) + else: + a = float(self.wall_snap_ema) + t.wall_edge_x = float(a * float(t.wall_edge_x) + (1.0 - a) * float(x_edge)) + + # Clamp to image if known + if img_wh is not None: + iw, ih = int(img_wh[0]), int(img_wh[1]) + cur_box[0] = np.float32(np.clip(float(cur_box[0]), 0.0, float(iw - 1))) + cur_box[2] = np.float32(np.clip(float(cur_box[2]), 0.0, float(iw))) + cur_box[1] = np.float32(np.clip(float(cur_box[1]), 0.0, float(ih - 1))) + cur_box[3] = np.float32(np.clip(float(cur_box[3]), 0.0, float(ih))) + + # Occlusion latch: if we have any occlusion signal this frame, latch an occluded box for TTL frames. + # This is the "make a box and keep it for N seconds" behavior. + if bool(self.full_body_bbox) and ttl_frames > 0: + occ_signal = str(getattr(t, "occlusion_dir", "none")) != "none" + if occ_signal: + # If we have a wall boundary estimate, ensure sideways direction matches it + # before we compute the latched expansion. + wx = getattr(t, "wall_edge_x", None) + od0 = str(getattr(t, "occlusion_dir", "none")) + if wx is not None and np.isfinite(float(wx)) and od0.startswith("sideways_"): + wl = abs(float(wx) - float(raw_box[0])) + wr = abs(float(wx) - float(raw_box[2])) + t.occlusion_dir = "sideways_left" if wl < wr else "sideways_right" + + t.occ_latch = int(ttl_frames) + # Prefer snapped/frozen box if available, else current expanded box, else full_ref + box_latch = None + if getattr(t, "freeze_active", False) and getattr(t, "freeze_box_xyxy", None) is not None: + box_latch = t.freeze_box_xyxy.astype(np.float32).copy() + elif cur_box is not None: + box_latch = cur_box.astype(np.float32).copy() + elif getattr(t, "full_ref_cxcywh", None) is not None: + box_latch = _cxcywh_to_xyxy(t.full_ref_cxcywh).astype(np.float32) + + # Expand latched box in occlusion direction (simple, practical heuristic). + if box_latch is not None: + bx1, by1, bx2, by2 = [float(v) for v in box_latch.tolist()] + # base size from full-body reference if available (more stable) + if getattr(t, "full_ref_cxcywh", None) is not None: + ref_xyxy = _cxcywh_to_xyxy(t.full_ref_cxcywh).astype(np.float32) + ref_w = float(max(1.0, float(ref_xyxy[2] - ref_xyxy[0]))) + ref_h = float(max(1.0, float(ref_xyxy[3] - ref_xyxy[1]))) + else: + ref_w = float(max(1.0, bx2 - bx1)) + ref_h = float(max(1.0, by2 - by1)) + + ex = float(self.occ_latch_expand_x) * ref_w + ey = float(self.occ_latch_expand_y) * ref_h + od = str(getattr(t, "occlusion_dir", "none")) + if od == "sideways_left": + bx1 -= ex + elif od == "sideways_right": + bx2 += ex + elif od == "bottom_up": + by2 += ey + # If we have a wall edge, bias the expansion to go "into" the occluder. + if getattr(t, "wall_edge_x", None) is not None and od.startswith("sideways_"): + wx = float(getattr(t, "wall_edge_x")) + pad = float(self.wall_snap_pad_px) + if od == "sideways_left": + bx1 = min(bx1, wx - pad - ex) + else: + bx2 = max(bx2, wx + pad + ex) + + if img_wh is not None: + iw, ih = int(img_wh[0]), int(img_wh[1]) + bx1 = float(np.clip(bx1, 0.0, float(iw - 1))) + bx2 = float(np.clip(bx2, 0.0, float(iw))) + by1 = float(np.clip(by1, 0.0, float(ih - 1))) + by2 = float(np.clip(by2, 0.0, float(ih))) + box_latch = np.array([bx1, by1, bx2, by2], dtype=np.float32) + t.occ_latch_box_xyxy = box_latch + else: + # If the track is clearly visible again, release latch early. + # Heuristic: sufficient lower-body visibility and width not collapsed. + release = False + if cur_kpt is not None and getattr(t, "full_ref_cxcywh", None) is not None: + lower_vis_now = _count_visible_kpts(cur_kpt, _LOWER_KPT_IDXS, kpt_thr=self.kpt_thr) + ref_xyxy = _cxcywh_to_xyxy(t.full_ref_cxcywh).astype(np.float32) + ref_w = float(max(1.0, float(ref_xyxy[2] - ref_xyxy[0]))) + cur_w = float(max(1.0, float(cur_box[2] - cur_box[0]))) + if int(lower_vis_now) >= int(self.full_body_lower_update_min) and cur_w >= (0.9 * ref_w): + release = True + if release: + t.occ_latch = 0 + t.occ_latch_box_xyxy = None + + prev_box = t.box_smooth if t.box_smooth is not None else t.box_xyxy + prev_area = _area_xyxy(prev_box) + cur_area = _area_xyxy(cur_box) + if prev_area > 1.0 and cur_area / prev_area < self.min_area_ratio: + cur_box = prev_box.astype(np.float32) + + # Kalman measurement update (cxcywh) + meas = _xyxy_to_cxcywh(cur_box) + if t.kf is None: + t.kf = _KalmanFilter8D( + meas, + process_noise=self.process_noise, + measurement_noise=self.measurement_noise, + ) + else: + t.kf.update(meas) + + t.box_xyxy = cur_box + det_s = float(det_scores[di]) + a = float(self.score_ema) + t.score = float(a * float(t.score) + (1.0 - a) * det_s) + t.time_since_update = 0 + t.hits = int(t.hits) + 1 + if int(t.hits) >= int(self.confirm_hits): + t.confirmed = True + if det_pose_quality is not None and det_pose_quality.size > di: + t.pose_quality = float(det_pose_quality[di]) + t.last_meas_cxcywh = meas.astype(np.float32) + # Update full-body reference only when we believe lower body is actually visible. + # Otherwise bottom-up occlusion would overwrite footline and "full-body bbox" would collapse. + if cur_kpt is not None: + lower_vis_now = _count_visible_kpts(cur_kpt, _LOWER_KPT_IDXS, kpt_thr=self.kpt_thr) + else: + lower_vis_now = int(self.full_body_lower_update_min) + # Don't overwrite the full-body reference with a narrow sideways-occluded sliver. + allow_update_visible = (not self.full_body_bbox) or (int(lower_vis_now) >= int(self.full_body_lower_update_min)) + if self.full_body_bbox and t.full_ref_cxcywh is not None: + ref_xyxy = _cxcywh_to_xyxy(t.full_ref_cxcywh).astype(np.float32) + ref_w = float(max(1.0, float(ref_xyxy[2] - ref_xyxy[0]))) + cur_w = float(max(1.0, float(cur_box[2] - cur_box[0]))) + if cur_w < (self.full_body_min_width_ratio * ref_w): + allow_update_visible = False + if allow_update_visible: + t.y_foot_ref = float(cur_box[3]) # y2 in pixels + t.last_visible_box_xyxy = cur_box.astype(np.float32).copy() + t.freeze_active = False + t.freeze_box_xyxy = None + t.oof_count = 0 + t.occluded_conf = 0.0 + t.occluder_front_frac = 0.0 + + t.keypoints = cur_kpt if cur_kpt is not None else t.keypoints + + # depth measurement update (inverse depth proxy) + if self.depth_enabled and det_inv_depth is not None and int(det_inv_depth.size) > di: + z_meas = float(det_inv_depth[di]) + if np.isfinite(z_meas): + if t.zf is None: + t.zf = _KalmanFilter2D( + init_z=z_meas, + process_noise=self.depth_process_noise, + measurement_noise=self.depth_measurement_noise, + ) + else: + # regular measurement update + t.zf.update(z_meas, r_scale=1.0) + t.inv_depth = float(t.zf.get()) + t.inv_depth_vel = float(t.zf.get_vel()) + t.inv_depth_conf = float(np.clip(1.0 / (abs(float(t.inv_depth_vel)) + 1.0), 0.0, 1.0)) + + pq_ok = (self.posture_poseq_thr <= 0.0) or (float(t.pose_quality) >= float(self.posture_poseq_thr)) + probs = _posture_probs_from_kpts(cur_box, cur_kpt, kpt_thr=self.posture_kpt_thr) if pq_ok else None + if probs is not None and probs.size == 3: + t.posture_probs = probs + t.posture_missing = 0 + else: + t.posture_probs, t.posture_missing = _posture_decay_update( + t.posture_probs, t.posture_missing, self.posture_hold_frames, self.posture_decay + ) + + if self.smooth_boxes: + if t.box_smooth is None: + t.box_smooth = cur_box.copy() + else: + t.box_smooth = self._ema(t.box_smooth, cur_box) + else: + t.box_smooth = None + + if self.smooth_keypoints and t.keypoints is not None: + if t.kpt_smooth is None: + t.kpt_smooth = t.keypoints.copy() + else: + prev = t.kpt_smooth + cur = t.keypoints + sm_xy = self._ema(prev[..., :2], cur[..., :2]) + t.kpt_smooth = np.concatenate([sm_xy, cur[..., 2:3]], axis=-1) + else: + t.kpt_smooth = None + + # 4) spawn new tracks for unmatched detections + for di in sorted(unmatched_dets): + box = det_boxes_xyxy[di].astype(np.float32) + score = float(det_scores[di]) + if float(score) < float(self.new_track_score_thr): + continue + kpt = det_keypoints[di].astype(np.float32) if det_keypoints is not None else None + pq = float(det_pose_quality[di]) if det_pose_quality is not None and det_pose_quality.size > di else 0.0 + meas = _xyxy_to_cxcywh(box) + t = KalmanTrack( + track_id=self._next_id, + box_xyxy=box, + score=score, + keypoints=kpt, + pose_quality=pq, + kf=_KalmanFilter8D( + meas, + process_noise=self.process_noise, + measurement_noise=self.measurement_noise, + ), + ) + t.hits = 1 + t.confirmed = int(t.hits) >= int(self.confirm_hits) + t.last_meas_cxcywh = meas.astype(np.float32) + t.y_foot_ref = float(box[3]) # y2 in pixels (update only when visible) + t.last_visible_box_xyxy = box.astype(np.float32).copy() + t.freeze_active = False + t.freeze_box_xyxy = None + t.oof_count = 0 + t.occluded_conf = 0.0 + t.occluder_front_frac = 0.0 + pq_ok = (self.posture_poseq_thr <= 0.0) or (float(t.pose_quality) >= float(self.posture_poseq_thr)) + t.posture_probs = _posture_probs_from_kpts(box, kpt, kpt_thr=self.posture_kpt_thr) if pq_ok else None + t.posture_missing = 0 if (t.posture_probs is not None) else 1 + if self.smooth_boxes: + t.box_smooth = box.copy() + if self.smooth_keypoints and kpt is not None: + t.kpt_smooth = kpt.copy() + + if self.depth_enabled and det_inv_depth is not None and int(det_inv_depth.size) > di: + z_meas = float(det_inv_depth[di]) + if np.isfinite(z_meas): + t.zf = _KalmanFilter2D( + init_z=z_meas, + process_noise=self.depth_process_noise, + measurement_noise=self.depth_measurement_noise, + ) + t.inv_depth = float(t.zf.get()) + t.inv_depth_vel = float(t.zf.get_vel()) + t.inv_depth_conf = 1.0 + self._tracks.append(t) + self._next_id += 1 + + # occlusion handling for unmatched tracks (not updated this frame) + for ti in sorted(unmatched_tracks): + t = self._tracks[ti] + t.posture_probs, t.posture_missing = _posture_decay_update( + t.posture_probs, t.posture_missing, self.posture_hold_frames, self.posture_decay + ) + + # Update occluder evidence while unmatched (detector missed / partial occlusion) + if self.depth_enabled and inv_depth_map is not None and t.inv_depth is not None: + frac = occluder_front_fraction( + inv_depth_map, + t.box_xyxy, + inv_depth_expected=float(t.inv_depth), + region=self.depth_region, + margin_abs=self.depth_occ_margin_abs, + margin_rel=self.depth_occ_margin_rel, + ) + t.occluder_front_frac = float(frac) + occ_like = 1.0 if float(frac) >= float(self.depth_occ_frac_thr) else 0.0 + a = float(self.depth_occ_conf_ema) + t.occluded_conf = float(a * float(t.occluded_conf) + (1.0 - a) * float(occ_like)) + + # 5) prune dead + kept = [] + for t in self._tracks: + if int(t.time_since_update) > int(self.max_age): + continue + if (not bool(t.confirmed)) and int(t.time_since_update) > int(self.unconfirmed_max_age): + continue + kept.append(t) + self._tracks = kept + return self.tracks + diff --git a/segmentation_sivert/base_dfine/dfine_hgnetv2_x_obj2coco.yml b/segmentation_sivert/base_dfine/dfine_hgnetv2_x_obj2coco.yml new file mode 100644 index 00000000..6efb3714 --- /dev/null +++ b/segmentation_sivert/base_dfine/dfine_hgnetv2_x_obj2coco.yml @@ -0,0 +1,61 @@ +__include__: [ + '../../configs/dataset/coco_detection.yml', + '../../configs/runtime.yml', + '../../configs/dfine/include/dataloader.yml', + '../../configs/dfine/include/optimizer.yml', + '../../configs/dfine/include/dfine_hgnetv2.yml', +] + +output_dir: ./output/dfine_hgnetv2_x_obj2coco + + +DFINE: + backbone: HGNetv2 + +HGNetv2: + name: 'B5' + return_idx: [1, 2, 3] + freeze_stem_only: True + freeze_at: 0 + freeze_norm: True + +HybridEncoder: + # intra + hidden_dim: 384 + dim_feedforward: 2048 + +DFINETransformer: + feat_channels: [384, 384, 384] + reg_scale: 8 + +optimizer: + type: AdamW + params: + - + params: '^(?=.*backbone)(?!.*norm|bn).*$' + lr: 0.0000025 + - + params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$' + weight_decay: 0. + + lr: 0.00025 + betas: [0.9, 0.999] + weight_decay: 0.000125 + + +epochs: 36 # Early stop +train_dataloader: + dataset: + transforms: + policy: + epoch: 30 + collate_fn: + stop_epoch: 30 + ema_restart_decay: 0.9999 + base_size_repeat: 3 + +ema: + warmups: 0 + +lr_warmup_scheduler: + warmup_duration: 0 diff --git a/segmentation_sivert/configs/advanced.yaml b/segmentation_sivert/configs/advanced.yaml new file mode 100644 index 00000000..3c52d25e --- /dev/null +++ b/segmentation_sivert/configs/advanced.yaml @@ -0,0 +1,169 @@ +# Advanced Configuration for DFINE Segmentation +# Maximum accuracy and robustness - ideal for high-end GPU setups + +# Inherit from base configuration +base: "base.yaml" + +# Model Tier +tier: "advanced" + +# Model-specific Configuration +model: + segmentation_head: + type: "advanced" + feature_dim: 384 # Increased from 256 + dropout_rate: 0.15 # Increased for regularization + + # Robust FPN settings + fpn: + enable_attention_gates: true + enable_self_attention: true + enable_fusion_weights: true + + # Advanced ASPP settings + aspp: + dilations: [1, 6, 12, 18] # Full dilation set + enable_cbam: true + enable_global_pool: true + + # Advanced decoder settings + decoder: + stages: 3 # Multi-stage decoder + deep_supervision: true + boundary_refinement: true + enable_cbam: true + +# Training Configuration Overrides +training: + batch_size: 4 # Smaller batch due to larger model + epochs: 120 # More epochs for convergence + learning_rate: 0.0003 # 3e-4 in decimal # Lower LR for stability + + # Advanced tier specific settings + multi_scale_training: true + augmentation_strength: "strong" + + # Progressive training + progressive_resizing: true + + # Advanced unfreezing strategy + unfreeze_epoch1: 30 # First backbone stage + unfreeze_epoch2: 45 # Second backbone stage + backbone_lr_factor: 0.05 # Very low backbone LR + +# Loss Configuration for Advanced Tier +loss: + type: "advanced" + focal_alpha: 0.25 + focal_gamma: 2.0 + dice_weight: 0.3 + tversky_weight: 0.2 # Additional loss component + ce_weight: 0.1 + boundary_weight: 0.3 + size_weight: 0.1 # For distant objects + +# Advanced Training Techniques +advanced_training: + # Deep supervision + deep_supervision: + aux_loss_weight: 0.4 + num_aux_outputs: 2 + + # Attention mechanisms + attention: + self_attention: true + channel_attention: true + spatial_attention: true + attention_gates: true + + # Multi-scale inference + multiscale_inference: + scales: [0.75, 1.0, 1.25] + ensemble_method: "average" + +# Performance Targets +targets: + accuracy: + miou: 0.80 # Highest accuracy target + pixel_accuracy: 0.90 + class_consistency: 0.85 + + speed: + min_fps: 15 # Lower FPS acceptable for max accuracy + max_latency_ms: 70 # Higher latency acceptable + + memory: + max_gpu_mb: 8000 # High memory usage acceptable + max_model_mb: 400 # Large model size + +# Output Configuration +output: + base_dir: "outputs/advanced" + save_intermediate: true # Save auxiliary outputs + +# Logging Configuration +logging: + wandb: + project: "dfine-advanced-segmentation" + tags: ["advanced", "high-accuracy", "robust"] + + detailed_metrics: true + save_attention_maps: true + +# TensorRT Configuration for Advanced Tier +tensorrt: + fp16: true + int8: true + workspace_gb: 8 # Large workspace for optimization + + # Advanced tier optimization + optimization_level: 5 # Maximum optimization + enable_dla: false # DLA not needed for high-end GPUs + + # Larger input size for accuracy + input_shape: [1, 3, 768, 768] # Increased from 640x640 + + # Multi-scale TensorRT + multiscale_engines: true + scales: [640, 768, 896] + + calibration: + mode: "mixed" # Mixed dataset for robustness + max_images: 1000 # More calibration images + pascal_ratio: 0.6 # Balanced mix + +# Platform Configuration +platform: + high_end_gpu: + rtx_4090: true + rtx_3090: true + tesla_v100: true + a100: true + + server: + multi_gpu: true + distributed_training: true + + workstation: + content_creation: true + research: true + +# Advanced Features +features: + # Robustness enhancements + robustness: + weather_adaptation: true + lighting_adaptation: true + distance_optimization: true + + # Quality enhancements + quality: + boundary_refinement: true + temporal_consistency: true + multi_scale_fusion: true + + # Analysis features + analysis: + uncertainty_estimation: true + failure_detection: true + confidence_calibration: true \ No newline at end of file diff --git a/segmentation_sivert/configs/base.yaml b/segmentation_sivert/configs/base.yaml new file mode 100644 index 00000000..6a10ba64 --- /dev/null +++ b/segmentation_sivert/configs/base.yaml @@ -0,0 +1,94 @@ +# Base Configuration for DFINE Segmentation +# Common settings shared across all model tiers + +# Dataset Configuration +dataset: + name: "pascal_person_parts" + root_dir: "datasets/pascal_person_parts" + num_classes: 7 + class_names: ["background", "head", "torso", "arms", "hands", "legs", "feet"] + image_size: 640 + +# DFINE Model Configuration +dfine: + config_path: "base_dfine/dfine_hgnetv2_x_obj2coco.yml" + checkpoint_path: "base_dfine/dfine_x_obj2coco.pth" + +# Training Configuration +training: + batch_size: 8 + epochs: 250 + learning_rate: 0.0005 # 5e-4 in decimal + weight_decay: 0.0001 # 1e-4 in decimal + + # Scheduler settings + scheduler: + type: "cosine_warmup" + T_0: 15 + T_mult: 2 + eta_min: 0.000001 # 1e-6 in decimal + + # Optimization settings + gradient_clipping: 1.0 + freeze_detection: true + unfreeze_epoch: 40 + + # Augmentation settings + multi_scale_training: false + augmentation_strength: "medium" + +# Loss Configuration +loss: + focal_alpha: 0.25 + focal_gamma: 2.0 + dice_weight: 0.4 + ce_weight: 0.2 + boundary_weight: 0.3 + ignore_index: 255 + +# System Configuration +system: + num_workers: 4 + pin_memory: true + device: "auto" # auto, cuda, cpu + +# Output Configuration +output: + base_dir: "outputs" + save_every: 10 + save_best: true + +# Logging Configuration +logging: + wandb: + project: "dfine-segmentation" + entity: null + tags: [] + + tensorboard: false + + log_every: 20 + print_every: 1 + +# Evaluation Configuration +evaluation: + metrics: + - "miou" + - "pixel_accuracy" + - "class_iou" + - "class_accuracy" + + save_predictions: false + visualize_samples: 5 + +# TensorRT Configuration +tensorrt: + fp16: true + int8: false + workspace_gb: 4 + max_batch_size: 1 + + calibration: + mode: "mixed" # pascal, coco, mixed + max_images: 500 + pascal_ratio: 0.7 \ No newline at end of file diff --git a/segmentation_sivert/configs/camo_finetune.yaml b/segmentation_sivert/configs/camo_finetune.yaml new file mode 100644 index 00000000..eeac751a --- /dev/null +++ b/segmentation_sivert/configs/camo_finetune.yaml @@ -0,0 +1,49 @@ +# CAMO Fine-tune Config (Phase 2) +# Fine-tune from COD10K pretrained checkpoint on CAMO dataset +# Run: python train.py --tier advanced --config-file configs/camo_finetune.yaml \ +# --output-dir outputs/camo_finetune \ +# --resume outputs/cod10k_pretrain/best_model.pth \ +# --no-wandb + +base: "advanced.yaml" + +dataset: + name: "camo" + root_dir: "data/camo" + num_classes: 2 + class_names: ["background", "camouflaged"] + image_size: 640 + +training: + batch_size: 4 + epochs: 60 + learning_rate: 0.00005 # 10x lower than pretraining + weight_decay: 0.0001 + + scheduler: + type: "cosine_warmup" + T_0: 10 + T_mult: 2 + eta_min: 0.000001 + + # Unfreeze backbone sooner — already adapted from COD10K + unfreeze_epoch: 10 + backbone_lr_factor: 0.05 + multi_scale_training: true + +loss: + focal_alpha: 0.25 + focal_gamma: 2.0 + dice_weight: 0.5 + ce_weight: 0.1 + boundary_weight: 0.5 + size_weight: 0.0 + ignore_index: 255 + +output: + base_dir: "outputs/camo_finetune" + +logging: + wandb: + project: "dfine-camouflage" + tags: ["camo", "finetune", "advanced"] diff --git a/segmentation_sivert/configs/cod10k_advanced.yaml b/segmentation_sivert/configs/cod10k_advanced.yaml new file mode 100644 index 00000000..51280b46 --- /dev/null +++ b/segmentation_sivert/configs/cod10k_advanced.yaml @@ -0,0 +1,34 @@ +# COD10K Pretraining Config (Phase 1) +# Advanced tier, binary camouflage segmentation +# Run: python train.py --tier advanced --config-dir configs --dataset data/cod10k --output-dir outputs/cod10k_pretrain + +base: "advanced.yaml" + +dataset: + name: "cod10k" + root_dir: "data/cod10k" + num_classes: 2 + class_names: ["background", "camouflaged"] + image_size: 640 + +training: + batch_size: 4 + epochs: 120 + learning_rate: 0.0003 + +loss: + focal_alpha: 0.25 + focal_gamma: 2.0 + dice_weight: 0.5 # Higher — binary task with class imbalance + ce_weight: 0.1 + boundary_weight: 0.5 # Higher — camouflage edges are subtle + size_weight: 0.0 # Not relevant for camouflage + ignore_index: 255 + +output: + base_dir: "outputs/cod10k_pretrain" + +logging: + wandb: + project: "dfine-camouflage" + tags: ["cod10k", "pretrain", "advanced"] diff --git a/segmentation_sivert/configs/lightweight.yaml b/segmentation_sivert/configs/lightweight.yaml new file mode 100644 index 00000000..24e8d35d --- /dev/null +++ b/segmentation_sivert/configs/lightweight.yaml @@ -0,0 +1,109 @@ +# Lightweight Configuration for DFINE Segmentation +# Optimized for speed and low memory usage - ideal for edge devices and mobile + +# Inherit from base configuration +base: "base.yaml" + +# Model Tier +tier: "lightweight" + +# Model-specific Configuration +model: + segmentation_head: + type: "lightweight" + feature_dim: 128 # Reduced from 256 + dropout_rate: 0.05 # Reduced from 0.1 + + # Lightweight FPN settings + fpn: + use_depthwise_separable: true + reduced_channels: true + + # Ultra-lightweight ASPP settings + aspp: + branches: 3 # Reduced from 4 + use_depthwise_separable: true + enable_global_pool: true + + # Minimal decoder settings + decoder: + stages: 1 # Minimal processing + enable_batch_norm: true + +# Training Configuration Overrides +training: + batch_size: 12 # Can use larger batch due to smaller model + epochs: 150 # Fewer epochs due to simpler model + learning_rate: 0.0007 # 7e-4 in decimal # Slightly higher LR + + # Lightweight tier specific settings + multi_scale_training: True # Disabled for speed + augmentation_strength: "strong" + + # Simplified unfreezing strategy + unfreeze_epoch: 30 + backbone_lr_factor: 0.1 + +# Loss Configuration for Lightweight Tier +loss: + type: "lightweight" + focal_alpha: 0.25 + focal_gamma: 2.0 + dice_weight: 0.3 # Reduced complexity + ce_weight: 0.1 + # No boundary loss for speed + +# Performance Targets +targets: + accuracy: + miou: 0.60 # Lower than standard but still good + pixel_accuracy: 0.80 + + speed: + min_fps: 60 # High FPS target + max_latency_ms: 16 # Very low latency + + memory: + max_gpu_mb: 2000 # Very low memory usage + max_model_mb: 50 # Small model size + +# Output Configuration +output: + base_dir: "outputs/lightweight" + +# Logging Configuration +logging: + wandb: + project: "dfine-lightweight-segmentation" + tags: ["lightweight", "edge", "fast"] + +# TensorRT Configuration for Lightweight Tier +tensorrt: + fp16: true + int8: true + workspace_gb: 2 # Reduced workspace + + # Lightweight tier optimization + optimization_level: 5 # Maximum optimization + enable_dla: true # Enable DLA for edge devices + + # Smaller input size for speed + input_shape: [1, 3, 512, 512] # Reduced from 640x640 + + calibration: + mode: "pascal" # Pascal-only for consistency + max_images: 300 # Fewer calibration images + +# Platform-specific optimizations +platform: + edge_devices: + jetson_nano: true + jetson_xavier: true + raspberry_pi: false # Too resource constrained + + mobile: + android: true + ios: true + + embedded: + custom_asic: true \ No newline at end of file diff --git a/segmentation_sivert/configs/mcs1k_finetune.yaml b/segmentation_sivert/configs/mcs1k_finetune.yaml new file mode 100644 index 00000000..026af591 --- /dev/null +++ b/segmentation_sivert/configs/mcs1k_finetune.yaml @@ -0,0 +1,49 @@ +# MCS1K Fine-tune Config (Phase 3) +# Fine-tune from CAMO checkpoint on MCS1K (military camouflage soldiers) +# Run: python train.py --tier advanced --config-file configs/mcs1k_finetune.yaml \ +# --output-dir outputs/mcs1k_finetune \ +# --resume outputs/camo_finetune/best_model.pth \ +# --finetune --no-wandb + +base: "advanced.yaml" + +dataset: + name: "mcs1k" + root_dir: "data/mcs1k" + num_classes: 2 + class_names: ["background", "camouflaged"] + image_size: 640 + +training: + batch_size: 4 + epochs: 80 + learning_rate: 0.00002 # 2.5x lower than CAMO fine-tune + weight_decay: 0.0001 + + scheduler: + type: "cosine_warmup" + T_0: 15 + T_mult: 2 + eta_min: 0.0000005 + + # Backbone already well-adapted — unfreeze quickly + unfreeze_epoch: 5 + backbone_lr_factor: 0.03 + multi_scale_training: true + +loss: + focal_alpha: 0.25 + focal_gamma: 2.0 + dice_weight: 0.6 + ce_weight: 0.1 + boundary_weight: 0.6 + size_weight: 0.0 + ignore_index: 255 + +output: + base_dir: "outputs/mcs1k_finetune" + +logging: + wandb: + project: "dfine-camouflage" + tags: ["mcs1k", "finetune", "military", "advanced"] diff --git a/segmentation_sivert/configs/standard.yaml b/segmentation_sivert/configs/standard.yaml new file mode 100644 index 00000000..7f73ac82 --- /dev/null +++ b/segmentation_sivert/configs/standard.yaml @@ -0,0 +1,193 @@ +# Standard Configuration for DFINE Segmentation +# Balanced performance and accuracy - Updated for unified training + +# Inherit from base configuration +base: "base.yaml" + +# Model Tier +tier: "standard" + +# DFINE Model Configuration (overrides base if needed) +dfine: + config_path: "base_dfine/dfine_hgnetv2_x_obj2coco.yml" + checkpoint_path: "base_dfine/dfine_x_obj2coco.pth" + +# Dataset Configuration (overrides base if needed) +dataset: + name: "pascal_person_parts" + root_dir: "datasets/pascal_person_parts" + num_classes: 7 + class_names: ["background", "head", "torso", "arms", "hands", "legs", "feet"] + image_size: 640 + +# Model-specific Configuration +model: + segmentation_head: + type: "standard" + feature_dim: 256 + dropout_rate: 0.1 + + # Standard FPN settings + fpn: + enable_fusion_weights: true + + # Standard ASPP settings + aspp: + dilations: [1, 6, 12] + enable_global_pool: true + + # Standard decoder settings + decoder: + stages: 3 + enable_batch_norm: true + +# Training Configuration (overrides base) +training: + batch_size: 16 + epochs: 150 + learning_rate: 0.0005 # 5e-4 in decimal + weight_decay: 0.0001 # 1e-4 in decimal + + # Scheduler settings + scheduler: + type: "cosine_warmup" + T_0: 15 + T_mult: 2 + eta_min: 0.000001 # 1e-6 in decimal + + # Optimization settings + gradient_clipping: 1.0 + freeze_detection: true + unfreeze_epoch: 35 + backbone_lr_factor: 0.1 + + # Standard tier specific settings + multi_scale_training: true + augmentation_strength: "medium" + +# Loss Configuration +loss: + type: "standard" + focal_alpha: 0.25 + focal_gamma: 2.0 + dice_weight: 0.4 + ce_weight: 0.2 + boundary_weight: 0.3 + ignore_index: 255 + +# Performance Targets +targets: + accuracy: + miou: 0.70 # Target 70% mIoU + pixel_accuracy: 0.85 + class_consistency: 0.75 + + speed: + min_fps: 30 # Target 30+ FPS on RTX 3080 + max_latency_ms: 35 + + memory: + max_gpu_mb: 4000 # Max 4GB GPU memory + max_model_mb: 100 # Max 100MB model size + +# System Configuration (overrides base if needed) +system: + num_workers: 4 + pin_memory: true + device: "auto" + +# Output Configuration +output: + base_dir: "outputs/standard" + save_every: 10 + save_best: true + +# Logging Configuration +logging: + wandb: + project: "dfine-standard-segmentation" + entity: null + tags: ["standard", "balanced", "unified"] + + tensorboard: false + log_every: 20 + print_every: 1 + +# Evaluation Configuration +evaluation: + metrics: + - "miou" + - "pixel_accuracy" + - "class_iou" + - "class_accuracy" + - "class_precision" + - "f1_score" + + save_predictions: false + visualize_samples: 5 + +# TensorRT Configuration for deployment +tensorrt: + fp16: true + int8: true + workspace_gb: 4 + max_batch_size: 1 + + # Standard tier optimization + optimization_level: 3 + enable_dla: false + + # Input configuration + input_shape: [1, 3, 640, 640] + + calibration: + mode: "mixed" # pascal, coco, mixed + max_images: 500 + pascal_ratio: 0.7 + +# Platform Configuration +platform: + general_purpose: + rtx_3070: true + rtx_3080: true + rtx_4070: true + gtx_1080: true + + server: + single_gpu: true + multi_gpu: false + + workstation: + content_creation: true + development: true + +# Features enabled for standard tier +features: + # Performance features + performance: + multi_scale_training: true + progressive_unfreezing: true + gradient_clipping: true + + # Quality features + quality: + weighted_feature_fusion: true + global_context_pooling: true + batch_normalization: true + + # Analysis features + analysis: + detailed_metrics: true + class_wise_evaluation: true + +# Comments for clarity +comments: | + Standard tier configuration provides: + - Balanced accuracy and speed performance + - Good for general purpose applications + - Target: ~70% mIoU at 30+ FPS + - Memory efficient: <4GB GPU, <100MB model + - Multi-scale training for robustness + - Progressive backbone unfreezing + - Suitable for most deployment scenarios + - Integrates with unified training script \ No newline at end of file diff --git a/segmentation_sivert/core/__init__.py b/segmentation_sivert/core/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/segmentation_sivert/core/calibrators.py b/segmentation_sivert/core/calibrators.py new file mode 100644 index 00000000..48ce6648 --- /dev/null +++ b/segmentation_sivert/core/calibrators.py @@ -0,0 +1,350 @@ +#!/usr/bin/env python3 +""" +Core TensorRT Calibrators for DFINE Segmentation +Reusable INT8 calibration classes for different datasets +""" + +import os +import numpy as np +import random +from typing import List, Optional +import tensorrt as trt +import pycuda.driver as cuda +import pycuda.autoinit +import cv2 +from PIL import Image +import torchvision.transforms as transforms +from abc import ABC, abstractmethod + + +class BaseCalibrator(trt.IInt8EntropyCalibrator2, ABC): + """Base class for INT8 calibrators""" + + def __init__(self, + batch_size: int = 1, + cache_file: str = "calibration.cache", + max_calibration_images: int = 500, + input_shape: tuple = (640, 640)): + super().__init__() + + self.batch_size = batch_size + self.cache_file = cache_file + self.max_calibration_images = max_calibration_images + self.input_shape = input_shape + self.current_index = 0 + + # Image preprocessing + self.transform = transforms.Compose([ + transforms.Resize(input_shape), + transforms.ToTensor(), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + ]) + + # Load image paths + self.image_files = self._load_image_paths() + print(f"📊 Loaded {len(self.image_files)} images for calibration") + + # Allocate GPU memory + self._allocate_memory() + + @abstractmethod + def _load_image_paths(self) -> List[str]: + """Load paths to calibration images""" + pass + + def _allocate_memory(self): + """Allocate GPU memory for calibration""" + # Input image memory + image_size = self.batch_size * 3 * self.input_shape[0] * self.input_shape[1] * 4 # float32 + self.device_input_image = cuda.mem_alloc(image_size) + + # Input sizes memory (for DFINE model) + sizes_size = self.batch_size * 2 * 8 # int64 + self.device_input_sizes = cuda.mem_alloc(sizes_size) + + # Host memory for sizes + self.host_sizes = np.array([[self.input_shape[1], self.input_shape[0]] for _ in range(self.batch_size)], dtype=np.int64) + + # CUDA stream + self.stream = cuda.Stream() + + def get_batch_size(self): + """Return batch size""" + return self.batch_size + + def get_batch(self, names): + """Get next batch for calibration""" + if self.current_index + self.batch_size > len(self.image_files): + print(f"📊 Calibration complete. Processed {self.current_index} images.") + return None + + batch_progress = self.current_index // self.batch_size + 1 + total_batches = len(self.image_files) // self.batch_size + 1 + print(f"📊 Calibration batch {batch_progress}/{total_batches}") + + # Prepare batch data + batch_data = np.zeros((self.batch_size, 3, *self.input_shape), dtype=np.float32) + + valid_images = 0 + for i in range(self.batch_size): + if self.current_index + i >= len(self.image_files): + break + + image_path = self.image_files[self.current_index + i] + + try: + pil_image = Image.open(image_path).convert('RGB') + processed_image = self.transform(pil_image) + batch_data[i] = processed_image.numpy() + valid_images += 1 + + if i == 0: + print(f" Processing: {os.path.basename(image_path)}") + + except Exception as e: + print(f" Error processing {image_path}: {e}") + continue + + if valid_images == 0: + self.current_index += self.batch_size + return self.get_batch(names) + + # Copy to GPU + cuda.memcpy_htod_async(self.device_input_image, batch_data.ravel(), self.stream) + cuda.memcpy_htod_async(self.device_input_sizes, self.host_sizes.ravel(), self.stream) + self.stream.synchronize() + + self.current_index += self.batch_size + + # Return bindings based on input names + bindings = [] + for name in names: + if name in ["images", "input"]: + bindings.append(int(self.device_input_image)) + elif name in ["orig_target_sizes", "sizes"]: + bindings.append(int(self.device_input_sizes)) + else: + print(f"⚠️ Unknown binding name: {name}") + bindings.append(0) + + return bindings + + def read_calibration_cache(self): + """Read calibration cache if exists""" + if os.path.exists(self.cache_file): + with open(self.cache_file, "rb") as f: + cache_data = f.read() + print(f"📁 Read {len(cache_data)} bytes from calibration cache: {self.cache_file}") + return cache_data + return None + + def write_calibration_cache(self, cache): + """Write calibration cache""" + with open(self.cache_file, "wb") as f: + f.write(cache) + print(f"💾 Wrote {len(cache)} bytes to calibration cache: {self.cache_file}") + + +class PascalPersonPartsCalibrator(BaseCalibrator): + """Pascal Person Parts calibrator for human segmentation optimization""" + + def __init__(self, + pascal_data_dir: str, + batch_size: int = 1, + cache_file: str = "pascal_calibration.cache", + max_calibration_images: int = 500, + input_shape: tuple = (640, 640)): + + self.pascal_data_dir = pascal_data_dir + super().__init__(batch_size, cache_file, max_calibration_images, input_shape) + + def _load_image_paths(self) -> List[str]: + """Load Pascal Person Parts image paths""" + image_files = [] + + # Search in multiple possible directories + search_dirs = [ + os.path.join(self.pascal_data_dir, 'images', 'val'), + os.path.join(self.pascal_data_dir, 'images', 'train'), + os.path.join(self.pascal_data_dir, 'images'), + self.pascal_data_dir + ] + + for img_dir in search_dirs: + if os.path.exists(img_dir): + print(f"🔍 Searching Pascal images in: {img_dir}") + for root, dirs, files in os.walk(img_dir): + for file in files: + if file.lower().endswith(('.jpg', '.jpeg', '.png')): + image_files.append(os.path.join(root, file)) + + if len(image_files) >= self.max_calibration_images: + break + + if len(image_files) == 0: + raise ValueError(f"No Pascal images found in {self.pascal_data_dir}") + + # Limit and shuffle + if len(image_files) > self.max_calibration_images: + random.shuffle(image_files) + image_files = image_files[:self.max_calibration_images] + + return image_files + + +class COCOCalibrator(BaseCalibrator): + """COCO calibrator for general detection optimization""" + + def __init__(self, + coco_data_dir: str, + batch_size: int = 1, + cache_file: str = "coco_calibration.cache", + max_calibration_images: int = 500, + input_shape: tuple = (640, 640)): + + self.coco_data_dir = coco_data_dir + super().__init__(batch_size, cache_file, max_calibration_images, input_shape) + + def _load_image_paths(self) -> List[str]: + """Load COCO image paths""" + image_files = [] + + if os.path.exists(self.coco_data_dir): + print(f"🔍 Searching COCO images in: {self.coco_data_dir}") + for root, dirs, files in os.walk(self.coco_data_dir): + for file in files: + if file.lower().endswith(('.jpg', '.jpeg', '.png')): + image_files.append(os.path.join(root, file)) + + if len(image_files) == 0: + raise ValueError(f"No COCO images found in {self.coco_data_dir}") + + # Limit and shuffle + if len(image_files) > self.max_calibration_images: + random.shuffle(image_files) + image_files = image_files[:self.max_calibration_images] + + return image_files + + +class MixedDatasetCalibrator(BaseCalibrator): + """Mixed Pascal + COCO calibrator for balanced optimization""" + + def __init__(self, + pascal_data_dir: str, + coco_data_dir: str, + pascal_ratio: float = 0.7, + batch_size: int = 1, + cache_file: str = "mixed_calibration.cache", + max_calibration_images: int = 500, + input_shape: tuple = (640, 640)): + + self.pascal_data_dir = pascal_data_dir + self.coco_data_dir = coco_data_dir + self.pascal_ratio = pascal_ratio + super().__init__(batch_size, cache_file, max_calibration_images, input_shape) + + def _load_image_paths(self) -> List[str]: + """Load mixed dataset image paths""" + # Collect Pascal images + pascal_images = [] + pascal_search_dirs = [ + os.path.join(self.pascal_data_dir, 'images', 'val'), + os.path.join(self.pascal_data_dir, 'images', 'train'), + os.path.join(self.pascal_data_dir, 'images'), + self.pascal_data_dir + ] + + for img_dir in pascal_search_dirs: + if os.path.exists(img_dir): + print(f"🔍 Searching Pascal images in: {img_dir}") + for root, dirs, files in os.walk(img_dir): + for file in files: + if file.lower().endswith(('.jpg', '.jpeg', '.png')): + pascal_images.append(os.path.join(root, file)) + break + + # Collect COCO images + coco_images = [] + if os.path.exists(self.coco_data_dir): + print(f"🔍 Searching COCO images in: {self.coco_data_dir}") + for root, dirs, files in os.walk(self.coco_data_dir): + for file in files: + if file.lower().endswith(('.jpg', '.jpeg', '.png')): + coco_images.append(os.path.join(root, file)) + + # Mix datasets according to ratio + num_pascal = min(len(pascal_images), int(self.max_calibration_images * self.pascal_ratio)) + num_coco = min(len(coco_images), self.max_calibration_images - num_pascal) if coco_images else 0 + + # Sample images + selected_pascal = [] + if pascal_images: + random.shuffle(pascal_images) + selected_pascal = pascal_images[:num_pascal] + + selected_coco = [] + if coco_images and num_coco > 0: + random.shuffle(coco_images) + selected_coco = coco_images[:num_coco] + + # Combine and shuffle + image_files = selected_pascal + selected_coco + random.shuffle(image_files) + + print(f"📊 Mixed calibration dataset:") + print(f" Pascal images: {len(selected_pascal)}") + print(f" COCO images: {len(selected_coco)}") + print(f" Total: {len(image_files)}") + + if len(image_files) == 0: + raise ValueError("No calibration images found in either dataset") + + return image_files + + +# Factory functions for creating calibrators +def create_pascal_calibrator(pascal_data_dir: str, **kwargs) -> PascalPersonPartsCalibrator: + """Create Pascal Person Parts calibrator""" + return PascalPersonPartsCalibrator(pascal_data_dir, **kwargs) + + +def create_coco_calibrator(coco_data_dir: str, **kwargs) -> COCOCalibrator: + """Create COCO calibrator""" + return COCOCalibrator(coco_data_dir, **kwargs) + + +def create_mixed_calibrator(pascal_data_dir: str, coco_data_dir: str, **kwargs) -> MixedDatasetCalibrator: + """Create mixed dataset calibrator""" + return MixedDatasetCalibrator(pascal_data_dir, coco_data_dir, **kwargs) + + +def create_calibrator(mode: str, + pascal_data_dir: Optional[str] = None, + coco_data_dir: Optional[str] = None, + **kwargs): + """Factory function to create calibrators based on mode""" + + if mode == 'pascal': + if not pascal_data_dir: + raise ValueError("pascal_data_dir required for Pascal mode") + return create_pascal_calibrator(pascal_data_dir, **kwargs) + + elif mode == 'coco': + if not coco_data_dir: + raise ValueError("coco_data_dir required for COCO mode") + return create_coco_calibrator(coco_data_dir, **kwargs) + + elif mode == 'mixed': + if not pascal_data_dir or not coco_data_dir: + raise ValueError("Both pascal_data_dir and coco_data_dir required for mixed mode") + return create_mixed_calibrator(pascal_data_dir, coco_data_dir, **kwargs) + + else: + raise ValueError(f"Unknown calibration mode: {mode}") + + +# Legacy aliases for backward compatibility +class COCOOnlyCalibrator(COCOCalibrator): + """Legacy alias""" + pass \ No newline at end of file diff --git a/segmentation_sivert/core/datasets.py b/segmentation_sivert/core/datasets.py new file mode 100644 index 00000000..d65afa68 --- /dev/null +++ b/segmentation_sivert/core/datasets.py @@ -0,0 +1,329 @@ +#!/usr/bin/env python3 +""" +Core Dataset Classes for DFINE Segmentation +Base classes for dataset loading with tier-specific augmentations +""" + +import os +import random +import numpy as np +import torch +import torch.nn.functional as F +from torch.utils.data import Dataset +import cv2 +import albumentations as A +from albumentations.pytorch import ToTensorV2 +from abc import ABC, abstractmethod +from typing import Tuple, List, Optional, Dict, Any +from pathlib import Path + + +class BaseSegmentationDataset(Dataset, ABC): + """Abstract base class for segmentation datasets""" + + def __init__(self, + root_dir: str, + split: str = 'train', + image_size: int = 640, + num_classes: int = 7, + tier: str = 'standard'): + """Initialize base dataset + + Args: + root_dir: Root directory of dataset + split: Dataset split ('train' or 'val') + image_size: Target image size + num_classes: Number of segmentation classes + tier: Model tier ('lightweight', 'standard', 'advanced') + """ + self.root_dir = root_dir + self.split = split + self.image_size = image_size + self.num_classes = num_classes + self.tier = tier + + # Paths + self.img_dir = os.path.join(root_dir, 'images', split) + self.mask_dir = os.path.join(root_dir, 'masks', split) + + # Load image list + self.image_names = self._load_image_list() + + # Normalization constants + self.mean = torch.tensor([0.485, 0.456, 0.406]) + self.std = torch.tensor([0.229, 0.224, 0.225]) + + # Setup augmentations based on tier + self.setup_augmentations() + + print(f"📊 Loaded {len(self.image_names)} {split} samples for {tier} tier") + + @abstractmethod + def _load_image_list(self) -> List[str]: + """Load list of image filenames""" + pass + + @abstractmethod + def _load_image_and_mask(self, idx: int) -> Tuple[np.ndarray, np.ndarray]: + """Load image and mask for given index""" + pass + + def setup_augmentations(self): + """Setup augmentations based on tier""" + if self.split != 'train': + self.augmentations = None + return + + if self.tier == 'lightweight': + self.augmentations = self._get_lightweight_augmentations() + elif self.tier == 'standard': + self.augmentations = self._get_standard_augmentations() + elif self.tier == 'advanced': + self.augmentations = self._get_advanced_augmentations() + else: + self.augmentations = self._get_standard_augmentations() + + def _get_lightweight_augmentations(self): + """Lightweight augmentations for fast training""" + return A.Compose([ + A.HorizontalFlip(p=0.5), + A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.6), + A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1, rotate_limit=10, p=0.5), + A.GaussNoise(var_limit=(10.0, 30.0), p=0.3), + ], additional_targets={'mask': 'mask'}) + + def _get_standard_augmentations(self): + """Standard augmentations for balanced training""" + return A.Compose([ + A.HorizontalFlip(p=0.5), + A.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=0.8), + A.HueSaturationValue(hue_shift_limit=15, sat_shift_limit=25, val_shift_limit=15, p=0.6), + A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=15, p=0.7), + A.GaussNoise(var_limit=(10.0, 40.0), p=0.3), + A.Blur(blur_limit=3, p=0.2), + A.CLAHE(clip_limit=3.0, tile_grid_size=(8, 8), p=0.3), + ], additional_targets={'mask': 'mask'}) + + def _get_advanced_augmentations(self): + """Advanced augmentations for robust training""" + return A.Compose([ + A.HorizontalFlip(p=0.5), + A.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=0.8), + A.HueSaturationValue(hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20, p=0.6), + A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=15, p=0.7), + A.Perspective(scale=(0.05, 0.15), p=0.3), + A.ElasticTransform(alpha=1, sigma=50, alpha_affine=50, p=0.2), + A.GaussNoise(var_limit=(10.0, 50.0), p=0.3), + A.ISONoise(color_shift=(0.01, 0.05), intensity=(0.1, 0.5), p=0.2), + A.Blur(blur_limit=3, p=0.2), + A.MotionBlur(blur_limit=5, p=0.2), + A.CLAHE(clip_limit=4.0, tile_grid_size=(8, 8), p=0.4), + A.RandomShadow(shadow_roi=(0, 0.5, 1, 1), num_shadows_lower=1, num_shadows_upper=2, p=0.3), + A.ImageCompression(quality_lower=60, quality_upper=100, p=0.3), + ], additional_targets={'mask': 'mask'}) + + def apply_augmentations(self, image: np.ndarray, mask: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """Apply augmentations to image and mask""" + if self.augmentations is None: + return image, mask + + transformed = self.augmentations(image=image, mask=mask) + return transformed['image'], transformed['mask'] + + def preprocess_image(self, image: np.ndarray) -> torch.Tensor: + """Convert image to tensor and normalize""" + # Resize + image = cv2.resize(image, (self.image_size, self.image_size)) + + # Convert to tensor and normalize + image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0 + image = (image - self.mean.view(3, 1, 1)) / self.std.view(3, 1, 1) + + return image + + def preprocess_mask(self, mask: np.ndarray) -> torch.Tensor: + """Convert mask to tensor""" + mask = cv2.resize(mask, (self.image_size, self.image_size), interpolation=cv2.INTER_NEAREST) + return torch.from_numpy(mask).long() + + def __len__(self) -> int: + return len(self.image_names) + + def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]: + """Get item by index""" + # Load image and mask + image, mask = self._load_image_and_mask(idx) + + # Apply augmentations if training + if self.split == 'train': + image, mask = self.apply_augmentations(image, mask) + + # Preprocess + image_tensor = self.preprocess_image(image) + mask_tensor = self.preprocess_mask(mask) + + return image_tensor, mask_tensor + + +class PascalPersonPartsDataset(BaseSegmentationDataset): + """Pascal Person Parts dataset implementation""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def _load_image_list(self) -> List[str]: + """Load Pascal Person Parts image list""" + if not os.path.exists(self.img_dir): + raise ValueError(f"Image directory not found: {self.img_dir}") + + image_names = [f for f in os.listdir(self.img_dir) if f.endswith(('.jpg', '.jpeg', '.png'))] + + if len(image_names) == 0: + raise ValueError(f"No images found in {self.img_dir}") + + return sorted(image_names) + + def _load_image_and_mask(self, idx: int) -> Tuple[np.ndarray, np.ndarray]: + """Load Pascal Person Parts image and mask""" + img_name = self.image_names[idx] + + # Determine mask name (assuming .jpg -> .png conversion) + mask_name = img_name.replace('.jpg', '.png').replace('.jpeg', '.png') + + # Load image + img_path = os.path.join(self.img_dir, img_name) + image = cv2.imread(img_path) + if image is None: + raise ValueError(f"Failed to load image: {img_path}") + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + + # Load mask + mask_path = os.path.join(self.mask_dir, mask_name) + mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) + if mask is None: + raise ValueError(f"Failed to load mask: {mask_path}") + + return image, mask + + +class COD10KDataset(BaseSegmentationDataset): + """COD10K camouflaged object detection dataset (binary: background / camouflaged)""" + + def __init__(self, *args, **kwargs): + kwargs.setdefault('num_classes', 2) + super().__init__(*args, **kwargs) + + def _load_image_list(self) -> List[str]: + if not os.path.exists(self.img_dir): + raise ValueError(f"Image directory not found: {self.img_dir}") + images = [f for f in os.listdir(self.img_dir) if f.endswith(('.jpg', '.jpeg', '.png'))] + if not images: + raise ValueError(f"No images found in {self.img_dir}") + return sorted(images) + + def _load_image_and_mask(self, idx: int) -> Tuple[np.ndarray, np.ndarray]: + img_name = self.image_names[idx] + mask_name = Path(img_name).stem + '.png' + + img_path = os.path.join(self.img_dir, img_name) + image = cv2.imread(img_path) + if image is None: + raise ValueError(f"Failed to load image: {img_path}") + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + + mask_path = os.path.join(self.mask_dir, mask_name) + mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) + if mask is None: + raise ValueError(f"Failed to load mask: {mask_path}") + + return image, mask + + def _get_advanced_augmentations(self): + """Camouflage-specific augmentations: preserve texture cues, add field conditions""" + return A.Compose([ + A.HorizontalFlip(p=0.5), + A.VerticalFlip(p=0.1), + A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.7), + A.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=15, val_shift_limit=10, p=0.5), + A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=15, p=0.7), + A.Perspective(scale=(0.05, 0.1), p=0.3), + A.ElasticTransform(alpha=1, sigma=50, alpha_affine=50, p=0.2), + A.GridDistortion(num_steps=5, distort_limit=0.3, p=0.2), + A.GaussNoise(var_limit=(10.0, 40.0), p=0.3), + A.ISONoise(color_shift=(0.01, 0.04), intensity=(0.1, 0.4), p=0.2), + A.MotionBlur(blur_limit=5, p=0.2), + A.CLAHE(clip_limit=4.0, tile_grid_size=(8, 8), p=0.5), + A.RandomFog(fog_coef_lower=0.1, fog_coef_upper=0.3, p=0.2), + A.RandomShadow(shadow_roi=(0, 0.5, 1, 1), num_shadows_lower=1, num_shadows_upper=2, p=0.3), + A.ImageCompression(quality_lower=70, quality_upper=100, p=0.2), + ], additional_targets={'mask': 'mask'}) + + +class CAMODataset(COD10KDataset): + """CAMO camouflaged object dataset — same format as COD10K.""" + pass + + +class MultiScaleDataset(torch.utils.data.Dataset): + """Multi-scale training dataset wrapper""" + + def __init__(self, base_dataset: BaseSegmentationDataset, scales: List[float] = None): + self.base_dataset = base_dataset + self.scales = scales or [0.75, 0.85, 1.0, 1.15, 1.3] + + def __len__(self) -> int: + return len(self.base_dataset) + + def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]: + # Get base item + image, mask = self.base_dataset._load_image_and_mask(idx) + + # Apply multi-scale if training + if self.base_dataset.split == 'train': + scale = random.choice(self.scales) + h, w = image.shape[:2] + new_h, new_w = int(h * scale), int(w * scale) + + image = cv2.resize(image, (new_w, new_h)) + mask = cv2.resize(mask, (new_w, new_h), interpolation=cv2.INTER_NEAREST) + + # Apply augmentations + image, mask = self.base_dataset.apply_augmentations(image, mask) + + # Preprocess to final size + image_tensor = self.base_dataset.preprocess_image(image) + mask_tensor = self.base_dataset.preprocess_mask(mask) + + return image_tensor, mask_tensor + + +def create_dataset(dataset_name: str, + root_dir: str, + split: str, + image_size: int, + tier: str, + multi_scale: bool = False) -> BaseSegmentationDataset: + """Factory function to create datasets""" + + dataset_map = { + 'pascal_person_parts': PascalPersonPartsDataset, + 'cod10k': COD10KDataset, + 'camo': CAMODataset, + 'mcs1k': COD10KDataset, + } + + cls = dataset_map.get(dataset_name.lower()) + if cls is None: + raise ValueError(f"Unknown dataset: {dataset_name}. Available: {list(dataset_map)}") + + base_dataset = cls(root_dir=root_dir, split=split, image_size=image_size, tier=tier) + + if multi_scale and split == 'train': + return MultiScaleDataset(base_dataset) + return base_dataset + + +# Export functions for convenience +def create_pascal_dataset(root_dir: str, split: str, image_size: int, tier: str, multi_scale: bool = False): + """Create Pascal Person Parts dataset""" + return create_dataset('pascal_person_parts', root_dir, split, image_size, tier, multi_scale) \ No newline at end of file diff --git a/segmentation_sivert/core/engines.py b/segmentation_sivert/core/engines.py new file mode 100644 index 00000000..2c413e61 --- /dev/null +++ b/segmentation_sivert/core/engines.py @@ -0,0 +1,453 @@ +#!/usr/bin/env python3 +""" +Core TensorRT Engine Classes for DFINE Segmentation +Reusable TensorRT engine management and inference +""" + +import os +import time +import numpy as np +import cv2 +import tensorrt as trt +import pycuda.driver as cuda +import pycuda.autoinit +from typing import List, Dict, Tuple, Optional, Any +from abc import ABC, abstractmethod +import logging + + +class BaseTensorRTEngine(ABC): + """Base class for TensorRT engines with common functionality""" + + def __init__(self, engine_path: str, logger_level: int = trt.Logger.WARNING): + self.engine_path = engine_path + self.logger = trt.Logger(logger_level) + self.runtime = None + self.engine = None + self.context = None + self.stream = None + + # Memory management + self.host_inputs = [] + self.host_outputs = [] + self.device_inputs = [] + self.device_outputs = [] + self.input_shapes = {} + self.output_shapes = {} + self.input_names = [] + self.output_names = [] + + # Performance tracking + self.inference_times = [] + + # Load and setup engine + self._load_engine() + self._allocate_memory() + + def _load_engine(self): + """Load TensorRT engine from file""" + print(f"🔧 Loading TensorRT engine: {self.engine_path}") + + if not os.path.exists(self.engine_path): + raise FileNotFoundError(f"Engine file not found: {self.engine_path}") + + self.runtime = trt.Runtime(self.logger) + + with open(self.engine_path, 'rb') as f: + engine_bytes = f.read() + + self.engine = self.runtime.deserialize_cuda_engine(engine_bytes) + if not self.engine: + raise RuntimeError(f"Failed to load TensorRT engine from {self.engine_path}") + + self.context = self.engine.create_execution_context() + if not self.context: + raise RuntimeError("Failed to create TensorRT execution context") + + self.stream = cuda.Stream() + print(f"✅ Engine loaded successfully") + + def _allocate_memory(self): + """Pre-allocate host and device memory""" + print(f"🧠 Allocating memory buffers...") + + # Check TensorRT version for API compatibility + use_new_api = hasattr(self.engine, 'num_io_tensors') + + if use_new_api: + self._allocate_memory_new_api() + else: + self._allocate_memory_legacy_api() + + def _allocate_memory_new_api(self): + """Allocate memory using TensorRT 10.x+ API""" + for i in range(self.engine.num_io_tensors): + tensor_name = self.engine.get_tensor_name(i) + shape = self.engine.get_tensor_shape(tensor_name) + dtype = self.engine.get_tensor_dtype(tensor_name) + is_input = self.engine.get_tensor_mode(tensor_name) == trt.TensorIOMode.INPUT + + # Handle dynamic shapes + if -1 in shape: + actual_shape = tuple(1 if dim == -1 else dim for dim in shape) + self.context.set_input_shape(tensor_name, actual_shape) + shape = actual_shape + + # Convert dtype and calculate size + np_dtype, element_size = self._get_numpy_dtype(dtype) + size = trt.volume(shape) + + # Allocate memory + host_mem = cuda.pagelocked_empty(size, np_dtype) + device_mem = cuda.mem_alloc(size * element_size) + + if is_input: + self.host_inputs.append(host_mem) + self.device_inputs.append(device_mem) + self.input_shapes[tensor_name] = shape + self.input_names.append(tensor_name) + else: + self.host_outputs.append(host_mem) + self.device_outputs.append(device_mem) + self.output_shapes[tensor_name] = shape + self.output_names.append(tensor_name) + + # Set tensor address + self.context.set_tensor_address(tensor_name, int(device_mem)) + + print(f" {'📥' if is_input else '📤'} {tensor_name}: {shape} ({size * element_size / 1024 / 1024:.1f} MB)") + + def _allocate_memory_legacy_api(self): + """Allocate memory using legacy TensorRT API""" + self.bindings = [] + + for i in range(self.engine.num_bindings): + binding_name = self.engine.get_binding_name(i) + shape = self.engine.get_binding_shape(i) + dtype = self.engine.get_binding_dtype(i) + is_input = self.engine.binding_is_input(i) + + # Handle dynamic shapes + if -1 in shape: + actual_shape = tuple(1 if dim == -1 else dim for dim in shape) + self.context.set_binding_shape(i, actual_shape) + shape = actual_shape + + # Convert dtype and calculate size + np_dtype, element_size = self._get_numpy_dtype(dtype) + size = trt.volume(shape) + + # Allocate memory + host_mem = cuda.pagelocked_empty(size, np_dtype) + device_mem = cuda.mem_alloc(size * element_size) + + if is_input: + self.host_inputs.append(host_mem) + self.device_inputs.append(device_mem) + self.input_shapes[binding_name] = shape + self.input_names.append(binding_name) + else: + self.host_outputs.append(host_mem) + self.device_outputs.append(device_mem) + self.output_shapes[binding_name] = shape + self.output_names.append(binding_name) + + self.bindings.append(int(device_mem)) + + print(f" {'📥' if is_input else '📤'} {binding_name}: {shape} ({size * element_size / 1024 / 1024:.1f} MB)") + + def _get_numpy_dtype(self, trt_dtype) -> Tuple[np.dtype, int]: + """Convert TensorRT dtype to numpy dtype and element size""" + if trt_dtype == trt.DataType.FLOAT: + return np.float32, 4 + elif trt_dtype == trt.DataType.HALF: + return np.float16, 2 + elif trt_dtype == trt.DataType.INT32: + return np.int32, 4 + elif trt_dtype == trt.DataType.INT64: + return np.int64, 8 + elif trt_dtype == trt.DataType.INT8: + return np.int8, 1 + else: + return np.float32, 4 # Default fallback + + def infer(self, inputs: List[np.ndarray]) -> List[np.ndarray]: + """Perform inference with timing""" + start_time = time.perf_counter() + + # Copy inputs to device + for i, input_data in enumerate(inputs): + np.copyto(self.host_inputs[i][:input_data.size], input_data.ravel()) + cuda.memcpy_htod_async(self.device_inputs[i], self.host_inputs[i], self.stream) + + # Execute inference + if hasattr(self.context, 'execute_async_v3'): + success = self.context.execute_async_v3(stream_handle=self.stream.handle) + elif hasattr(self.context, 'execute_async_v2'): + success = self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle) + else: + success = self.context.execute_async(bindings=self.bindings, stream_handle=self.stream.handle) + + if not success: + raise RuntimeError("TensorRT inference failed") + + # Copy outputs from device + outputs = [] + for i, (host_output, device_output, shape) in enumerate( + zip(self.host_outputs, self.device_outputs, self.output_shapes.values()) + ): + cuda.memcpy_dtoh_async(host_output, device_output, self.stream) + self.stream.synchronize() + output = host_output[:np.prod(shape)].reshape(shape).copy() + outputs.append(output) + + # Track inference time + inference_time = time.perf_counter() - start_time + self.inference_times.append(inference_time) + + return outputs + + def get_engine_info(self) -> Dict[str, Any]: + """Get engine information""" + info = { + 'engine_path': self.engine_path, + 'num_inputs': len(self.input_names), + 'num_outputs': len(self.output_names), + 'input_shapes': self.input_shapes, + 'output_shapes': self.output_shapes, + 'input_names': self.input_names, + 'output_names': self.output_names + } + + if self.inference_times: + info.update({ + 'avg_inference_time_ms': np.mean(self.inference_times) * 1000, + 'min_inference_time_ms': np.min(self.inference_times) * 1000, + 'max_inference_time_ms': np.max(self.inference_times) * 1000, + 'avg_fps': 1.0 / np.mean(self.inference_times) + }) + + return info + + def print_engine_info(self): + """Print engine information""" + info = self.get_engine_info() + print(f"📊 TensorRT Engine Info:") + print(f" Engine: {os.path.basename(info['engine_path'])}") + print(f" Inputs: {info['num_inputs']}") + print(f" Outputs: {info['num_outputs']}") + + if 'avg_inference_time_ms' in info: + print(f" Avg inference time: {info['avg_inference_time_ms']:.2f}ms") + print(f" Avg FPS: {info['avg_fps']:.1f}") + + @abstractmethod + def preprocess(self, *args, **kwargs) -> List[np.ndarray]: + """Preprocess inputs for inference""" + pass + + @abstractmethod + def postprocess(self, outputs: List[np.ndarray], *args, **kwargs) -> Any: + """Postprocess inference outputs""" + pass + + +class SegmentationTensorRTEngine(BaseTensorRTEngine): + """TensorRT engine for segmentation models""" + + def __init__(self, engine_path: str, target_size: Tuple[int, int] = (640, 640)): + self.target_size = target_size + super().__init__(engine_path) + + # Validate expected outputs for segmentation + self._validate_segmentation_outputs() + + def _validate_segmentation_outputs(self): + """Validate that engine has expected segmentation outputs""" + expected_outputs = ['labels', 'boxes', 'scores', 'seg_probs', 'seg_preds'] + + if len(self.output_names) != 5: + print(f"⚠️ Warning: Expected 5 outputs for segmentation model, found {len(self.output_names)}") + print(f" Expected: {expected_outputs}") + print(f" Found: {self.output_names}") + else: + print(f"✅ Segmentation model outputs validated") + + def preprocess(self, frame: np.ndarray) -> List[np.ndarray]: + """Preprocess frame for segmentation inference""" + # Resize frame + resized = cv2.resize(frame, self.target_size, interpolation=cv2.INTER_LINEAR) + + # Convert to RGB and normalize + rgb_frame = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB) + normalized = rgb_frame.astype(np.float32) / 255.0 + + # Apply ImageNet normalization + mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) + std = np.array([0.229, 0.224, 0.225], dtype=np.float32) + normalized = (normalized - mean) / std + + # Convert to NCHW format + input_tensor = np.transpose(normalized, (2, 0, 1)) + input_tensor = np.expand_dims(input_tensor, axis=0) # Add batch dimension + + # Target size tensor + target_size_tensor = np.array([[self.target_size[0], self.target_size[1]]], dtype=np.int64) + + return [input_tensor, target_size_tensor] + + def postprocess(self, outputs: List[np.ndarray]) -> Dict[str, np.ndarray]: + """Postprocess segmentation inference outputs""" + if len(outputs) != 5: + raise ValueError(f"Expected 5 outputs, got {len(outputs)}") + + # Extract outputs + labels = outputs[0].flatten() # Detection labels + boxes = outputs[1][0] # Detection boxes (N, 4) + scores = outputs[2].flatten() # Detection scores + seg_probs = outputs[3][0] # Segmentation probabilities (7, 640, 640) + seg_preds = outputs[4][0] # Segmentation predictions (640, 640) + + return { + 'labels': labels, + 'boxes': boxes, + 'scores': scores, + 'seg_probs': seg_probs, + 'seg_preds': seg_preds + } + + def infer_frame(self, frame: np.ndarray) -> Dict[str, np.ndarray]: + """Complete inference pipeline for a single frame""" + # Preprocess + inputs = self.preprocess(frame) + + # Inference + outputs = self.infer(inputs) + + # Postprocess + results = self.postprocess(outputs) + + return results + + +class LightweightSegmentationEngine(SegmentationTensorRTEngine): + """Lightweight segmentation engine optimized for speed""" + + def __init__(self, engine_path: str, target_size: Tuple[int, int] = (512, 512)): + # Use smaller input size for lightweight model + super().__init__(engine_path, target_size) + print(f"🚀 Lightweight segmentation engine ready") + print(f" Target size: {target_size}") + print(f" Optimized for: Speed and low memory") + + +class StandardSegmentationEngine(SegmentationTensorRTEngine): + """Standard segmentation engine with balanced performance""" + + def __init__(self, engine_path: str, target_size: Tuple[int, int] = (640, 640)): + super().__init__(engine_path, target_size) + print(f"🚀 Standard segmentation engine ready") + print(f" Target size: {target_size}") + print(f" Optimized for: Balanced performance") + + +class AdvancedSegmentationEngine(SegmentationTensorRTEngine): + """Advanced segmentation engine with maximum accuracy""" + + def __init__(self, engine_path: str, target_size: Tuple[int, int] = (768, 768)): + # Use larger input size for advanced model + super().__init__(engine_path, target_size) + print(f"🚀 Advanced segmentation engine ready") + print(f" Target size: {target_size}") + print(f" Optimized for: Maximum accuracy") + + def infer_multiscale(self, frame: np.ndarray, scales: List[float] = [0.75, 1.0, 1.25]) -> Dict[str, np.ndarray]: + """Multi-scale inference for better accuracy""" + all_results = [] + + for scale in scales: + # Scale frame + h, w = frame.shape[:2] + scaled_h, scaled_w = int(h * scale), int(w * scale) + scaled_frame = cv2.resize(frame, (scaled_w, scaled_h)) + + # Infer on scaled frame + result = self.infer_frame(scaled_frame) + all_results.append(result) + + # Ensemble results (simple averaging for seg_probs) + ensemble_seg_probs = np.mean([r['seg_probs'] for r in all_results], axis=0) + ensemble_seg_preds = np.argmax(ensemble_seg_probs, axis=0) + + # Use results from scale 1.0 for detection + base_result = all_results[1] if len(all_results) > 1 else all_results[0] + + return { + 'labels': base_result['labels'], + 'boxes': base_result['boxes'], + 'scores': base_result['scores'], + 'seg_probs': ensemble_seg_probs, + 'seg_preds': ensemble_seg_preds + } + + +# Factory functions +def create_segmentation_engine(tier: str, engine_path: str, **kwargs) -> SegmentationTensorRTEngine: + """Factory function to create segmentation engines""" + + if tier == 'lightweight': + return LightweightSegmentationEngine(engine_path, **kwargs) + elif tier == 'standard': + return StandardSegmentationEngine(engine_path, **kwargs) + elif tier == 'advanced': + return AdvancedSegmentationEngine(engine_path, **kwargs) + else: + raise ValueError(f"Unknown engine tier: {tier}") + + +# Utility functions +def benchmark_engine(engine: BaseTensorRTEngine, num_iterations: int = 100) -> Dict[str, float]: + """Benchmark engine performance""" + print(f"🏃 Benchmarking engine for {num_iterations} iterations...") + + # Create dummy input + dummy_inputs = [] + for shape in engine.input_shapes.values(): + dummy_input = np.random.randn(*shape).astype(np.float32) + dummy_inputs.append(dummy_input) + + # Warmup + for _ in range(10): + _ = engine.infer(dummy_inputs) + + # Benchmark + start_time = time.perf_counter() + for _ in range(num_iterations): + _ = engine.infer(dummy_inputs) + + total_time = time.perf_counter() - start_time + avg_time = total_time / num_iterations + fps = 1.0 / avg_time + + print(f"✅ Benchmark completed:") + print(f" Average inference time: {avg_time*1000:.2f}ms") + print(f" Average FPS: {fps:.1f}") + + return { + 'avg_inference_time_ms': avg_time * 1000, + 'avg_fps': fps, + 'total_time_s': total_time, + 'iterations': num_iterations + } + + +# Export classes +__all__ = [ + 'BaseTensorRTEngine', + 'SegmentationTensorRTEngine', + 'LightweightSegmentationEngine', + 'StandardSegmentationEngine', + 'AdvancedSegmentationEngine', + 'create_segmentation_engine', + 'benchmark_engine' +] \ No newline at end of file diff --git a/segmentation_sivert/core/losses.py b/segmentation_sivert/core/losses.py new file mode 100644 index 00000000..6e960a89 --- /dev/null +++ b/segmentation_sivert/core/losses.py @@ -0,0 +1,322 @@ +#!/usr/bin/env python3 +""" +Fixed Core Loss Functions for DFINE Segmentation +Includes the original AdvancedSegmentationLoss functionality +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +import cv2 +from typing import Dict, Any, Optional, Tuple + + +class FocalLoss(nn.Module): + """Focal Loss for handling class imbalance""" + + def __init__(self, alpha: float = 0.25, gamma: float = 2.0, ignore_index: int = 255): + super().__init__() + self.alpha = alpha + self.gamma = gamma + self.ignore_index = ignore_index + + def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: + ce_loss = F.cross_entropy(inputs, targets, ignore_index=self.ignore_index, reduction='none') + pt = torch.exp(-ce_loss) + focal_loss = self.alpha * (1 - pt) ** self.gamma * ce_loss + return focal_loss.mean() + + +class DiceLoss(nn.Module): + """Dice Loss for better boundary handling""" + + def __init__(self, smooth: float = 1.0, ignore_index: int = 255): + super().__init__() + self.smooth = smooth + self.ignore_index = ignore_index + + def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: + # Create mask for valid pixels + valid_mask = (targets != self.ignore_index) + + # Apply softmax to get probabilities + inputs_soft = F.softmax(inputs, dim=1) + + # One-hot encode targets + targets_one_hot = F.one_hot(targets * valid_mask.long(), num_classes=inputs.shape[1]) + targets_one_hot = targets_one_hot.permute(0, 3, 1, 2).float() + + # Apply valid mask + inputs_soft = inputs_soft * valid_mask.unsqueeze(1).float() + targets_one_hot = targets_one_hot * valid_mask.unsqueeze(1).float() + + # Calculate dice coefficient + intersection = (inputs_soft * targets_one_hot).sum(dim=(2, 3)) + union = inputs_soft.sum(dim=(2, 3)) + targets_one_hot.sum(dim=(2, 3)) + + dice = (2 * intersection + self.smooth) / (union + self.smooth) + return 1 - dice.mean() + + +class AdvancedSegmentationLoss(nn.Module): + """Advanced loss combining multiple components for robust training - ORIGINAL VERSION""" + + def __init__(self, num_classes=7, focal_alpha=0.25, focal_gamma=2.0, ignore_index=255, + dice_weight=0.4, ce_weight=0.2, boundary_weight=0.3, size_weight=0.1): + super().__init__() + + self.focal_loss = FocalLoss(focal_alpha, focal_gamma, ignore_index) + self.dice_loss = DiceLoss(ignore_index=ignore_index) + + self.dice_weight = dice_weight + self.ce_weight = ce_weight + self.boundary_weight = boundary_weight + self.size_weight = size_weight + self.ignore_index = ignore_index + + def size_sensitive_loss(self, pred, target): + """Loss that gives more weight to smaller objects (distant people)""" + # Calculate object sizes + unique_labels = torch.unique(target) + size_weights = torch.ones_like(target, dtype=torch.float32) + + for label in unique_labels: + if label == self.ignore_index: + continue + + mask = (target == label) + size = mask.sum().float() + + # Inverse size weighting - smaller objects get higher weight + if size > 0: + weight = 1.0 / (torch.sqrt(size) + 1e-6) + size_weights[mask] = weight + + # Apply size weights to cross entropy loss + ce_loss = F.cross_entropy(pred, target, ignore_index=self.ignore_index, reduction='none') + weighted_ce = (ce_loss * size_weights).mean() + + return weighted_ce + + def forward(self, pred, target, boundary_pred=None): + # Focal loss (main loss) + focal = self.focal_loss(pred, target) + + # Dice loss (boundary preservation) + dice = self.dice_loss(pred, target) + + # Standard CE loss + ce = F.cross_entropy(pred, target, ignore_index=self.ignore_index) + + # Size-sensitive loss for distant objects + size_loss = self.size_sensitive_loss(pred, target) + + # Boundary loss if boundary prediction is provided + boundary_loss_val = 0 + if boundary_pred is not None: + # Create boundary target from segmentation mask + boundary_target = self.create_boundary_target(target) + boundary_loss_val = F.binary_cross_entropy(boundary_pred.squeeze(1), boundary_target) + + # Combine losses + total_loss = (focal + + self.dice_weight * dice + + self.ce_weight * ce + + self.size_weight * size_loss + + self.boundary_weight * boundary_loss_val) + + return total_loss, { + 'focal': focal.item(), + 'dice': dice.item(), + 'ce': ce.item(), + 'size': size_loss.item(), + 'boundary': boundary_loss_val.item() if isinstance(boundary_loss_val, torch.Tensor) else boundary_loss_val, + 'total': total_loss.item() + } + + def create_boundary_target(self, target): + """Create boundary target from segmentation mask""" + # Use morphological operations to detect boundaries + target_np = target.cpu().numpy() + boundary_targets = [] + + for i in range(target_np.shape[0]): + mask = target_np[i] + + # Detect edges using morphological gradient + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)) + boundary = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_GRADIENT, kernel) + boundary = (boundary > 0).astype(np.float32) + + boundary_targets.append(torch.from_numpy(boundary)) + + return torch.stack(boundary_targets).to(target.device) + + +class BaseLoss(nn.Module): + """Base loss class for different model tiers""" + + def __init__(self, + num_classes: int = 7, + ignore_index: int = 255, + tier: str = 'standard'): + super().__init__() + self.num_classes = num_classes + self.ignore_index = ignore_index + self.tier = tier + + # Configure loss based on tier + self._configure_loss() + + def _configure_loss(self): + """Configure loss components based on tier""" + if self.tier == 'lightweight': + self._configure_lightweight_loss() + elif self.tier == 'standard': + self._configure_standard_loss() + elif self.tier == 'advanced': + self._configure_advanced_loss() + else: + self._configure_standard_loss() + + def _configure_lightweight_loss(self): + """Configure lightweight loss (simple, fast)""" + self.focal_loss = FocalLoss(alpha=0.25, gamma=2.0, ignore_index=self.ignore_index) + self.dice_loss = DiceLoss(ignore_index=self.ignore_index) + + self.weights = { + 'focal': 1.0, + 'dice': 0.3, + 'ce': 0.1 + } + + def _configure_standard_loss(self): + """Configure standard loss (enhanced to match original)""" + self.focal_loss = FocalLoss(alpha=0.25, gamma=2.0, ignore_index=self.ignore_index) + self.dice_loss = DiceLoss(ignore_index=self.ignore_index) + + # Use the same weights as the original AdvancedSegmentationLoss + self.weights = { + 'focal': 1.0, + 'dice': 0.4, # Same as original + 'ce': 0.2, # Same as original + 'boundary': 0.3, # Same as original + 'size': 0.1 # Same as original + } + + def _configure_advanced_loss(self): + """Configure advanced loss (comprehensive)""" + self.focal_loss = FocalLoss(alpha=0.25, gamma=2.0, ignore_index=self.ignore_index) + self.dice_loss = DiceLoss(ignore_index=self.ignore_index) + + self.weights = { + 'focal': 1.0, + 'dice': 0.4, + 'ce': 0.2, + 'boundary': 0.3, + 'size': 0.1 + } + + def size_sensitive_loss(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """Loss that gives more weight to smaller objects (distant people)""" + # Calculate object sizes + unique_labels = torch.unique(target) + size_weights = torch.ones_like(target, dtype=torch.float32) + + for label in unique_labels: + if label == self.ignore_index: + continue + + mask = (target == label) + size = mask.sum().float() + + # Inverse size weighting - smaller objects get higher weight + if size > 0: + weight = 1.0 / (torch.sqrt(size) + 1e-6) + size_weights[mask] = weight + + # Apply size weights to cross entropy loss + ce_loss = F.cross_entropy(pred, target, ignore_index=self.ignore_index, reduction='none') + weighted_ce = (ce_loss * size_weights).mean() + + return weighted_ce + + def create_boundary_target(self, target: torch.Tensor) -> torch.Tensor: + """Create boundary target from segmentation mask""" + target_np = target.cpu().numpy() + boundary_targets = [] + + for i in range(target_np.shape[0]): + mask = target_np[i] + + # Detect edges using morphological gradient + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)) + boundary = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_GRADIENT, kernel) + boundary = (boundary > 0).astype(np.float32) + + boundary_targets.append(torch.from_numpy(boundary)) + + return torch.stack(boundary_targets).to(target.device) + + def forward(self, + pred: torch.Tensor, + target: torch.Tensor, + boundary_pred: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, Dict[str, float]]: + """Forward pass of loss computation""" + + loss_dict = {} + total_loss = 0.0 + + # Focal loss (main loss) + if 'focal' in self.weights: + focal = self.focal_loss(pred, target) + loss_dict['focal'] = focal.item() + total_loss += self.weights['focal'] * focal + + # Dice loss + if 'dice' in self.weights: + dice = self.dice_loss(pred, target) + loss_dict['dice'] = dice.item() + total_loss += self.weights['dice'] * dice + + # Standard CE loss + if 'ce' in self.weights: + ce = F.cross_entropy(pred, target, ignore_index=self.ignore_index) + loss_dict['ce'] = ce.item() + total_loss += self.weights['ce'] * ce + + # Size-sensitive loss + if 'size' in self.weights: + size_loss = self.size_sensitive_loss(pred, target) + loss_dict['size'] = size_loss.item() + total_loss += self.weights['size'] * size_loss + + # Boundary loss + if 'boundary' in self.weights: + if boundary_pred is not None: + # Use provided boundary prediction + boundary_target = self.create_boundary_target(target) + boundary_loss_val = F.binary_cross_entropy(boundary_pred.squeeze(1), boundary_target) + else: + # No boundary prediction, skip boundary loss + boundary_loss_val = torch.tensor(0.0, device=pred.device) + + loss_dict['boundary'] = boundary_loss_val.item() + total_loss += self.weights['boundary'] * boundary_loss_val + + loss_dict['total'] = total_loss.item() + + return total_loss, loss_dict + + +# Factory function for creating losses +def create_loss(tier: str, num_classes: int = 7, ignore_index: int = 255, **kwargs) -> BaseLoss: + """Factory function to create loss based on tier""" + return BaseLoss(num_classes=num_classes, ignore_index=ignore_index, tier=tier) + + +# Keep the original AdvancedSegmentationLoss for backward compatibility +def create_advanced_loss(num_classes=7, **kwargs): + """Create the original advanced loss""" + return AdvancedSegmentationLoss(num_classes=num_classes, **kwargs) \ No newline at end of file diff --git a/segmentation_sivert/core/metrics.py b/segmentation_sivert/core/metrics.py new file mode 100644 index 00000000..5e06896e --- /dev/null +++ b/segmentation_sivert/core/metrics.py @@ -0,0 +1,355 @@ +#!/usr/bin/env python3 +""" +Core Metrics for DFINE Segmentation - GPU Enhanced +Evaluation metrics for segmentation performance with GPU optimizations +""" + +import torch +import torch.nn.functional as F +import numpy as np +from typing import Dict, List, Optional, Tuple +from collections import defaultdict + + +class SegmentationMetrics: + """Comprehensive segmentation metrics calculation - GPU Enhanced""" + + def __init__(self, num_classes: int = 7, ignore_index: int = 255, class_names: Optional[List[str]] = None): + self.num_classes = num_classes + self.ignore_index = ignore_index + self.class_names = class_names or [f'class_{i}' for i in range(num_classes)] + + # Default Pascal Person Parts class names + if num_classes == 7 and class_names is None: + self.class_names = [ + 'background', 'head', 'torso', 'arms', 'hands', 'legs', 'feet' + ] + + # Determine device + self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + + self.reset() + + def reset(self): + """Reset all metrics - GPU optimized""" + # Keep confusion matrix on GPU for speed + self.confusion_matrix = torch.zeros((self.num_classes, self.num_classes), + dtype=torch.long, device=self.device) + self.pixel_counts = torch.zeros(self.num_classes, dtype=torch.long, device=self.device) + self.total_pixels = 0 + + def update(self, pred: torch.Tensor, target: torch.Tensor): + """Update metrics with new predictions and targets - GPU OPTIMIZED + + Args: + pred: Predicted logits [N, C, H, W] or class predictions [N, H, W] + target: Ground truth labels [N, H, W] + """ + # Convert logits to class predictions if needed + if pred.dim() == 4: # [N, C, H, W] + pred = torch.argmax(pred, dim=1) + + # Ensure tensors are on GPU + pred = pred.to(self.device) + target = target.to(self.device) + + # Create mask for valid pixels (GPU operation) + valid_mask = (target != self.ignore_index) + + # Get valid predictions and targets (stay on GPU) + pred_valid = pred[valid_mask] + target_valid = target[valid_mask] + + # GPU-optimized confusion matrix update using bincount + # This replaces the SLOW Python loop! + n_valid = pred_valid.numel() + if n_valid > 0: + # Ensure values are within valid range + valid_pred_mask = (pred_valid >= 0) & (pred_valid < self.num_classes) + valid_target_mask = (target_valid >= 0) & (target_valid < self.num_classes) + final_mask = valid_pred_mask & valid_target_mask + + if final_mask.sum() > 0: + pred_final = pred_valid[final_mask] + target_final = target_valid[final_mask] + + # Create indices: target * num_classes + pred + indices = target_final * self.num_classes + pred_final + + # Vectorized confusion matrix update (GPU operation) + bincount = torch.bincount(indices, minlength=self.num_classes**2) + bincount = bincount.reshape(self.num_classes, self.num_classes) + + # Update confusion matrix and pixel counts (GPU operations) + self.confusion_matrix += bincount + + # Update pixel counts + pixel_bincount = torch.bincount(target_final, minlength=self.num_classes) + self.pixel_counts += pixel_bincount + + self.total_pixels += pred_final.numel() + + def compute_miou(self) -> float: + """Compute mean Intersection over Union""" + # Convert to numpy for final computation (only transfer small matrix) + cm = self.confusion_matrix.cpu().numpy() + + ious = [] + for i in range(self.num_classes): + true_positive = cm[i, i] + false_positive = cm[:, i].sum() - true_positive + false_negative = cm[i, :].sum() - true_positive + + union = true_positive + false_positive + false_negative + if union > 0: + iou = true_positive / union + ious.append(iou) + + return np.mean(ious) if ious else 0.0 + + def compute_class_iou(self) -> Dict[str, float]: + """Compute per-class IoU""" + cm = self.confusion_matrix.cpu().numpy() + + class_ious = {} + for i in range(self.num_classes): + true_positive = cm[i, i] + false_positive = cm[:, i].sum() - true_positive + false_negative = cm[i, :].sum() - true_positive + + union = true_positive + false_positive + false_negative + if union > 0: + iou = true_positive / union + else: + iou = 0.0 + + class_ious[self.class_names[i]] = iou + + return class_ious + + def compute_pixel_accuracy(self) -> float: + """Compute overall pixel accuracy""" + cm = self.confusion_matrix.cpu().numpy() + correct_pixels = np.diag(cm).sum() + return correct_pixels / self.total_pixels if self.total_pixels > 0 else 0.0 + + def compute_class_accuracy(self) -> Dict[str, float]: + """Compute per-class accuracy (recall)""" + cm = self.confusion_matrix.cpu().numpy() + + class_accuracies = {} + for i in range(self.num_classes): + true_positive = cm[i, i] + total_true = cm[i, :].sum() + + if total_true > 0: + accuracy = true_positive / total_true + else: + accuracy = 0.0 + + class_accuracies[self.class_names[i]] = accuracy + + return class_accuracies + + def compute_class_precision(self) -> Dict[str, float]: + """Compute per-class precision""" + cm = self.confusion_matrix.cpu().numpy() + + class_precisions = {} + for i in range(self.num_classes): + true_positive = cm[i, i] + total_pred = cm[:, i].sum() + + if total_pred > 0: + precision = true_positive / total_pred + else: + precision = 0.0 + + class_precisions[self.class_names[i]] = precision + + return class_precisions + + def compute_f1_score(self) -> Dict[str, float]: + """Compute per-class F1 score""" + class_f1s = {} + accuracies = self.compute_class_accuracy() + precisions = self.compute_class_precision() + + for class_name in self.class_names: + acc = accuracies[class_name] + prec = precisions[class_name] + + if acc + prec > 0: + f1 = 2 * (acc * prec) / (acc + prec) + else: + f1 = 0.0 + + class_f1s[class_name] = f1 + + return class_f1s + + def compute_frequency_weighted_iou(self) -> float: + """Compute frequency weighted IoU""" + class_ious = list(self.compute_class_iou().values()) + pixel_counts = self.pixel_counts.cpu().numpy() + class_frequencies = pixel_counts / self.total_pixels if self.total_pixels > 0 else np.zeros(self.num_classes) + + fwiou = 0.0 + for i in range(self.num_classes): + fwiou += class_frequencies[i] * class_ious[i] + + return fwiou + + def get_summary(self) -> Dict[str, float]: + """Get comprehensive metrics summary""" + summary = { + 'miou': self.compute_miou(), + 'pixel_accuracy': self.compute_pixel_accuracy(), + 'fwiou': self.compute_frequency_weighted_iou() + } + + # Add per-class metrics + class_ious = self.compute_class_iou() + class_accs = self.compute_class_accuracy() + class_precs = self.compute_class_precision() + class_f1s = self.compute_f1_score() + + for class_name in self.class_names: + summary[f'{class_name}_iou'] = class_ious[class_name] + summary[f'{class_name}_acc'] = class_accs[class_name] + summary[f'{class_name}_prec'] = class_precs[class_name] + summary[f'{class_name}_f1'] = class_f1s[class_name] + + return summary + + def print_summary(self): + """Print metrics summary""" + summary = self.get_summary() + + print("📊 Segmentation Metrics Summary:") + print(f" mIoU: {summary['miou']:.4f}") + print(f" Pixel Accuracy: {summary['pixel_accuracy']:.4f}") + print(f" Frequency Weighted IoU: {summary['fwiou']:.4f}") + + print("\n📋 Per-Class Metrics:") + for class_name in self.class_names: + iou = summary[f'{class_name}_iou'] + acc = summary[f'{class_name}_acc'] + prec = summary[f'{class_name}_prec'] + f1 = summary[f'{class_name}_f1'] + print(f" {class_name:>12}: IoU={iou:.3f}, Acc={acc:.3f}, Prec={prec:.3f}, F1={f1:.3f}") + + +def compute_miou(pred: torch.Tensor, target: torch.Tensor, num_classes: int = 7, ignore_index: int = 255) -> float: + """Quick mIoU computation for training loops - Already GPU optimized""" + if pred.dim() == 4: + pred = torch.argmax(pred, dim=1) + + ious = [] + for cls in range(num_classes): + pred_cls = (pred == cls) + target_cls = (target == cls) + + # Exclude ignore_index pixels + if ignore_index is not None: + valid_mask = (target != ignore_index) + pred_cls = pred_cls & valid_mask + target_cls = target_cls & valid_mask + + intersection = (pred_cls & target_cls).sum().float() + union = (pred_cls | target_cls).sum().float() + + if union > 0: + ious.append(intersection / union) + + return torch.tensor(ious).mean().item() if ious else 0.0 + + +def compute_pixel_accuracy(pred: torch.Tensor, target: torch.Tensor, ignore_index: int = 255) -> float: + """Quick pixel accuracy computation - Already GPU optimized""" + if pred.dim() == 4: + pred = torch.argmax(pred, dim=1) + + if ignore_index is not None: + valid_mask = (target != ignore_index) + correct = (pred == target) & valid_mask + total = valid_mask.sum() + else: + correct = (pred == target) + total = torch.numel(pred) + + return (correct.sum().float() / total.float()).item() if total > 0 else 0.0 + + +def compute_class_iou(pred: torch.Tensor, target: torch.Tensor, num_classes: int = 7, ignore_index: int = 255) -> List[float]: + """Compute per-class IoU - Already GPU optimized""" + if pred.dim() == 4: + pred = torch.argmax(pred, dim=1) + + class_ious = [] + for cls in range(num_classes): + pred_cls = (pred == cls) + target_cls = (target == cls) + + # Exclude ignore_index pixels + if ignore_index is not None: + valid_mask = (target != ignore_index) + pred_cls = pred_cls & valid_mask + target_cls = target_cls & valid_mask + + intersection = (pred_cls & target_cls).sum().float() + union = (pred_cls | target_cls).sum().float() + + if union > 0: + class_ious.append((intersection / union).item()) + else: + class_ious.append(0.0) + + return class_ious + + +class BatchMetricsTracker: + """Track metrics across training batches - Already fast""" + + def __init__(self, num_classes: int = 7, ignore_index: int = 255): + self.num_classes = num_classes + self.ignore_index = ignore_index + self.reset() + + def reset(self): + """Reset tracked metrics""" + self.total_miou = 0.0 + self.total_pixel_acc = 0.0 + self.batch_count = 0 + self.class_ious = [0.0] * self.num_classes + + def update(self, pred: torch.Tensor, target: torch.Tensor): + """Update with batch predictions""" + miou = compute_miou(pred, target, self.num_classes, self.ignore_index) + pixel_acc = compute_pixel_accuracy(pred, target, self.ignore_index) + class_ious = compute_class_iou(pred, target, self.num_classes, self.ignore_index) + + self.total_miou += miou + self.total_pixel_acc += pixel_acc + self.batch_count += 1 + + # Update class IoUs + for i, iou in enumerate(class_ious): + self.class_ious[i] += iou + + def get_averages(self) -> Dict[str, float]: + """Get average metrics""" + if self.batch_count == 0: + return {'miou': 0.0, 'pixel_accuracy': 0.0} + + return { + 'miou': self.total_miou / self.batch_count, + 'pixel_accuracy': self.total_pixel_acc / self.batch_count, + 'class_ious': [iou / self.batch_count for iou in self.class_ious] + } + + +# Factory function - SAME NAME +def create_metrics_tracker(num_classes: int = 7, ignore_index: int = 255, class_names: Optional[List[str]] = None) -> SegmentationMetrics: + """Create metrics tracker - GPU enhanced""" + return SegmentationMetrics(num_classes, ignore_index, class_names) \ No newline at end of file diff --git a/segmentation_sivert/core/models.py b/segmentation_sivert/core/models.py new file mode 100644 index 00000000..416ddaf8 --- /dev/null +++ b/segmentation_sivert/core/models.py @@ -0,0 +1,322 @@ +#!/usr/bin/env python3 +""" +Core Model Classes for DFINE Segmentation +Base classes for segmentation heads and model combinations +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from abc import ABC, abstractmethod +from typing import List, Dict, Any, Optional, Tuple +import sys +import os + +# Add src to path for DFINE imports (src is at parent level) +current_dir = os.path.dirname(os.path.abspath(__file__)) +project_root = os.path.dirname(os.path.dirname(current_dir)) # Go up from core/ to project root +src_path = os.path.join(project_root, 'src') +if os.path.exists(src_path) and src_path not in sys.path: + sys.path.insert(0, src_path) + sys.path.insert(0, project_root) + + +class BaseSegmentationHead(nn.Module, ABC): + """Abstract base class for segmentation heads""" + + def __init__(self, + in_channels_list: List[int], + num_classes: int = 7, + feature_dim: int = 256, + dropout_rate: float = 0.1, + tier: str = 'standard'): + """Initialize base segmentation head + + Args: + in_channels_list: List of input channel counts from backbone + num_classes: Number of segmentation classes + feature_dim: Feature dimension for processing + dropout_rate: Dropout rate + tier: Model tier ('lightweight', 'standard', 'advanced') + """ + super().__init__() + + self.in_channels_list = in_channels_list + self.num_classes = num_classes + self.feature_dim = feature_dim + self.dropout_rate = dropout_rate + self.tier = tier + + # Build the segmentation head + self._build_head() + + # Initialize weights + self._init_weights() + + print(f"🏗️ Created {tier} segmentation head with channels: {in_channels_list}") + + @abstractmethod + def _build_head(self): + """Build the segmentation head architecture""" + pass + + @abstractmethod + def forward(self, features: List[torch.Tensor]) -> torch.Tensor: + """Forward pass through segmentation head""" + pass + + def _init_weights(self): + """Initialize weights using standard initialization""" + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, nn.BatchNorm2d): + nn.init.ones_(m.weight) + nn.init.zeros_(m.bias) + elif isinstance(m, nn.Linear): + nn.init.normal_(m.weight, 0, 0.01) + if m.bias is not None: + nn.init.zeros_(m.bias) + + def get_model_info(self) -> Dict[str, Any]: + """Get model information""" + total_params = sum(p.numel() for p in self.parameters()) + trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad) + + return { + 'tier': self.tier, + 'num_classes': self.num_classes, + 'feature_dim': self.feature_dim, + 'dropout_rate': self.dropout_rate, + 'total_params': total_params, + 'trainable_params': trainable_params, + 'model_size_mb': total_params * 4 / (1024 * 1024) # Assuming float32 + } + + def print_model_info(self): + """Print model information""" + info = self.get_model_info() + print(f"📊 {info['tier'].capitalize()} Segmentation Head Info:") + print(f" Classes: {info['num_classes']}") + print(f" Feature dim: {info['feature_dim']}") + print(f" Dropout: {info['dropout_rate']}") + print(f" Total params: {info['total_params']:,}") + print(f" Model size: {info['model_size_mb']:.2f} MB") + + +class BaseCombinedModel(nn.Module, ABC): + """Base class for DFINE + Segmentation combination""" + + def __init__(self, + dfine_model: nn.Module, + seg_head: BaseSegmentationHead, + freeze_detection: bool = True): + """Initialize combined model + + Args: + dfine_model: Pretrained DFINE detection model + seg_head: Segmentation head + freeze_detection: Whether to freeze detection weights + """ + super().__init__() + + self.dfine_model = dfine_model + self.seg_head = seg_head + self.freeze_detection = freeze_detection + + if freeze_detection: + self._freeze_detection() + + def _freeze_detection(self): + """Freeze detection model parameters""" + frozen_params = 0 + total_params = 0 + + for name, param in self.named_parameters(): + total_params += param.numel() + if 'seg_head' not in name: + param.requires_grad = False + frozen_params += param.numel() + + trainable_params = total_params - frozen_params + + print(f"🔒 Frozen detection components:") + print(f" 📊 Total parameters: {total_params:,}") + print(f" ❄️ Frozen parameters: {frozen_params:,} ({frozen_params/total_params*100:.1f}%)") + print(f" 🎯 Trainable parameters: {trainable_params:,} ({trainable_params/total_params*100:.1f}%)") + + def unfreeze_detection(self): + """Unfreeze detection model parameters""" + unfrozen_count = 0 + for name, param in self.dfine_model.named_parameters(): + # Only unfreeze parameters that can require gradients (floating point and complex dtypes) + if param.dtype.is_floating_point or param.dtype.is_complex: + param.requires_grad = True + unfrozen_count += 1 + + self.freeze_detection = False + print(f"🔓 Unfrozen detection components: {unfrozen_count} parameters") + + def forward(self, x: torch.Tensor, targets: Optional[Any] = None) -> Dict[str, torch.Tensor]: + """Forward pass through combined model""" + # Get backbone features (these will have gradients if backbone is unfrozen) + backbone_features = self.dfine_model.backbone(x) + + outputs = {} + + # Detection branch - always keep in inference mode for segmentation training + # Even when backbone is unfrozen, we don't want to trigger detection training + with torch.no_grad(): + # Temporarily set to eval mode to avoid denoising training + was_training = self.dfine_model.training + if was_training: + self.dfine_model.eval() + + det_outputs = self.dfine_model(x) + outputs.update(det_outputs) + + # Restore training mode + if was_training: + self.dfine_model.train() + + # Segmentation branch (uses backbone_features which may have gradients) + seg_outputs = self.seg_head(backbone_features) + + # Handle different segmentation head output formats + if isinstance(seg_outputs, tuple): + # Advanced head returns (main_logits, aux1_logits, aux2_logits, boundary_map) + seg_logits = seg_outputs[0] # Use main output + # Store auxiliary outputs for potential use in loss calculation + if len(seg_outputs) >= 4: + outputs['aux1_logits'] = seg_outputs[1] + outputs['aux2_logits'] = seg_outputs[2] + outputs['boundary_map'] = seg_outputs[3] + else: + # Standard/lightweight heads return single tensor + seg_logits = seg_outputs + + # Upsample to input resolution + seg_logits = F.interpolate( + seg_logits, size=x.shape[-2:], + mode='bilinear', align_corners=False + ) + + outputs['segmentation'] = seg_logits + + return outputs + + def get_model_info(self) -> Dict[str, Any]: + """Get combined model information""" + det_params = sum(p.numel() for p in self.dfine_model.parameters()) + seg_params = sum(p.numel() for p in self.seg_head.parameters()) + total_params = det_params + seg_params + + trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad) + + seg_info = self.seg_head.get_model_info() + + return { + 'tier': seg_info['tier'], + 'detection_params': det_params, + 'segmentation_params': seg_params, + 'total_params': total_params, + 'trainable_params': trainable_params, + 'freeze_detection': self.freeze_detection, + 'model_size_mb': total_params * 4 / (1024 * 1024) + } + + +def get_actual_backbone_channels(model: nn.Module, input_size: Tuple[int, int] = (640, 640)) -> List[int]: + """Get actual backbone output channels by running a forward pass""" + model.eval() + dummy_input = torch.randn(1, 3, *input_size) + + with torch.no_grad(): + try: + backbone_features = model.backbone(dummy_input) + actual_channels = [feat.shape[1] for feat in backbone_features] + return actual_channels + except Exception as e: + print(f"Warning: Could not determine backbone channels: {e}") + # Return default channels for HGNetv2 + return [64, 128, 256, 512] + + +def load_pretrained_dfine(config_path: str, checkpoint_path: str) -> nn.Module: + """Load pretrained DFINE model""" + print(f"🚀 Loading pretrained DFINE from {checkpoint_path}") + + try: + from src.core import YAMLConfig + cfg = YAMLConfig(config_path) + model = cfg.model + + checkpoint = torch.load(checkpoint_path, map_location='cpu') + + # Extract state dict from various checkpoint formats + if 'ema' in checkpoint and 'module' in checkpoint['ema']: + state_dict = checkpoint['ema']['module'] + elif 'model' in checkpoint: + state_dict = checkpoint['model'] + else: + state_dict = checkpoint + + model.load_state_dict(state_dict, strict=False) + print(f"✅ Loaded pretrained DFINE model") + return model + + except Exception as e: + print(f"❌ Failed to load DFINE model: {e}") + raise + + +def create_combined_model(dfine_model: nn.Module, + seg_head: BaseSegmentationHead, + freeze_detection: bool = True) -> BaseCombinedModel: + """Create combined DFINE + Segmentation model""" + + class CombinedModel(BaseCombinedModel): + """Concrete implementation of combined model""" + pass + + return CombinedModel(dfine_model, seg_head, freeze_detection) + + +# Utility functions for model management +def count_parameters(model: nn.Module) -> Dict[str, int]: + """Count model parameters""" + total_params = sum(p.numel() for p in model.parameters()) + trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + + return { + 'total': total_params, + 'trainable': trainable_params, + 'frozen': total_params - trainable_params + } + + +def get_model_memory_usage(model: nn.Module, input_size: Tuple[int, int, int, int] = (1, 3, 640, 640)) -> Dict[str, float]: + """Estimate model memory usage""" + # Count parameters + param_count = sum(p.numel() for p in model.parameters()) + param_memory = param_count * 4 / (1024 * 1024) # MB, assuming float32 + + # Estimate activation memory with a forward pass + model.eval() + dummy_input = torch.randn(*input_size) + + with torch.no_grad(): + try: + _ = model(dummy_input) + # This is a rough estimate + activation_memory = torch.cuda.max_memory_allocated() / (1024 * 1024) if torch.cuda.is_available() else 0 + except: + activation_memory = 0 + + return { + 'parameters_mb': param_memory, + 'estimated_activation_mb': activation_memory, + 'estimated_total_mb': param_memory + activation_memory + } \ No newline at end of file diff --git a/segmentation_sivert/infer_video.py b/segmentation_sivert/infer_video.py new file mode 100644 index 00000000..2ac3b067 --- /dev/null +++ b/segmentation_sivert/infer_video.py @@ -0,0 +1,196 @@ +#!/usr/bin/env python3 +""" +Video inference with trained camouflage segmentation model. +Overlays segmentation mask on each frame and writes output video. + +Usage: + python infer_video.py --input "Sekvens 5.mp4" --checkpoint outputs/camo_finetune/best_model.pth +""" + +import os +import sys +import argparse +import numpy as np +import torch +import cv2 +from pathlib import Path +from tqdm import tqdm + +current_dir = os.path.dirname(os.path.abspath(__file__)) +if current_dir not in sys.path: + sys.path.insert(0, current_dir) +sys.path.append('src') + +from core.models import load_pretrained_dfine, get_actual_backbone_channels, create_combined_model +from models import create_segmentation_head + + +def load_model(checkpoint_path: str, device: torch.device): + checkpoint = torch.load(checkpoint_path, map_location=device) + + # Read config from checkpoint + config = checkpoint['config'] + tier = checkpoint['args']['tier'] + num_classes = config['dataset']['num_classes'] + + seg_head_config = config.get('model', {}).get('segmentation_head', {}) + feature_dim = seg_head_config.get('feature_dim', 384) + dropout_rate = seg_head_config.get('dropout_rate', 0.15) + + dfine_config = config['dfine']['config_path'] + dfine_checkpoint = config['dfine']['checkpoint_path'] + + dfine_model = load_pretrained_dfine(dfine_config, dfine_checkpoint) + backbone_channels = get_actual_backbone_channels(dfine_model) + + seg_head = create_segmentation_head( + tier=tier, + in_channels_list=backbone_channels, + num_classes=num_classes, + feature_dim=feature_dim, + dropout_rate=dropout_rate + ) + + model = create_combined_model(dfine_model=dfine_model, seg_head=seg_head, freeze_detection=False) + model.load_state_dict(checkpoint['model_state_dict']) + model = model.to(device) + model.eval() + + print(f"Loaded {tier} model — {num_classes} classes, mIoU at save: {checkpoint.get('best_miou', 'N/A'):.4f}") + return model, num_classes + + +def preprocess_frame(frame: np.ndarray, image_size: int = 640): + rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + resized = cv2.resize(rgb, (image_size, image_size)) + tensor = torch.from_numpy(resized).permute(2, 0, 1).float() / 255.0 + mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1) + std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1) + tensor = (tensor - mean) / std + return tensor.unsqueeze(0) + + +def filter_person_blobs(mask: np.ndarray, frame_h: int, frame_w: int, + min_area_frac: float = 0.002, + max_aspect: float = 6.0) -> np.ndarray: + """ + Keep only blobs that are plausibly person-shaped: + - Large enough (> min_area_frac of frame) + - Not too elongated (aspect ratio < max_aspect) — removes sticks/branches + """ + min_area = int(frame_h * frame_w * min_area_frac) + out = np.zeros_like(mask) + + num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8) + + for i in range(1, num_labels): + area = stats[i, cv2.CC_STAT_AREA] + bw = stats[i, cv2.CC_STAT_WIDTH] + bh = stats[i, cv2.CC_STAT_HEIGHT] + + if area < min_area: + continue + + aspect = max(bw, bh) / max(min(bw, bh), 1) + if aspect > max_aspect: + continue + + out[labels == i] = 1 + + return out + + +def overlay_mask(frame: np.ndarray, mask: np.ndarray, alpha: float = 0.5, color=(0, 255, 80)): + """Overlay binary segmentation mask on frame with colored highlight.""" + h, w = frame.shape[:2] + mask_resized = cv2.resize(mask.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST) + + overlay = frame.copy() + overlay[mask_resized == 1] = ( + overlay[mask_resized == 1] * (1 - alpha) + + np.array(color, dtype=np.float32) * alpha + ).astype(np.uint8) + + # Draw contour for crisp edge + contours, _ = cv2.findContours(mask_resized, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + cv2.drawContours(overlay, contours, -1, color, 2) + + return overlay + + +def run_inference(args): + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + print(f"Device: {device}") + + model, num_classes = load_model(args.checkpoint, device) + + cap = cv2.VideoCapture(args.input) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + fps = cap.get(cv2.CAP_PROP_FPS) + w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + print(f"Video: {w}x{h} @ {fps:.1f}fps — {total_frames} frames") + + # Output path + input_path = Path(args.input) + out_path = args.output or str(input_path.parent / (input_path.stem + '_segmented.mp4')) + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + writer = cv2.VideoWriter(out_path, fourcc, fps, (w, h)) + + smooth_prob = None # running average of per-pixel camouflage probability + + with torch.no_grad(): + for _ in tqdm(range(total_frames), desc="Inferring"): + ret, frame = cap.read() + if not ret: + break + + inp = preprocess_frame(frame, image_size=args.image_size).to(device) + outputs = model(inp) + seg_logits = outputs['segmentation'] # [1, C, H, W] + + # Softmax prob for class 1 (camouflaged), at model resolution + prob = torch.softmax(seg_logits, dim=1)[0, 1].cpu().numpy() # [H_model, W_model] + + # Temporal smoothing: EMA across frames + if smooth_prob is None: + smooth_prob = prob + else: + smooth_prob = args.temporal_alpha * prob + (1 - args.temporal_alpha) * smooth_prob + + # Threshold smoothed probability + mask = (smooth_prob > args.threshold).astype(np.uint8) + + # Filter to person-shaped blobs only + mask = filter_person_blobs(mask, mask.shape[0], mask.shape[1]) + + result = overlay_mask(frame, mask, alpha=args.alpha) + + if args.heatmap: + prob_resized = cv2.resize(smooth_prob, (w, h)) + heatmap = cv2.applyColorMap((prob_resized * 255).astype(np.uint8), cv2.COLORMAP_JET) + result = cv2.addWeighted(result, 0.7, heatmap, 0.3, 0) + + writer.write(result) + + cap.release() + writer.release() + print(f"\nDone. Output saved to: {out_path}") + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument('--input', required=True, help='Input video path') + parser.add_argument('--checkpoint', default='outputs/camo_finetune/best_model.pth') + parser.add_argument('--output', default=None, help='Output video path (default: input_segmented.mp4)') + parser.add_argument('--image-size', type=int, default=640) + parser.add_argument('--alpha', type=float, default=0.5, help='Mask overlay transparency') + parser.add_argument('--threshold', type=float, default=0.5, help='Probability threshold for mask') + parser.add_argument('--temporal-alpha', type=float, default=0.4, + help='EMA weight for current frame (0=fully smooth, 1=no smoothing)') + parser.add_argument('--heatmap', action='store_true', help='Also overlay confidence heatmap') + return parser.parse_args() + + +if __name__ == '__main__': + run_inference(parse_args()) diff --git a/segmentation_sivert/legacy/build_engine.py b/segmentation_sivert/legacy/build_engine.py new file mode 100644 index 00000000..ba8975dc --- /dev/null +++ b/segmentation_sivert/legacy/build_engine.py @@ -0,0 +1,322 @@ +#!/usr/bin/env python3 +""" +DFINE Segmentation TensorRT Engine Builder + +Converts ONNX segmentation model to TensorRT engine for faster inference. +Usage: + python tensorrt/build_segmentation_engine.py -i model_segmentation.onnx -o model_segmentation.engine +""" + +import os +import sys +import argparse +import time +import tensorrt as trt +import numpy as np +import pycuda.driver as cuda +import pycuda.autoinit + +def parse_args(): + """Parse command line arguments""" + parser = argparse.ArgumentParser(description='DFINE Segmentation TensorRT Engine Builder') + parser.add_argument('-i', '--input', type=str, required=True, + help='Path to input ONNX segmentation model') + parser.add_argument('-o', '--output', type=str, default=None, + help='Path to output TensorRT engine (default: input.onnx -> input.engine)') + parser.add_argument('--fp16', action='store_true', + help='Enable FP16 precision (default: FP32)') + parser.add_argument('--workspace', type=int, default=2, + help='Maximum workspace size in GB (default: 2GB for segmentation)') + parser.add_argument('--max-batch-size', type=int, default=1, + help='Maximum batch size (default: 1)') + parser.add_argument('--verbose', action='store_true', + help='Enable verbose logging') + parser.add_argument('--min-timing-iterations', type=int, default=2, + help='Minimum number of timing iterations') + parser.add_argument('--avg-timing-iterations', type=int, default=1, + help='Average number of timing iterations') + parser.add_argument('--force', action='store_true', + help='Force engine building even if output file exists') + return parser.parse_args() + +def get_engine_info(): + """Print TensorRT version information""" + print(f"TensorRT version: {trt.__version__}") + + # Get CUDA version + try: + cuda_version = cuda.get_version() + print(f"CUDA driver version: {cuda_version}") + except Exception as e: + print(f"Error getting CUDA version: {e}") + + # Check if FP16 is supported + logger = trt.Logger(trt.Logger.WARNING) + builder = trt.Builder(logger) + has_fp16 = builder.platform_has_fast_fp16 + print(f"Platform has fast FP16: {has_fp16}") + +def build_segmentation_engine(args): + """Build TensorRT engine from ONNX segmentation model""" + # Set logger level + if args.verbose: + logger = trt.Logger(trt.Logger.VERBOSE) + else: + logger = trt.Logger(trt.Logger.WARNING) + + # Check if output engine file exists + output_file = args.output if args.output else args.input.replace('.onnx', '.engine') + if os.path.exists(output_file) and not args.force: + print(f"Engine file {output_file} already exists. Use --force to overwrite.") + return output_file + + # Print build information + print(f"Building TensorRT segmentation engine from {args.input} to {output_file}") + print(f"Precision: {'FP16' if args.fp16 else 'FP32'}") + print(f"Workspace size: {args.workspace} GB") + print(f"Maximum batch size: {args.max_batch_size}") + + start_time = time.time() + + # Create builder and config + builder = trt.Builder(logger) + config = builder.create_builder_config() + + # Set max workspace size (in bytes) - UPDATED for TensorRT 10.x + config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, args.workspace * (1 << 30)) + + # Set precision flags + if args.fp16: + if builder.platform_has_fast_fp16: + print("Enabling FP16 mode") + config.set_flag(trt.BuilderFlag.FP16) + else: + print("Warning: Platform doesn't support fast FP16, falling back to FP32") + + # Set optimization flags if available in this TensorRT version + try: + config.builder_optimization_level = 3 # Maximum optimization + except: + print("Note: Builder optimization level not set (may not be available in this TensorRT version)") + + # Create network definition with explicit batch flag + explicit_batch = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) + network = builder.create_network(explicit_batch) + + # Parse ONNX model + parser = trt.OnnxParser(network, logger) + + # Read the ONNX model file + print(f"Parsing ONNX segmentation model: {args.input}") + with open(args.input, 'rb') as model: + model_bytes = model.read() + if not parser.parse(model_bytes): + print("Failed to parse ONNX model:") + for error in range(parser.num_errors): + print(f" {parser.get_error(error)}") + return None + + # Print network information + print(f"Network inputs: {network.num_inputs}") + for i in range(network.num_inputs): + tensor = network.get_input(i) + print(f" Input {i}: {tensor.name}, shape: {tensor.shape}, dtype: {tensor.dtype}") + + print(f"Network outputs: {network.num_outputs}") + expected_outputs = ['labels', 'boxes', 'scores', 'seg_probs', 'seg_preds'] + for i in range(network.num_outputs): + tensor = network.get_output(i) + expected_name = expected_outputs[i] if i < len(expected_outputs) else f"output_{i}" + print(f" Output {i}: {tensor.name} (expected: {expected_name}), shape: {tensor.shape}, dtype: {tensor.dtype}") + + if network.num_outputs != 5: + print(f"Warning: Expected 5 outputs for segmentation model, found {network.num_outputs}") + print("Expected outputs: labels, boxes, scores, seg_probs, seg_preds") + + # Create optimization profile for dynamic shapes + print("Creating optimization profile for dynamic input shapes...") + profile = builder.create_optimization_profile() + + # Set shapes for each input tensor based on its name + for i in range(network.num_inputs): + input_tensor = network.get_input(i) + input_name = input_tensor.name + input_shape = input_tensor.shape + + # Check if this input has dynamic dimensions + if -1 in input_shape: + # Create a shape with actual values (replacing -1 with concrete values) + min_shape = [] + opt_shape = [] + max_shape = [] + + for dim in input_shape: + if dim == -1: + # This is a dynamic dimension, typically batch size + min_shape.append(1) # Minimum batch size (1) + opt_shape.append(1) # Optimal batch size (1 for video processing) + max_shape.append(args.max_batch_size) # Maximum batch size + else: + # Fixed dimension + min_shape.append(dim) + opt_shape.append(dim) + max_shape.append(dim) + + # Convert to tuple + min_shape = tuple(min_shape) + opt_shape = tuple(opt_shape) + max_shape = tuple(max_shape) + + print(f" Setting profile for {input_name}: min={min_shape}, opt={opt_shape}, max={max_shape}") + profile.set_shape(input_name, min_shape, opt_shape, max_shape) + + # Add the optimization profile to the config + config.add_optimization_profile(profile) + + # Mark output tensors if not done in ONNX + if network.num_outputs == 0: + print("Warning: No outputs detected in ONNX model. Marking last layer's outputs.") + for i in range(network.num_layers): + layer = network.get_layer(i) + for j in range(layer.num_outputs): + network.mark_output(layer.get_output(j)) + + # Build engine - UPDATED for TensorRT 10.x + print("Building TensorRT segmentation engine (this may take a while)...") + plan = builder.build_serialized_network(network, config) + if not plan: + print("Failed to build TensorRT engine!") + return None + + # Save engine to file + with open(output_file, 'wb') as f: + f.write(plan) + + build_time = time.time() - start_time + print(f"Segmentation engine built successfully in {build_time:.2f} seconds") + print(f"Engine saved to: {output_file}") + + return output_file + +def verify_segmentation_engine(engine_file): + """Verify the TensorRT segmentation engine by loading it""" + print(f"Verifying segmentation engine: {engine_file}") + + logger = trt.Logger(trt.Logger.WARNING) + runtime = trt.Runtime(logger) + + try: + with open(engine_file, 'rb') as f: + engine_bytes = f.read() + + engine = runtime.deserialize_cuda_engine(engine_bytes) + if not engine: + print("Failed to deserialize engine!") + return False + + print("Segmentation engine verification successful!") + + # Print engine information using TensorRT 10.x API + try: + # Try new API first (TensorRT 10.x) + input_count = 0 + output_count = 0 + for i in range(engine.num_io_tensors): + tensor_name = engine.get_tensor_name(i) + if engine.get_tensor_mode(tensor_name) == trt.TensorIOMode.INPUT: + input_count += 1 + else: + output_count += 1 + + print(f"Engine inputs: {input_count}") + print(f"Engine outputs: {output_count}") + + # Check that we have the expected number of outputs for segmentation + if output_count == 5: + print("✅ Correct number of outputs for segmentation model (5)") + else: + print(f"⚠️ Warning: Expected 5 outputs, found {output_count}") + + for i in range(engine.num_io_tensors): + tensor_name = engine.get_tensor_name(i) + if engine.get_tensor_mode(tensor_name) == trt.TensorIOMode.INPUT: + shape = engine.get_tensor_shape(tensor_name) + dtype = engine.get_tensor_dtype(tensor_name) + print(f" Input: {tensor_name}, shape: {shape}, dtype: {dtype}") + else: + shape = engine.get_tensor_shape(tensor_name) + dtype = engine.get_tensor_dtype(tensor_name) + # Try to identify the output type + output_type = "unknown" + if 'label' in tensor_name.lower() or 'class' in tensor_name.lower(): + output_type = "detection_labels" + elif 'box' in tensor_name.lower(): + output_type = "detection_boxes" + elif 'score' in tensor_name.lower(): + output_type = "detection_scores" + elif 'seg_prob' in tensor_name.lower(): + output_type = "segmentation_probabilities" + elif 'seg_pred' in tensor_name.lower(): + output_type = "segmentation_predictions" + + print(f" Output: {tensor_name} ({output_type}), shape: {shape}, dtype: {dtype}") + except: + # Fall back to old API (just in case) + print("Note: Using legacy API for engine inspection") + input_count = sum(1 for i in range(engine.num_bindings) if engine.binding_is_input(i)) + output_count = engine.num_bindings - input_count + print(f"Engine inputs: {input_count}") + print(f"Engine outputs: {output_count}") + + if output_count == 5: + print("✅ Correct number of outputs for segmentation model (5)") + else: + print(f"⚠️ Warning: Expected 5 outputs, found {output_count}") + + for i in range(engine.num_bindings): + name = engine.get_binding_name(i) + shape = engine.get_binding_shape(i) + dtype = engine.get_binding_dtype(i) + if engine.binding_is_input(i): + print(f" Input {i}: {name}, shape: {shape}, dtype: {dtype}") + else: + print(f" Output {i}: {name}, shape: {shape}, dtype: {dtype}") + + return True + except Exception as e: + print(f"Engine verification failed: {e}") + return False + +def main(): + """Main function""" + args = parse_args() + + # Print TensorRT information + get_engine_info() + + # Check if input file exists + if not os.path.isfile(args.input): + print(f"Error: Input file {args.input} not found!") + sys.exit(1) + + # Set default output path if not provided + if not args.output: + args.output = args.input.replace('.onnx', '.engine') + + # Build engine + engine_file = build_segmentation_engine(args) + if not engine_file: + print("Engine building failed!") + sys.exit(1) + + # Verify engine + if not verify_segmentation_engine(engine_file): + print("Engine verification failed!") + sys.exit(1) + + print(f"TensorRT segmentation engine successfully built and saved to {engine_file}") + print("\nTo use this engine with the DFINE segmentation model:") + print(f" python tensorrt/infer_segmentation_trt.py -e {engine_file} -i input_video.mp4 -o output_video.mp4") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/segmentation_sivert/legacy/build_int8_engine.py b/segmentation_sivert/legacy/build_int8_engine.py new file mode 100644 index 00000000..6d56ec22 --- /dev/null +++ b/segmentation_sivert/legacy/build_int8_engine.py @@ -0,0 +1,531 @@ +#!/usr/bin/env python3 +""" +Complete DFINE Segmentation TensorRT INT8 Engine Builder + +Builds optimized INT8 TensorRT engines for DFINE + Segmentation models with multiple calibration options. + +Usage Examples: + # Pascal Person Parts only (recommended for segmentation) + python build_complete_segmentation_int8.py -i model.onnx --pascal-data datasets/pascal_person_parts --mode pascal + + # Mixed Pascal + COCO (balanced detection + segmentation) + python build_complete_segmentation_int8.py -i model.onnx --pascal-data datasets/pascal_person_parts --coco-data datasets/coco_val2017 --mode mixed + + # COCO only (detection-focused) + python build_complete_segmentation_int8.py -i model.onnx --coco-data datasets/coco_val2017 --mode coco +""" + +import os +import sys +import argparse +import time +import tensorrt as trt +import numpy as np +import pycuda.driver as cuda +import pycuda.autoinit + +# Import our calibrators +from int8_calibrator import ( + PascalPersonPartsCalibrator, + MixedDatasetCalibrator, + COCOOnlyCalibrator +) + +def parse_args(): + parser = argparse.ArgumentParser(description='Complete DFINE Segmentation INT8 Engine Builder') + + # Required arguments + parser.add_argument('-i', '--input', type=str, required=True, + help='Path to input ONNX segmentation model') + + # Dataset paths + parser.add_argument('--pascal-data', type=str, default='datasets/pascal_person_parts', + help='Path to Pascal Person Parts dataset (default: datasets/pascal_person_parts)') + parser.add_argument('--coco-data', type=str, default='datasets/coco_val2017', + help='Path to COCO validation dataset (default: datasets/coco_val2017)') + + # Calibration mode selection + parser.add_argument('--mode', type=str, choices=['pascal', 'mixed', 'coco'], default='mixed', + help='Calibration mode: pascal (Pascal-only), mixed (Pascal+COCO), coco (COCO-only)') + + # Optional arguments + parser.add_argument('-o', '--output', type=str, default=None, + help='Path to output TensorRT engine (default: auto-generated)') + parser.add_argument('--workspace', type=int, default=4, + help='Maximum workspace size in GB (default: 4)') + parser.add_argument('--max-calibration-images', type=int, default=5000, + help='Maximum number of images for calibration (default: 500)') + parser.add_argument('--pascal-ratio', type=float, default=0.7, + help='Pascal ratio in mixed mode (default: 0.7 = 70% Pascal, 30% COCO)') + + # Build options + parser.add_argument('--force', action='store_true', + help='Force engine building even if output file exists') + parser.add_argument('--fp16-fallback', action='store_true', default=False, + help='Use FP16 for segmentation layers (recommended, but disable if build fails)') + parser.add_argument('--int8-only', action='store_true', + help='Use pure INT8 without FP16 fallback (faster build, may reduce accuracy)') + parser.add_argument('--verbose', action='store_true', + help='Enable verbose logging') + + return parser.parse_args() + +def get_engine_info(): + """Print TensorRT and system information""" + print("="*60) + print("TENSORRT SYSTEM INFORMATION") + print("="*60) + print(f"TensorRT version: {trt.__version__}") + + try: + cuda_version = cuda.get_version() + print(f"CUDA driver version: {cuda_version}") + except Exception as e: + print(f"Error getting CUDA version: {e}") + + logger = trt.Logger(trt.Logger.WARNING) + builder = trt.Builder(logger) + print(f"Platform has fast INT8: {builder.platform_has_fast_int8}") + print(f"Platform has fast FP16: {builder.platform_has_fast_fp16}") + + if not builder.platform_has_fast_int8: + print("⚠️ WARNING: Your GPU may not support efficient INT8. Consider using FP16 instead.") + if not builder.platform_has_fast_fp16: + print("⚠️ WARNING: Your GPU may not support efficient FP16. Mixed precision may be limited.") + + print("="*60) + +def identify_segmentation_layers(network): + """Identify layers that should use FP16 instead of INT8 for better accuracy""" + + fp16_layers = [] + + for layer_idx in range(network.num_layers): + layer = network.get_layer(layer_idx) + layer_name = layer.name + layer_type = layer.type + + should_use_fp16 = False + + # EXCLUDE layers that CANNOT use FP16 (these must stay as INT64/INT32) + excluded_layer_types = [ + trt.LayerType.CONSTANT, # Constant layers with INT weights + trt.LayerType.SHAPE, # Shape layers must be INT64 + trt.LayerType.GATHER, # Index operations + trt.LayerType.CAST, # Type casting operations + trt.LayerType.SLICE, # Indexing operations + ] + + # Skip if this is an excluded layer type + if layer_type in excluded_layer_types: + continue + + # Skip shape-related layers (these have integer outputs and can't be FP16) + shape_keywords = [ + "constant", "shape", "gather", "cast", "slice", + "elementwise", "concat", "split", "unsqueeze" + ] + + if any(keyword in layer_name.lower() for keyword in shape_keywords): + continue + + # Only target specific segmentation-related COMPUTATION layers + segmentation_patterns = [ + "seg_head", # Segmentation head layers + "/aspp/", # ASPP module layers + "/fpn/", # FPN module layers + "/decoder/", # Decoder layers + ] + + # Check if this is a segmentation computation layer + for pattern in segmentation_patterns: + if pattern in layer_name.lower(): + # Additional check: only if it's a computation layer type + computation_layer_types = [ + trt.LayerType.CONVOLUTION, + trt.LayerType.DECONVOLUTION, + trt.LayerType.ACTIVATION, + trt.LayerType.POOLING, + trt.LayerType.SCALE, + trt.LayerType.SOFTMAX, + trt.LayerType.RESIZE, + ] + + if layer_type in computation_layer_types: + should_use_fp16 = True + break + + # Also target final output layers (last 5 layers only) + if layer_idx >= network.num_layers - 5: + # But only if they're computation layers, not shape layers + if layer_type in [trt.LayerType.CONVOLUTION, trt.LayerType.SOFTMAX, trt.LayerType.ACTIVATION]: + should_use_fp16 = True + + if should_use_fp16: + fp16_layers.append((layer_idx, layer_name, layer_type)) + + return fp16_layers + +def create_calibrator(args): + """Create the appropriate calibrator based on mode""" + + if args.mode == 'pascal': + print("🎯 CALIBRATION MODE: Pascal Person Parts Only") + print(" - Optimized for human segmentation") + print(" - Best for segmentation accuracy") + print(" - Detection works well on humans") + print(f" - Using dataset: {args.pascal_data}") + + if not os.path.exists(args.pascal_data): + raise ValueError(f"Pascal dataset not found: {args.pascal_data}") + + return PascalPersonPartsCalibrator( + pascal_data_dir=args.pascal_data, + batch_size=1, + cache_file="pascal_segmentation_calibration.cache", + max_calibration_images=args.max_calibration_images + ) + + elif args.mode == 'mixed': + print("🎯 CALIBRATION MODE: Mixed Pascal + COCO") + print(f" - Pascal ratio: {args.pascal_ratio*100:.0f}%") + print(f" - COCO ratio: {(1-args.pascal_ratio)*100:.0f}%") + print(" - Balanced detection + segmentation") + print(" - Good general performance") + print(f" - Pascal dataset: {args.pascal_data}") + print(f" - COCO dataset: {args.coco_data}") + + if not os.path.exists(args.pascal_data): + raise ValueError(f"Pascal dataset not found: {args.pascal_data}") + if not os.path.exists(args.coco_data): + raise ValueError(f"COCO dataset not found: {args.coco_data}") + + return MixedDatasetCalibrator( + pascal_data_dir=args.pascal_data, + coco_data_dir=args.coco_data, + pascal_ratio=args.pascal_ratio, + batch_size=1, + cache_file="mixed_segmentation_calibration.cache", + max_calibration_images=args.max_calibration_images + ) + + elif args.mode == 'coco': + print("🎯 CALIBRATION MODE: COCO Only") + print(" - Optimized for general detection") + print(" - Broad object category coverage") + print(" - May impact segmentation accuracy") + print(f" - Using dataset: {args.coco_data}") + + if not os.path.exists(args.coco_data): + raise ValueError(f"COCO dataset not found: {args.coco_data}") + + return COCOOnlyCalibrator( + coco_data_dir=args.coco_data, + batch_size=1, + cache_file="coco_segmentation_calibration.cache", + max_calibration_images=args.max_calibration_images + ) + + else: + raise ValueError(f"Unknown calibration mode: {args.mode}") + +def build_segmentation_int8_engine(args): + """Build optimized INT8 TensorRT engine""" + + # Set logger level + logger = trt.Logger(trt.Logger.INFO if args.verbose else trt.Logger.WARNING) + + # Generate output filename if not provided + if not args.output: + base_name = os.path.splitext(args.input)[0] + mode_suffix = f"_{args.mode}" + + if args.int8_only: + precision_suffix = "_int8_pure" + elif args.fp16_fallback: + precision_suffix = "_int8_mixed" + else: + precision_suffix = "_int8" + + args.output = f"{base_name}_segmentation{mode_suffix}{precision_suffix}.engine" + + # Check if output exists + if os.path.exists(args.output) and not args.force: + print(f"❌ Engine file {args.output} already exists.") + print(" Use --force to overwrite, or specify different --output path.") + return None + + print("\n" + "="*60) + print("BUILDING INT8 SEGMENTATION ENGINE") + print("="*60) + print(f"📥 Input ONNX: {args.input}") + print(f"📤 Output Engine: {args.output}") + print(f"💾 Workspace: {args.workspace} GB") + print(f"🖼️ Max calibration images: {args.max_calibration_images}") + + if args.int8_only: + precision_desc = "Pure INT8" + elif use_mixed_precision: + precision_desc = "INT8 + FP16 mixed precision" + else: + precision_desc = "INT8 (FP16 disabled due to platform limitations)" + + print(f"⚡ Precision: {precision_desc}") + print("="*60) + + start_time = time.time() + + # Create builder and config + builder = trt.Builder(logger) + config = builder.create_builder_config() + + # Set workspace size + config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, args.workspace * (1 << 30)) + + # Enable INT8 + config.set_flag(trt.BuilderFlag.INT8) + print("✅ Enabled INT8 precision") + + # Enable FP16 for mixed precision + if args.fp16_fallback: + config.set_flag(trt.BuilderFlag.FP16) + print("✅ Enabled FP16 fallback for critical layers") + + # Create calibrator + print("\n📊 SETTING UP CALIBRATION...") + calibrator = create_calibrator(args) + config.int8_calibrator = calibrator + + # Create network + explicit_batch = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) + network = builder.create_network(explicit_batch) + + # Parse ONNX model + print(f"\n🔍 PARSING ONNX MODEL: {args.input}") + parser = trt.OnnxParser(network, logger) + with open(args.input, 'rb') as model: + model_bytes = model.read() + if not parser.parse(model_bytes): + print("❌ Failed to parse ONNX model:") + for error in range(parser.num_errors): + print(f" Error {error}: {parser.get_error(error)}") + return None + + # Print network info + print(f"✅ Parsed ONNX model successfully") + print(f" Inputs: {network.num_inputs}") + for i in range(network.num_inputs): + tensor = network.get_input(i) + print(f" {i}: {tensor.name} {tensor.shape} ({tensor.dtype})") + + print(f" Outputs: {network.num_outputs}") + expected_outputs = ['labels', 'boxes', 'scores', 'seg_probs', 'seg_preds'] + for i in range(network.num_outputs): + tensor = network.get_output(i) + expected = expected_outputs[i] if i < len(expected_outputs) else f"output_{i}" + print(f" {i}: {tensor.name} ({expected}) {tensor.shape} ({tensor.dtype})") + + if network.num_outputs != 5: + print(f"⚠️ Warning: Expected 5 outputs for segmentation model, found {network.num_outputs}") + + # Configure mixed precision + if args.fp16_fallback: + print(f"\n🔧 CONFIGURING MIXED PRECISION...") + fp16_layers = identify_segmentation_layers(network) + print(f" Found {len(fp16_layers)} layers to use FP16 precision") + + fp16_count = 0 + for layer_idx, layer_name, layer_type in fp16_layers: + try: + layer = network.get_layer(layer_idx) + layer.precision = trt.float16 + + # Set output precision + for i in range(layer.num_outputs): + layer.set_output_type(i, trt.float16) + + fp16_count += 1 + if args.verbose: + print(f" Layer {layer_idx}: '{layer_name}' -> FP16") + + except Exception as e: + if args.verbose: + print(f" Warning: Could not set precision for '{layer_name}': {e}") + + print(f"✅ Set {fp16_count}/{len(fp16_layers)} layers to FP16 precision") + + # Create optimization profile + print(f"\n⚙️ CREATING OPTIMIZATION PROFILE...") + profile = builder.create_optimization_profile() + + for i in range(network.num_inputs): + input_tensor = network.get_input(i) + input_name = input_tensor.name + input_shape = input_tensor.shape + + if -1 in input_shape: + min_shape = [] + opt_shape = [] + max_shape = [] + + for dim in input_shape: + if dim == -1: + min_shape.append(1) + opt_shape.append(1) + max_shape.append(1) + else: + min_shape.append(dim) + opt_shape.append(dim) + max_shape.append(dim) + + min_shape = tuple(min_shape) + opt_shape = tuple(opt_shape) + max_shape = tuple(max_shape) + + print(f" {input_name}: min={min_shape}, opt={opt_shape}, max={max_shape}") + profile.set_shape(input_name, min_shape, opt_shape, max_shape) + + config.add_optimization_profile(profile) + + # Set optimization level + try: + config.builder_optimization_level = 5 + print("✅ Set optimization level to 5 (maximum)") + except: + try: + config.builder_optimization_level = 3 + print("✅ Set optimization level to 3") + except: + print("ℹ️ Using default optimization level") + + # Build engine + print(f"\n🏗️ BUILDING INT8 ENGINE...") + print(" This will take a while due to INT8 calibration...") + print(" Calibration determines optimal quantization scales for your model.") + + plan = builder.build_serialized_network(network, config) + if not plan: + print("❌ Failed to build INT8 engine!") + return None + + # Save engine + print(f"\n💾 SAVING ENGINE...") + with open(args.output, 'wb') as f: + f.write(plan) + + build_time = time.time() - start_time + print(f"✅ Engine built successfully in {build_time:.1f} seconds") + print(f"📁 Saved to: {args.output}") + + return args.output, use_mixed_precision + +def verify_engine(engine_file): + """Verify the built engine""" + print(f"\n🔍 VERIFYING ENGINE: {engine_file}") + + logger = trt.Logger(trt.Logger.WARNING) + runtime = trt.Runtime(logger) + + try: + with open(engine_file, 'rb') as f: + engine_bytes = f.read() + + engine = runtime.deserialize_cuda_engine(engine_bytes) + if not engine: + print("❌ Failed to deserialize engine!") + return False + + print("✅ Engine verification successful!") + + # Get engine info + input_count = 0 + output_count = 0 + + for i in range(engine.num_io_tensors): + tensor_name = engine.get_tensor_name(i) + if engine.get_tensor_mode(tensor_name) == trt.TensorIOMode.INPUT: + input_count += 1 + else: + output_count += 1 + + print(f" Inputs: {input_count}") + print(f" Outputs: {output_count}") + + if output_count == 5: + print("✅ Correct outputs for segmentation model") + else: + print(f"⚠️ Expected 5 outputs, found {output_count}") + + # Print tensor details + for i in range(engine.num_io_tensors): + tensor_name = engine.get_tensor_name(i) + shape = engine.get_tensor_shape(tensor_name) + dtype = engine.get_tensor_dtype(tensor_name) + + if engine.get_tensor_mode(tensor_name) == trt.TensorIOMode.INPUT: + print(f" Input: {tensor_name} {shape} ({dtype})") + else: + print(f" Output: {tensor_name} {shape} ({dtype})") + + return True + + except Exception as e: + print(f"❌ Engine verification failed: {e}") + return False + +def print_summary(args, engine_file, build_time): + """Print build summary""" + print("\n" + "="*60) + print("🎉 BUILD SUMMARY") + print("="*60) + print(f"✅ INT8 segmentation engine built successfully!") + print(f"📁 Engine file: {engine_file}") + print(f"⏱️ Build time: {build_time:.1f} seconds") + print(f"🎯 Calibration mode: {args.mode.upper()}") + print(f"🖼️ Calibration images: {args.max_calibration_images}") + print(f"⚡ Precision: INT8 + {'FP16 fallback' if args.fp16_fallback else 'INT8 only'}") + + print(f"\n📈 EXPECTED PERFORMANCE:") + print(f" • Speed: ~3-4x faster than FP32") + print(f" • Speed: ~2x faster than FP16") + print(f" • Memory: ~75% reduction") + print(f" • Quality: Good (with mixed precision)") + + print(f"\n🚀 USAGE:") + print(f" python tensorrt/infer_segmentation_trt.py \\") + print(f" -e {engine_file} \\") + print(f" -i input_video.mp4 \\") + print(f" -o output_segmented.mp4") + + print("="*60) + +def main(): + args = parse_args() + + # Print system info + get_engine_info() + + # Validate inputs + if not os.path.isfile(args.input): + print(f"❌ Input ONNX file not found: {args.input}") + sys.exit(1) + + # Build engine + start_time = time.time() + engine_file, use_mixed_precision = build_segmentation_int8_engine(args) + + if not engine_file: + print("❌ Engine building failed!") + sys.exit(1) + + # Verify engine + if not verify_engine(engine_file): + print("❌ Engine verification failed!") + sys.exit(1) + + # Print summary + build_time = time.time() - start_time + print_summary(args, engine_file, build_time, use_mixed_precision) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/segmentation_sivert/legacy/export_onnx.py b/segmentation_sivert/legacy/export_onnx.py new file mode 100644 index 00000000..24ee4a5a --- /dev/null +++ b/segmentation_sivert/legacy/export_onnx.py @@ -0,0 +1,400 @@ +#!/usr/bin/env python3 +""" +Standard DFINE Segmentation ONNX Export Script +Refactored to use base classes and modular structure +""" + +import os +import sys +import argparse +import torch +import torch.nn as nn +import torch.nn.functional as F + +# Add paths for imports +current_dir = os.path.dirname(os.path.abspath(__file__)) +segmentation_root = os.path.dirname(os.path.dirname(current_dir)) # Go up to segmentation_sivert/ +project_root = os.path.dirname(segmentation_root) # Go up to D-FINE-NEWBRINGER/ + +# Add segmentation_sivert to path for core imports +if segmentation_root not in sys.path: + sys.path.insert(0, segmentation_root) + +# Add project root to path for src imports +src_path = os.path.join(project_root, 'src') +if os.path.exists(src_path) and src_path not in sys.path: + sys.path.insert(0, src_path) + sys.path.insert(0, project_root) + +# Import core modules +from core.models import load_pretrained_dfine, get_actual_backbone_channels, create_combined_model + +# Import model +from models.standard import create_standard_segmentation_head + + +def parse_args(): + """Parse command line arguments""" + parser = argparse.ArgumentParser(description='Standard DFINE Segmentation ONNX Export') + + # Model configuration + parser.add_argument('-c', '--config', type=str, default="configs/dfine_hgnetv2_x_obj2coco.yml", + help='Path to original DFINE configuration YAML') + parser.add_argument('-r', '--resume', type=str, required=True, + help='Path to trained segmentation model weights (.pth)') + parser.add_argument('--original-weights', default="models/dfine_x_obj2coco.pth", type=str, + help='Path to original DFINE model weights (.pth)') + + # Export configuration + parser.add_argument('-o', '--output', type=str, default=None, + help='Path to save the ONNX model (default: auto-generated)') + parser.add_argument('--input-shape', type=str, default='1,3,640,640', + help='Input shape for the model (default: 1,3,640,640)') + parser.add_argument('--opset', type=int, default=17, + help='ONNX opset version (default: 17)') + + # Model configuration + parser.add_argument('--num-classes', type=int, default=7, + help='Number of segmentation classes (default: 7 for Pascal Person Parts)') + parser.add_argument('--feature-dim', type=int, default=256, + help='Feature dimension for segmentation head') + parser.add_argument('--dropout-rate', type=float, default=0.1, + help='Dropout rate') + + # Options + parser.add_argument('--check', action='store_true', default=True, + help='Check the exported ONNX model') + parser.add_argument('--simplify', action='store_true', default=True, + help='Simplify the ONNX model') + parser.add_argument('--device', type=str, default='cpu', + choices=['cpu', 'cuda'], + help='Device for export (default: cpu)') + + return parser.parse_args() + + +def load_hyperparameters_from_checkpoint(checkpoint_path: str) -> dict: + """Load hyperparameters from saved checkpoint""" + print(f"📚 Loading hyperparameters from {checkpoint_path}") + + checkpoint = torch.load(checkpoint_path, map_location='cpu') + + # Extract hyperparameters + if 'hyperparameters' in checkpoint: + hyperparams = checkpoint['hyperparameters'] + print(f"✅ Found saved hyperparameters:") + for key, value in hyperparams.items(): + print(f" {key}: {value}") + return hyperparams + + elif 'model_info' in checkpoint: + # Extract from model info + model_info = checkpoint['model_info'] + hyperparams = { + 'feature_dim': model_info.get('fpn_feature_dim', 256), + 'dropout_rate': 0.1, # Default + 'tier': model_info.get('tier', 'standard') + } + print(f"✅ Extracted hyperparameters from model_info:") + for key, value in hyperparams.items(): + print(f" {key}: {value}") + return hyperparams + + else: + print("⚠️ No hyperparameters found in checkpoint, using defaults") + return { + 'feature_dim': 256, + 'dropout_rate': 0.1, + 'tier': 'standard' + } + + +def create_segmentation_model(args) -> nn.Module: + """Create segmentation model and load trained weights""" + + # Load hyperparameters from checkpoint + hyperparams = load_hyperparameters_from_checkpoint(args.resume) + + # Override with command line arguments if provided + feature_dim = args.feature_dim if args.feature_dim != 256 else hyperparams.get('feature_dim', 256) + dropout_rate = args.dropout_rate if args.dropout_rate != 0.1 else hyperparams.get('dropout_rate', 0.1) + + print(f"🏗️ Creating standard segmentation model:") + print(f" Feature dim: {feature_dim}") + print(f" Dropout rate: {dropout_rate}") + print(f" Num classes: {args.num_classes}") + + # Load original DFINE model + dfine_model = load_pretrained_dfine(args.config, args.original_weights) + backbone_channels = get_actual_backbone_channels(dfine_model) + + print(f" Backbone channels: {backbone_channels}") + + # Create standard segmentation head with correct hyperparameters + seg_head = create_standard_segmentation_head( + in_channels_list=backbone_channels, + num_classes=args.num_classes, + feature_dim=feature_dim, + dropout_rate=dropout_rate + ) + + # Create combined model + model = create_combined_model( + dfine_model=dfine_model, + seg_head=seg_head, + freeze_detection=False # Don't freeze during inference + ) + + # Load trained weights + print(f"📚 Loading trained segmentation weights from {args.resume}") + checkpoint = torch.load(args.resume, map_location='cpu') + + if 'model_state_dict' in checkpoint: + state_dict = checkpoint['model_state_dict'] + else: + state_dict = checkpoint + + # Load state dict + try: + model.load_state_dict(state_dict, strict=True) + print(f"✅ Loaded trained segmentation model successfully!") + except RuntimeError as e: + print(f"⚠️ Strict loading failed: {e}") + # Try non-strict loading + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + print(f"✅ Loaded with non-strict mode") + if missing_keys: + print(f" Missing keys: {len(missing_keys)}") + if unexpected_keys: + print(f" Unexpected keys: {len(unexpected_keys)}") + + return model, hyperparams + + +class StandardSegmentationModelWrapper(nn.Module): + """Wrapper model for ONNX export with proper postprocessing""" + + def __init__(self, seg_model, postprocessor): + super().__init__() + self.seg_model = seg_model.eval() + self.postprocessor = postprocessor + + def forward(self, images, orig_target_sizes): + # Run the segmentation model + outputs = self.seg_model(images) + + # Extract detection outputs for postprocessing + det_outputs = {k: v for k, v in outputs.items() if k != 'segmentation'} + + # Process detection outputs + processed_det = self.postprocessor(det_outputs, orig_target_sizes) + + # Get segmentation outputs + seg_logits = outputs['segmentation'] + + # Apply softmax to get probabilities + seg_probs = torch.softmax(seg_logits, dim=1) + + # Get segmentation predictions + seg_preds = torch.argmax(seg_logits, dim=1) + + # Return standard segmentation outputs + return processed_det[0], processed_det[1], processed_det[2], seg_probs, seg_preds + + +def export_onnx(args): + """Export the standard segmentation model to ONNX format""" + + # Set device + device = torch.device(args.device) + + # Create the segmentation model + model, hyperparams = create_segmentation_model(args) + model = model.to(device) + + # Create wrapper model for ONNX export + try: + from src.core import YAMLConfig + cfg = YAMLConfig(args.config) + postprocessor = cfg.postprocessor.deploy() + except Exception as e: + print(f"⚠️ Could not load postprocessor: {e}") + print(" Using identity postprocessor") + postprocessor = lambda x, y: (x.get('pred_logits', torch.empty(0)), + x.get('pred_boxes', torch.empty(0)), + torch.empty(0)) + + wrapper_model = StandardSegmentationModelWrapper(model, postprocessor) + wrapper_model.eval() + + # Parse input shape + input_shape = [int(x) for x in args.input_shape.split(",")] + print(f"📐 Using input shape: {input_shape}") + + # Create dummy input + data = torch.randn(*input_shape).to(device) + size = torch.tensor([[input_shape[3], input_shape[2]]]).to(device) # width, height format + + # Run test forward pass + print("🧪 Running test forward pass...") + with torch.no_grad(): + try: + test_outputs = wrapper_model(data, size) + print(f"✅ Test forward pass successful") + print(f" Outputs: {len(test_outputs)} tensors") + for i, out in enumerate(test_outputs): + print(f" Output {i}: shape {out.shape}, dtype {out.dtype}") + except Exception as e: + print(f"❌ Test forward pass failed: {e}") + raise + + # Define dynamic axes for ONNX export + dynamic_axes = { + "images": {0: "N"}, # Batch dimension + "orig_target_sizes": {0: "N"}, # Batch dimension + "labels": {0: "N"}, + "boxes": {0: "N"}, + "scores": {0: "N"}, + "seg_probs": {0: "N"}, + "seg_preds": {0: "N"}, + } + + # Determine output file name + if args.output: + output_file = args.output + else: + # Create descriptive filename + base_name = os.path.splitext(os.path.basename(args.resume))[0] + feature_dim = hyperparams.get('feature_dim', 256) + output_file = f"{base_name}_standard_fd{feature_dim}.onnx" + + print(f"📤 Exporting standard segmentation model to ONNX: {output_file}") + print(f" Model: {hyperparams.get('tier', 'standard')} tier") + print(f" Feature dim: {hyperparams.get('feature_dim', 256)}") + print(f" ONNX opset: {args.opset}") + + # Export to ONNX + try: + torch.onnx.export( + wrapper_model, + (data, size), + output_file, + input_names=["images", "orig_target_sizes"], + output_names=["labels", "boxes", "scores", "seg_probs", "seg_preds"], + dynamic_axes=dynamic_axes, + opset_version=args.opset, + verbose=False, + do_constant_folding=True, + export_params=True + ) + print("✅ ONNX export successful!") + except Exception as e: + print(f"❌ ONNX export failed: {e}") + raise + + # Check the ONNX model + if args.check: + try: + import onnx + print("🔍 Checking ONNX model...") + onnx_model = onnx.load(output_file) + onnx.checker.check_model(onnx_model) + print("✅ ONNX model check successful!") + + # Print model info + print(f"📊 ONNX Model Info:") + print(f" Inputs: {len(onnx_model.graph.input)}") + print(f" Outputs: {len(onnx_model.graph.output)}") + print(f" Nodes: {len(onnx_model.graph.node)}") + + except ImportError: + print("⚠️ ONNX package not found, skipping model check") + print(" Install with: pip install onnx") + except Exception as e: + print(f"❌ ONNX model check failed: {e}") + + # Simplify the ONNX model + if args.simplify: + try: + import onnx + import onnxsim + print("🔧 Simplifying ONNX model...") + + # Create input shapes dictionary + input_shapes = { + "images": input_shape, + "orig_target_sizes": [input_shape[0], 2] + } + + # Simplify + onnx_model = onnx.load(output_file) + model_simp, check = onnxsim.simplify( + onnx_model, + test_input_shapes=input_shapes, + dynamic_input_shape=True + ) + + if check: + onnx.save(model_simp, output_file) + print(f"✅ ONNX model simplified and saved") + else: + print("⚠️ ONNX model simplification could not be validated") + + except ImportError: + print("⚠️ onnx-simplifier package not found, skipping simplification") + print(" Install with: pip install onnx-simplifier") + except Exception as e: + print(f"❌ ONNX model simplification failed: {e}") + + # Print summary + file_size_mb = os.path.getsize(output_file) / (1024 * 1024) + print(f"\n🎉 Export completed successfully!") + print(f"📁 Output file: {output_file}") + print(f"📊 File size: {file_size_mb:.2f} MB") + print(f"🎯 Tier: Standard (balanced performance)") + print(f"\n📝 Next steps:") + print(f" 1. Build TensorRT engine:") + print(f" python tensorrt/standard/build_engine.py -i {output_file}") + print(f" 2. Or build INT8 engine:") + print(f" python tensorrt/standard/build_int8_engine.py -i {output_file}") + + return output_file + + +def main(): + """Main function""" + args = parse_args() + + # Validate inputs + if not os.path.isfile(args.resume): + print(f"❌ Model checkpoint not found: {args.resume}") + sys.exit(1) + + if not os.path.isfile(args.original_weights): + print(f"❌ Original DFINE weights not found: {args.original_weights}") + sys.exit(1) + + if not os.path.isfile(args.config): + print(f"❌ DFINE config not found: {args.config}") + sys.exit(1) + + print("=" * 60) + print("🚀 STANDARD DFINE SEGMENTATION ONNX EXPORT") + print("=" * 60) + print(f"📁 Checkpoint: {args.resume}") + print(f"📁 Original weights: {args.original_weights}") + print(f"📁 Config: {args.config}") + print(f"🎯 Target: Balanced performance (standard tier)") + print("=" * 60) + + try: + export_onnx(args) + except Exception as e: + print(f"\n❌ Export failed: {e}") + import traceback + traceback.print_exc() + sys.exit(1) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/segmentation_sivert/legacy/infer_trt.py b/segmentation_sivert/legacy/infer_trt.py new file mode 100644 index 00000000..5afe45fd --- /dev/null +++ b/segmentation_sivert/legacy/infer_trt.py @@ -0,0 +1,461 @@ +#!/usr/bin/env python3 +""" +Standard DFINE Segmentation TensorRT Inference Script +Refactored to use base classes and modular structure +""" + +import os +import sys +import argparse +import time +import cv2 +import numpy as np +import subprocess +import tempfile +from typing import Dict, List, Tuple, Optional + +# Add paths for imports +current_dir = os.path.dirname(os.path.abspath(__file__)) +segmentation_root = os.path.dirname(os.path.dirname(current_dir)) # Go up to segmentation_sivert/ + +# Add segmentation_sivert to path for core imports +if segmentation_root not in sys.path: + sys.path.insert(0, segmentation_root) +from core.engines import create_segmentation_engine, benchmark_engine + + +def parse_args(): + """Parse command line arguments""" + parser = argparse.ArgumentParser(description='Standard DFINE Segmentation TensorRT Inference') + + # Required arguments + parser.add_argument('-e', '--engine', type=str, required=True, + help='Path to TensorRT engine file') + parser.add_argument('-i', '--input', type=str, required=True, + help='Path to input video file') + + # Output configuration + parser.add_argument('-o', '--output', type=str, + help='Path to output video file (required for video processing mode)') + + # Processing modes + parser.add_argument('--benchmark', action='store_true', + help='Run pure inference benchmark mode') + parser.add_argument('--parallel-streams', type=int, default=1, + help='Number of parallel streams for benchmark (default: 1)') + + # Processing options + parser.add_argument('--max-frames', type=int, + help='Maximum number of frames to process') + parser.add_argument('--resize-factor', type=float, default=1.0, + help='Resize input frames by this factor') + parser.add_argument('--confidence-threshold', type=float, default=0.5, + help='Detection confidence threshold') + parser.add_argument('--seg-alpha', type=float, default=0.6, + help='Segmentation overlay alpha') + + # Video output options + parser.add_argument('--fps', type=float, default=None, + help='Output FPS (default: same as input)') + parser.add_argument('--crf', type=int, default=23, + help='CRF value for video encoding (default: 23)') + parser.add_argument('--preset', type=str, default='medium', + choices=['ultrafast', 'superfast', 'veryfast', 'faster', 'fast', 'medium', 'slow'], + help='x264 encoding preset') + + return parser.parse_args() + + +def get_body_part_colors() -> Dict[int, Tuple[int, int, int]]: + """Get colors for each body part class""" + return { + 0: (0, 0, 0), # background - black (transparent) + 1: (255, 0, 0), # head - red + 2: (0, 255, 0), # torso - green + 3: (0, 0, 255), # arms - blue + 4: (255, 255, 0), # hands - yellow + 5: (255, 0, 255), # legs - magenta + 6: (0, 255, 255), # feet - cyan + } + + +def get_body_part_names() -> Dict[int, str]: + """Get names for each body part class""" + return { + 0: 'background', + 1: 'head', + 2: 'torso', + 3: 'arms', + 4: 'hands', + 5: 'legs', + 6: 'feet' + } + + +def create_segmentation_overlay(frame: np.ndarray, seg_mask: np.ndarray, alpha: float = 0.6) -> np.ndarray: + """Create segmentation overlay on frame""" + colors = get_body_part_colors() + height, width = frame.shape[:2] + + # Resize segmentation mask to frame size + seg_resized = cv2.resize(seg_mask, (width, height), interpolation=cv2.INTER_NEAREST) + + # Create colored overlay + overlay = np.zeros((height, width, 3), dtype=np.uint8) + for class_id, color in colors.items(): + if class_id == 0: # Skip background + continue + mask = (seg_resized == class_id) + overlay[mask] = color + + # Blend with original frame + result = cv2.addWeighted(frame, 1-alpha, overlay, alpha, 0) + return result + + +def draw_detections(frame: np.ndarray, boxes: np.ndarray, scores: np.ndarray, labels: np.ndarray, + confidence_threshold: float = 0.5) -> np.ndarray: + """Draw detection boxes on frame""" + height, width = frame.shape[:2] + result = frame.copy() + + # COCO class names (simplified) + class_names = { + 0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane', + 5: 'bus', 6: 'train', 7: 'truck', 8: 'boat', 9: 'traffic light' + } + + colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255), + (0, 255, 255), (128, 0, 128), (255, 165, 0), (0, 128, 128), (128, 128, 0)] + + detection_count = 0 + + for i in range(len(scores)): + if scores[i] < confidence_threshold: + continue + + # Convert normalized coordinates to pixel coordinates + x1 = int(boxes[i][0] * width) + y1 = int(boxes[i][1] * height) + x2 = int(boxes[i][2] * width) + y2 = int(boxes[i][3] * height) + + label = int(labels[i]) + color = colors[label % len(colors)] + + # Draw bounding box + cv2.rectangle(result, (x1, y1), (x2, y2), color, 2) + + # Draw label + class_name = class_names.get(label, f'Class {label}') + label_text = f'{class_name}: {scores[i]:.2f}' + + (text_w, text_h), _ = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1) + cv2.rectangle(result, (x1, y1 - text_h - 10), (x1 + text_w, y1), color, -1) + cv2.putText(result, label_text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1) + + detection_count += 1 + + return result, detection_count + + +def process_video_with_visualization(engine, input_video: str, output_video: str, args) -> Dict: + """Process video with segmentation visualization""" + print(f"\n🎬 Processing video with visualization") + print(f"📥 Input: {input_video}") + print(f"📤 Output: {output_video}") + + # Open input video + cap = cv2.VideoCapture(input_video) + if not cap.isOpened(): + raise ValueError(f"Could not open input video: {input_video}") + + # Get video properties + video_fps = cap.get(cv2.CAP_PROP_FPS) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + if args.max_frames: + total_frames = min(total_frames, args.max_frames) + + # Calculate output dimensions + if args.resize_factor != 1.0: + output_width = int(width * args.resize_factor) + output_height = int(height * args.resize_factor) + else: + output_width, output_height = width, height + + output_fps = args.fps if args.fps else video_fps + + print(f"📊 Video: {width}x{height} @ {video_fps:.1f} FPS") + print(f"📊 Output: {output_width}x{output_height} @ {output_fps:.1f} FPS") + print(f"📊 Processing {total_frames:,} frames") + + # Create temporary directory for frames + temp_dir = tempfile.mkdtemp() + frames_dir = os.path.join(temp_dir, 'frames') + os.makedirs(frames_dir) + + # Performance tracking + frame_count = 0 + total_detections = 0 + total_body_parts = 0 + start_time = time.perf_counter() + + try: + while True: + ret, frame = cap.read() + if not ret or (args.max_frames and frame_count >= args.max_frames): + break + + frame_start_time = time.perf_counter() + + # Resize frame if needed + if args.resize_factor != 1.0: + frame = cv2.resize(frame, (output_width, output_height)) + + # Run inference + results = engine.infer_frame(frame) + + # Create visualization + vis_frame = frame.copy() + + # Add segmentation overlay + seg_preds = results['seg_preds'].astype(np.uint8) + vis_frame = create_segmentation_overlay(vis_frame, seg_preds, args.seg_alpha) + + # Add detection boxes + boxes = results['boxes'] + scores = results['scores'] + labels = results['labels'] + + vis_frame, detection_count = draw_detections( + vis_frame, boxes, scores, labels, args.confidence_threshold + ) + + # Count active body parts + unique_parts = np.unique(seg_preds) + active_parts = len([p for p in unique_parts if p != 0]) + + # Add performance info + frame_time = time.perf_counter() - frame_start_time + current_fps = 1.0 / frame_time if frame_time > 0 else 0 + + # Performance overlay + perf_text = [ + f"Frame: {frame_count+1}/{total_frames}", + f"FPS: {current_fps:.1f}", + f"Detections: {detection_count}", + f"Body parts: {active_parts}", + f"Standard Tier" + ] + + y_offset = 30 + for text in perf_text: + cv2.putText(vis_frame, text, (10, y_offset), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) + y_offset += 30 + + # Save frame + frame_path = os.path.join(frames_dir, f"frame_{frame_count:06d}.png") + cv2.imwrite(frame_path, vis_frame) + + # Update counters + frame_count += 1 + total_detections += detection_count + total_body_parts += active_parts + + # Progress reporting + if frame_count % max(1, total_frames // 20) == 0: + elapsed = time.perf_counter() - start_time + processing_fps = frame_count / elapsed + eta = (total_frames - frame_count) / processing_fps if processing_fps > 0 else 0 + print(f" Processing: {frame_count:4d}/{total_frames} ({frame_count/total_frames*100:5.1f}%) | " + f"Speed: {processing_fps:6.1f} FPS | ETA: {eta:4.1f}s") + + finally: + cap.release() + + total_time = time.perf_counter() - start_time + + # Create video using ffmpeg + print(f"🎬 Creating output video...") + ffmpeg_cmd = [ + 'ffmpeg', '-y', # Overwrite output file + '-framerate', str(output_fps), + '-i', os.path.join(frames_dir, 'frame_%06d.png'), + '-c:v', 'libx264', + '-pix_fmt', 'yuv420p', + '-crf', str(args.crf), + '-preset', args.preset, + output_video + ] + + try: + subprocess.run(ffmpeg_cmd, check=True, capture_output=True) + print(f"✅ Video saved to: {output_video}") + except subprocess.CalledProcessError as e: + print(f"❌ FFmpeg failed: {e}") + print(f"FFmpeg stderr: {e.stderr.decode()}") + raise + finally: + # Cleanup temporary files + import shutil + shutil.rmtree(temp_dir) + + # Calculate results + overall_fps = frame_count / total_time + + return { + 'frames_processed': frame_count, + 'total_time': total_time, + 'processing_fps': overall_fps, + 'total_detections': total_detections, + 'total_body_parts': total_body_parts, + 'avg_detections_per_frame': total_detections / frame_count if frame_count > 0 else 0, + 'avg_body_parts_per_frame': total_body_parts / frame_count if frame_count > 0 else 0 + } + + +def run_benchmark(engine, input_video: str, args) -> Dict: + """Run inference benchmark""" + print(f"\n🚀 Running inference benchmark") + + # Open video + cap = cv2.VideoCapture(input_video) + if not cap.isOpened(): + raise ValueError(f"Could not open input video: {input_video}") + + # Get video info + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + if args.max_frames: + total_frames = min(total_frames, args.max_frames) + + print(f"📊 Benchmarking {total_frames:,} frames") + + # Preprocess all frames + frames = [] + frame_count = 0 + + print("📥 Loading frames...") + while frame_count < total_frames: + ret, frame = cap.read() + if not ret: + break + + if args.resize_factor != 1.0: + h, w = frame.shape[:2] + new_h, new_w = int(h * args.resize_factor), int(w * args.resize_factor) + frame = cv2.resize(frame, (new_w, new_h)) + + frames.append(frame) + frame_count += 1 + + if frame_count % max(1, total_frames // 10) == 0: + print(f" Loaded {frame_count}/{total_frames} frames") + + cap.release() + + # Benchmark inference + print("🏃 Running benchmark...") + start_time = time.perf_counter() + + for i, frame in enumerate(frames): + _ = engine.infer_frame(frame) + + if (i + 1) % max(1, len(frames) // 20) == 0: + elapsed = time.perf_counter() - start_time + current_fps = (i + 1) / elapsed + eta = (len(frames) - i - 1) / current_fps if current_fps > 0 else 0 + print(f" Progress: {i + 1:4d}/{len(frames)} | FPS: {current_fps:6.1f} | ETA: {eta:4.1f}s") + + total_time = time.perf_counter() - start_time + + # Get engine performance info + engine_info = engine.get_engine_info() + + return { + 'frames_processed': len(frames), + 'total_time': total_time, + 'pure_inference_fps': len(frames) / total_time, + 'avg_inference_time_ms': engine_info.get('avg_inference_time_ms', 0), + 'engine_info': engine_info + } + + +def main(): + """Main function""" + args = parse_args() + + # Validate arguments + if not args.benchmark and not args.output: + print("❌ Error: --output is required when not in benchmark mode") + sys.exit(1) + + if not os.path.isfile(args.engine): + print(f"❌ Error: Engine file not found: {args.engine}") + sys.exit(1) + + if not os.path.isfile(args.input): + print(f"❌ Error: Input video not found: {args.input}") + sys.exit(1) + + print("=" * 80) + print("🚀 STANDARD DFINE SEGMENTATION TENSORRT INFERENCE") + print("=" * 80) + print(f"📁 Engine: {args.engine}") + print(f"📁 Input: {args.input}") + + if args.benchmark: + print(f"⚡ Mode: INFERENCE BENCHMARK") + else: + print(f"📁 Output: {args.output}") + print(f"🎬 Mode: VIDEO PROCESSING") + + if args.max_frames: + print(f"🔢 Max frames: {args.max_frames:,}") + + print("🎯 Tier: Standard (balanced performance)") + print("=" * 80) + + try: + # Create TensorRT engine + engine = create_segmentation_engine('standard', args.engine) + + # Print engine info + engine.print_engine_info() + + # Run processing + if args.benchmark: + # Benchmark mode + results = run_benchmark(engine, args.input, args) + + print(f"\n🎉 BENCHMARK RESULTS:") + print(f" Frames processed: {results['frames_processed']:,}") + print(f" Total time: {results['total_time']:.2f}s") + print(f" Pure inference FPS: {results['pure_inference_fps']:.1f}") + print(f" Average inference time: {results['avg_inference_time_ms']:.2f}ms") + + else: + # Video processing mode + results = process_video_with_visualization(engine, args.input, args.output, args) + + print(f"\n🎉 VIDEO PROCESSING RESULTS:") + print(f" Frames processed: {results['frames_processed']:,}") + print(f" Processing FPS: {results['processing_fps']:.1f}") + print(f" Total detections: {results['total_detections']:,}") + print(f" Avg detections/frame: {results['avg_detections_per_frame']:.1f}") + print(f" Avg body parts/frame: {results['avg_body_parts_per_frame']:.1f}") + print(f" Real-time capable: {'✅ Yes' if results['processing_fps'] >= 30 else '⚠️ No'}") + + except Exception as e: + print(f"❌ Processing failed: {e}") + import traceback + traceback.print_exc() + sys.exit(1) + + print(f"\n✅ Standard tier inference completed successfully!") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/segmentation_sivert/legacy/infer_trt_real_time.py b/segmentation_sivert/legacy/infer_trt_real_time.py new file mode 100644 index 00000000..105cc6aa --- /dev/null +++ b/segmentation_sivert/legacy/infer_trt_real_time.py @@ -0,0 +1,577 @@ +#!/usr/bin/env python3 +""" +Single-Threaded Optimized Webcam Inference Server +Avoids CUDA context issues by handling clients sequentially. + +Usage: + python optimized_server.py --engine model.engine --port 65432 +""" + +import os +import sys +import argparse +import time +import socket +import struct +import numpy as np +import cv2 +import tensorrt as trt +import pycuda.driver as cuda +import pycuda.autoinit +import traceback + +class TensorRTEngine: + """TensorRT engine wrapper - single threaded""" + + def __init__(self, engine_path: str): + self.engine_path = engine_path + self.logger = trt.Logger(trt.Logger.WARNING) + self.runtime = None + self.engine = None + self.context = None + self.stream = None + + # Memory management + self.host_inputs = [] + self.host_outputs = [] + self.device_inputs = [] + self.device_outputs = [] + self.input_shapes = {} + self.output_shapes = {} + + self._load_engine() + self._allocate_memory() + + def _load_engine(self): + """Load TensorRT engine""" + print(f"🔧 Loading TensorRT engine: {self.engine_path}") + + self.runtime = trt.Runtime(self.logger) + + with open(self.engine_path, 'rb') as f: + engine_bytes = f.read() + + self.engine = self.runtime.deserialize_cuda_engine(engine_bytes) + if not self.engine: + raise RuntimeError(f"Failed to load TensorRT engine from {self.engine_path}") + + self.context = self.engine.create_execution_context() + if not self.context: + raise RuntimeError("Failed to create TensorRT execution context") + + self.stream = cuda.Stream() + print(f"✅ Engine loaded successfully") + + def _allocate_memory(self): + """Pre-allocate host and device memory""" + print(f"🧠 Allocating memory buffers...") + + # Check TensorRT version + use_new_api = hasattr(self.engine, 'num_io_tensors') + + if use_new_api: + # TensorRT 10.x+ API + for i in range(self.engine.num_io_tensors): + tensor_name = self.engine.get_tensor_name(i) + shape = self.engine.get_tensor_shape(tensor_name) + dtype = self.engine.get_tensor_dtype(tensor_name) + is_input = self.engine.get_tensor_mode(tensor_name) == trt.TensorIOMode.INPUT + + # Handle dynamic shapes - set to batch size 1 + if -1 in shape: + actual_shape = tuple(1 if dim == -1 else dim for dim in shape) + self.context.set_input_shape(tensor_name, actual_shape) + shape = actual_shape + + # Convert dtype and calculate size + if dtype == trt.DataType.FLOAT: + np_dtype, element_size = np.float32, 4 + elif dtype == trt.DataType.HALF: + np_dtype, element_size = np.float16, 2 + elif dtype == trt.DataType.INT32: + np_dtype, element_size = np.int32, 4 + elif dtype == trt.DataType.INT64: + np_dtype, element_size = np.int64, 8 + else: + np_dtype, element_size = np.float32, 4 + + size = trt.volume(shape) + + # Allocate memory + host_mem = cuda.pagelocked_empty(size, np_dtype) + device_mem = cuda.mem_alloc(size * element_size) + + if is_input: + self.host_inputs.append(host_mem) + self.device_inputs.append(device_mem) + self.input_shapes[tensor_name] = shape + else: + self.host_outputs.append(host_mem) + self.device_outputs.append(device_mem) + self.output_shapes[tensor_name] = shape + + # Set tensor address + self.context.set_tensor_address(tensor_name, int(device_mem)) + + print(f" {'📥' if is_input else '📤'} {tensor_name}: {shape} ({size * element_size / 1024 / 1024:.1f} MB)") + + else: + # Legacy TensorRT API + self.bindings = [] + for i in range(self.engine.num_bindings): + binding_name = self.engine.get_binding_name(i) + shape = self.engine.get_binding_shape(i) + dtype = self.engine.get_binding_dtype(i) + is_input = self.engine.binding_is_input(i) + + if -1 in shape: + actual_shape = tuple(1 if dim == -1 else dim for dim in shape) + self.context.set_binding_shape(i, actual_shape) + shape = actual_shape + + if dtype == trt.DataType.FLOAT: + np_dtype, element_size = np.float32, 4 + elif dtype == trt.DataType.INT64: + np_dtype, element_size = np.int64, 8 + else: + np_dtype, element_size = np.float32, 4 + + size = trt.volume(shape) + host_mem = cuda.pagelocked_empty(size, np_dtype) + device_mem = cuda.mem_alloc(size * element_size) + + if is_input: + self.host_inputs.append(host_mem) + self.device_inputs.append(device_mem) + self.input_shapes[binding_name] = shape + else: + self.host_outputs.append(host_mem) + self.device_outputs.append(device_mem) + self.output_shapes[binding_name] = shape + + self.bindings.append(int(device_mem)) + + print(f" {'📥' if is_input else '📤'} {binding_name}: {shape} ({size * element_size / 1024 / 1024:.1f} MB)") + + def preprocess_frame(self, frame: np.ndarray) -> tuple: + """Fast preprocessing""" + # Resize frame + resized = cv2.resize(frame, (640, 640), interpolation=cv2.INTER_LINEAR) + + # Convert to RGB and normalize (simplified) + rgb_frame = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB) + normalized = rgb_frame.astype(np.float32) / 255.0 + + # Apply ImageNet normalization (keep this if your model needs it) + mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) + std = np.array([0.229, 0.224, 0.225], dtype=np.float32) + normalized = (normalized - mean) / std + + # Convert to NCHW format + input_tensor = np.transpose(normalized, (2, 0, 1)) + input_tensor = np.expand_dims(input_tensor, axis=0) # Add batch dimension + + # Target size tensor + target_size_tensor = np.array([[640, 640]], dtype=np.int64) + + return input_tensor, target_size_tensor + + def infer(self, inputs: list) -> list: + """Perform inference""" + + # Copy inputs to device + for i, input_data in enumerate(inputs): + np.copyto(self.host_inputs[i][:input_data.size], input_data.ravel()) + cuda.memcpy_htod_async(self.device_inputs[i], self.host_inputs[i], self.stream) + + # Execute inference + if hasattr(self.context, 'execute_async_v3'): + success = self.context.execute_async_v3(stream_handle=self.stream.handle) + elif hasattr(self.context, 'execute_async_v2'): + success = self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle) + else: + success = self.context.execute_async(bindings=self.bindings, stream_handle=self.stream.handle) + + if not success: + raise RuntimeError("TensorRT inference failed") + + # Copy outputs from device + outputs = [] + for i, (host_output, device_output, shape) in enumerate( + zip(self.host_outputs, self.device_outputs, self.output_shapes.values()) + ): + cuda.memcpy_dtoh_async(host_output, device_output, self.stream) + self.stream.synchronize() + output = host_output[:np.prod(shape)].reshape(shape).copy() + outputs.append(output) + + return outputs + +def create_segmentation_overlay(frame: np.ndarray, seg_mask: np.ndarray, alpha: float = 0.6) -> np.ndarray: + """Create segmentation overlay on frame""" + + # Pascal Person Parts color map (7 classes) + colors = [ + [0, 0, 0], # 0: background - black + [128, 0, 0], # 1: head - dark red + [255, 0, 0], # 2: torso - red + [0, 128, 0], # 3: upper arms - dark green + [0, 255, 0], # 4: lower arms - green + [0, 0, 128], # 5: upper legs - dark blue + [0, 0, 255], # 6: lower legs - blue + ] + + height, width = frame.shape[:2] + + # Resize segmentation mask to frame size + seg_resized = cv2.resize(seg_mask, (width, height), interpolation=cv2.INTER_NEAREST) + + # Create colored overlay + overlay = np.zeros((height, width, 3), dtype=np.uint8) + for class_id, color in enumerate(colors): + mask = (seg_resized == class_id) + overlay[mask] = color + + # Blend with original frame + result = cv2.addWeighted(frame, 1-alpha, overlay, alpha, 0) + + return result + +def draw_detections(frame: np.ndarray, boxes: np.ndarray, scores: np.ndarray, labels: np.ndarray, + confidence_threshold: float = 0.5) -> np.ndarray: + """Draw detection boxes on frame""" + + height, width = frame.shape[:2] + result = frame.copy() + + # COCO class names (simplified) + class_names = { + 0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane', + 5: 'bus', 6: 'train', 7: 'truck', 8: 'boat', 9: 'traffic light' + } + + colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255), + (0, 255, 255), (128, 0, 128), (255, 165, 0), (0, 128, 128), (128, 128, 0)] + + for i in range(len(scores)): + if scores[i] < confidence_threshold: + continue + + # Convert normalized coordinates to pixel coordinates + x1 = int(boxes[i][0] * width) + y1 = int(boxes[i][1] * height) + x2 = int(boxes[i][2] * width) + y2 = int(boxes[i][3] * height) + + label = int(labels[i]) + color = colors[label % len(colors)] + + # Draw bounding box + cv2.rectangle(result, (x1, y1), (x2, y2), color, 2) + + # Draw label + class_name = class_names.get(label, f'Class {label}') + label_text = f'{class_name}: {scores[i]:.2f}' + + (text_w, text_h), _ = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1) + cv2.rectangle(result, (x1, y1 - text_h - 10), (x1 + text_w, y1), color, -1) + cv2.putText(result, label_text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1) + + return result + +class SingleThreadedWebcamServer: + """Single-threaded webcam inference server (avoids CUDA context issues)""" + + def __init__(self, engine_path: str, port: int = 65432): + self.engine_path = engine_path + self.port = port + + # Server components + self.server_socket = None + self.running = False + + # Inference engine (single instance, single thread) + self.engine = None + + # Performance tracking + self.frame_count = 0 + self.start_time = time.perf_counter() + + def initialize_engine(self): + """Initialize the inference engine""" + print("🚀 Initializing TensorRT engine (single-threaded)...") + try: + # Initialize CUDA context in main thread + cuda.init() + + self.engine = TensorRTEngine(self.engine_path) + + # Warmup + print("🔥 Warming up...") + dummy_frame = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8) + for _ in range(3): + input_tensor, target_size_tensor = self.engine.preprocess_frame(dummy_frame) + _ = self.engine.infer([input_tensor, target_size_tensor]) + + print("✅ Engine ready!") + return True + + except Exception as e: + print(f"❌ Engine initialization failed: {e}") + traceback.print_exc() + return False + + def handle_client_session(self, client_socket, client_address): + """Handle a complete client session""" + print(f"🔗 Client connected: {client_address}") + + try: + # Set socket timeouts to avoid hanging + client_socket.settimeout(30.0) + + while self.running: + try: + # Receive frame size + size_data = self._recv_exact(client_socket, 4) + if not size_data: + print("❌ Client disconnected (no size data)") + break + + frame_size = struct.unpack("!I", size_data)[0] + + # Sanity check + if frame_size <= 0 or frame_size > 2 * 1024 * 1024: # Max 2MB + print(f"⚠️ Invalid frame size: {frame_size}") + continue + + # Receive compressed frame + compressed_data = self._recv_exact(client_socket, frame_size) + if not compressed_data: + print("❌ Client disconnected (no frame data)") + break + + # Decompress frame + frame_array = np.frombuffer(compressed_data, dtype=np.uint8) + frame = cv2.imdecode(frame_array, cv2.IMREAD_COLOR) + + if frame is None: + print("⚠️ Failed to decode frame") + # Send back error frame + error_frame = np.zeros((480, 640, 3), dtype=np.uint8) + cv2.putText(error_frame, "Decode Error", (10, 50), + cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2) + self._send_frame(client_socket, error_frame) + continue + + except socket.timeout: + print("⚠️ Client timeout - closing connection") + break + except Exception as e: + print(f"❌ Frame receive error: {e}") + break + + # Process frame (single-threaded, no CUDA context issues) + try: + process_start = time.perf_counter() + + # Preprocess and infer + input_tensor, target_size_tensor = self.engine.preprocess_frame(frame) + outputs = self.engine.infer([input_tensor, target_size_tensor]) + + # Extract outputs and create visualization + if len(outputs) >= 5: + labels = outputs[0].flatten() # Detection labels + boxes = outputs[1][0] # Detection boxes + scores = outputs[2].flatten() # Detection scores + seg_probs = outputs[3][0] # Segmentation probabilities + seg_preds = outputs[4][0] # Segmentation predictions + + # Create processed frame + result_frame = frame.copy() + + # Add segmentation overlay + result_frame = create_segmentation_overlay(result_frame, seg_preds.astype(np.uint8), alpha=0.6) + + # Add detection boxes + result_frame = draw_detections(result_frame, boxes, scores, labels, confidence_threshold=0.5) + + else: + print("⚠️ Unexpected number of outputs, using original frame") + result_frame = frame.copy() + + # Add performance info + process_time = time.perf_counter() - process_start + self.frame_count += 1 + + elapsed = time.perf_counter() - self.start_time + fps = self.frame_count / elapsed if elapsed > 0 else 0 + + cv2.putText(result_frame, f"Server FPS: {fps:.1f}", (10, 30), + cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) + cv2.putText(result_frame, f"Process: {process_time*1000:.1f}ms", (10, 60), + cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) + + # Send result back to client + if not self._send_frame(client_socket, result_frame): + print("❌ Failed to send result") + break + + except Exception as e: + print(f"❌ Processing error: {e}") + traceback.print_exc() + + # Send original frame back as fallback + try: + error_frame = frame.copy() + cv2.putText(error_frame, f"Process Error: {str(e)[:30]}", (10, 50), + cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2) + self._send_frame(client_socket, error_frame) + except: + break + + except Exception as e: + print(f"❌ Client session error: {e}") + traceback.print_exc() + finally: + try: + client_socket.close() + except: + pass + print(f"👋 Client disconnected: {client_address}") + + def _recv_exact(self, sock, size: int) -> bytes: + """Receive exactly 'size' bytes""" + data = b'' + while len(data) < size: + try: + chunk = sock.recv(size - len(data)) + if not chunk: + return b'' + data += chunk + except Exception as e: + return b'' + return data + + def _send_frame(self, sock, frame: np.ndarray) -> bool: + """Send frame as JPEG to client""" + try: + # Compress frame + _, compressed = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 75]) + + # Send size and data + size_data = struct.pack("!I", len(compressed)) + sock.sendall(size_data) + sock.sendall(compressed.tobytes()) + return True + + except Exception as e: + print(f"❌ Send error: {e}") + return False + + def start_server(self): + """Start the server""" + try: + print(f"🚀 Starting single-threaded server on port {self.port}...") + + self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self.server_socket.bind(('0.0.0.0', self.port)) + self.server_socket.listen(1) # Only accept 1 client at a time + + self.running = True + self.start_time = time.perf_counter() + + print(f"✅ Server listening on port {self.port}") + print("📝 Note: Server handles one client at a time (avoids CUDA issues)") + return True + + except Exception as e: + print(f"❌ Failed to start server: {e}") + return False + + def run(self): + """Main server loop""" + print("🏃 Server running (single-threaded)... Press Ctrl+C to stop") + + try: + while self.running: + try: + print("⏳ Waiting for client connection...") + client_socket, client_address = self.server_socket.accept() + + # Handle this client completely before accepting another + self.handle_client_session(client_socket, client_address) + + except Exception as e: + if self.running: + print(f"❌ Accept error: {e}") + + except KeyboardInterrupt: + print("\n👋 Server shutdown requested") + finally: + self.stop() + + def stop(self): + """Stop the server""" + print("🛑 Stopping server...") + self.running = False + + if self.server_socket: + try: + self.server_socket.close() + except: + pass + + print("✅ Server stopped") + +def parse_args(): + """Parse command line arguments""" + parser = argparse.ArgumentParser(description='Single-Threaded Optimized Webcam Inference Server') + parser.add_argument('-e', '--engine', type=str, required=True, + help='Path to TensorRT engine file') + parser.add_argument('-p', '--port', type=int, default=65432, + help='Server port (default: 65432)') + return parser.parse_args() + +def main(): + """Main function""" + args = parse_args() + + if not os.path.isfile(args.engine): + print(f"❌ Engine file not found: {args.engine}") + sys.exit(1) + + print("=" * 70) + print("🚀 SINGLE-THREADED OPTIMIZED WEBCAM INFERENCE SERVER") + print("=" * 70) + print(f"📁 Engine: {args.engine}") + print(f"🌐 Port: {args.port}") + print("🎯 Focus: No CUDA context issues + fast inference") + print("⚠️ Note: Handles one client at a time") + print("=" * 70) + + # System info + if cuda.Device.count() > 0: + device = cuda.Device(0) + print(f"🎮 GPU: {device.name()} ({device.total_memory() / 1024**3:.1f} GB)") + + server = SingleThreadedWebcamServer(args.engine, args.port) + + try: + if not server.initialize_engine(): + print("❌ Failed to initialize engine") + sys.exit(1) + + if not server.start_server(): + print("❌ Failed to start server") + sys.exit(1) + + server.run() + + except Exception as e: + print(f"❌ Server error: {e}") + traceback.print_exc() + sys.exit(1) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/segmentation_sivert/legacy/int8_calibrator.py b/segmentation_sivert/legacy/int8_calibrator.py new file mode 100644 index 00000000..59fb3881 --- /dev/null +++ b/segmentation_sivert/legacy/int8_calibrator.py @@ -0,0 +1,380 @@ +import os +import numpy as np +import pycuda.driver as cuda +import pycuda.autoinit +import tensorrt as trt +import cv2 +from PIL import Image +import torchvision.transforms as transforms +import random + +class PascalPersonPartsCalibrator(trt.IInt8EntropyCalibrator2): + """Pascal Person Parts only calibrator - optimized for human segmentation""" + + def __init__(self, pascal_data_dir, batch_size=1, cache_file="pascal_calibration.cache", + max_calibration_images=500): + super(PascalPersonPartsCalibrator, self).__init__() + + # Image preprocessing - match your training preprocessing + self.transform = transforms.Compose([ + transforms.Resize((640, 640)), + transforms.ToTensor(), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + ]) + + # Find Pascal Person Parts images + self.image_files = [] + + # Look in datasets/pascal_person_parts/images/ + pascal_image_dirs = [ + os.path.join(pascal_data_dir, 'images', 'val'), # Validation set (preferred) + os.path.join(pascal_data_dir, 'images', 'train'), # Training set if val is small + os.path.join(pascal_data_dir, 'images'), # Direct images folder + pascal_data_dir # Root directory fallback + ] + + for img_dir in pascal_image_dirs: + if os.path.exists(img_dir): + print(f"Searching Pascal images in: {img_dir}") + for root, dirs, files in os.walk(img_dir): + for file in files: + if file.lower().endswith(('.jpg', '.jpeg', '.png')): + self.image_files.append(os.path.join(root, file)) + + if len(self.image_files) >= max_calibration_images: + break + + # Limit and shuffle for better representation + if len(self.image_files) > max_calibration_images: + random.shuffle(self.image_files) + self.image_files = self.image_files[:max_calibration_images] + + print(f"Pascal-only calibration: Found {len(self.image_files)} images") + if len(self.image_files) == 0: + raise ValueError(f"No Pascal Person Parts images found in {pascal_data_dir}") + + self.batch_size = batch_size + self.current_index = 0 + self.cache_file = cache_file + + # Allocate GPU memory + self.device_input_image = cuda.mem_alloc(self.batch_size * 3 * 640 * 640 * 4) + self.device_input_sizes = cuda.mem_alloc(self.batch_size * 2 * 8) + self.host_sizes = np.array([[640, 640] for _ in range(batch_size)], dtype=np.int64) + self.stream = cuda.Stream() + + def get_batch_size(self): + return self.batch_size + + def get_batch(self, names): + if self.current_index + self.batch_size > len(self.image_files): + print(f"Pascal calibration complete. Processed {self.current_index} images.") + return None + + print(f"Pascal calibration batch {self.current_index//self.batch_size + 1}/{len(self.image_files)//self.batch_size + 1}") + + batch_data = np.zeros((self.batch_size, 3, 640, 640), dtype=np.float32) + + valid_images = 0 + for i in range(self.batch_size): + if self.current_index + i >= len(self.image_files): + break + + image_path = self.image_files[self.current_index + i] + + try: + pil_image = Image.open(image_path).convert('RGB') + processed_image = self.transform(pil_image) + batch_data[i] = processed_image.numpy() + valid_images += 1 + + if i == 0: + print(f" Processing Pascal image: {os.path.basename(image_path)}") + + except Exception as e: + print(f"Error processing {image_path}: {e}") + continue + + if valid_images == 0: + self.current_index += self.batch_size + return self.get_batch(names) + + # Copy to GPU + cuda.memcpy_htod_async(self.device_input_image, batch_data.ravel(), self.stream) + cuda.memcpy_htod_async(self.device_input_sizes, self.host_sizes.ravel(), self.stream) + self.stream.synchronize() + + self.current_index += self.batch_size + + # Return bindings + bindings = [] + for name in names: + if name == "images": + bindings.append(int(self.device_input_image)) + elif name == "orig_target_sizes": + bindings.append(int(self.device_input_sizes)) + else: + print(f"Warning: Unknown binding name: {name}") + bindings.append(0) + + return bindings + + def read_calibration_cache(self): + if os.path.exists(self.cache_file): + with open(self.cache_file, "rb") as f: + cache_data = f.read() + print(f"Read {len(cache_data)} bytes from Pascal calibration cache") + return cache_data + return None + + def write_calibration_cache(self, cache): + with open(self.cache_file, "wb") as f: + f.write(cache) + print(f"Wrote {len(cache)} bytes to Pascal calibration cache") + + +class MixedDatasetCalibrator(trt.IInt8EntropyCalibrator2): + """Mixed Pascal Person Parts + COCO calibrator for balanced detection + segmentation""" + + def __init__(self, pascal_data_dir, coco_data_dir, pascal_ratio=0.7, batch_size=1, + cache_file="mixed_calibration.cache", max_calibration_images=500): + super(MixedDatasetCalibrator, self).__init__() + + self.transform = transforms.Compose([ + transforms.Resize((640, 640)), + transforms.ToTensor(), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + ]) + + # Collect Pascal Person Parts images + pascal_images = [] + pascal_image_dirs = [ + os.path.join(pascal_data_dir, 'images', 'val'), + os.path.join(pascal_data_dir, 'images', 'train'), + os.path.join(pascal_data_dir, 'images'), + pascal_data_dir + ] + + for img_dir in pascal_image_dirs: + if os.path.exists(img_dir): + print(f"Searching Pascal images in: {img_dir}") + for root, dirs, files in os.walk(img_dir): + for file in files: + if file.lower().endswith(('.jpg', '.jpeg', '.png')): + pascal_images.append(os.path.join(root, file)) + break # Use first found directory + + # Collect COCO images from datasets/coco_val2017/ + coco_images = [] + if os.path.exists(coco_data_dir): + print(f"Searching COCO images in: {coco_data_dir}") + for root, dirs, files in os.walk(coco_data_dir): + for file in files: + if file.lower().endswith(('.jpg', '.jpeg', '.png')): + coco_images.append(os.path.join(root, file)) + else: + print(f"Warning: COCO directory {coco_data_dir} not found, using Pascal-only") + + # Mix datasets according to ratio + num_pascal = min(len(pascal_images), int(max_calibration_images * pascal_ratio)) + num_coco = min(len(coco_images), max_calibration_images - num_pascal) if coco_images else 0 + + # Randomly sample images + selected_pascal = [] + if pascal_images: + random.shuffle(pascal_images) + selected_pascal = pascal_images[:num_pascal] + + selected_coco = [] + if coco_images and num_coco > 0: + random.shuffle(coco_images) + selected_coco = coco_images[:num_coco] + + # Combine and shuffle final dataset + self.image_files = selected_pascal + selected_coco + random.shuffle(self.image_files) + + print(f"Mixed calibration dataset:") + print(f" Pascal Person Parts images: {len(selected_pascal)}") + print(f" COCO images: {len(selected_coco)}") + print(f" Total calibration images: {len(self.image_files)}") + + if len(self.image_files) == 0: + raise ValueError("No calibration images found. Check Pascal and COCO dataset paths.") + + self.batch_size = batch_size + self.current_index = 0 + self.cache_file = cache_file + + # Allocate GPU memory + self.device_input_image = cuda.mem_alloc(self.batch_size * 3 * 640 * 640 * 4) + self.device_input_sizes = cuda.mem_alloc(self.batch_size * 2 * 8) + self.host_sizes = np.array([[640, 640] for _ in range(batch_size)], dtype=np.int64) + self.stream = cuda.Stream() + + def get_batch_size(self): + return self.batch_size + + def get_batch(self, names): + if self.current_index + self.batch_size > len(self.image_files): + print(f"Mixed calibration complete. Processed {self.current_index} images.") + return None + + print(f"Mixed calibration batch {self.current_index//self.batch_size + 1}/{len(self.image_files)//self.batch_size + 1}") + + batch_data = np.zeros((self.batch_size, 3, 640, 640), dtype=np.float32) + + for i in range(self.batch_size): + if self.current_index + i >= len(self.image_files): + break + + image_path = self.image_files[self.current_index + i] + + try: + pil_image = Image.open(image_path).convert('RGB') + processed_image = self.transform(pil_image) + batch_data[i] = processed_image.numpy() + + if i == 0: + dataset_type = "Pascal" if "pascal" in image_path.lower() else "COCO" + print(f" Processing {dataset_type} image: {os.path.basename(image_path)}") + + except Exception as e: + print(f"Error processing {image_path}: {e}") + continue + + # Copy to GPU + cuda.memcpy_htod_async(self.device_input_image, batch_data.ravel(), self.stream) + cuda.memcpy_htod_async(self.device_input_sizes, self.host_sizes.ravel(), self.stream) + self.stream.synchronize() + + self.current_index += self.batch_size + + # Return bindings + bindings = [] + for name in names: + if name == "images": + bindings.append(int(self.device_input_image)) + elif name == "orig_target_sizes": + bindings.append(int(self.device_input_sizes)) + else: + bindings.append(0) + + return bindings + + def read_calibration_cache(self): + if os.path.exists(self.cache_file): + with open(self.cache_file, "rb") as f: + cache_data = f.read() + print(f"Read {len(cache_data)} bytes from mixed calibration cache") + return cache_data + return None + + def write_calibration_cache(self, cache): + with open(self.cache_file, "wb") as f: + f.write(cache) + print(f"Wrote {len(cache)} bytes to mixed calibration cache") + + +class COCOOnlyCalibrator(trt.IInt8EntropyCalibrator2): + """COCO-only calibrator for detection-focused optimization""" + + def __init__(self, coco_data_dir, batch_size=1, cache_file="coco_calibration.cache", + max_calibration_images=500): + super(COCOOnlyCalibrator, self).__init__() + + self.transform = transforms.Compose([ + transforms.Resize((640, 640)), + transforms.ToTensor(), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + ]) + + # Find COCO images in datasets/coco_val2017/ + self.image_files = [] + if os.path.exists(coco_data_dir): + print(f"Searching COCO images in: {coco_data_dir}") + for root, dirs, files in os.walk(coco_data_dir): + for file in files: + if file.lower().endswith(('.jpg', '.jpeg', '.png')): + self.image_files.append(os.path.join(root, file)) + + # Limit and shuffle + if len(self.image_files) > max_calibration_images: + random.shuffle(self.image_files) + self.image_files = self.image_files[:max_calibration_images] + + print(f"COCO-only calibration: Found {len(self.image_files)} images") + if len(self.image_files) == 0: + raise ValueError(f"No COCO images found in {coco_data_dir}") + + self.batch_size = batch_size + self.current_index = 0 + self.cache_file = cache_file + + # Allocate GPU memory + self.device_input_image = cuda.mem_alloc(self.batch_size * 3 * 640 * 640 * 4) + self.device_input_sizes = cuda.mem_alloc(self.batch_size * 2 * 8) + self.host_sizes = np.array([[640, 640] for _ in range(batch_size)], dtype=np.int64) + self.stream = cuda.Stream() + + def get_batch_size(self): + return self.batch_size + + def get_batch(self, names): + if self.current_index + self.batch_size > len(self.image_files): + print(f"COCO calibration complete. Processed {self.current_index} images.") + return None + + print(f"COCO calibration batch {self.current_index//self.batch_size + 1}/{len(self.image_files)//self.batch_size + 1}") + + batch_data = np.zeros((self.batch_size, 3, 640, 640), dtype=np.float32) + + for i in range(self.batch_size): + if self.current_index + i >= len(self.image_files): + break + + image_path = self.image_files[self.current_index + i] + + try: + pil_image = Image.open(image_path).convert('RGB') + processed_image = self.transform(pil_image) + batch_data[i] = processed_image.numpy() + + if i == 0: + print(f" Processing COCO image: {os.path.basename(image_path)}") + + except Exception as e: + print(f"Error processing {image_path}: {e}") + continue + + # Copy to GPU + cuda.memcpy_htod_async(self.device_input_image, batch_data.ravel(), self.stream) + cuda.memcpy_htod_async(self.device_input_sizes, self.host_sizes.ravel(), self.stream) + self.stream.synchronize() + + self.current_index += self.batch_size + + # Return bindings + bindings = [] + for name in names: + if name == "images": + bindings.append(int(self.device_input_image)) + elif name == "orig_target_sizes": + bindings.append(int(self.device_input_sizes)) + else: + bindings.append(0) + + return bindings + + def read_calibration_cache(self): + if os.path.exists(self.cache_file): + with open(self.cache_file, "rb") as f: + cache_data = f.read() + print(f"Read {len(cache_data)} bytes from COCO calibration cache") + return cache_data + return None + + def write_calibration_cache(self, cache): + with open(self.cache_file, "wb") as f: + f.write(cache) + print(f"Wrote {len(cache)} bytes to COCO calibration cache") \ No newline at end of file diff --git a/segmentation_sivert/legacy/test_standard_model.py b/segmentation_sivert/legacy/test_standard_model.py new file mode 100644 index 00000000..13a92962 --- /dev/null +++ b/segmentation_sivert/legacy/test_standard_model.py @@ -0,0 +1,788 @@ +#!/usr/bin/env python3 +""" +Enhanced D-FINE Segmentation Video Testing Script +Complete working version that matches the robust training architecture +""" + +import os +import sys +import argparse +import time +import subprocess +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision.transforms as T +from PIL import Image, ImageDraw +import cv2 + +# Add src to path +sys.path.append('src') + +def parse_args(): + """Parse command line arguments""" + parser = argparse.ArgumentParser(description='Enhanced D-FINE Segmentation Video Testing') + parser.add_argument('-c', '--config', type=str, default="configs/dfine_hgnetv2_x_obj2coco.yml", + help='Path to original DFINE configuration YAML') + parser.add_argument('-m', '--model', type=str, default='outputs/dfine_segmentation_robust/best_model.pth', + help='Path to enhanced segmentation model checkpoint') + parser.add_argument('-i', '--input', type=str, required=True, + help='Path to input video') + parser.add_argument('-o', '--output', type=str, default='enhanced_segmentation_result.mp4', + help='Path to output video') + parser.add_argument('-t', '--threshold', type=float, default=0.6, + help='Detection confidence threshold (default: 0.6)') + parser.add_argument('--seg-alpha', type=float, default=0.6, + help='Segmentation overlay alpha (default: 0.6)') + parser.add_argument('-d', '--device', type=str, default='cuda', + choices=['cuda', 'cpu'], + help='Device to run inference on') + parser.add_argument('--resize-factor', type=float, default=1.0, + help='Resize input frames by this factor') + parser.add_argument('--fps', type=float, default=None, + help='Output FPS (default: same as input)') + parser.add_argument('--crf', type=int, default=23, + help='CRF value for video encoding (0-51)') + parser.add_argument('--preset', type=str, default='medium', + choices=['ultrafast', 'superfast', 'veryfast', 'faster', 'fast', 'medium', 'slow', 'slower', 'veryslow'], + help='x264 encoding preset') + return parser.parse_args() + + +# Import core components +from src.core import YAMLConfig + + +class AttentionGate(nn.Module): + """Attention gate for better feature selection""" + + def __init__(self, F_g, F_l, F_int): + super().__init__() + self.W_g = nn.Sequential( + nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=True), + nn.BatchNorm2d(F_int) + ) + + self.W_x = nn.Sequential( + nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=True), + nn.BatchNorm2d(F_int) + ) + + self.psi = nn.Sequential( + nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=True), + nn.BatchNorm2d(1), + nn.Sigmoid() + ) + + self.relu = nn.ReLU(inplace=True) + + def forward(self, g, x): + g1 = self.W_g(g) + x1 = self.W_x(x) + psi = self.relu(g1 + x1) + psi = self.psi(psi) + return x * psi + + +class EnhancedASPP(nn.Module): + """Enhanced ASPP with attention and multi-scale context""" + + def __init__(self, in_channels, out_channels, dilations=[1, 6, 12, 18]): + super().__init__() + + self.convs = nn.ModuleList() + for dilation in dilations: + if dilation == 1: + conv = nn.Conv2d(in_channels, out_channels//len(dilations), 1, bias=False) + else: + conv = nn.Conv2d(in_channels, out_channels//len(dilations), 3, + padding=dilation, dilation=dilation, bias=False) + self.convs.append(nn.Sequential( + conv, + nn.BatchNorm2d(out_channels//len(dilations)), + nn.ReLU(inplace=True) + )) + + # Global average pooling branch + self.global_pool = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(in_channels, out_channels//len(dilations), 1, bias=False), + nn.BatchNorm2d(out_channels//len(dilations)), + nn.ReLU(inplace=True) + ) + + # Channel attention + total_channels = out_channels//len(dilations) * (len(dilations) + 1) + self.channel_attention = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(total_channels, total_channels//8, 1), + nn.ReLU(inplace=True), + nn.Conv2d(total_channels//8, total_channels, 1), + nn.Sigmoid() + ) + + # Final projection + self.project = nn.Sequential( + nn.Conv2d(total_channels, out_channels, 1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(inplace=True), + nn.Dropout2d(0.1) + ) + + def forward(self, x): + h, w = x.shape[-2:] + features = [] + + # Apply dilated convolutions + for conv in self.convs: + features.append(conv(x)) + + # Global pooling branch + global_feat = self.global_pool(x) + global_feat = F.interpolate(global_feat, size=(h, w), mode='bilinear', align_corners=False) + features.append(global_feat) + + # Combine features + combined = torch.cat(features, dim=1) + + # Apply channel attention + att = self.channel_attention(combined) + combined = combined * att + + return self.project(combined) + + +class RobustFPN(nn.Module): + """Robust FPN with attention gates and multi-scale fusion""" + + def __init__(self, in_channels_list, out_channels=256): + super().__init__() + + # Lateral connections + self.lateral_convs = nn.ModuleList([ + nn.Conv2d(in_ch, out_channels, 1, bias=False) + for in_ch in in_channels_list + ]) + + # Output convolutions + self.fpn_convs = nn.ModuleList([ + nn.Sequential( + nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(inplace=True) + ) + for _ in in_channels_list + ]) + + # Attention gates for better feature selection + self.attention_gates = nn.ModuleList([ + AttentionGate(out_channels, out_channels, out_channels//2) + for _ in range(len(in_channels_list)-1) + ]) + + # Multi-scale fusion weights + self.fusion_weights = nn.Parameter(torch.ones(len(in_channels_list))) + + def forward(self, features): + # Build laterals + laterals = [conv(feat) for conv, feat in zip(self.lateral_convs, features)] + + # Build top-down pathway with attention + for i in range(len(laterals) - 1, 0, -1): + upsampled = F.interpolate( + laterals[i], size=laterals[i-1].shape[-2:], + mode='bilinear', align_corners=False + ) + + # Apply attention gate + att_feat = self.attention_gates[i-1](upsampled, laterals[i-1]) + laterals[i-1] = laterals[i-1] + att_feat + + # Apply final convs + outputs = [conv(lateral) for conv, lateral in zip(self.fpn_convs, laterals)] + + # Multi-scale weighted feature fusion + target_size = outputs[0].shape[-2:] + fused_features = [] + + for i, feat in enumerate(outputs): + if feat.shape[-2:] != target_size: + feat = F.interpolate(feat, size=target_size, mode='bilinear', align_corners=False) + # Apply learnable weights + weighted_feat = feat * torch.sigmoid(self.fusion_weights[i]) + fused_features.append(weighted_feat) + + return sum(fused_features) + + +class SuperiorSegmentationHead(nn.Module): + """Superior segmentation head matching training architecture""" + + def __init__(self, in_channels_list, num_classes=7, feature_dim=256, dropout_rate=0.1): + super().__init__() + self.num_classes = num_classes + + # Robust FPN with attention + self.fpn = RobustFPN(in_channels_list, feature_dim) + + # Enhanced ASPP for multi-scale context + self.aspp = EnhancedASPP(feature_dim, feature_dim) + + # Multi-scale decoder for better detail preservation + self.decoder = nn.Sequential( + # First stage - preserve details + nn.Conv2d(feature_dim, feature_dim, 3, padding=1, bias=False), + nn.BatchNorm2d(feature_dim), + nn.ReLU(inplace=True), + + # Second stage - refine features + nn.Conv2d(feature_dim, feature_dim//2, 3, padding=1, bias=False), + nn.BatchNorm2d(feature_dim//2), + nn.ReLU(inplace=True), + + # Third stage - final refinement + nn.Conv2d(feature_dim//2, feature_dim//4, 3, padding=1, bias=False), + nn.BatchNorm2d(feature_dim//4), + nn.ReLU(inplace=True), + + nn.Dropout2d(dropout_rate), + nn.Conv2d(feature_dim//4, num_classes, 1) + ) + + # Boundary refinement branch + self.boundary_head = nn.Sequential( + nn.Conv2d(feature_dim, 64, 3, padding=1, bias=False), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True), + nn.Conv2d(64, 1, 1), + nn.Sigmoid() + ) + + self._init_weights() + + def _init_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, nn.BatchNorm2d): + nn.init.ones_(m.weight) + nn.init.zeros_(m.bias) + + def forward(self, features): + # Multi-scale feature fusion + fused_feature = self.fpn(features) + + # Apply enhanced ASPP + context_feature = self.aspp(fused_feature) + + # Generate segmentation + seg_logits = self.decoder(context_feature) + + # Generate boundary map for auxiliary loss + boundary_map = self.boundary_head(context_feature) + + return seg_logits, boundary_map + + +class DFineWithRobustSegmentation(nn.Module): + """D-FINE with robust segmentation matching training architecture""" + + def __init__(self, dfine_model, seg_head): + super().__init__() + self.dfine_model = dfine_model + self.seg_head = seg_head + + # Set to eval mode for inference + self.eval() + + def forward(self, x): + # Get backbone features + backbone_features = self.dfine_model.backbone(x) + + outputs = {} + + # Detection branch + det_outputs = self.dfine_model(x) + outputs.update(det_outputs) + + # Segmentation branch + seg_logits, boundary_map = self.seg_head(backbone_features) + + # Upsample to input resolution + seg_logits = F.interpolate( + seg_logits, size=x.shape[-2:], + mode='bilinear', align_corners=False + ) + + boundary_map = F.interpolate( + boundary_map, size=x.shape[-2:], + mode='bilinear', align_corners=False + ) + + outputs['segmentation'] = seg_logits + outputs['boundary'] = boundary_map + + return outputs + + +def get_actual_backbone_channels(model): + """Get actual backbone output channels""" + model.eval() + dummy_input = torch.randn(1, 3, 640, 640) + + with torch.no_grad(): + backbone_features = model.backbone(dummy_input) + + actual_channels = [feat.shape[1] for feat in backbone_features] + return actual_channels + + +def load_enhanced_segmentation_model(config_file, checkpoint_path, device): + """Load the enhanced segmentation model with correct architecture""" + print(f"Loading enhanced D-FINE model from config: {config_file}") + + # Load base D-FINE model + cfg = YAMLConfig(config_file) + base_model = cfg.model + + # Detect actual backbone channels + actual_channels = get_actual_backbone_channels(base_model) + print(f"Detected backbone channels: {actual_channels}") + + # Load checkpoint + print(f"Loading enhanced segmentation weights from: {checkpoint_path}") + checkpoint = torch.load(checkpoint_path, map_location='cpu') + + if 'model_state_dict' in checkpoint: + state_dict = checkpoint['model_state_dict'] + hyperparams = checkpoint.get('hyperparameters', {}) + else: + state_dict = checkpoint + hyperparams = {} + + # Get hyperparameters from checkpoint or use defaults from best results + feature_dim = hyperparams.get('feature_dim', 448) # From your best training results + dropout_rate = hyperparams.get('dropout_rate', 0.05) + + print(f"Using hyperparameters: feature_dim={feature_dim}, dropout_rate={dropout_rate}") + + # Create superior segmentation head with correct architecture + seg_head = SuperiorSegmentationHead( + in_channels_list=actual_channels, + num_classes=7, + feature_dim=feature_dim, + dropout_rate=dropout_rate + ) + + # Create enhanced combined model + model = DFineWithRobustSegmentation( + dfine_model=base_model, + seg_head=seg_head + ) + + # Load the state dict + try: + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=True) + print("✅ Model loaded successfully with strict=True") + except RuntimeError as e: + print(f"⚠️ Strict loading failed, using filtered loading") + # Filter state dict to only include matching keys with correct shapes + model_state_dict = model.state_dict() + filtered_state_dict = {} + + for key, value in state_dict.items(): + if key in model_state_dict and model_state_dict[key].shape == value.shape: + filtered_state_dict[key] = value + + missing_keys, unexpected_keys = model.load_state_dict(filtered_state_dict, strict=False) + print(f"✅ Model loaded with {len(filtered_state_dict)}/{len(state_dict)} weights") + + if missing_keys: + print(f"⚠️ Missing keys: {len(missing_keys)}") + if unexpected_keys: + print(f"⚠️ Unexpected keys: {len(unexpected_keys)}") + + model = model.to(device) + model.eval() + + print(f"Enhanced segmentation model loaded successfully on {device}") + return model + + +def get_body_part_colors(): + """Get colors for each body part class""" + return { + 0: (0, 0, 0), # background - black (transparent) + 1: (255, 0, 0), # head - red + 2: (0, 255, 0), # torso - green + 3: (0, 0, 255), # arms - blue + 4: (255, 255, 0), # hands - yellow + 5: (255, 0, 255), # legs - magenta + 6: (0, 255, 255), # feet - cyan + } + + +def get_body_part_names(): + """Get names for each body part class""" + return { + 0: 'background', + 1: 'head', + 2: 'torso', + 3: 'arms', + 4: 'hands', + 5: 'legs', + 6: 'feet' + } + + +def process_frame(model, frame_bgr, device, transforms_val, threshold=0.6): + """Process a single frame with the enhanced segmentation model""" + # Convert BGR to RGB and then to PIL + rgb_frame = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) + pil_image = Image.fromarray(rgb_frame) + + # Get original dimensions for detection postprocessing + w, h = pil_image.size + orig_size = torch.tensor([[w, h]]).to(device) + + # Transform image + im_data = transforms_val(pil_image).unsqueeze(0).to(device) + + with torch.no_grad(): + outputs = model(im_data) + + # Get segmentation results + seg_logits = outputs['segmentation'][0] # Remove batch dimension + seg_pred = torch.argmax(seg_logits, dim=0).cpu().numpy() + + # Get detection results if available + labels, boxes, scores = None, None, None + if 'pred_logits' in outputs and 'pred_boxes' in outputs: + # Process detection outputs (simplified, might need adjustment based on your postprocessor) + pred_logits = outputs['pred_logits'][0] + pred_boxes = outputs['pred_boxes'][0] + + # Apply softmax and get scores + pred_scores = torch.softmax(pred_logits, dim=-1) + max_scores, pred_labels = torch.max(pred_scores, dim=-1) + + # Filter by threshold and person class (assuming class 0 is person) + mask = (max_scores > threshold) & (pred_labels == 0) + + if mask.any(): + filtered_boxes = pred_boxes[mask] + filtered_scores = max_scores[mask] + filtered_labels = pred_labels[mask] + + # Convert boxes from normalized to pixel coordinates + filtered_boxes = filtered_boxes * torch.tensor([w, h, w, h]).to(device) + # Convert from cxcywh to xyxy + boxes_xyxy = torch.zeros_like(filtered_boxes) + boxes_xyxy[:, 0] = filtered_boxes[:, 0] - filtered_boxes[:, 2] / 2 # x1 + boxes_xyxy[:, 1] = filtered_boxes[:, 1] - filtered_boxes[:, 3] / 2 # y1 + boxes_xyxy[:, 2] = filtered_boxes[:, 0] + filtered_boxes[:, 2] / 2 # x2 + boxes_xyxy[:, 3] = filtered_boxes[:, 1] + filtered_boxes[:, 3] / 2 # y2 + + labels = [filtered_labels.cpu()] + boxes = [boxes_xyxy.cpu()] + scores = [filtered_scores.cpu()] + else: + labels, boxes, scores = [[]], [[]], [[]] + + return pil_image, seg_pred, labels, boxes, scores + + +def create_segmentation_overlay(pil_image, seg_pred, alpha=0.6): + """Create segmentation overlay on PIL image""" + colors = get_body_part_colors() + + # Create colored segmentation mask + h, w = seg_pred.shape + colored_mask = np.zeros((h, w, 3), dtype=np.uint8) + + for class_id, color in colors.items(): + if class_id == 0: # Skip background + continue + mask = seg_pred == class_id + colored_mask[mask] = color + + # Convert to PIL and resize to match original image + mask_pil = Image.fromarray(colored_mask).resize(pil_image.size, Image.NEAREST) + + # Create overlay + overlay = Image.blend(pil_image, mask_pil, alpha) + + return overlay + + +def draw_detections_and_segmentation(pil_image, seg_pred, labels, boxes, scores, + threshold=0.6, seg_alpha=0.6): + """Draw both detection boxes and segmentation overlay""" + + # First create segmentation overlay + result_image = create_segmentation_overlay(pil_image, seg_pred, seg_alpha) + + # Then draw detection boxes + draw = ImageDraw.Draw(result_image) + person_count = 0 + + if labels and boxes and scores: + # Process detection results (assuming batch size 1) + scores_per_image = scores[0] if scores[0] is not None and len(scores[0]) > 0 else [] + boxes_per_image = boxes[0] if boxes[0] is not None and len(boxes[0]) > 0 else [] + + for i in range(len(scores_per_image)): + score = scores_per_image[i].item() + if score > threshold: + person_count += 1 + box = boxes_per_image[i].round().int().tolist() + + # Draw bounding box + draw.rectangle(box, outline="red", width=3) + + # Draw label + text = f"Person: {score:.2f}" + text_y_pos = box[1] - 15 if box[1] > 20 else box[1] + 2 + draw.text((box[0] + 2, text_y_pos), text, fill="white") + + # Draw segmentation info + unique_parts = np.unique(seg_pred) + part_names = get_body_part_names() + active_parts = [part_names[part_id] for part_id in unique_parts if part_id != 0] + + if active_parts: + parts_text = f"Body parts: {', '.join(active_parts)}" + draw.text((10, result_image.size[1] - 30), parts_text, fill="white") + + # Enhanced info display + info_text = f"Robust Enhanced Model | People: {person_count} | Parts: {len(active_parts)}" + draw.text((10, 10), info_text, fill="white") + + return result_image, person_count, len(active_parts) + + +def check_ffmpeg(): + """Check if ffmpeg is available""" + try: + subprocess.run(["ffmpeg", "-version"], stdout=subprocess.PIPE, + stderr=subprocess.PIPE, check=True) + return True + except (FileNotFoundError, subprocess.CalledProcessError): + print("WARNING: ffmpeg not found. Please install ffmpeg.") + return False + + +def process_video(model, input_path, output_path, device, threshold, seg_alpha, + resize_factor, output_fps_target, crf, preset): + """Process video with enhanced segmentation model""" + + if not check_ffmpeg(): + raise RuntimeError("ffmpeg is required for video output") + + cap = cv2.VideoCapture(input_path) + if not cap.isOpened(): + raise ValueError(f"Could not open video: {input_path}") + + original_fps = cap.get(cv2.CAP_PROP_FPS) + original_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + original_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + + # Calculate output dimensions + if resize_factor != 1.0: + output_width = int(original_width * resize_factor) + output_height = int(original_height * resize_factor) + else: + output_width = original_width + output_height = original_height + + # Model transforms + transforms_val = T.Compose([ + T.Resize((640, 640)), + T.ToTensor(), + T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + ]) + + output_fps = output_fps_target if output_fps_target is not None else original_fps + if output_fps <= 0: + output_fps = 30 + + print(f"Input: {input_path} ({original_width}x{original_height} @ {original_fps:.2f} FPS)") + print(f"Output: {output_path} ({output_width}x{output_height} @ {output_fps:.2f} FPS)") + print(f"Device: {device}, Detection threshold: {threshold}, Seg alpha: {seg_alpha}") + print(f"Using ROBUST ENHANCED architecture with superior segmentation head") + + # Setup ffmpeg + command = [ + 'ffmpeg', '-y', '-f', 'rawvideo', '-vcodec', 'rawvideo', + '-s', f'{output_width}x{output_height}', '-pix_fmt', 'bgr24', + '-r', str(output_fps), '-i', '-', '-c:v', 'libx264', + '-pix_fmt', 'yuv420p', '-crf', str(crf), '-preset', preset, output_path, + ] + + ffmpeg_process = subprocess.Popen(command, stdin=subprocess.PIPE, + stdout=subprocess.PIPE, stderr=subprocess.PIPE) + + frame_idx = 0 + total_people = 0 + total_parts_detected = 0 + processing_start_time = time.time() + stdin_broken = False + + try: + while True: + ret, frame_bgr_original = cap.read() + if not ret: + break + + loop_start_time = time.time() + + # Resize frame if needed + if resize_factor != 1.0: + frame_to_process = cv2.resize(frame_bgr_original, + (output_width, output_height)) + else: + frame_to_process = frame_bgr_original + + # Process frame with enhanced model + pil_image, seg_pred, labels, boxes, scores = process_frame( + model, frame_to_process, device, transforms_val, threshold + ) + + # Draw results with enhanced visualization + result_pil, person_count, parts_count = draw_detections_and_segmentation( + pil_image, seg_pred, labels, boxes, scores, threshold, seg_alpha + ) + + total_people += person_count + total_parts_detected += parts_count + + # Convert back to BGR for ffmpeg + output_frame_bgr = cv2.cvtColor(np.array(result_pil), cv2.COLOR_RGB2BGR) + + # Ensure correct dimensions + if (output_frame_bgr.shape[1] != output_width or + output_frame_bgr.shape[0] != output_height): + output_frame_bgr = cv2.resize(output_frame_bgr, + (output_width, output_height)) + + # Write to ffmpeg + try: + if ffmpeg_process.stdin and not ffmpeg_process.stdin.closed: + ffmpeg_process.stdin.write(output_frame_bgr.tobytes()) + else: + if not stdin_broken: + print("\nffmpeg stdin closed") + stdin_broken = True + break + except (BrokenPipeError, IOError) as e: + if not stdin_broken: + print(f"\nffmpeg stdin error: {e}") + stdin_broken = True + break + + frame_idx += 1 + + # Progress update + if frame_idx % max(1, int(output_fps)) == 0 or frame_idx == 1: + elapsed = time.time() - processing_start_time + fps = frame_idx / elapsed if elapsed > 0 else 0 + eta = (total_frames - frame_idx) / fps if fps > 0 else 0 + + progress = (frame_idx / total_frames * 100) if total_frames > 0 else 0 + + sys.stdout.write( + f"\rFrame: {frame_idx}/{total_frames} ({progress:.1f}%) | " + f"FPS: {fps:.1f} | ETA: {eta:.0f}s | " + f"People: {person_count} | Parts: {parts_count} " + ) + sys.stdout.flush() + + finally: + cap.release() + sys.stdout.write("\n") + + # Robust ffmpeg cleanup + if ffmpeg_process.stdin: + if not ffmpeg_process.stdin.closed and not stdin_broken: + try: + ffmpeg_process.stdin.close() + except (BrokenPipeError, IOError) as e: + print(f"Error closing ffmpeg stdin: {e}") + + stdout_data, stderr_data = None, None + try: + stdout_data, stderr_data = ffmpeg_process.communicate(timeout=30) + except subprocess.TimeoutExpired: + print("ffmpeg process timed out. Forcing termination.") + ffmpeg_process.kill() + stdout_data, stderr_data = ffmpeg_process.communicate() + except ValueError as e: + print(f"ValueError during ffmpeg.communicate(): {e}") + if ffmpeg_process.poll() is None: + ffmpeg_process.wait(timeout=5) + except Exception as e: + print(f"Error during ffmpeg.communicate(): {e}") + if ffmpeg_process.poll() is None: + ffmpeg_process.kill() + ffmpeg_process.wait(timeout=5) + + processing_duration = time.time() - processing_start_time + final_fps = frame_idx / processing_duration if processing_duration > 0 else 0 + + print(f"Processed {frame_idx} frames in {processing_duration:.2f}s ({final_fps:.2f} FPS)") + print(f"Total people detected: {total_people}") + if frame_idx > 0: + print(f"Average body parts per frame: {total_parts_detected/frame_idx:.1f}") + + if ffmpeg_process.returncode is not None and ffmpeg_process.returncode == 0: + print("ffmpeg process completed successfully.") + elif ffmpeg_process.returncode is not None: + print(f"ffmpeg process exited with code {ffmpeg_process.returncode}") + + if os.path.exists(output_path): + size_mb = os.path.getsize(output_path) / 1024 / 1024 + print(f"✅ Robust enhanced segmentation output saved: {output_path} ({size_mb:.2f} MB)") + else: + print(f"❌ Output file not found: {output_path}") + + +def main(): + args = parse_args() + + if args.device == 'cuda' and not torch.cuda.is_available(): + print("CUDA not available, using CPU") + device = torch.device('cpu') + else: + device = torch.device(args.device) + + try: + if not os.path.isfile(args.input): + raise FileNotFoundError(f"Input file not found: {args.input}") + + if not os.path.isfile(args.model): + raise FileNotFoundError(f"Model checkpoint not found: {args.model}") + + model = load_enhanced_segmentation_model(args.config, args.model, device) + + process_video( + model=model, input_path=args.input, output_path=args.output, + device=device, threshold=args.threshold, seg_alpha=args.seg_alpha, + resize_factor=args.resize_factor, output_fps_target=args.fps, + crf=args.crf, preset=args.preset + ) + + print("\n🎉 Robust enhanced segmentation processing complete!") + + except Exception as e: + print(f"\nError: {e}") + import traceback + traceback.print_exc() + sys.exit(1) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/segmentation_sivert/legacy/train_standard_model.py b/segmentation_sivert/legacy/train_standard_model.py new file mode 100644 index 00000000..f0b489a3 --- /dev/null +++ b/segmentation_sivert/legacy/train_standard_model.py @@ -0,0 +1,780 @@ +#!/usr/bin/env python3 +""" +Fixed Enhanced D-FINE Segmentation Training Script +Stable multi-scale training without batch collation issues +""" + +import os +import sys +import yaml +import argparse +import random +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader +import torch.nn.functional as F +from tqdm import tqdm +import wandb +import numpy as np +from pathlib import Path + +# Add src to path +sys.path.append('src') + +from src.core import YAMLConfig + + +class PascalPersonPartsDataset(torch.utils.data.Dataset): + """Fixed Pascal Person Parts Dataset with stable multi-scale training""" + + def __init__(self, root_dir, split='train', image_size=640, multi_scale=False): + self.root_dir = root_dir + self.split = split + self.base_image_size = image_size + self.multi_scale = multi_scale + + # Get image paths + self.img_dir = os.path.join(root_dir, 'images', split) + self.mask_dir = os.path.join(root_dir, 'masks', split) + + self.image_names = [f for f in os.listdir(self.img_dir) if f.endswith('.jpg')] + + print(f"📊 Loaded {len(self.image_names)} {split} samples") + + # Normalization + self.mean = torch.tensor([0.485, 0.456, 0.406]) + self.std = torch.tensor([0.229, 0.224, 0.225]) + + # Multi-scale options - but we'll resize everything to base_image_size for batching + self.scales = [0.75, 0.85, 1.0, 1.15, 1.3] if multi_scale else [1.0] + + def __len__(self): + return len(self.image_names) + + def __getitem__(self, idx): + img_name = self.image_names[idx] + mask_name = img_name.replace('.jpg', '.png') + + # Load image and mask + img_path = os.path.join(self.img_dir, img_name) + mask_path = os.path.join(self.mask_dir, mask_name) + + # Load image + import cv2 + image = cv2.imread(img_path) + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + + # Load mask + mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) + + # Multi-scale training with random crop/resize + if self.split == 'train' and self.multi_scale: + # Choose random scale + scale = random.choice(self.scales) + scaled_size = int(self.base_image_size * scale) + + # Resize to scaled size first + image = cv2.resize(image, (scaled_size, scaled_size)) + mask = cv2.resize(mask, (scaled_size, scaled_size), interpolation=cv2.INTER_NEAREST) + + # Random crop/pad to base_image_size for consistent batching + if scaled_size > self.base_image_size: + # Random crop + start_x = random.randint(0, scaled_size - self.base_image_size) + start_y = random.randint(0, scaled_size - self.base_image_size) + image = image[start_y:start_y+self.base_image_size, start_x:start_x+self.base_image_size] + mask = mask[start_y:start_y+self.base_image_size, start_x:start_x+self.base_image_size] + elif scaled_size < self.base_image_size: + # Center pad + pad_x = (self.base_image_size - scaled_size) // 2 + pad_y = (self.base_image_size - scaled_size) // 2 + image = cv2.copyMakeBorder( + image, pad_y, self.base_image_size-scaled_size-pad_y, + pad_x, self.base_image_size-scaled_size-pad_x, cv2.BORDER_REFLECT + ) + mask = cv2.copyMakeBorder( + mask, pad_y, self.base_image_size-scaled_size-pad_y, + pad_x, self.base_image_size-scaled_size-pad_x, cv2.BORDER_CONSTANT, value=0 + ) + else: + # Standard resize for val or non-multi-scale + image = cv2.resize(image, (self.base_image_size, self.base_image_size)) + mask = cv2.resize(mask, (self.base_image_size, self.base_image_size), interpolation=cv2.INTER_NEAREST) + + # Training augmentations + if self.split == 'train': + # Color augmentations + if random.random() < 0.6: + alpha = random.uniform(0.8, 1.2) # contrast + beta = random.randint(-20, 20) # brightness + image = cv2.convertScaleAbs(image, alpha=alpha, beta=beta) + + # Horizontal flip + if random.random() < 0.5: + image = cv2.flip(image, 1) + mask = cv2.flip(mask, 1) + + # Gaussian blur (helps with distance) + if random.random() < 0.3: + kernel_size = random.choice([3, 5]) + image = cv2.GaussianBlur(image, (kernel_size, kernel_size), 0) + + # Random rotation (small angles) + if random.random() < 0.4: + angle = random.uniform(-10, 10) + h, w = image.shape[:2] + center = (w//2, h//2) + M = cv2.getRotationMatrix2D(center, angle, 1.0) + image = cv2.warpAffine(image, M, (w, h), borderMode=cv2.BORDER_REFLECT) + mask = cv2.warpAffine(mask, M, (w, h), borderMode=cv2.BORDER_CONSTANT, borderValue=0) + + # Convert to tensor and normalize + image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0 + image = (image - self.mean.view(3, 1, 1)) / self.std.view(3, 1, 1) + + mask = torch.from_numpy(mask).long() + + # Now all tensors have consistent size for batching + return image, mask + + +class SimplifiedASPP(nn.Module): + """Simplified ASPP for better speed/accuracy balance""" + + def __init__(self, in_channels, out_channels, dilations=[1, 6, 12]): + super().__init__() + + self.convs = nn.ModuleList() + for dilation in dilations: + if dilation == 1: + conv = nn.Conv2d(in_channels, out_channels//len(dilations), 1, bias=False) + else: + conv = nn.Conv2d(in_channels, out_channels//len(dilations), 3, + padding=dilation, dilation=dilation, bias=False) + self.convs.append(nn.Sequential( + conv, + nn.BatchNorm2d(out_channels//len(dilations)), + nn.ReLU(inplace=True) + )) + + # Global average pooling branch + self.global_pool = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(in_channels, out_channels//len(dilations), 1, bias=False), + nn.BatchNorm2d(out_channels//len(dilations)), + nn.ReLU(inplace=True) + ) + + # Final projection + total_channels = out_channels//len(dilations) * (len(dilations) + 1) + self.project = nn.Sequential( + nn.Conv2d(total_channels, out_channels, 1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(inplace=True), + nn.Dropout2d(0.1) + ) + + def forward(self, x): + h, w = x.shape[-2:] + features = [] + + # Apply dilated convolutions + for conv in self.convs: + features.append(conv(x)) + + # Global pooling branch + global_feat = self.global_pool(x) + global_feat = F.interpolate(global_feat, size=(h, w), mode='bilinear', align_corners=False) + features.append(global_feat) + + # Combine and project + combined = torch.cat(features, dim=1) + return self.project(combined) + + +class ImprovedFPN(nn.Module): + """Improved FPN with better feature fusion""" + + def __init__(self, in_channels_list, out_channels=256): + super().__init__() + + # Lateral connections + self.lateral_convs = nn.ModuleList([ + nn.Conv2d(in_ch, out_channels, 1, bias=False) + for in_ch in in_channels_list + ]) + + # Output convolutions + self.fpn_convs = nn.ModuleList([ + nn.Sequential( + nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(inplace=True) + ) + for _ in in_channels_list + ]) + + # Feature fusion weights + self.fusion_weights = nn.Parameter(torch.ones(len(in_channels_list))) + + def forward(self, features): + # Build laterals + laterals = [conv(feat) for conv, feat in zip(self.lateral_convs, features)] + + # Build top-down pathway + for i in range(len(laterals) - 1, 0, -1): + upsampled = F.interpolate( + laterals[i], size=laterals[i-1].shape[-2:], + mode='bilinear', align_corners=False + ) + laterals[i-1] = laterals[i-1] + upsampled + + # Apply final convs + outputs = [conv(lateral) for conv, lateral in zip(self.fpn_convs, laterals)] + + # Weighted feature fusion + target_size = outputs[0].shape[-2:] + fused_features = [] + + for i, feat in enumerate(outputs): + if feat.shape[-2:] != target_size: + feat = F.interpolate(feat, size=target_size, mode='bilinear', align_corners=False) + fused_features.append(feat * self.fusion_weights[i]) + + # Return highest resolution feature with weighted fusion + return sum(fused_features) + + +class EnhancedSegmentationHead(nn.Module): + """Enhanced but stable segmentation head""" + + def __init__(self, in_channels_list, num_classes=7, feature_dim=384, dropout_rate=0.1): + super().__init__() + self.num_classes = num_classes + + print(f"🏗️ Creating enhanced segmentation head with channels: {in_channels_list}") + + # Improved FPN + self.fpn = ImprovedFPN(in_channels_list, feature_dim) + + # Simplified ASPP for multi-scale context + self.aspp = SimplifiedASPP(feature_dim, feature_dim) + + # Enhanced decoder + self.decoder = nn.Sequential( + nn.Conv2d(feature_dim, feature_dim, 3, padding=1, bias=False), + nn.BatchNorm2d(feature_dim), + nn.ReLU(inplace=True), + + nn.Conv2d(feature_dim, feature_dim//2, 3, padding=1, bias=False), + nn.BatchNorm2d(feature_dim//2), + nn.ReLU(inplace=True), + + nn.Dropout2d(dropout_rate), + nn.Conv2d(feature_dim//2, num_classes, 1) + ) + + self._init_weights() + + def _init_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, nn.BatchNorm2d): + nn.init.ones_(m.weight) + nn.init.zeros_(m.bias) + + def forward(self, features): + # Multi-scale feature fusion + fused_feature = self.fpn(features) + + # Apply simplified ASPP + context_feature = self.aspp(fused_feature) + + # Generate segmentation + seg_logits = self.decoder(context_feature) + + return seg_logits + + +class DFineWithSegmentation(nn.Module): + """D-FINE with enhanced segmentation""" + + def __init__(self, dfine_model, seg_head, freeze_detection=True): + super().__init__() + self.dfine_model = dfine_model + self.seg_head = seg_head + self.freeze_detection = freeze_detection + + if freeze_detection: + self._freeze_detection() + + def _freeze_detection(self): + frozen_params = 0 + total_params = 0 + + for name, param in self.named_parameters(): + total_params += param.numel() + if 'seg_head' not in name: + param.requires_grad = False + frozen_params += param.numel() + + trainable_params = total_params - frozen_params + + print(f"🔒 Frozen detection components:") + print(f" 📊 Total parameters: {total_params:,}") + print(f" ❄️ Frozen parameters: {frozen_params:,} ({frozen_params/total_params*100:.1f}%)") + print(f" 🎯 Trainable parameters: {trainable_params:,} ({trainable_params/total_params*100:.1f}%)") + + def forward(self, x, targets=None): + # Get backbone features + backbone_features = self.dfine_model.backbone(x) + + outputs = {} + + # Detection branch (frozen during training) + if not self.training or not self.freeze_detection: + with torch.no_grad() if self.freeze_detection else torch.enable_grad(): + det_outputs = self.dfine_model(x) + outputs.update(det_outputs) + + # Segmentation branch + seg_logits = self.seg_head(backbone_features) + + # Upsample to input resolution + seg_logits = F.interpolate( + seg_logits, size=x.shape[-2:], + mode='bilinear', align_corners=False + ) + + outputs['segmentation'] = seg_logits + + return outputs + + +class FocalLoss(nn.Module): + """Focal Loss for handling class imbalance""" + + def __init__(self, alpha=0.25, gamma=2.0, ignore_index=255): + super().__init__() + self.alpha = alpha + self.gamma = gamma + self.ignore_index = ignore_index + + def forward(self, inputs, targets): + ce_loss = F.cross_entropy(inputs, targets, ignore_index=self.ignore_index, reduction='none') + pt = torch.exp(-ce_loss) + focal_loss = self.alpha * (1 - pt) ** self.gamma * ce_loss + return focal_loss.mean() + + +class DiceLoss(nn.Module): + """Dice Loss for better boundary handling""" + + def __init__(self, smooth=1.0, ignore_index=255): + super().__init__() + self.smooth = smooth + self.ignore_index = ignore_index + + def forward(self, inputs, targets): + # Create mask for valid pixels + valid_mask = (targets != self.ignore_index) + + # Apply softmax to get probabilities + inputs_soft = F.softmax(inputs, dim=1) + + # One-hot encode targets + targets_one_hot = F.one_hot(targets * valid_mask.long(), num_classes=inputs.shape[1]) + targets_one_hot = targets_one_hot.permute(0, 3, 1, 2).float() + + # Apply valid mask + inputs_soft = inputs_soft * valid_mask.unsqueeze(1).float() + targets_one_hot = targets_one_hot * valid_mask.unsqueeze(1).float() + + # Calculate dice coefficient + intersection = (inputs_soft * targets_one_hot).sum(dim=(2, 3)) + union = inputs_soft.sum(dim=(2, 3)) + targets_one_hot.sum(dim=(2, 3)) + + dice = (2 * intersection + self.smooth) / (union + self.smooth) + return 1 - dice.mean() + + +class AdvancedSegmentationLoss(nn.Module): + """Combined advanced loss for better segmentation""" + + def __init__(self, num_classes=7, focal_alpha=0.25, focal_gamma=2.0, ignore_index=255): + super().__init__() + + self.focal_loss = FocalLoss(focal_alpha, focal_gamma, ignore_index) + self.dice_loss = DiceLoss(ignore_index=ignore_index) + + def forward(self, pred, target): + # Focal loss (main loss) + focal = self.focal_loss(pred, target) + + # Dice loss (boundary preservation) + dice = self.dice_loss(pred, target) + + # Standard CE loss + ce = F.cross_entropy(pred, target, ignore_index=255) + + # Combine losses + total_loss = focal + 0.4 * dice + 0.2 * ce + + return total_loss, { + 'focal': focal.item(), + 'dice': dice.item(), + 'ce': ce.item(), + 'total': total_loss.item() + } + + +def compute_miou(pred, target, num_classes=7): + """Compute mean IoU""" + pred = torch.argmax(pred, dim=1) + + ious = [] + for cls in range(num_classes): + pred_cls = (pred == cls) + target_cls = (target == cls) + + intersection = (pred_cls & target_cls).sum().float() + union = (pred_cls | target_cls).sum().float() + + if union > 0: + ious.append(intersection / union) + + return torch.tensor(ious).mean() if ious else torch.tensor(0.0) + + +def get_actual_backbone_channels(model): + """Get actual backbone output channels""" + model.eval() + dummy_input = torch.randn(1, 3, 640, 640) + + with torch.no_grad(): + backbone_features = model.backbone(dummy_input) + + actual_channels = [feat.shape[1] for feat in backbone_features] + return actual_channels + + +def load_pretrained_dfine(config_path, checkpoint_path): + """Load pretrained D-FINE model""" + print(f"🚀 Loading pretrained D-FINE from {checkpoint_path}") + + cfg = YAMLConfig(config_path) + model = cfg.model + + checkpoint = torch.load(checkpoint_path, map_location='cpu') + + if 'ema' in checkpoint and 'module' in checkpoint['ema']: + state_dict = checkpoint['ema']['module'] + elif 'model' in checkpoint: + state_dict = checkpoint['model'] + else: + state_dict = checkpoint + + model.load_state_dict(state_dict, strict=False) + print(f"✅ Loaded pretrained D-FINE model") + return model + + +def create_enhanced_segmentation_model(config_path, checkpoint_path, num_classes=7): + """Create enhanced segmentation model""" + + det_model = load_pretrained_dfine(config_path, checkpoint_path) + backbone_channels = get_actual_backbone_channels(det_model) + + print(f"🏗️ Using backbone channels: {backbone_channels}") + + # Create enhanced segmentation head + seg_head = EnhancedSegmentationHead( + in_channels_list=backbone_channels, + num_classes=num_classes, + feature_dim=256 + ) + + model = DFineWithSegmentation( + dfine_model=det_model, + seg_head=seg_head, + freeze_detection=True + ) + + return model + + +def unfreeze_backbone_gradually(model, optimizer, epoch): + """Gradually unfreeze backbone layers""" + + if epoch == 35: + print("🔓 Unfreezing backbone stage 4 (highest level features)") + new_params = [] + for name, param in model.dfine_model.backbone.named_parameters(): + if 'stages.3' in name: # Stage 4 of HGNetv2 + param.requires_grad = True + new_params.append(param) + + if new_params: + optimizer.add_param_group({ + 'params': new_params, + 'lr': optimizer.param_groups[0]['lr'] * 0.1 # 10x lower LR + }) + print(f" Added {len(new_params)} parameters to optimizer") + + elif epoch == 45: + print("🔓 Unfreezing backbone stage 3") + new_params = [] + for name, param in model.dfine_model.backbone.named_parameters(): + if 'stages.2' in name: + param.requires_grad = True + new_params.append(param) + + if new_params: + optimizer.add_param_group({ + 'params': new_params, + 'lr': optimizer.param_groups[0]['lr'] * 0.1 + }) + print(f" Added {len(new_params)} parameters to optimizer") + + +def train_epoch(model, dataloader, criterion, optimizer, device, epoch): + """Enhanced training epoch""" + model.train() + + total_loss = 0 + total_miou = 0 + loss_components = {'focal': 0, 'dice': 0, 'ce': 0, 'total': 0} + num_batches = len(dataloader) + + pbar = tqdm(dataloader, desc=f'Epoch {epoch}') + + for batch_idx, (images, masks) in enumerate(pbar): + images = images.to(device) + masks = masks.to(device) + + optimizer.zero_grad() + + # Forward pass + outputs = model(images) + seg_pred = outputs['segmentation'] + + # Calculate loss + loss, loss_dict = criterion(seg_pred, masks) + + # Backward pass + loss.backward() + + # Gradient clipping for stability + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) + + optimizer.step() + + # Calculate metrics + with torch.no_grad(): + miou = compute_miou(seg_pred, masks) + + total_loss += loss.item() + total_miou += miou.item() + + # Accumulate loss components + for key in loss_components: + loss_components[key] += loss_dict.get(key, 0) + + # Update progress bar + pbar.set_postfix({ + 'Loss': f'{loss.item():.4f}', + 'mIoU': f'{miou.item():.3f}', + 'FL': f'{loss_dict.get("focal", 0):.3f}', + 'DL': f'{loss_dict.get("dice", 0):.3f}' + }) + + # Log to wandb + if batch_idx % 20 == 0: + wandb.log({ + 'train/batch_loss': loss.item(), + 'train/batch_miou': miou.item(), + 'train/focal_loss': loss_dict.get('focal', 0), + 'train/dice_loss': loss_dict.get('dice', 0), + 'train/ce_loss': loss_dict.get('ce', 0), + 'train/lr': optimizer.param_groups[0]['lr'] + }) + + # Average metrics + avg_loss = total_loss / num_batches + avg_miou = total_miou / num_batches + avg_loss_components = {k: v / num_batches for k, v in loss_components.items()} + + return avg_loss, avg_miou, avg_loss_components + + +def validate_epoch(model, dataloader, criterion, device): + """Enhanced validation""" + model.eval() + + total_loss = 0 + total_miou = 0 + loss_components = {'focal': 0, 'dice': 0, 'ce': 0, 'total': 0} + num_batches = len(dataloader) + + with torch.no_grad(): + for images, masks in tqdm(dataloader, desc='Validation'): + images = images.to(device) + masks = masks.to(device) + + outputs = model(images) + seg_pred = outputs['segmentation'] + + loss, loss_dict = criterion(seg_pred, masks) + miou = compute_miou(seg_pred, masks) + + total_loss += loss.item() + total_miou += miou.item() + + for key in loss_components: + loss_components[key] += loss_dict.get(key, 0) + + avg_loss = total_loss / num_batches + avg_miou = total_miou / num_batches + avg_loss_components = {k: v / num_batches for k, v in loss_components.items()} + + return avg_loss, avg_miou, avg_loss_components + + +def main(): + parser = argparse.ArgumentParser(description='Fixed Enhanced D-FINE Segmentation Training') + parser.add_argument('--config', default='configs/dfine_hgnetv2_x_obj2coco.yml') + parser.add_argument('--checkpoint', default='models/dfine_x_obj2coco.pth') + parser.add_argument('--dataset', default='datasets/pascal_person_parts') + parser.add_argument('--batch_size', type=int, default=8) # Reduced for stability + parser.add_argument('--epochs', type=int, default=80) + parser.add_argument('--lr', type=float, default=5e-4) + parser.add_argument('--image_size', type=int, default=640) + parser.add_argument('--num_workers', type=int, default=4) # Reduced for stability + parser.add_argument('--output_dir', default='outputs/dfine_segmentation_enhanced') + + args = parser.parse_args() + + os.makedirs(args.output_dir, exist_ok=True) + + # Initialize wandb + wandb.init( + project='dfine-enhanced-segmentation-fixed', + config=vars(args), + name=f'fixed_enhanced_dfine_seg_{args.epochs}ep' + ) + + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + print(f"🚀 Using device: {device}") + + # Create enhanced model + print("🏗️ Creating enhanced segmentation model...") + model = create_enhanced_segmentation_model( + args.config, args.checkpoint, num_classes=7 + ) + model = model.to(device) + + # Enhanced datasets with FIXED multi-scale training + print("📊 Loading enhanced datasets...") + train_dataset = PascalPersonPartsDataset( + args.dataset, split='train', image_size=args.image_size, multi_scale=True + ) + val_dataset = PascalPersonPartsDataset( + args.dataset, split='val', image_size=args.image_size, multi_scale=False + ) + + train_loader = DataLoader( + train_dataset, batch_size=args.batch_size, shuffle=True, + num_workers=args.num_workers, pin_memory=True + ) + val_loader = DataLoader( + val_dataset, batch_size=args.batch_size, shuffle=False, + num_workers=args.num_workers, pin_memory=True + ) + + # Advanced loss and optimizer + criterion = AdvancedSegmentationLoss(num_classes=7) + + seg_params = [p for p in model.seg_head.parameters() if p.requires_grad] + optimizer = optim.AdamW(seg_params, lr=args.lr, weight_decay=1e-4) + + # Advanced scheduler + scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts( + optimizer, T_0=15, T_mult=2, eta_min=1e-6 + ) + + print(f"🎯 Training {sum(p.numel() for p in seg_params):,} segmentation parameters") + + # Training loop + best_miou = 0 + + for epoch in range(args.epochs): + print(f"\n📅 Epoch {epoch+1}/{args.epochs}") + + # Gradual backbone unfreezing + unfreeze_backbone_gradually(model, optimizer, epoch) + + # Train + train_loss, train_miou, train_loss_components = train_epoch( + model, train_loader, criterion, optimizer, device, epoch + ) + + # Validate + val_loss, val_miou, val_loss_components = validate_epoch( + model, val_loader, criterion, device + ) + + # Update scheduler + scheduler.step() + + # Log comprehensive metrics + wandb.log({ + 'epoch': epoch, + 'train/loss': train_loss, + 'train/miou': train_miou, + 'train/focal_loss': train_loss_components['focal'], + 'train/dice_loss': train_loss_components['dice'], + 'train/ce_loss': train_loss_components['ce'], + 'val/loss': val_loss, + 'val/miou': val_miou, + 'val/focal_loss': val_loss_components['focal'], + 'val/dice_loss': val_loss_components['dice'], + 'val/ce_loss': val_loss_components['ce'], + 'lr': optimizer.param_groups[0]['lr'] + }) + + print(f"📊 Train - Loss: {train_loss:.4f}, mIoU: {train_miou:.3f}") + print(f"📊 Val - Loss: {val_loss:.4f}, mIoU: {val_miou:.3f}") + + # Save best model + if val_miou > best_miou: + best_miou = val_miou + torch.save({ + 'epoch': epoch, + 'model_state_dict': model.state_dict(), + 'optimizer_state_dict': optimizer.state_dict(), + 'scheduler_state_dict': scheduler.state_dict(), + 'best_miou': best_miou, + 'config': vars(args) + }, f'{args.output_dir}/best_model.pth') + print(f"💾 Saved best model with mIoU: {best_miou:.3f}") + + # Save checkpoint every 10 epochs + if (epoch + 1) % 10 == 0: + torch.save({ + 'epoch': epoch, + 'model_state_dict': model.state_dict(), + 'optimizer_state_dict': optimizer.state_dict(), + 'scheduler_state_dict': scheduler.state_dict(), + 'miou': val_miou, + }, f'{args.output_dir}/checkpoint_epoch_{epoch+1}.pth') + + print(f"🎉 Training completed! Best mIoU: {best_miou:.3f}") + wandb.finish() + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/segmentation_sivert/legacy/train_standard_model_optuna.py b/segmentation_sivert/legacy/train_standard_model_optuna.py new file mode 100644 index 00000000..a2dc6ba5 --- /dev/null +++ b/segmentation_sivert/legacy/train_standard_model_optuna.py @@ -0,0 +1,1345 @@ +#!/usr/bin/env python3 +""" +Enhanced D-FINE Segmentation Training Script with Superior Robustness +Optimized for distant objects, legs/feet segmentation, and overall robustness +Features advanced augmentations, multi-scale training, and progressive strategies +""" + +import os +import sys +import yaml +import argparse +import random +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader +import torch.nn.functional as F +from tqdm import tqdm +import wandb +import numpy as np +from pathlib import Path +import json +import optuna +from optuna.integration.wandb import WeightsAndBiasesCallback +import threading +import cv2 +import albumentations as A +from albumentations.pytorch import ToTensorV2 +import math + +# Add src to path +sys.path.append('src') + +from src.core import YAMLConfig + +# Global variables to track best across trials +best_global_miou = 0.0 +best_global_lock = threading.Lock() + + +def set_seed(seed=42): + """Set seed for reproducibility""" + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + os.environ['PYTHONHASHSEED'] = str(seed) + print(f"🌱 Set random seed to {seed}") + + +class AdvancedPascalPersonPartsDataset(torch.utils.data.Dataset): + """Advanced Pascal Person Parts Dataset with robust augmentations for distant objects and lower body""" + + def __init__(self, root_dir, split='train', image_size=640, multi_scale=False, progressive_epoch=0): + self.root_dir = root_dir + self.split = split + self.base_image_size = image_size + self.multi_scale = multi_scale + self.progressive_epoch = progressive_epoch + + # Get image paths + self.img_dir = os.path.join(root_dir, 'images', split) + self.mask_dir = os.path.join(root_dir, 'masks', split) + + self.image_names = [f for f in os.listdir(self.img_dir) if f.endswith('.jpg')] + + print(f"📊 Loaded {len(self.image_names)} {split} samples") + + # Normalization + self.mean = torch.tensor([0.485, 0.456, 0.406]) + self.std = torch.tensor([0.229, 0.224, 0.225]) + + # Enhanced multi-scale options for better distant object handling + self.scales = [0.5, 0.6, 0.75, 0.85, 1.0, 1.15, 1.3, 1.5] if multi_scale else [1.0] + + # Progressive training: start smaller, go bigger + self.current_size = self._get_progressive_size() + + # Advanced augmentations + self.setup_augmentations() + + def _get_progressive_size(self): + """Progressive resizing: start smaller for better distant object learning""" + if self.split != 'train': + return self.base_image_size + + # Progressive sizing over epochs + min_size = int(self.base_image_size * 0.75) # Start at 75% of target + max_size = self.base_image_size + + # Gradually increase size over first 40 epochs + progress = min(self.progressive_epoch / 40.0, 1.0) + current_size = int(min_size + (max_size - min_size) * progress) + + # Ensure divisible by 32 for model compatibility + return ((current_size + 31) // 32) * 32 + + def setup_augmentations(self): + """Setup advanced augmentations for robustness""" + if self.split == 'train': + self.spatial_transform = A.Compose([ + # Geometric augmentations for robustness + A.HorizontalFlip(p=0.5), + A.ShiftScaleRotate( + shift_limit=0.1, scale_limit=0.2, rotate_limit=15, + border_mode=cv2.BORDER_REFLECT, p=0.7 + ), + # Perspective transform for distant object simulation + A.Perspective(scale=(0.05, 0.15), p=0.3), + # Elastic transform for natural deformation + A.ElasticTransform( + alpha=1, sigma=50, alpha_affine=50, + border_mode=cv2.BORDER_REFLECT, p=0.2 + ), + # Grid distortion for robustness + A.GridDistortion(num_steps=5, distort_limit=0.1, p=0.2), + # Optical distortion + A.OpticalDistortion(distort_limit=0.1, shift_limit=0.05, p=0.2), + ], additional_targets={'mask': 'mask'}) + + self.color_transform = A.Compose([ + # Enhanced color augmentations + A.RandomBrightnessContrast( + brightness_limit=0.3, contrast_limit=0.3, p=0.8 + ), + A.HueSaturationValue( + hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20, p=0.6 + ), + A.CLAHE(clip_limit=4.0, tile_grid_size=(8, 8), p=0.4), + A.RandomGamma(gamma_limit=(80, 120), p=0.4), + A.ChannelShuffle(p=0.1), + # Weather and lighting simulation + A.RandomShadow( + shadow_roi=(0, 0.5, 1, 1), + num_shadows_lower=1, num_shadows_upper=2, p=0.3 + ), + A.RandomSunFlare( + flare_roi=(0, 0, 1, 0.5), + angle_lower=0, angle_upper=1, + num_flare_circles_lower=6, num_flare_circles_upper=10, p=0.1 + ), + ]) + + self.noise_transform = A.Compose([ + # Noise and blur for distant object simulation + A.GaussNoise(var_limit=(10.0, 50.0), p=0.3), + A.ISONoise(color_shift=(0.01, 0.05), intensity=(0.1, 0.5), p=0.2), + A.Blur(blur_limit=3, p=0.2), + A.MotionBlur(blur_limit=5, p=0.2), + A.GaussianBlur(blur_limit=3, sigma_limit=0, p=0.1), + # JPEG compression simulation + A.ImageCompression(quality_lower=60, quality_upper=100, p=0.3), + ]) + else: + # No augmentations for validation + self.spatial_transform = None + self.color_transform = None + self.noise_transform = None + + def __len__(self): + return len(self.image_names) + + def apply_mixup_cutmix(self, image, mask, alpha=0.2): + """Apply MixUp or CutMix for robustness""" + if random.random() > 0.3: # 30% chance + return image, mask + + # Get another random sample + idx2 = random.randint(0, len(self.image_names) - 1) + img_name2 = self.image_names[idx2] + mask_name2 = img_name2.replace('.jpg', '.png') + + img_path2 = os.path.join(self.img_dir, img_name2) + mask_path2 = os.path.join(self.mask_dir, mask_name2) + + image2 = cv2.imread(img_path2) + image2 = cv2.cvtColor(image2, cv2.COLOR_BGR2RGB) + mask2 = cv2.imread(mask_path2, cv2.IMREAD_GRAYSCALE) + + # Resize to match current image + h, w = image.shape[:2] + image2 = cv2.resize(image2, (w, h)) + mask2 = cv2.resize(mask2, (w, h), interpolation=cv2.INTER_NEAREST) + + if random.random() < 0.5: # MixUp + lam = np.random.beta(alpha, alpha) + mixed_image = (lam * image + (1 - lam) * image2).astype(np.uint8) + # For masks, choose based on lambda threshold + mixed_mask = np.where(np.random.random((h, w)) < lam, mask, mask2) + else: # CutMix + lam = np.random.beta(alpha, alpha) + bbx1, bby1, bbx2, bby2 = self.rand_bbox(w, h, lam) + + mixed_image = image.copy() + mixed_image[bby1:bby2, bbx1:bbx2] = image2[bby1:bby2, bbx1:bbx2] + + mixed_mask = mask.copy() + mixed_mask[bby1:bby2, bbx1:bbx2] = mask2[bby1:bby2, bbx1:bbx2] + + return mixed_image, mixed_mask + + def rand_bbox(self, W, H, lam): + """Generate random bounding box for CutMix""" + cut_rat = np.sqrt(1. - lam) + cut_w = int(W * cut_rat) + cut_h = int(H * cut_rat) + + cx = np.random.randint(W) + cy = np.random.randint(H) + + bbx1 = np.clip(cx - cut_w // 2, 0, W) + bby1 = np.clip(cy - cut_h // 2, 0, H) + bbx2 = np.clip(cx + cut_w // 2, 0, W) + bby2 = np.clip(cy + cut_h // 2, 0, H) + + return bbx1, bby1, bbx2, bby2 + + def __getitem__(self, idx): + img_name = self.image_names[idx] + mask_name = img_name.replace('.jpg', '.png') + + # Load image and mask + img_path = os.path.join(self.img_dir, img_name) + mask_path = os.path.join(self.mask_dir, mask_name) + + image = cv2.imread(img_path) + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) + + # Enhanced multi-scale training with focus on distant objects + if self.split == 'train' and self.multi_scale: + # Bias towards smaller scales for better distant object learning + scale_weights = [0.2, 0.15, 0.15, 0.15, 0.15, 0.1, 0.05, 0.05] # Favor smaller scales + scale = np.random.choice(self.scales, p=scale_weights) + scaled_size = int(self.current_size * scale) + + # Resize to scaled size first + image = cv2.resize(image, (scaled_size, scaled_size)) + mask = cv2.resize(mask, (scaled_size, scaled_size), interpolation=cv2.INTER_NEAREST) + + # Smart crop/pad to current_size for consistent batching + if scaled_size > self.current_size: + # Random crop with bias towards center (where people usually are) + margin = scaled_size - self.current_size + # Center bias: 70% chance to crop near center + if random.random() < 0.7: + start_x = random.randint(margin//4, 3*margin//4) + start_y = random.randint(margin//4, 3*margin//4) + else: + start_x = random.randint(0, margin) + start_y = random.randint(0, margin) + + image = image[start_y:start_y+self.current_size, start_x:start_x+self.current_size] + mask = mask[start_y:start_y+self.current_size, start_x:start_x+self.current_size] + elif scaled_size < self.current_size: + # Center pad with reflection + pad_x = (self.current_size - scaled_size) // 2 + pad_y = (self.current_size - scaled_size) // 2 + image = cv2.copyMakeBorder( + image, pad_y, self.current_size-scaled_size-pad_y, + pad_x, self.current_size-scaled_size-pad_x, cv2.BORDER_REFLECT + ) + mask = cv2.copyMakeBorder( + mask, pad_y, self.current_size-scaled_size-pad_y, + pad_x, self.current_size-scaled_size-pad_x, cv2.BORDER_CONSTANT, value=0 + ) + else: + # Standard resize + target_size = self.current_size if self.split == 'train' else self.base_image_size + image = cv2.resize(image, (target_size, target_size)) + mask = cv2.resize(mask, (target_size, target_size), interpolation=cv2.INTER_NEAREST) + + # Apply advanced augmentations for training + if self.split == 'train': + # Spatial augmentations + if self.spatial_transform: + transformed = self.spatial_transform(image=image, mask=mask) + image, mask = transformed['image'], transformed['mask'] + + # MixUp/CutMix for robustness + image, mask = self.apply_mixup_cutmix(image, mask) + + # Color augmentations + if self.color_transform: + image = self.color_transform(image=image)['image'] + + # Noise and blur + if self.noise_transform: + image = self.noise_transform(image=image)['image'] + + # Convert to tensor and normalize + image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0 + image = (image - self.mean.view(3, 1, 1)) / self.std.view(3, 1, 1) + + mask = torch.from_numpy(mask).long() + + return image, mask + + def update_progressive_epoch(self, epoch): + """Update progressive training epoch""" + self.progressive_epoch = epoch + old_size = self.current_size + self.current_size = self._get_progressive_size() + if old_size != self.current_size: + print(f"📏 Progressive resize: {old_size} -> {self.current_size}") + + +class AttentionGate(nn.Module): + """Attention gate for better feature selection""" + + def __init__(self, F_g, F_l, F_int): + super().__init__() + self.W_g = nn.Sequential( + nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=True), + nn.BatchNorm2d(F_int) + ) + + self.W_x = nn.Sequential( + nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=True), + nn.BatchNorm2d(F_int) + ) + + self.psi = nn.Sequential( + nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=True), + nn.BatchNorm2d(1), + nn.Sigmoid() + ) + + self.relu = nn.ReLU(inplace=True) + + def forward(self, g, x): + g1 = self.W_g(g) + x1 = self.W_x(x) + psi = self.relu(g1 + x1) + psi = self.psi(psi) + return x * psi + + +class EnhancedASPP(nn.Module): + """Enhanced ASPP with attention and multi-scale context""" + + def __init__(self, in_channels, out_channels, dilations=[1, 6, 12, 18]): + super().__init__() + + self.convs = nn.ModuleList() + for dilation in dilations: + if dilation == 1: + conv = nn.Conv2d(in_channels, out_channels//len(dilations), 1, bias=False) + else: + conv = nn.Conv2d(in_channels, out_channels//len(dilations), 3, + padding=dilation, dilation=dilation, bias=False) + self.convs.append(nn.Sequential( + conv, + nn.BatchNorm2d(out_channels//len(dilations)), + nn.ReLU(inplace=True) + )) + + # Global average pooling branch + self.global_pool = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(in_channels, out_channels//len(dilations), 1, bias=False), + nn.BatchNorm2d(out_channels//len(dilations)), + nn.ReLU(inplace=True) + ) + + # Channel attention + total_channels = out_channels//len(dilations) * (len(dilations) + 1) + self.channel_attention = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(total_channels, total_channels//8, 1), + nn.ReLU(inplace=True), + nn.Conv2d(total_channels//8, total_channels, 1), + nn.Sigmoid() + ) + + # Final projection + self.project = nn.Sequential( + nn.Conv2d(total_channels, out_channels, 1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(inplace=True), + nn.Dropout2d(0.1) + ) + + def forward(self, x): + h, w = x.shape[-2:] + features = [] + + # Apply dilated convolutions + for conv in self.convs: + features.append(conv(x)) + + # Global pooling branch + global_feat = self.global_pool(x) + global_feat = F.interpolate(global_feat, size=(h, w), mode='bilinear', align_corners=False) + features.append(global_feat) + + # Combine features + combined = torch.cat(features, dim=1) + + # Apply channel attention + att = self.channel_attention(combined) + combined = combined * att + + return self.project(combined) + + +class RobustFPN(nn.Module): + """Robust FPN with attention gates and multi-scale fusion""" + + def __init__(self, in_channels_list, out_channels=256): + super().__init__() + + # Lateral connections + self.lateral_convs = nn.ModuleList([ + nn.Conv2d(in_ch, out_channels, 1, bias=False) + for in_ch in in_channels_list + ]) + + # Output convolutions + self.fpn_convs = nn.ModuleList([ + nn.Sequential( + nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(inplace=True) + ) + for _ in in_channels_list + ]) + + # Attention gates for better feature selection + self.attention_gates = nn.ModuleList([ + AttentionGate(out_channels, out_channels, out_channels//2) + for _ in range(len(in_channels_list)-1) + ]) + + # Multi-scale fusion weights + self.fusion_weights = nn.Parameter(torch.ones(len(in_channels_list))) + + def forward(self, features): + # Build laterals + laterals = [conv(feat) for conv, feat in zip(self.lateral_convs, features)] + + # Build top-down pathway with attention + for i in range(len(laterals) - 1, 0, -1): + upsampled = F.interpolate( + laterals[i], size=laterals[i-1].shape[-2:], + mode='bilinear', align_corners=False + ) + + # Apply attention gate + att_feat = self.attention_gates[i-1](upsampled, laterals[i-1]) + laterals[i-1] = laterals[i-1] + att_feat + + # Apply final convs + outputs = [conv(lateral) for conv, lateral in zip(self.fpn_convs, laterals)] + + # Multi-scale weighted feature fusion + target_size = outputs[0].shape[-2:] + fused_features = [] + + for i, feat in enumerate(outputs): + if feat.shape[-2:] != target_size: + feat = F.interpolate(feat, size=target_size, mode='bilinear', align_corners=False) + # Apply learnable weights + weighted_feat = feat * torch.sigmoid(self.fusion_weights[i]) + fused_features.append(weighted_feat) + + return sum(fused_features) + + +class SuperiorSegmentationHead(nn.Module): + """Superior segmentation head optimized for distant objects and lower body parts""" + + def __init__(self, in_channels_list, num_classes=7, feature_dim=256, dropout_rate=0.1): + super().__init__() + self.num_classes = num_classes + + print(f"🏗️ Creating superior segmentation head with channels: {in_channels_list}") + + # Robust FPN with attention + self.fpn = RobustFPN(in_channels_list, feature_dim) + + # Enhanced ASPP for multi-scale context + self.aspp = EnhancedASPP(feature_dim, feature_dim) + + # Multi-scale decoder for better detail preservation + self.decoder = nn.Sequential( + # First stage - preserve details + nn.Conv2d(feature_dim, feature_dim, 3, padding=1, bias=False), + nn.BatchNorm2d(feature_dim), + nn.ReLU(inplace=True), + + # Second stage - refine features + nn.Conv2d(feature_dim, feature_dim//2, 3, padding=1, bias=False), + nn.BatchNorm2d(feature_dim//2), + nn.ReLU(inplace=True), + + # Third stage - final refinement + nn.Conv2d(feature_dim//2, feature_dim//4, 3, padding=1, bias=False), + nn.BatchNorm2d(feature_dim//4), + nn.ReLU(inplace=True), + + nn.Dropout2d(dropout_rate), + nn.Conv2d(feature_dim//4, num_classes, 1) + ) + + # Boundary refinement branch + self.boundary_head = nn.Sequential( + nn.Conv2d(feature_dim, 64, 3, padding=1, bias=False), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True), + nn.Conv2d(64, 1, 1), + nn.Sigmoid() + ) + + self._init_weights() + + def _init_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, nn.BatchNorm2d): + nn.init.ones_(m.weight) + nn.init.zeros_(m.bias) + + def forward(self, features): + # Multi-scale feature fusion + fused_feature = self.fpn(features) + + # Apply enhanced ASPP + context_feature = self.aspp(fused_feature) + + # Generate segmentation + seg_logits = self.decoder(context_feature) + + # Generate boundary map for auxiliary loss + boundary_map = self.boundary_head(context_feature) + + return seg_logits, boundary_map + + +class DFineWithRobustSegmentation(nn.Module): + """D-FINE with robust segmentation optimized for distant objects""" + + def __init__(self, dfine_model, seg_head, freeze_detection=True): + super().__init__() + self.dfine_model = dfine_model + self.seg_head = seg_head + self.freeze_detection = freeze_detection + + if freeze_detection: + self._freeze_detection() + + def _freeze_detection(self): + frozen_params = 0 + total_params = 0 + + for name, param in self.named_parameters(): + total_params += param.numel() + if 'seg_head' not in name: + param.requires_grad = False + frozen_params += param.numel() + + trainable_params = total_params - frozen_params + + print(f"🔒 Frozen detection components:") + print(f" 📊 Total parameters: {total_params:,}") + print(f" ❄️ Frozen parameters: {frozen_params:,} ({frozen_params/total_params*100:.1f}%)") + print(f" 🎯 Trainable parameters: {trainable_params:,} ({trainable_params/total_params*100:.1f}%)") + + def forward(self, x, targets=None): + # Get backbone features + backbone_features = self.dfine_model.backbone(x) + + outputs = {} + + # Detection branch (frozen during training) + if not self.training or not self.freeze_detection: + with torch.no_grad() if self.freeze_detection else torch.enable_grad(): + det_outputs = self.dfine_model(x) + outputs.update(det_outputs) + + # Segmentation branch + seg_logits, boundary_map = self.seg_head(backbone_features) + + # Multi-scale inference for better distant object detection + if not self.training: + seg_logits = self.multi_scale_inference(x, seg_logits) + + # Upsample to input resolution + seg_logits = F.interpolate( + seg_logits, size=x.shape[-2:], + mode='bilinear', align_corners=False + ) + + boundary_map = F.interpolate( + boundary_map, size=x.shape[-2:], + mode='bilinear', align_corners=False + ) + + outputs['segmentation'] = seg_logits + outputs['boundary'] = boundary_map + + return outputs + + def multi_scale_inference(self, x, base_logits): + """Multi-scale test-time inference for better results""" + scales = [0.75, 1.25] # Additional scales for inference + multi_scale_logits = [base_logits] + + for scale in scales: + # Resize input + h, w = x.shape[-2:] + new_h, new_w = int(h * scale), int(w * scale) + scaled_x = F.interpolate(x, size=(new_h, new_w), mode='bilinear', align_corners=False) + + # Forward pass + with torch.no_grad(): + scaled_features = self.dfine_model.backbone(scaled_x) + scaled_logits, _ = self.seg_head(scaled_features) + + # Resize back to original scale + scaled_logits = F.interpolate( + scaled_logits, size=base_logits.shape[-2:], + mode='bilinear', align_corners=False + ) + multi_scale_logits.append(scaled_logits) + + # Average multi-scale predictions + return torch.stack(multi_scale_logits).mean(dim=0) + + +class BoundaryLoss(nn.Module): + """Boundary-aware loss for better edge preservation""" + + def __init__(self, theta0=3, theta=5): + super().__init__() + self.theta0 = theta0 + self.theta = theta + + def forward(self, pred, target): + """ + pred: [N, C, H, W] - predicted logits + target: [N, H, W] - ground truth labels + """ + # Get prediction probabilities + pred_soft = torch.softmax(pred, dim=1) + + # Create one-hot target + target_one_hot = F.one_hot(target, num_classes=pred.shape[1]).permute(0, 3, 1, 2).float() + + # Calculate gradients for boundary detection + def get_boundaries(tensor): + # Sobel operator for edge detection + sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32, device=tensor.device) + sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32, device=tensor.device) + + sobel_x = sobel_x.view(1, 1, 3, 3) + sobel_y = sobel_y.view(1, 1, 3, 3) + + grad_x = F.conv2d(tensor, sobel_x, padding=1) + grad_y = F.conv2d(tensor, sobel_y, padding=1) + + return torch.sqrt(grad_x**2 + grad_y**2) + + # Calculate boundary regions + pred_boundaries = get_boundaries(pred_soft.max(dim=1)[0].unsqueeze(1)) + target_boundaries = get_boundaries(target_one_hot.max(dim=1)[0].unsqueeze(1)) + + # Boundary loss + boundary_loss = F.mse_loss(pred_boundaries, target_boundaries) + + return boundary_loss + + +class AdvancedSegmentationLoss(nn.Module): + """Advanced loss combining multiple components for robust training""" + + def __init__(self, num_classes=7, focal_alpha=0.25, focal_gamma=2.0, ignore_index=255, + dice_weight=0.4, ce_weight=0.2, boundary_weight=0.3, size_weight=0.1): + super().__init__() + + self.focal_loss = FocalLoss(focal_alpha, focal_gamma, ignore_index) + self.dice_loss = DiceLoss(ignore_index=ignore_index) + self.boundary_loss = BoundaryLoss() + + self.dice_weight = dice_weight + self.ce_weight = ce_weight + self.boundary_weight = boundary_weight + self.size_weight = size_weight + self.ignore_index = ignore_index + + def size_sensitive_loss(self, pred, target): + """Loss that gives more weight to smaller objects (distant people)""" + # Calculate object sizes + unique_labels = torch.unique(target) + size_weights = torch.ones_like(target, dtype=torch.float32) + + for label in unique_labels: + if label == self.ignore_index: + continue + + mask = (target == label) + size = mask.sum().float() + + # Inverse size weighting - smaller objects get higher weight + if size > 0: + weight = 1.0 / (torch.sqrt(size) + 1e-6) + size_weights[mask] = weight + + # Apply size weights to cross entropy loss + ce_loss = F.cross_entropy(pred, target, ignore_index=self.ignore_index, reduction='none') + weighted_ce = (ce_loss * size_weights).mean() + + return weighted_ce + + def forward(self, pred, target, boundary_pred=None): + # Focal loss (main loss) + focal = self.focal_loss(pred, target) + + # Dice loss (boundary preservation) + dice = self.dice_loss(pred, target) + + # Standard CE loss + ce = F.cross_entropy(pred, target, ignore_index=self.ignore_index) + + # Size-sensitive loss for distant objects + size_loss = self.size_sensitive_loss(pred, target) + + # Boundary loss if boundary prediction is provided + boundary_loss_val = 0 + if boundary_pred is not None: + # Create boundary target from segmentation mask + boundary_target = self.create_boundary_target(target) + boundary_loss_val = F.binary_cross_entropy(boundary_pred.squeeze(1), boundary_target) + + # Combine losses + total_loss = (focal + + self.dice_weight * dice + + self.ce_weight * ce + + self.size_weight * size_loss + + self.boundary_weight * boundary_loss_val) + + return total_loss, { + 'focal': focal.item(), + 'dice': dice.item(), + 'ce': ce.item(), + 'size': size_loss.item(), + 'boundary': boundary_loss_val.item() if isinstance(boundary_loss_val, torch.Tensor) else boundary_loss_val, + 'total': total_loss.item() + } + + def create_boundary_target(self, target): + """Create boundary target from segmentation mask""" + # Use morphological operations to detect boundaries + target_np = target.cpu().numpy() + boundary_targets = [] + + for i in range(target_np.shape[0]): + mask = target_np[i] + + # Detect edges using morphological gradient + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)) + boundary = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_GRADIENT, kernel) + boundary = (boundary > 0).astype(np.float32) + + boundary_targets.append(torch.from_numpy(boundary)) + + return torch.stack(boundary_targets).to(target.device) + + +class FocalLoss(nn.Module): + """Focal Loss for handling class imbalance""" + + def __init__(self, alpha=0.25, gamma=2.0, ignore_index=255): + super().__init__() + self.alpha = alpha + self.gamma = gamma + self.ignore_index = ignore_index + + def forward(self, inputs, targets): + ce_loss = F.cross_entropy(inputs, targets, ignore_index=self.ignore_index, reduction='none') + pt = torch.exp(-ce_loss) + focal_loss = self.alpha * (1 - pt) ** self.gamma * ce_loss + return focal_loss.mean() + + +class DiceLoss(nn.Module): + """Dice Loss for better boundary handling""" + + def __init__(self, smooth=1.0, ignore_index=255): + super().__init__() + self.smooth = smooth + self.ignore_index = ignore_index + + def forward(self, inputs, targets): + # Create mask for valid pixels + valid_mask = (targets != self.ignore_index) + + # Apply softmax to get probabilities + inputs_soft = F.softmax(inputs, dim=1) + + # One-hot encode targets + targets_one_hot = F.one_hot(targets * valid_mask.long(), num_classes=inputs.shape[1]) + targets_one_hot = targets_one_hot.permute(0, 3, 1, 2).float() + + # Apply valid mask + inputs_soft = inputs_soft * valid_mask.unsqueeze(1).float() + targets_one_hot = targets_one_hot * valid_mask.unsqueeze(1).float() + + # Calculate dice coefficient + intersection = (inputs_soft * targets_one_hot).sum(dim=(2, 3)) + union = inputs_soft.sum(dim=(2, 3)) + targets_one_hot.sum(dim=(2, 3)) + + dice = (2 * intersection + self.smooth) / (union + self.smooth) + return 1 - dice.mean() + + +def compute_miou(pred, target, num_classes=7): + """Compute mean IoU""" + pred = torch.argmax(pred, dim=1) + + ious = [] + for cls in range(num_classes): + pred_cls = (pred == cls) + target_cls = (target == cls) + + intersection = (pred_cls & target_cls).sum().float() + union = (pred_cls | target_cls).sum().float() + + if union > 0: + ious.append(intersection / union) + + return torch.tensor(ious).mean() if ious else torch.tensor(0.0) + + +def get_actual_backbone_channels(model): + """Get actual backbone output channels""" + model.eval() + dummy_input = torch.randn(1, 3, 640, 640) + + with torch.no_grad(): + backbone_features = model.backbone(dummy_input) + + actual_channels = [feat.shape[1] for feat in backbone_features] + return actual_channels + + +def load_pretrained_dfine(config_path, checkpoint_path): + """Load pretrained D-FINE model""" + print(f"🚀 Loading pretrained D-FINE from {checkpoint_path}") + + cfg = YAMLConfig(config_path) + model = cfg.model + + checkpoint = torch.load(checkpoint_path, map_location='cpu') + + if 'ema' in checkpoint and 'module' in checkpoint['ema']: + state_dict = checkpoint['ema']['module'] + elif 'model' in checkpoint: + state_dict = checkpoint['model'] + else: + state_dict = checkpoint + + model.load_state_dict(state_dict, strict=False) + print(f"✅ Loaded pretrained D-FINE model") + return model + + +def create_robust_segmentation_model(config_path, checkpoint_path, num_classes=7, + feature_dim=256, dropout_rate=0.1): + """Create robust segmentation model""" + + det_model = load_pretrained_dfine(config_path, checkpoint_path) + backbone_channels = get_actual_backbone_channels(det_model) + + print(f"🏗️ Using backbone channels: {backbone_channels}") + + # Create superior segmentation head + seg_head = SuperiorSegmentationHead( + in_channels_list=backbone_channels, + num_classes=num_classes, + feature_dim=feature_dim, + dropout_rate=dropout_rate + ) + + model = DFineWithRobustSegmentation( + dfine_model=det_model, + seg_head=seg_head, + freeze_detection=True + ) + + return model + + +def unfreeze_backbone_gradually(model, optimizer, epoch, unfreeze_epoch1=35, unfreeze_epoch2=45): + """Gradually unfreeze backbone layers""" + + if epoch == unfreeze_epoch1: + print("🔓 Unfreezing backbone stage 4 (highest level features)") + new_params = [] + for name, param in model.dfine_model.backbone.named_parameters(): + if 'stages.3' in name: # Stage 4 of HGNetv2 + param.requires_grad = True + new_params.append(param) + + if new_params: + optimizer.add_param_group({ + 'params': new_params, + 'lr': optimizer.param_groups[0]['lr'] * 0.1 # 10x lower LR + }) + print(f" Added {len(new_params)} parameters to optimizer") + + elif epoch == unfreeze_epoch2: + print("🔓 Unfreezing backbone stage 3") + new_params = [] + for name, param in model.dfine_model.backbone.named_parameters(): + if 'stages.2' in name: + param.requires_grad = True + new_params.append(param) + + if new_params: + optimizer.add_param_group({ + 'params': new_params, + 'lr': optimizer.param_groups[0]['lr'] * 0.1 + }) + print(f" Added {len(new_params)} parameters to optimizer") + + +def train_epoch(model, dataloader, criterion, optimizer, device, epoch): + """Enhanced training epoch with progressive training""" + model.train() + + # Update progressive training + if hasattr(dataloader.dataset, 'update_progressive_epoch'): + dataloader.dataset.update_progressive_epoch(epoch) + + total_loss = 0 + total_miou = 0 + loss_components = {'focal': 0, 'dice': 0, 'ce': 0, 'size': 0, 'boundary': 0, 'total': 0} + num_batches = len(dataloader) + + pbar = tqdm(dataloader, desc=f'Epoch {epoch}', leave=False) + + for batch_idx, (images, masks) in enumerate(pbar): + images = images.to(device) + masks = masks.to(device) + + optimizer.zero_grad() + + # Forward pass + outputs = model(images) + seg_pred = outputs['segmentation'] + boundary_pred = outputs.get('boundary', None) + + # Calculate loss + loss, loss_dict = criterion(seg_pred, masks, boundary_pred) + + # Backward pass + loss.backward() + + # Gradient clipping for stability + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) + + optimizer.step() + + # Calculate metrics + with torch.no_grad(): + miou = compute_miou(seg_pred, masks) + + total_loss += loss.item() + total_miou += miou.item() + + # Accumulate loss components + for key in loss_components: + loss_components[key] += loss_dict.get(key, 0) + + # Update progress bar + pbar.set_postfix({ + 'Loss': f'{loss.item():.4f}', + 'mIoU': f'{miou.item():.3f}', + 'FL': f'{loss_dict.get("focal", 0):.3f}', + 'DL': f'{loss_dict.get("dice", 0):.3f}', + 'BL': f'{loss_dict.get("boundary", 0):.3f}' + }) + + # Average metrics + avg_loss = total_loss / num_batches + avg_miou = total_miou / num_batches + avg_loss_components = {k: v / num_batches for k, v in loss_components.items()} + + return avg_loss, avg_miou, avg_loss_components + + +def validate_epoch(model, dataloader, criterion, device): + """Enhanced validation""" + model.eval() + + total_loss = 0 + total_miou = 0 + loss_components = {'focal': 0, 'dice': 0, 'ce': 0, 'size': 0, 'boundary': 0, 'total': 0} + num_batches = len(dataloader) + + with torch.no_grad(): + for images, masks in tqdm(dataloader, desc='Validation', leave=False): + images = images.to(device) + masks = masks.to(device) + + outputs = model(images) + seg_pred = outputs['segmentation'] + boundary_pred = outputs.get('boundary', None) + + loss, loss_dict = criterion(seg_pred, masks, boundary_pred) + miou = compute_miou(seg_pred, masks) + + total_loss += loss.item() + total_miou += miou.item() + + for key in loss_components: + loss_components[key] += loss_dict.get(key, 0) + + avg_loss = total_loss / num_batches + avg_miou = total_miou / num_batches + avg_loss_components = {k: v / num_batches for k, v in loss_components.items()} + + return avg_loss, avg_miou, avg_loss_components + + +def save_best_trial_results(output_dir, trial_number, trial_miou, trial_params, model, optimizer, scheduler, args, epoch): + """Save best trial results and model""" + global best_global_miou, best_global_lock + + with best_global_lock: + if trial_miou > best_global_miou: + best_global_miou = trial_miou + + print(f"🏆 New best trial found! Trial {trial_number} with mIoU: {trial_miou:.4f}") + + # Save best model + model_save_path = f'{output_dir}/best_model.pth' + torch.save({ + 'epoch': epoch, + 'trial_number': trial_number, + 'model_state_dict': model.state_dict(), + 'optimizer_state_dict': optimizer.state_dict(), + 'scheduler_state_dict': scheduler.state_dict(), + 'best_miou': trial_miou, + 'hyperparameters': trial_params, + 'config': vars(args), + 'seed': args.seed + }, model_save_path) + + # Save best results JSON + best_results = { + 'best_miou': trial_miou, + 'best_trial_number': trial_number, + 'best_hyperparameters': trial_params, + 'best_epoch': epoch, + 'seed': args.seed, + 'config': vars(args), + 'timestamp': str(torch.cuda.Event().query() if torch.cuda.is_available() else 'cpu') + } + + results_save_path = f'{output_dir}/best_results.json' + with open(results_save_path, 'w') as f: + json.dump(best_results, f, indent=2) + + print(f"💾 Saved best model to: {model_save_path}") + print(f"💾 Saved best results to: {results_save_path}") + + return True # Indicates this was a new best + + return False # Not a new best + + +def objective(trial, args, device, train_loader, val_loader, backbone_channels): + """Optuna objective function for hyperparameter optimization""" + + # Set seed for each trial + trial_seed = args.seed + trial.number + set_seed(trial_seed) + + # Enhanced hyperparameters for robust training + hyperparams = { + 'lr': trial.suggest_categorical('lr', [1e-4, 2e-4, 3e-4, 5e-4, 7e-4, 1e-3]), + 'batch_size': trial.suggest_categorical('batch_size', [4, 6, 8, 10, 12]), # Smaller batches for stability + 'feature_dim': trial.suggest_categorical('feature_dim', [256, 320, 384, 448]), + 'dropout_rate': trial.suggest_categorical('dropout_rate', [0.05, 0.1, 0.15, 0.2]), + 'focal_alpha': trial.suggest_categorical('focal_alpha', [0.15, 0.2, 0.25, 0.3]), + 'focal_gamma': trial.suggest_categorical('focal_gamma', [1.5, 2.0, 2.5]), + 'dice_weight': trial.suggest_categorical('dice_weight', [0.3, 0.4, 0.5, 0.6]), + 'ce_weight': trial.suggest_categorical('ce_weight', [0.1, 0.15, 0.2]), + 'boundary_weight': trial.suggest_categorical('boundary_weight', [0.2, 0.3, 0.4]), + 'size_weight': trial.suggest_categorical('size_weight', [0.05, 0.1, 0.15, 0.2]), + 'weight_decay': trial.suggest_categorical('weight_decay', [1e-5, 5e-5, 1e-4, 2e-4]), + 'scheduler_t0': trial.suggest_categorical('scheduler_t0', [12, 15, 18, 20]), + 'unfreeze_epoch1': trial.suggest_categorical('unfreeze_epoch1', [25, 30, 35]), + 'unfreeze_epoch2': trial.suggest_categorical('unfreeze_epoch2', [35, 40, 45]), + } + + # Ensure unfreeze_epoch2 > unfreeze_epoch1 + if hyperparams['unfreeze_epoch2'] <= hyperparams['unfreeze_epoch1']: + hyperparams['unfreeze_epoch2'] = hyperparams['unfreeze_epoch1'] + 10 + + print(f"\n🔬 Trial {trial.number} hyperparameters:") + for key, value in hyperparams.items(): + print(f" {key}: {value}") + + # Initialize wandb for this trial + wandb.init( + project=f'{args.project_name}-robust', + config={**vars(args), **hyperparams, 'trial_number': trial.number}, + name=f'robust_trial_{trial.number}', + reinit=True + ) + + try: + # Create model with trial hyperparameters + model = create_robust_segmentation_model( + args.config, args.checkpoint, + num_classes=7, + feature_dim=hyperparams['feature_dim'], + dropout_rate=hyperparams['dropout_rate'] + ) + model = model.to(device) + + # Create advanced loss function + criterion = AdvancedSegmentationLoss( + num_classes=7, + focal_alpha=hyperparams['focal_alpha'], + focal_gamma=hyperparams['focal_gamma'], + dice_weight=hyperparams['dice_weight'], + ce_weight=hyperparams['ce_weight'], + boundary_weight=hyperparams['boundary_weight'], + size_weight=hyperparams['size_weight'] + ) + + # Create optimizer with trial hyperparameters + seg_params = [p for p in model.seg_head.parameters() if p.requires_grad] + optimizer = optim.AdamW( + seg_params, + lr=hyperparams['lr'], + weight_decay=hyperparams['weight_decay'] + ) + + # Create scheduler with trial hyperparameters + scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts( + optimizer, T_0=hyperparams['scheduler_t0'], T_mult=2, eta_min=1e-6 + ) + + best_trial_miou = 0 + best_trial_epoch = 0 + patience_counter = 0 + patience = 30 # Increased patience for robust training + + # Extended training for robust results + max_epochs = max(args.epochs, 100) # Minimum 100 epochs for robust training + + for epoch in range(max_epochs): + # Gradual backbone unfreezing + unfreeze_backbone_gradually( + model, optimizer, epoch, + hyperparams['unfreeze_epoch1'], + hyperparams['unfreeze_epoch2'] + ) + + # Train + train_loss, train_miou, train_loss_components = train_epoch( + model, train_loader, criterion, optimizer, device, epoch + ) + + # Validate + val_loss, val_miou, val_loss_components = validate_epoch( + model, val_loader, criterion, device + ) + + # Update scheduler + scheduler.step() + + # Log metrics + wandb.log({ + 'epoch': epoch, + 'train/loss': train_loss, + 'train/miou': train_miou, + 'val/loss': val_loss, + 'val/miou': val_miou, + 'lr': optimizer.param_groups[0]['lr'], + 'trial_number': trial.number, + 'progressive_size': getattr(train_loader.dataset, 'current_size', args.image_size) + }) + + # Report to Optuna + trial.report(val_miou, epoch) + + # Check if this is the best for this trial + if val_miou > best_trial_miou: + best_trial_miou = val_miou + best_trial_epoch = epoch + patience_counter = 0 + + # Save if this is globally the best + save_best_trial_results( + args.output_dir, trial.number, val_miou, hyperparams, + model, optimizer, scheduler, args, epoch + ) + else: + patience_counter += 1 + + # Prune unpromising trials + if trial.should_prune(): + wandb.finish() + raise optuna.exceptions.TrialPruned() + + # Early stopping + if patience_counter >= patience: + print(f"⏹️ Early stopping at epoch {epoch}") + break + + print(f"🏁 Trial {trial.number} completed. Best mIoU: {best_trial_miou:.4f} at epoch {best_trial_epoch}") + + wandb.finish() + return best_trial_miou + + except Exception as e: + print(f"❌ Trial {trial.number} failed: {e}") + wandb.finish() + raise e + + +def main(): + parser = argparse.ArgumentParser(description='Robust D-FINE Segmentation with Advanced Augmentations') + parser.add_argument('--config', default='configs/dfine_hgnetv2_x_obj2coco.yml') + parser.add_argument('--checkpoint', default='models/dfine_x_obj2coco.pth') + parser.add_argument('--dataset', default='datasets/pascal_person_parts') + parser.add_argument('--batch_size', type=int, default=32) # Smaller for stability + parser.add_argument('--epochs', type=int, default=250) # More epochs for robust training + parser.add_argument('--lr', type=float, default=3e-4) + parser.add_argument('--image_size', type=int, default=640) + parser.add_argument('--num_workers', type=int, default=12) + parser.add_argument('--output_dir', default='outputs/dfine_segmentation_robust') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') + parser.add_argument('--n_trials', type=int, default=25, help='Number of Optuna trials') + parser.add_argument('--project_name', default='dfine-segmentation-robust', help='Wandb project name') + parser.add_argument('--study_name', default='dfine_seg_robust_distant', help='Optuna study name') + + args = parser.parse_args() + + # Set global seed + set_seed(args.seed) + + os.makedirs(args.output_dir, exist_ok=True) + + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + print(f"🚀 Using device: {device}") + + # Load datasets with advanced augmentations + print("📊 Loading datasets with advanced augmentations...") + train_dataset = AdvancedPascalPersonPartsDataset( + args.dataset, split='train', image_size=args.image_size, multi_scale=True + ) + val_dataset = AdvancedPascalPersonPartsDataset( + args.dataset, split='val', image_size=args.image_size, multi_scale=False + ) + + train_loader = DataLoader( + train_dataset, batch_size=args.batch_size, shuffle=True, + num_workers=args.num_workers, pin_memory=True + ) + val_loader = DataLoader( + val_dataset, batch_size=args.batch_size, shuffle=False, + num_workers=args.num_workers, pin_memory=True + ) + + # Get backbone channels for model creation + temp_model = load_pretrained_dfine(args.config, args.checkpoint) + backbone_channels = get_actual_backbone_channels(temp_model) + del temp_model # Free memory + + # Create Optuna study + study = optuna.create_study( + direction='maximize', + study_name=args.study_name, + storage=f'sqlite:///{args.output_dir}/optuna_study.db', + load_if_exists=True, + pruner=optuna.pruners.MedianPruner(n_startup_trials=5, n_warmup_steps=30) + ) + + print(f"🔬 Starting robust optimization with {args.n_trials} trials") + print(f"📁 Output directory: {args.output_dir}") + print(f"🎯 Focus: Distant objects, legs/feet robustness, advanced augmentations") + print(f"⏱️ Using progressive training and multi-scale inference") + print(f"💾 Best model and results will be saved after each better trial") + + # Run optimization + study.optimize( + lambda trial: objective(trial, args, device, train_loader, val_loader, backbone_channels), + n_trials=args.n_trials, + timeout=None, + gc_after_trial=True + ) + + # Get best trial + best_trial = study.best_trial + best_params = best_trial.params + best_value = best_trial.value + + print(f"\n🏆 Robust optimization completed!") + print(f"📊 Best mIoU: {best_value:.4f}") + print(f"🎯 Best trial number: {best_trial.number}") + print(f"🎯 Best hyperparameters:") + for key, value in best_params.items(): + print(f" {key}: {value}") + + # Save final optimization summary + optimization_summary = { + 'best_miou': best_value, + 'best_trial_number': best_trial.number, + 'best_params': best_params, + 'total_trials': args.n_trials, + 'seed': args.seed, + 'optimization_config': vars(args), + 'features': [ + 'Progressive training', + 'Advanced augmentations', + 'Multi-scale inference', + 'Size-sensitive loss', + 'Boundary-aware loss', + 'Attention mechanisms' + ] + } + + with open(f'{args.output_dir}/optimization_summary.json', 'w') as f: + json.dump(optimization_summary, f, indent=2) + + print(f"\n✅ Robust optimization summary saved to: {args.output_dir}/optimization_summary.json") + print(f"🏆 Best model saved at: {args.output_dir}/best_model.pth") + print(f"📊 Best results saved at: {args.output_dir}/best_results.json") + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/segmentation_sivert/models/__init__.py b/segmentation_sivert/models/__init__.py new file mode 100644 index 00000000..365f7f75 --- /dev/null +++ b/segmentation_sivert/models/__init__.py @@ -0,0 +1,45 @@ +#!/usr/bin/env python3 +""" +Model tier implementations for DFINE Segmentation +Three specialized tiers: lightweight, standard, advanced +""" + +# Import tier-specific models +from .lightweight import create_lightweight_segmentation_head, LightweightSegmentationHead +from .standard import create_standard_segmentation_head, StandardSegmentationHead +from .advanced import create_advanced_segmentation_head, AdvancedSegmentationHead + +# Import base classes +from .base import SegmentationHeadInterface, ConvBlock, ASPP, LightweightASPP, SimpleFPN + +__version__ = "1.0.0" +__all__ = [ + # Factory functions + 'create_lightweight_segmentation_head', + 'create_standard_segmentation_head', + 'create_advanced_segmentation_head', + + # Model classes + 'LightweightSegmentationHead', + 'StandardSegmentationHead', + 'AdvancedSegmentationHead', + + # Base classes + 'SegmentationHeadInterface', + 'ConvBlock', + 'ASPP', + 'LightweightASPP', + 'SimpleFPN' +] + +def create_segmentation_head(tier: str, in_channels_list, **kwargs): + """Factory function to create segmentation head by tier""" + + if tier == 'lightweight': + return create_lightweight_segmentation_head(in_channels_list, **kwargs) + elif tier == 'standard': + return create_standard_segmentation_head(in_channels_list, **kwargs) + elif tier == 'advanced': + return create_advanced_segmentation_head(in_channels_list, **kwargs) + else: + raise ValueError(f"Unknown tier: {tier}") \ No newline at end of file diff --git a/segmentation_sivert/models/advanced.py b/segmentation_sivert/models/advanced.py new file mode 100644 index 00000000..08701196 --- /dev/null +++ b/segmentation_sivert/models/advanced.py @@ -0,0 +1,449 @@ +#!/usr/bin/env python3 +""" +Advanced Segmentation Head for DFINE +Maximum accuracy and robustness, suitable for high-end GPU setups +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import List, Optional, Tuple +import os +import sys + +# Add modules to path +sys.path.append(os.path.join(os.path.dirname(__file__), '..')) +from models.base import SegmentationHeadInterface, ConvBlock, ASPP + + +class AttentionGate(nn.Module): + """Attention gate for better feature selection""" + + def __init__(self, F_g: int, F_l: int, F_int: int): + super().__init__() + self.W_g = nn.Sequential( + nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=True), + nn.BatchNorm2d(F_int) + ) + + self.W_x = nn.Sequential( + nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=True), + nn.BatchNorm2d(F_int) + ) + + self.psi = nn.Sequential( + nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=True), + nn.BatchNorm2d(1), + nn.Sigmoid() + ) + + self.relu = nn.ReLU(inplace=True) + + def forward(self, g, x): + g1 = self.W_g(g) + x1 = self.W_x(x) + psi = self.relu(g1 + x1) + psi = self.psi(psi) + return x * psi + + +class SelfAttentionModule(nn.Module): + """Self-attention module for long-range dependencies""" + + def __init__(self, in_channels: int, reduction: int = 8): + super().__init__() + + self.query_conv = nn.Conv2d(in_channels, in_channels // reduction, 1) + self.key_conv = nn.Conv2d(in_channels, in_channels // reduction, 1) + self.value_conv = nn.Conv2d(in_channels, in_channels, 1) + + self.gamma = nn.Parameter(torch.zeros(1)) + self.softmax = nn.Softmax(dim=-1) + + def forward(self, x): + batch_size, C, H, W = x.size() + + # Generate query, key, value + proj_query = self.query_conv(x).view(batch_size, -1, H * W).permute(0, 2, 1) + proj_key = self.key_conv(x).view(batch_size, -1, H * W) + proj_value = self.value_conv(x).view(batch_size, -1, H * W) + + # Attention + energy = torch.bmm(proj_query, proj_key) + attention = self.softmax(energy) + + # Apply attention + out = torch.bmm(proj_value, attention.permute(0, 2, 1)) + out = out.view(batch_size, C, H, W) + + # Residual connection + out = self.gamma * out + x + return out + + +class ChannelAttention(nn.Module): + """Channel attention mechanism""" + + def __init__(self, in_channels: int, reduction: int = 16): + super().__init__() + + self.avg_pool = nn.AdaptiveAvgPool2d(1) + self.max_pool = nn.AdaptiveMaxPool2d(1) + + self.fc = nn.Sequential( + nn.Conv2d(in_channels, in_channels // reduction, 1, bias=False), + nn.ReLU(inplace=True), + nn.Conv2d(in_channels // reduction, in_channels, 1, bias=False) + ) + + self.sigmoid = nn.Sigmoid() + + def forward(self, x): + avg_out = self.fc(self.avg_pool(x)) + max_out = self.fc(self.max_pool(x)) + out = avg_out + max_out + return x * self.sigmoid(out) + + +class SpatialAttention(nn.Module): + """Spatial attention mechanism""" + + def __init__(self, kernel_size: int = 7): + super().__init__() + + self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size//2, bias=False) + self.sigmoid = nn.Sigmoid() + + def forward(self, x): + avg_out = torch.mean(x, dim=1, keepdim=True) + max_out, _ = torch.max(x, dim=1, keepdim=True) + out = torch.cat([avg_out, max_out], dim=1) + out = self.conv(out) + return x * self.sigmoid(out) + + +class CBAM(nn.Module): + """Convolutional Block Attention Module""" + + def __init__(self, in_channels: int, reduction: int = 16, kernel_size: int = 7): + super().__init__() + + self.channel_attention = ChannelAttention(in_channels, reduction) + self.spatial_attention = SpatialAttention(kernel_size) + + def forward(self, x): + out = self.channel_attention(x) + out = self.spatial_attention(out) + return out + + +class AdvancedASPP(nn.Module): + """Advanced ASPP with attention and multi-scale context""" + + def __init__(self, in_channels: int, out_channels: int, dilations: List[int] = [1, 6, 12, 18]): + super().__init__() + + self.convs = nn.ModuleList() + for dilation in dilations: + if dilation == 1: + conv = nn.Conv2d(in_channels, out_channels//len(dilations), 1, bias=False) + else: + conv = nn.Conv2d(in_channels, out_channels//len(dilations), 3, + padding=dilation, dilation=dilation, bias=False) + self.convs.append(nn.Sequential( + conv, + nn.BatchNorm2d(out_channels//len(dilations)), + nn.ReLU(inplace=True) + )) + + # Global average pooling branch + self.global_pool = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(in_channels, out_channels//len(dilations), 1, bias=False), + nn.BatchNorm2d(out_channels//len(dilations)), + nn.ReLU(inplace=True) + ) + + # Channel and spatial attention + total_channels = out_channels//len(dilations) * (len(dilations) + 1) + self.attention = CBAM(total_channels) + + # Final projection + self.project = nn.Sequential( + nn.Conv2d(total_channels, out_channels, 1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(inplace=True), + nn.Dropout2d(0.1) + ) + + def forward(self, x): + h, w = x.shape[-2:] + features = [] + + # Apply dilated convolutions + for conv in self.convs: + features.append(conv(x)) + + # Global pooling branch + global_feat = self.global_pool(x) + global_feat = F.interpolate(global_feat, size=(h, w), mode='bilinear', align_corners=False) + features.append(global_feat) + + # Combine features + combined = torch.cat(features, dim=1) + + # Apply attention + attended = self.attention(combined) + + return self.project(attended) + + +class RobustFPN(nn.Module): + """Robust FPN with attention gates and multi-scale fusion""" + + def __init__(self, in_channels_list: List[int], out_channels: int = 256): + super().__init__() + + # Lateral connections with batch normalization + self.lateral_convs = nn.ModuleList([ + nn.Sequential( + nn.Conv2d(in_ch, out_channels, 1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(inplace=True) + ) for in_ch in in_channels_list + ]) + + # Output convolutions with attention + self.fpn_convs = nn.ModuleList([ + nn.Sequential( + nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(inplace=True), + CBAM(out_channels) + ) for _ in in_channels_list + ]) + + # Attention gates for better feature selection + self.attention_gates = nn.ModuleList([ + AttentionGate(out_channels, out_channels, out_channels//2) + for _ in range(len(in_channels_list)-1) + ]) + + # Multi-scale fusion weights + self.fusion_weights = nn.Parameter(torch.ones(len(in_channels_list))) + + # Self-attention for global context + self.self_attention = SelfAttentionModule(out_channels) + + def forward(self, features: List[torch.Tensor]) -> torch.Tensor: + # Build laterals + laterals = [conv(feat) for conv, feat in zip(self.lateral_convs, features)] + + # Build top-down pathway with attention + for i in range(len(laterals) - 1, 0, -1): + upsampled = F.interpolate( + laterals[i], size=laterals[i-1].shape[-2:], + mode='bilinear', align_corners=False + ) + + # Apply attention gate + att_feat = self.attention_gates[i-1](upsampled, laterals[i-1]) + laterals[i-1] = laterals[i-1] + att_feat + + # Apply final convs with attention + outputs = [conv(lateral) for conv, lateral in zip(self.fpn_convs, laterals)] + + # Multi-scale weighted feature fusion + target_size = outputs[0].shape[-2:] + fused_features = [] + + for i, feat in enumerate(outputs): + if feat.shape[-2:] != target_size: + feat = F.interpolate(feat, size=target_size, mode='bilinear', align_corners=False) + # Apply learnable weights + weighted_feat = feat * torch.sigmoid(self.fusion_weights[i]) + fused_features.append(weighted_feat) + + # Combine with self-attention + fused = sum(fused_features) + attended = self.self_attention(fused) + + return attended + + +class AdvancedDecoder(nn.Module): + """Advanced decoder with skip connections and attention""" + + def __init__(self, in_channels: int, num_classes: int, dropout_rate: float = 0.1): + super().__init__() + + # Multi-stage decoder with progressive refinement + self.stage1 = nn.Sequential( + ConvBlock(in_channels, in_channels, 3, 1, 1), + CBAM(in_channels), + nn.Dropout2d(dropout_rate) + ) + + self.stage2 = nn.Sequential( + ConvBlock(in_channels, in_channels//2, 3, 1, 1), + CBAM(in_channels//2), + nn.Dropout2d(dropout_rate) + ) + + self.stage3 = nn.Sequential( + ConvBlock(in_channels//2, in_channels//4, 3, 1, 1), + CBAM(in_channels//4), + nn.Dropout2d(dropout_rate) + ) + + # Final classification + self.classifier = nn.Conv2d(in_channels//4, num_classes, 1) + + # Auxiliary outputs for deep supervision + self.aux_classifier1 = nn.Conv2d(in_channels, num_classes, 1) + self.aux_classifier2 = nn.Conv2d(in_channels//2, num_classes, 1) + + def forward(self, x): + # Progressive decoding with auxiliary outputs + feat1 = self.stage1(x) + aux1 = self.aux_classifier1(feat1) + + feat2 = self.stage2(feat1) + aux2 = self.aux_classifier2(feat2) + + feat3 = self.stage3(feat2) + main_out = self.classifier(feat3) + + return main_out, aux1, aux2 + + +class AdvancedSegmentationHead(SegmentationHeadInterface): + """Advanced segmentation head with maximum accuracy and robustness + + Incorporates state-of-the-art attention mechanisms, multi-scale processing, + and deep supervision for the highest possible segmentation quality. + """ + + def __init__(self, + in_channels_list: List[int], + num_classes: int = 7, + feature_dim: int = 384, # Increased for advanced model + dropout_rate: float = 0.15): # Slightly higher for regularization + """Initialize advanced segmentation head + + Args: + in_channels_list: List of input channel counts from backbone + num_classes: Number of segmentation classes (default: 7 for Pascal Person Parts) + feature_dim: Feature dimension for processing (default: 384, increased for capacity) + dropout_rate: Dropout rate (default: 0.15, higher for robustness) + """ + super().__init__( + in_channels_list=in_channels_list, + num_classes=num_classes, + feature_dim=feature_dim, + dropout_rate=dropout_rate, + tier='advanced' + ) + + def _build_head(self): + """Build advanced segmentation head architecture""" + + # Robust FPN with attention mechanisms + self.fpn = RobustFPN(self.in_channels_list, self.feature_dim) + + # Advanced ASPP with attention + self.aspp = AdvancedASPP(self.feature_dim, self.feature_dim, dilations=[1, 6, 12, 18]) + + # Advanced decoder with deep supervision + self.decoder = AdvancedDecoder(self.feature_dim, self.num_classes, self.dropout_rate) + + # Boundary refinement branch + self.boundary_head = nn.Sequential( + ConvBlock(self.feature_dim, 64, 3, 1, 1), + CBAM(64), + nn.Conv2d(64, 1, 1), + nn.Sigmoid() + ) + + print(f"🏗️ Built advanced segmentation head:") + print(f" Feature dim: {self.feature_dim} (increased)") + print(f" Dropout rate: {self.dropout_rate}") + print(f" Classes: {self.num_classes}") + print(f" 🎯 Features: Attention, multi-scale, deep supervision, boundary refinement") + + def forward(self, features: List[torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Forward pass through advanced segmentation head + + Args: + features: List of feature tensors from backbone [feat1, feat2, feat3, feat4] + + Returns: + Tuple of (main_logits, aux1_logits, aux2_logits, boundary_map) + """ + # Multi-scale feature fusion with attention + fused_feature = self.fpn(features) + + # Apply advanced ASPP with attention + context_feature = self.aspp(fused_feature) + + # Generate segmentation with deep supervision + main_logits, aux1_logits, aux2_logits = self.decoder(context_feature) + + # Generate boundary map for auxiliary loss + boundary_map = self.boundary_head(context_feature) + + return main_logits, aux1_logits, aux2_logits, boundary_map + + def forward_inference(self, features: List[torch.Tensor]) -> torch.Tensor: + """Forward pass for inference (returns only main output)""" + main_logits, _, _, _ = self.forward(features) + return main_logits + + def get_model_info(self) -> dict: + """Get detailed model information""" + base_info = super().get_model_info() + + # Add advanced-specific info + base_info.update({ + 'fpn_feature_dim': self.feature_dim, + 'aspp_dilations': [1, 6, 12, 18], + 'decoder_stages': 3, + 'uses_attention': True, + 'uses_self_attention': True, + 'uses_cbam': True, + 'deep_supervision': True, + 'boundary_refinement': True, + 'supports_multiscale': True, + 'optimized_for': 'maximum_accuracy', + 'target_platforms': ['high_end_gpu', 'server', 'workstation'], + 'estimated_accuracy': 'highest', + 'estimated_memory_mb': '200-400' + }) + + return base_info + + +def create_advanced_segmentation_head(in_channels_list: List[int], + num_classes: int = 7, + feature_dim: int = 384, + dropout_rate: float = 0.15) -> AdvancedSegmentationHead: + """Factory function to create advanced segmentation head""" + return AdvancedSegmentationHead( + in_channels_list=in_channels_list, + num_classes=num_classes, + feature_dim=feature_dim, + dropout_rate=dropout_rate + ) + + +# Export for use in other modules +__all__ = [ + 'AdvancedSegmentationHead', + 'RobustFPN', + 'AdvancedASPP', + 'AdvancedDecoder', + 'AttentionGate', + 'SelfAttentionModule', + 'CBAM', + 'create_advanced_segmentation_head' +] \ No newline at end of file diff --git a/segmentation_sivert/models/base.py b/segmentation_sivert/models/base.py new file mode 100644 index 00000000..f5d89dae --- /dev/null +++ b/segmentation_sivert/models/base.py @@ -0,0 +1,276 @@ +#!/usr/bin/env python3 +""" +Base Segmentation Head for DFINE +Abstract base class defining the interface for all segmentation heads +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from abc import ABC, abstractmethod +from typing import List, Dict, Any, Optional, Tuple +import sys +import os + +# Add core modules to path +sys.path.append(os.path.join(os.path.dirname(__file__), '..')) +from core.models import BaseSegmentationHead + + +class SegmentationHeadInterface(BaseSegmentationHead): + """Interface for all segmentation heads with common functionality""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + @abstractmethod + def _build_head(self): + """Build the specific segmentation head architecture""" + pass + + @abstractmethod + def forward(self, features: List[torch.Tensor]) -> torch.Tensor: + """Forward pass through segmentation head + + Args: + features: List of feature tensors from backbone + + Returns: + Segmentation logits tensor [N, num_classes, H, W] + """ + pass + + def get_complexity_info(self) -> Dict[str, Any]: + """Get model complexity information""" + info = self.get_model_info() + + # Estimate FLOPs (rough approximation) + sample_input = [torch.randn(1, ch, 64, 64) for ch in self.in_channels_list] + + # Count convolutions + conv_count = 0 + for module in self.modules(): + if isinstance(module, nn.Conv2d): + conv_count += 1 + + info.update({ + 'conv_layers': conv_count, + 'estimated_flops': conv_count * 64 * 64 * 1000, # Very rough estimate + 'memory_efficient': self.tier == 'lightweight' + }) + + return info + + +# Common building blocks for segmentation heads +class ConvBlock(nn.Module): + """Standard convolution block""" + + def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3, + stride: int = 1, padding: int = 1, bias: bool = False, + norm: bool = True, activation: bool = True): + super().__init__() + + self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias) + self.norm = nn.BatchNorm2d(out_channels) if norm else nn.Identity() + self.activation = nn.ReLU(inplace=True) if activation else nn.Identity() + + def forward(self, x): + return self.activation(self.norm(self.conv(x))) + + +class DepthwiseSeparableConv(nn.Module): + """Depthwise separable convolution for lightweight models""" + + def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3, + stride: int = 1, padding: int = 1, bias: bool = False): + super().__init__() + + self.depthwise = nn.Conv2d(in_channels, in_channels, kernel_size, stride, padding, + groups=in_channels, bias=bias) + self.pointwise = nn.Conv2d(in_channels, out_channels, 1, bias=bias) + self.bn1 = nn.BatchNorm2d(in_channels) + self.bn2 = nn.BatchNorm2d(out_channels) + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + x = self.relu(self.bn1(self.depthwise(x))) + x = self.relu(self.bn2(self.pointwise(x))) + return x + + +class SimpleFPN(nn.Module): + """Simple FPN for lightweight models""" + + def __init__(self, in_channels_list: List[int], out_channels: int = 256): + super().__init__() + + self.lateral_convs = nn.ModuleList([ + nn.Conv2d(in_ch, out_channels, 1, bias=False) + for in_ch in in_channels_list + ]) + + self.fpn_convs = nn.ModuleList([ + nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False) + for _ in in_channels_list + ]) + + def forward(self, features: List[torch.Tensor]) -> torch.Tensor: + # Build laterals + laterals = [conv(feat) for conv, feat in zip(self.lateral_convs, features)] + + # Build top-down pathway + for i in range(len(laterals) - 1, 0, -1): + upsampled = F.interpolate( + laterals[i], size=laterals[i-1].shape[-2:], + mode='bilinear', align_corners=False + ) + laterals[i-1] = laterals[i-1] + upsampled + + # Apply final convs and return highest resolution feature + outputs = [conv(lateral) for conv, lateral in zip(self.fpn_convs, laterals)] + return outputs[0] # Return highest resolution + + +class ASPP(nn.Module): + """Atrous Spatial Pyramid Pooling""" + + def __init__(self, in_channels: int, out_channels: int, dilations: List[int] = [1, 6, 12, 18]): + super().__init__() + + self.convs = nn.ModuleList() + for dilation in dilations: + if dilation == 1: + conv = nn.Conv2d(in_channels, out_channels//len(dilations), 1, bias=False) + else: + conv = nn.Conv2d(in_channels, out_channels//len(dilations), 3, + padding=dilation, dilation=dilation, bias=False) + self.convs.append(nn.Sequential( + conv, + nn.BatchNorm2d(out_channels//len(dilations)), + nn.ReLU(inplace=True) + )) + + # Global average pooling branch + self.global_pool = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(in_channels, out_channels//len(dilations), 1, bias=False), + nn.BatchNorm2d(out_channels//len(dilations)), + nn.ReLU(inplace=True) + ) + + # Final projection + total_channels = out_channels//len(dilations) * (len(dilations) + 1) + self.project = nn.Sequential( + nn.Conv2d(total_channels, out_channels, 1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(inplace=True), + nn.Dropout2d(0.1) + ) + + def forward(self, x): + h, w = x.shape[-2:] + features = [] + + # Apply dilated convolutions + for conv in self.convs: + features.append(conv(x)) + + # Global pooling branch + global_feat = self.global_pool(x) + global_feat = F.interpolate(global_feat, size=(h, w), mode='bilinear', align_corners=False) + features.append(global_feat) + + # Combine and project + combined = torch.cat(features, dim=1) + return self.project(combined) + + +class LightweightASPP(nn.Module): + """Lightweight ASPP using depthwise separable convolutions""" + + def __init__(self, in_channels: int, out_channels: int, dilations: List[int] = [1, 6, 12]): + super().__init__() + + self.convs = nn.ModuleList() + for dilation in dilations[:3]: # Limit to 3 dilations for efficiency + if dilation == 1: + conv = nn.Conv2d(in_channels, out_channels//len(dilations), 1, bias=False) + else: + conv = DepthwiseSeparableConv(in_channels, out_channels//len(dilations), 3, padding=dilation) + self.convs.append(conv) + + # Global pooling branch + self.global_pool = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(in_channels, out_channels//len(dilations), 1, bias=False), + nn.ReLU(inplace=True) + ) + + # Final projection + total_channels = out_channels//len(dilations) * (len(dilations) + 1) + self.project = nn.Sequential( + nn.Conv2d(total_channels, out_channels, 1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(inplace=True) + ) + + def forward(self, x): + h, w = x.shape[-2:] + features = [] + + # Apply dilated convolutions + for conv in self.convs: + features.append(conv(x)) + + # Global pooling branch + global_feat = self.global_pool(x) + global_feat = F.interpolate(global_feat, size=(h, w), mode='bilinear', align_corners=False) + features.append(global_feat) + + # Combine and project + combined = torch.cat(features, dim=1) + return self.project(combined) + + +def create_decoder(in_channels: int, num_classes: int, tier: str = 'standard', dropout_rate: float = 0.1) -> nn.Module: + """Create decoder based on tier""" + + if tier == 'lightweight': + return nn.Sequential( + ConvBlock(in_channels, in_channels//2, 3, 1, 1), + nn.Dropout2d(dropout_rate), + nn.Conv2d(in_channels//2, num_classes, 1) + ) + + elif tier == 'standard': + return nn.Sequential( + ConvBlock(in_channels, in_channels, 3, 1, 1), + ConvBlock(in_channels, in_channels//2, 3, 1, 1), + nn.Dropout2d(dropout_rate), + nn.Conv2d(in_channels//2, num_classes, 1) + ) + + elif tier == 'advanced': + return nn.Sequential( + ConvBlock(in_channels, in_channels, 3, 1, 1), + ConvBlock(in_channels, in_channels//2, 3, 1, 1), + ConvBlock(in_channels//2, in_channels//4, 3, 1, 1), + nn.Dropout2d(dropout_rate), + nn.Conv2d(in_channels//4, num_classes, 1) + ) + + else: + raise ValueError(f"Unknown tier: {tier}") + + +# Export commonly used components +__all__ = [ + 'SegmentationHeadInterface', + 'ConvBlock', + 'DepthwiseSeparableConv', + 'SimpleFPN', + 'ASPP', + 'LightweightASPP', + 'create_decoder' +] \ No newline at end of file diff --git a/segmentation_sivert/models/lightweight.py b/segmentation_sivert/models/lightweight.py new file mode 100644 index 00000000..3dc618e7 --- /dev/null +++ b/segmentation_sivert/models/lightweight.py @@ -0,0 +1,283 @@ +#!/usr/bin/env python3 +""" +Lightweight Segmentation Head for DFINE +Optimized for speed and low memory usage, suitable for edge devices and guns +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import List, Optional +import os +import sys + +# Add modules to path +sys.path.append(os.path.join(os.path.dirname(__file__), '..')) +from models.base import SegmentationHeadInterface, DepthwiseSeparableConv, LightweightASPP, SimpleFPN + + +class MobileNetV2Block(nn.Module): + """MobileNetV2 inverted residual block for lightweight processing""" + + def __init__(self, in_channels: int, out_channels: int, expansion_factor: int = 6, stride: int = 1): + super().__init__() + + hidden_dim = in_channels * expansion_factor + self.use_residual = stride == 1 and in_channels == out_channels + + layers = [] + + # Expansion phase + if expansion_factor != 1: + layers.extend([ + nn.Conv2d(in_channels, hidden_dim, 1, bias=False), + nn.BatchNorm2d(hidden_dim), + nn.ReLU6(inplace=True) + ]) + + # Depthwise convolution + layers.extend([ + nn.Conv2d(hidden_dim, hidden_dim, 3, stride=stride, padding=1, groups=hidden_dim, bias=False), + nn.BatchNorm2d(hidden_dim), + nn.ReLU6(inplace=True) + ]) + + # Projection phase + layers.extend([ + nn.Conv2d(hidden_dim, out_channels, 1, bias=False), + nn.BatchNorm2d(out_channels) + ]) + + self.conv = nn.Sequential(*layers) + + def forward(self, x): + if self.use_residual: + return x + self.conv(x) + else: + return self.conv(x) + + +class LightweightFPN(nn.Module): + """Lightweight Feature Pyramid Network using depthwise separable convolutions""" + + def __init__(self, in_channels_list: List[int], out_channels: int = 128): + super().__init__() + + # Use smaller output channels for lightweight model + self.out_channels = min(out_channels, 128) + + # Lightweight lateral connections + self.lateral_convs = nn.ModuleList([ + nn.Sequential( + nn.Conv2d(in_ch, self.out_channels, 1, bias=False), + nn.BatchNorm2d(self.out_channels), + nn.ReLU6(inplace=True) + ) for in_ch in in_channels_list + ]) + + # Lightweight output processing using depthwise separable convolutions + self.fpn_convs = nn.ModuleList([ + DepthwiseSeparableConv(self.out_channels, self.out_channels, 3, 1, 1) + for _ in in_channels_list + ]) + + def forward(self, features: List[torch.Tensor]) -> torch.Tensor: + # Build laterals + laterals = [conv(feat) for conv, feat in zip(self.lateral_convs, features)] + + # Build top-down pathway (simplified) + for i in range(len(laterals) - 1, 0, -1): + # Simple upsampling + addition + upsampled = F.interpolate( + laterals[i], size=laterals[i-1].shape[-2:], + mode='bilinear', align_corners=False + ) + laterals[i-1] = laterals[i-1] + upsampled + + # Apply lightweight processing and return highest resolution + output = self.fpn_convs[0](laterals[0]) + return output + + +class UltraLightweightASPP(nn.Module): + """Ultra-lightweight ASPP with minimal dilations and channels""" + + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + + # Reduce to only 2 dilations for maximum efficiency + branch_channels = out_channels // 3 + + # Point-wise convolution + self.branch1 = nn.Sequential( + nn.Conv2d(in_channels, branch_channels, 1, bias=False), + nn.BatchNorm2d(branch_channels), + nn.ReLU6(inplace=True) + ) + + # Single dilated convolution (standard conv with dilation) + self.branch2 = nn.Sequential( + nn.Conv2d(in_channels, branch_channels, 3, padding=6, dilation=6, bias=False), + nn.BatchNorm2d(branch_channels), + nn.ReLU6(inplace=True) + ) + + # Global average pooling + self.branch3 = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(in_channels, branch_channels, 1, bias=False), + nn.BatchNorm2d(branch_channels), + nn.ReLU6(inplace=True) + ) + + # Final fusion + self.conv_out = nn.Sequential( + nn.Conv2d(branch_channels * 3, out_channels, 1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU6(inplace=True) + ) + + def forward(self, x): + h, w = x.shape[-2:] + + # Three branches + branch1 = self.branch1(x) + branch2 = self.branch2(x) + branch3 = F.interpolate(self.branch3(x), size=(h, w), mode='bilinear', align_corners=False) + + # Concatenate and fuse + out = torch.cat([branch1, branch2, branch3], dim=1) + return self.conv_out(out) + + +class LightweightSegmentationHead(SegmentationHeadInterface): + """Lightweight segmentation head optimized for speed and low memory usage + + Designed for edge devices, AI chips in guns, and real-time applications. + Prioritizes speed and efficiency over maximum accuracy. + """ + + def __init__(self, + in_channels_list: List[int], + num_classes: int = 7, + feature_dim: int = 128, # Reduced from 256 + dropout_rate: float = 0.05): # Reduced dropout + """Initialize lightweight segmentation head + + Args: + in_channels_list: List of input channel counts from backbone + num_classes: Number of segmentation classes (default: 7 for Pascal Person Parts) + feature_dim: Feature dimension for processing (default: 128, reduced for efficiency) + dropout_rate: Dropout rate (default: 0.05, reduced for efficiency) + """ + super().__init__( + in_channels_list=in_channels_list, + num_classes=num_classes, + feature_dim=feature_dim, + dropout_rate=dropout_rate, + tier='lightweight' + ) + + def _build_head(self): + """Build lightweight segmentation head architecture""" + + # Lightweight Feature Pyramid Network + self.fpn = LightweightFPN(self.in_channels_list, self.feature_dim) + + # Ultra-lightweight ASPP + self.aspp = UltraLightweightASPP(self.feature_dim, self.feature_dim) + + # Minimal decoder for fastest processing + self.decoder = nn.Sequential( + # Single processing stage + DepthwiseSeparableConv(self.feature_dim, self.feature_dim//2, 3, 1, 1), + + # Minimal dropout + nn.Dropout2d(self.dropout_rate), + + # Final classification + nn.Conv2d(self.feature_dim//2, self.num_classes, 1) + ) + + print(f"🏗️ Built lightweight segmentation head:") + print(f" Feature dim: {self.feature_dim} (reduced)") + print(f" Dropout rate: {self.dropout_rate} (minimal)") + print(f" Classes: {self.num_classes}") + print(f" 🎯 Optimized for: Speed and low memory usage") + + def forward(self, features: List[torch.Tensor]) -> torch.Tensor: + """Forward pass through lightweight segmentation head + + Args: + features: List of feature tensors from backbone [feat1, feat2, feat3, feat4] + + Returns: + Segmentation logits [N, num_classes, H, W] where H, W are 1/4 of input size + """ + # Lightweight multi-scale feature fusion + fused_feature = self.fpn(features) + + # Apply ultra-lightweight ASPP + context_feature = self.aspp(fused_feature) + + # Generate segmentation logits with minimal processing + seg_logits = self.decoder(context_feature) + + return seg_logits + + def get_model_info(self) -> dict: + """Get detailed model information""" + base_info = super().get_model_info() + + # Add lightweight-specific info + base_info.update({ + 'fpn_feature_dim': self.feature_dim, + 'aspp_branches': 3, # Reduced from 4 + 'decoder_stages': 1, # Minimal decoder + 'uses_depthwise_separable': True, + 'supports_multiscale': False, # Simplified for speed + 'optimized_for': 'speed_and_memory', + 'target_platforms': ['edge_devices', 'mobile', 'embedded', 'gun_ai_chip'], + 'estimated_fps': '60+', # On typical edge hardware + 'estimated_memory_mb': '<50' + }) + + return base_info + + def get_complexity_info(self) -> dict: + """Get lightweight-specific complexity information""" + info = super().get_complexity_info() + + # Add lightweight-specific metrics + info.update({ + 'depthwise_conv_ratio': 0.8, # 80% of convs are depthwise separable + 'parameter_reduction': 0.6, # ~60% fewer parameters than standard + 'flops_reduction': 0.7, # ~70% fewer FLOPs than standard + 'memory_efficient': True, + 'edge_device_ready': True + }) + + return info + + +def create_lightweight_segmentation_head(in_channels_list: List[int], + num_classes: int = 7, + feature_dim: int = 128, + dropout_rate: float = 0.05) -> LightweightSegmentationHead: + """Factory function to create lightweight segmentation head""" + return LightweightSegmentationHead( + in_channels_list=in_channels_list, + num_classes=num_classes, + feature_dim=feature_dim, + dropout_rate=dropout_rate + ) + + +# Export for use in other modules +__all__ = [ + 'LightweightSegmentationHead', + 'LightweightFPN', + 'UltraLightweightASPP', + 'MobileNetV2Block', + 'create_lightweight_segmentation_head' +] \ No newline at end of file diff --git a/segmentation_sivert/models/standard.py b/segmentation_sivert/models/standard.py new file mode 100644 index 00000000..09a166bf --- /dev/null +++ b/segmentation_sivert/models/standard.py @@ -0,0 +1,243 @@ +#!/usr/bin/env python3 +""" +Standard Segmentation Head for DFINE +Refactored from the original implementation with improved structure +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import List, Optional +import os +import sys + +# Add modules to path +sys.path.append(os.path.join(os.path.dirname(__file__), '..')) +from models.base import SegmentationHeadInterface, ConvBlock, ASPP, create_decoder + + +class StandardFPN(nn.Module): + """Standard FPN with feature fusion for balanced performance""" + + def __init__(self, in_channels_list: List[int], out_channels: int = 256): + super().__init__() + + # Lateral connections + self.lateral_convs = nn.ModuleList([ + nn.Conv2d(in_ch, out_channels, 1, bias=False) + for in_ch in in_channels_list + ]) + + # Output convolutions + self.fpn_convs = nn.ModuleList([ + nn.Sequential( + nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(inplace=True) + ) + for _ in in_channels_list + ]) + + # Feature fusion weights (learnable) + self.fusion_weights = nn.Parameter(torch.ones(len(in_channels_list))) + + def forward(self, features: List[torch.Tensor]) -> torch.Tensor: + # Build laterals + laterals = [conv(feat) for conv, feat in zip(self.lateral_convs, features)] + + # Build top-down pathway + for i in range(len(laterals) - 1, 0, -1): + upsampled = F.interpolate( + laterals[i], size=laterals[i-1].shape[-2:], + mode='bilinear', align_corners=False + ) + laterals[i-1] = laterals[i-1] + upsampled + + # Apply final convs + outputs = [conv(lateral) for conv, lateral in zip(self.fpn_convs, laterals)] + + # Weighted feature fusion + target_size = outputs[0].shape[-2:] + fused_features = [] + + for i, feat in enumerate(outputs): + if feat.shape[-2:] != target_size: + feat = F.interpolate(feat, size=target_size, mode='bilinear', align_corners=False) + fused_features.append(feat * self.fusion_weights[i]) + + # Return weighted sum of features + return sum(fused_features) + + +class StandardASPP(nn.Module): + """Standard ASPP for multi-scale context""" + + def __init__(self, in_channels: int, out_channels: int, dilations: List[int] = [1, 6, 12]): + super().__init__() + + self.convs = nn.ModuleList() + for dilation in dilations: + if dilation == 1: + conv = nn.Conv2d(in_channels, out_channels//len(dilations), 1, bias=False) + else: + conv = nn.Conv2d(in_channels, out_channels//len(dilations), 3, + padding=dilation, dilation=dilation, bias=False) + self.convs.append(nn.Sequential( + conv, + nn.BatchNorm2d(out_channels//len(dilations)), + nn.ReLU(inplace=True) + )) + + # Global average pooling branch + self.global_pool = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(in_channels, out_channels//len(dilations), 1, bias=False), + nn.BatchNorm2d(out_channels//len(dilations)), + nn.ReLU(inplace=True) + ) + + # Final projection + total_channels = out_channels//len(dilations) * (len(dilations) + 1) + self.project = nn.Sequential( + nn.Conv2d(total_channels, out_channels, 1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(inplace=True), + nn.Dropout2d(0.1) + ) + + def forward(self, x): + h, w = x.shape[-2:] + features = [] + + # Apply dilated convolutions + for conv in self.convs: + features.append(conv(x)) + + # Global pooling branch + global_feat = self.global_pool(x) + global_feat = F.interpolate(global_feat, size=(h, w), mode='bilinear', align_corners=False) + features.append(global_feat) + + # Combine and project + combined = torch.cat(features, dim=1) + return self.project(combined) + + +class StandardSegmentationHead(SegmentationHeadInterface): + """Standard segmentation head with balanced performance and accuracy + + This is the refactored version of the original segmentation head, + providing good performance with reasonable computational cost. + """ + + def __init__(self, + in_channels_list: List[int], + num_classes: int = 7, + feature_dim: int = 256, + dropout_rate: float = 0.1): + """Initialize standard segmentation head + + Args: + in_channels_list: List of input channel counts from backbone + num_classes: Number of segmentation classes (default: 7 for Pascal Person Parts) + feature_dim: Feature dimension for processing (default: 256) + dropout_rate: Dropout rate (default: 0.1) + """ + super().__init__( + in_channels_list=in_channels_list, + num_classes=num_classes, + feature_dim=feature_dim, + dropout_rate=dropout_rate, + tier='standard' + ) + + def _build_head(self): + """Build standard segmentation head architecture""" + + # Feature Pyramid Network for multi-scale feature fusion + self.fpn = StandardFPN(self.in_channels_list, self.feature_dim) + + # Atrous Spatial Pyramid Pooling for multi-scale context + self.aspp = StandardASPP(self.feature_dim, self.feature_dim) + + # Decoder for final segmentation prediction + self.decoder = nn.Sequential( + # First stage - maintain feature richness + ConvBlock(self.feature_dim, self.feature_dim, 3, 1, 1), + + # Second stage - reduce dimensions + ConvBlock(self.feature_dim, self.feature_dim//2, 3, 1, 1), + + # Dropout for regularization + nn.Dropout2d(self.dropout_rate), + + # Final classification layer + nn.Conv2d(self.feature_dim//2, self.num_classes, 1) + ) + + print(f"🏗️ Built standard segmentation head:") + print(f" Feature dim: {self.feature_dim}") + print(f" Dropout rate: {self.dropout_rate}") + print(f" Classes: {self.num_classes}") + + def forward(self, features: List[torch.Tensor]) -> torch.Tensor: + """Forward pass through standard segmentation head + + Args: + features: List of feature tensors from backbone [feat1, feat2, feat3, feat4] + + Returns: + Segmentation logits [N, num_classes, H, W] where H, W are 1/4 of input size + """ + # Multi-scale feature fusion with FPN + fused_feature = self.fpn(features) + + # Apply ASPP for multi-scale context + context_feature = self.aspp(fused_feature) + + # Generate segmentation logits + seg_logits = self.decoder(context_feature) + + return seg_logits + + def get_model_info(self) -> dict: + """Get detailed model information""" + base_info = super().get_model_info() + + # Add standard-specific info + base_info.update({ + 'fpn_feature_dim': self.feature_dim, + 'aspp_dilations': [1, 6, 12], + 'decoder_stages': 3, + 'supports_multiscale': True, + 'optimized_for': 'balanced_performance' + }) + + return base_info + + +# Backward compatibility - alias for the original name +EnhancedSegmentationHead = StandardSegmentationHead + + +def create_standard_segmentation_head(in_channels_list: List[int], + num_classes: int = 7, + feature_dim: int = 256, + dropout_rate: float = 0.1) -> StandardSegmentationHead: + """Factory function to create standard segmentation head""" + return StandardSegmentationHead( + in_channels_list=in_channels_list, + num_classes=num_classes, + feature_dim=feature_dim, + dropout_rate=dropout_rate + ) + + +# Export for use in other modules +__all__ = [ + 'StandardSegmentationHead', + 'StandardFPN', + 'StandardASPP', + 'EnhancedSegmentationHead', # Backward compatibility + 'create_standard_segmentation_head' +] \ No newline at end of file diff --git a/segmentation_sivert/prepare_camo.py b/segmentation_sivert/prepare_camo.py new file mode 100644 index 00000000..c1449e97 --- /dev/null +++ b/segmentation_sivert/prepare_camo.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +""" +Prepare CAMO dataset for the segmentation training framework. + +Source structure: + CAMO-V.1.0-CVIU2019/Images/Train/ (1000 .jpg) + CAMO-V.1.0-CVIU2019/Images/Test/ (250 .jpg) + CAMO-V.1.0-CVIU2019/GT/ (1250 .png, 0/255 binary) + +Output: + camo/images/train/ + camo/masks/train/ + camo/images/val/ + camo/masks/val/ + +Masks converted from 0/255 → 0/1. +""" + +import os +import argparse +import numpy as np +import cv2 +from pathlib import Path +from tqdm import tqdm + + +CAMO_ROOT = "/home/berna/.cache/kagglehub/datasets/ivanomelchenkoim11/camo-dataset/versions/1/CAMO-V.1.0-CVIU2019" + + +def prepare_split(img_dir, gt_dir, dst_img_dir, dst_mask_dir): + dst_img_dir.mkdir(parents=True, exist_ok=True) + dst_mask_dir.mkdir(parents=True, exist_ok=True) + + images = sorted([f for f in os.listdir(img_dir) if f.endswith(('.jpg', '.jpeg', '.png'))]) + print(f" Processing {len(images)} images...") + + missing = 0 + for img_name in tqdm(images): + stem = Path(img_name).stem + mask_name = stem + '.png' + src_mask = gt_dir / mask_name + + if not src_mask.exists(): + missing += 1 + continue + + # Symlink image with absolute path + dst_img = dst_img_dir / img_name + if dst_img.exists() or dst_img.is_symlink(): + dst_img.unlink() + os.symlink((img_dir / img_name).resolve(), dst_img) + + # Convert mask 0/255 → 0/1 + mask = cv2.imread(str(src_mask), cv2.IMREAD_GRAYSCALE) + mask_binary = (mask > 127).astype(np.uint8) + cv2.imwrite(str(dst_mask_dir / mask_name), mask_binary) + + if missing: + print(f" Warning: {missing} images had no matching mask and were skipped.") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--src', default=CAMO_ROOT) + parser.add_argument('--dst', default='data/camo') + args = parser.parse_args() + + src = Path(args.src) + dst = Path(args.dst) + gt_dir = src / 'GT' + + print("Preparing train split...") + prepare_split(src / 'Images' / 'Train', gt_dir, dst / 'images' / 'train', dst / 'masks' / 'train') + + print("Preparing val split...") + prepare_split(src / 'Images' / 'Test', gt_dir, dst / 'images' / 'val', dst / 'masks' / 'val') + + train_count = len(list((dst / 'images' / 'train').iterdir())) + val_count = len(list((dst / 'images' / 'val').iterdir())) + print(f"\nDone. {train_count} train / {val_count} val images ready at: {dst}") + + +if __name__ == '__main__': + main() diff --git a/segmentation_sivert/prepare_cod10k.py b/segmentation_sivert/prepare_cod10k.py new file mode 100644 index 00000000..361c30d0 --- /dev/null +++ b/segmentation_sivert/prepare_cod10k.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python3 +""" +Prepare COD10K-v3 dataset for the segmentation training framework. + +Converts from: + COD10K-v3/Train/Image/ + COD10K-v3/Train/GT_Object/ + COD10K-v3/Test/Image/ + COD10K-v3/Test/GT_Object/ + +To: + cod10k/images/train/ + cod10k/masks/train/ + cod10k/images/val/ + cod10k/masks/val/ + +Masks are converted from 0/255 binary to 0/1 (0=background, 1=camouflaged). +""" + +import os +import argparse +import numpy as np +import cv2 +from pathlib import Path +from tqdm import tqdm + + +def prepare_split(src_img_dir, src_mask_dir, dst_img_dir, dst_mask_dir): + dst_img_dir.mkdir(parents=True, exist_ok=True) + dst_mask_dir.mkdir(parents=True, exist_ok=True) + + images = sorted([f for f in os.listdir(src_img_dir) if f.endswith(('.jpg', '.jpeg', '.png'))]) + print(f" Processing {len(images)} images from {src_img_dir.name}...") + + missing_masks = 0 + for img_name in tqdm(images): + mask_name = Path(img_name).stem + '.png' + src_mask = src_mask_dir / mask_name + if not src_mask.exists(): + missing_masks += 1 + continue + + # Symlink image using absolute path so symlink resolves correctly + dst_img = dst_img_dir / img_name + if dst_img.exists() or dst_img.is_symlink(): + dst_img.unlink() + os.symlink((src_img_dir / img_name).resolve(), dst_img) + + # Convert mask 0/255 → 0/1 + mask = cv2.imread(str(src_mask), cv2.IMREAD_GRAYSCALE) + mask_binary = (mask > 127).astype(np.uint8) + cv2.imwrite(str(dst_mask_dir / mask_name), mask_binary) + + if missing_masks: + print(f" Warning: {missing_masks} images had no matching mask and were skipped.") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--src', default='data/COD10K-v3', + help='Path to extracted COD10K-v3 directory') + parser.add_argument('--dst', default='data/cod10k', + help='Output directory for prepared dataset') + args = parser.parse_args() + + src = Path(args.src) + dst = Path(args.dst) + + assert (src / 'Train' / 'Image').exists(), f"Not found: {src}/Train/Image" + assert (src / 'Test' / 'Image').exists(), f"Not found: {src}/Test/Image" + + print("Preparing train split...") + prepare_split( + src / 'Train' / 'Image', + src / 'Train' / 'GT_Object', + dst / 'images' / 'train', + dst / 'masks' / 'train', + ) + + print("Preparing val split...") + prepare_split( + src / 'Test' / 'Image', + src / 'Test' / 'GT_Object', + dst / 'images' / 'val', + dst / 'masks' / 'val', + ) + + train_count = len(list((dst / 'images' / 'train').iterdir())) + val_count = len(list((dst / 'images' / 'val').iterdir())) + print(f"\nDone. {train_count} train / {val_count} val images ready at: {dst}") + + +if __name__ == '__main__': + main() diff --git a/segmentation_sivert/prepare_mcs1k.py b/segmentation_sivert/prepare_mcs1k.py new file mode 100644 index 00000000..1e23d407 --- /dev/null +++ b/segmentation_sivert/prepare_mcs1k.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +""" +Prepare MCS1K dataset for the segmentation training framework. + +Source: + dataset-splitM/Training/images/ + dataset-splitM/Training/GT/ + dataset-splitM/Testing/images/ + dataset-splitM/Testing/GT/ + +Output: + mcs1k/images/train/ + mcs1k/masks/train/ + mcs1k/images/val/ + mcs1k/masks/val/ + +Masks thresholded at 127 → binary 0/1. +""" + +import os +import argparse +import numpy as np +import cv2 +from pathlib import Path +from tqdm import tqdm + +MCS1K_ROOT = "/home/berna/.cache/kagglehub/datasets/aalihhiader/military-camouflage-soldiers-dataset-mcs1k/versions/1/dataset-splitM" + + +def prepare_split(img_dir, gt_dir, dst_img_dir, dst_mask_dir): + dst_img_dir.mkdir(parents=True, exist_ok=True) + dst_mask_dir.mkdir(parents=True, exist_ok=True) + + images = sorted([f for f in os.listdir(img_dir) if f.endswith(('.jpg', '.jpeg', '.png'))]) + print(f" Processing {len(images)} images...") + + missing = 0 + for img_name in tqdm(images): + stem = Path(img_name).stem + # GT has same filename as image + mask_name = stem + '.jpg' + src_mask = gt_dir / mask_name + if not src_mask.exists(): + # Try png + src_mask = gt_dir / (stem + '.png') + if not src_mask.exists(): + missing += 1 + continue + + dst_img = dst_img_dir / img_name + if dst_img.exists() or dst_img.is_symlink(): + dst_img.unlink() + os.symlink((img_dir / img_name).resolve(), dst_img) + + # Threshold to binary 0/1 + mask = cv2.imread(str(src_mask), cv2.IMREAD_GRAYSCALE) + mask_binary = (mask > 127).astype(np.uint8) + cv2.imwrite(str(dst_mask_dir / (stem + '.png')), mask_binary) + + if missing: + print(f" Warning: {missing} images had no matching mask and were skipped.") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--src', default=MCS1K_ROOT) + parser.add_argument('--dst', default='data/mcs1k') + args = parser.parse_args() + + src = Path(args.src) + dst = Path(args.dst) + + print("Preparing train split...") + prepare_split(src / 'Training' / 'images', src / 'Training' / 'GT', + dst / 'images' / 'train', dst / 'masks' / 'train') + + print("Preparing val split...") + prepare_split(src / 'Testing' / 'images', src / 'Testing' / 'GT', + dst / 'images' / 'val', dst / 'masks' / 'val') + + train_count = len(list((dst / 'images' / 'train').iterdir())) + val_count = len(list((dst / 'images' / 'val').iterdir())) + print(f"\nDone. {train_count} train / {val_count} val images ready at: {dst}") + + +if __name__ == '__main__': + main() diff --git a/segmentation_sivert/train.py b/segmentation_sivert/train.py new file mode 100644 index 00000000..e2d35415 --- /dev/null +++ b/segmentation_sivert/train.py @@ -0,0 +1,746 @@ +#!/usr/bin/env python3 +""" +Unified DFINE Segmentation Training Script +Single script supporting lightweight, standard, and advanced tiers +Uses tier-specific YAML configurations for all parameters +""" + +import os +import sys +import argparse +import yaml +import torch +import torch.optim as optim +from torch.utils.data import DataLoader +try: + import wandb + WANDB_AVAILABLE = True +except ImportError: + WANDB_AVAILABLE = False +from tqdm import tqdm +from pathlib import Path + +# Add paths for imports +current_dir = os.path.dirname(os.path.abspath(__file__)) +if current_dir not in sys.path: + sys.path.insert(0, current_dir) +sys.path.append('src') + +# Import core modules +from core.datasets import create_dataset +from core.losses import create_loss +from core.metrics import create_metrics_tracker +from core.models import load_pretrained_dfine, get_actual_backbone_channels, create_combined_model + +# Import tier-specific models +from models import create_segmentation_head + + +def parse_args(): + """Parse command line arguments""" + parser = argparse.ArgumentParser(description='Unified DFINE Segmentation Training') + + # Required tier selection + parser.add_argument('--tier', type=str, required=True, + choices=['lightweight', 'standard', 'advanced'], + help='Model tier to train (lightweight/standard/advanced)') + + # Optional overrides + parser.add_argument('--config-dir', default='configs', + help='Directory containing tier configuration files') + parser.add_argument('--config-file', default=None, + help='Path to a specific config YAML file (overrides --config-dir tier lookup)') + parser.add_argument('--dataset', default=None, + help='Override dataset path from config') + parser.add_argument('--output-dir', default=None, + help='Override output directory from config') + parser.add_argument('--resume', default=None, + help='Resume training from checkpoint') + parser.add_argument('--finetune', action='store_true', + help='Load only model weights from checkpoint (not optimizer/scheduler)') + + # System configuration + parser.add_argument('--device', type=str, default='auto', + choices=['auto', 'cuda', 'cpu'], + help='Device to use for training') + parser.add_argument('--num-workers', type=int, default=None, + help='Override number of data loading workers') + + # Logging + parser.add_argument('--no-wandb', action='store_true', + help='Disable Weights & Biases logging') + parser.add_argument('--wandb-project', default=None, + help='Override wandb project name') + parser.add_argument('--wandb-run-name', default=None, + help='Override wandb run name') + + # Debug options + parser.add_argument('--dry-run', action='store_true', + help='Test configuration loading without training') + parser.add_argument('--verbose', action='store_true', + help='Enable verbose logging') + + return parser.parse_args() + + +def parse_numeric_values(config): + """Recursively parse string numeric values in config to proper types""" + if isinstance(config, dict): + for key, value in config.items(): + if isinstance(value, str): + # Try to parse scientific notation + try: + if 'e-' in value.lower() or 'e+' in value.lower(): + config[key] = float(value) + elif value.replace('.', '').replace('-', '').isdigit(): + if '.' in value: + config[key] = float(value) + else: + config[key] = int(value) + except (ValueError, AttributeError): + pass # Keep as string if parsing fails + elif isinstance(value, dict): + parse_numeric_values(value) + elif isinstance(value, list): + for i, item in enumerate(value): + if isinstance(item, dict): + parse_numeric_values(item) + elif isinstance(item, str): + try: + if 'e-' in item.lower() or 'e+' in item.lower(): + value[i] = float(item) + elif item.replace('.', '').replace('-', '').isdigit(): + if '.' in item: + value[i] = float(item) + else: + value[i] = int(item) + except (ValueError, AttributeError): + pass + + +def load_tier_config(tier: str, config_dir: str, config_file: str = None) -> dict: + """Load tier-specific configuration""" + config_path = config_file if config_file else os.path.join(config_dir, f"{tier}.yaml") + + if not os.path.exists(config_path): + raise FileNotFoundError(f"Configuration file not found: {config_path}") + + print(f"📋 Loading {tier} tier configuration from: {config_path}") + + with open(config_path, 'r') as f: + config = yaml.safe_load(f) + + # Recursively resolve base configs + if 'base' in config and config['base']: + base_filename = config['base'] + if not base_filename.startswith('/'): + base_path = os.path.join(os.path.dirname(os.path.abspath(config_path)), base_filename) + else: + base_path = base_filename + + print(f"📋 Loading base configuration from: {base_path}") + base_config = load_tier_config(tier, config_dir, config_file=base_path) + + merged_config = deep_merge_dict(base_config, config) + parse_numeric_values(merged_config) + return merged_config + + parse_numeric_values(config) + return config + + +def deep_merge_dict(base_dict: dict, override_dict: dict) -> dict: + """Deep merge two dictionaries""" + result = base_dict.copy() + + for key, value in override_dict.items(): + if key in result and isinstance(result[key], dict) and isinstance(value, dict): + result[key] = deep_merge_dict(result[key], value) + else: + result[key] = value + + return result + + +def setup_device(device_arg: str) -> torch.device: + """Setup compute device""" + if device_arg == 'auto': + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + else: + device = torch.device(device_arg) + + print(f"🚀 Using device: {device}") + if device.type == 'cuda': + print(f" GPU: {torch.cuda.get_device_name()}") + print(f" Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB") + + return device + + +def create_model(config: dict, tier: str, device: torch.device): + """Create the segmentation model based on tier""" + print(f"🏗️ Creating {tier} tier segmentation model...") + + # Load pretrained DFINE model + dfine_config = config['dfine']['config_path'] + dfine_checkpoint = config['dfine']['checkpoint_path'] + + dfine_model = load_pretrained_dfine(dfine_config, dfine_checkpoint) + + # Get backbone channels + backbone_channels = get_actual_backbone_channels(dfine_model) + print(f" Detected backbone channels: {backbone_channels}") + + # Get model-specific parameters + model_config = config.get('model', {}) + seg_head_config = model_config.get('segmentation_head', {}) + + # Create tier-specific segmentation head + seg_head = create_segmentation_head( + tier=tier, + in_channels_list=backbone_channels, + num_classes=config['dataset']['num_classes'], + feature_dim=seg_head_config.get('feature_dim', 256), + dropout_rate=seg_head_config.get('dropout_rate', 0.1) + ) + + # Create combined model + model = create_combined_model( + dfine_model=dfine_model, + seg_head=seg_head, + freeze_detection=config['training'].get('freeze_detection', True) + ) + + model = model.to(device) + + # Print model info + model_info = model.get_model_info() + print(f"📊 Model Information:") + print(f" Tier: {model_info['tier']}") + print(f" Total parameters: {model_info['total_params']:,}") + print(f" Trainable parameters: {model_info['trainable_params']:,}") + print(f" Model size: {model_info['model_size_mb']:.2f} MB") + + return model + + +def create_datasets(config: dict, tier: str, dataset_override: str = None): + """Create training and validation datasets""" + print("📊 Creating datasets...") + + dataset_config = config['dataset'] + training_config = config['training'] + + # Use override or config dataset path + dataset_root = dataset_override or dataset_config['root_dir'] + + dataset_name = dataset_config.get('name', 'pascal_person_parts') + + train_dataset = create_dataset( + dataset_name=dataset_name, + root_dir=dataset_root, + split='train', + image_size=dataset_config['image_size'], + tier=tier, + multi_scale=training_config.get('multi_scale_training', False) + ) + + val_dataset = create_dataset( + dataset_name=dataset_name, + root_dir=dataset_root, + split='val', + image_size=dataset_config['image_size'], + tier=tier, + multi_scale=False + ) + + # Create data loaders + system_config = config.get('system', {}) + + train_loader = DataLoader( + train_dataset, + batch_size=training_config['batch_size'], + shuffle=True, + num_workers=system_config.get('num_workers', 4), + pin_memory=system_config.get('pin_memory', True), + drop_last=True + ) + + val_loader = DataLoader( + val_dataset, + batch_size=training_config['batch_size'], + shuffle=False, + num_workers=system_config.get('num_workers', 4), + pin_memory=system_config.get('pin_memory', True) + ) + + print(f" Train samples: {len(train_dataset)}") + print(f" Val samples: {len(val_dataset)}") + print(f" Train batches: {len(train_loader)}") + print(f" Val batches: {len(val_loader)}") + + return train_loader, val_loader + + +def create_optimizer_and_scheduler(model, config: dict): + """Create optimizer and learning rate scheduler""" + training_config = config['training'] + + # Get trainable parameters + trainable_params = [p for p in model.parameters() if p.requires_grad] + + # Create optimizer + optimizer = optim.AdamW( + trainable_params, + lr=training_config['learning_rate'], + weight_decay=training_config.get('weight_decay', 1e-4) + ) + + # Create scheduler + scheduler_config = training_config.get('scheduler', {}) + scheduler_type = scheduler_config.get('type', 'cosine_warmup') + + if scheduler_type == 'cosine_warmup': + scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts( + optimizer, + T_0=scheduler_config.get('T_0', 15), + T_mult=scheduler_config.get('T_mult', 2), + eta_min=scheduler_config.get('eta_min', 1e-6) + ) + else: + # Default to cosine scheduler + scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts( + optimizer, T_0=15, T_mult=2, eta_min=1e-6 + ) + + print(f"🎯 Optimizer: AdamW with {len(trainable_params)} parameter groups") + print(f"📈 Scheduler: {scheduler_type}") + + return optimizer, scheduler + + +def train_epoch(model, train_loader, criterion, optimizer, device, epoch, config, args): + """Train for one epoch""" + model.train() + + total_loss = 0.0 + total_miou = 0.0 + num_batches = len(train_loader) + + pbar = tqdm(train_loader, desc=f'Epoch {epoch}', leave=False) + + for batch_idx, (images, masks) in enumerate(pbar): + images = images.to(device, non_blocking=True) + masks = masks.to(device, non_blocking=True) + + # Forward pass + optimizer.zero_grad() + outputs = model(images) + seg_pred = outputs['segmentation'] + + # Calculate loss (handle different loss function signatures) + if hasattr(criterion, 'forward') and 'boundary' in str(criterion.forward.__code__.co_varnames): + # Advanced loss with boundary prediction + boundary_pred = outputs.get('boundary', None) + loss, loss_dict = criterion(seg_pred, masks, boundary_pred) + else: + # Standard loss + loss, loss_dict = criterion(seg_pred, masks) + + # Backward pass + loss.backward() + + # Gradient clipping + gradient_clipping = config['training'].get('gradient_clipping', 1.0) + if gradient_clipping > 0: + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=gradient_clipping) + + optimizer.step() + + # Calculate metrics + with torch.no_grad(): + from core.metrics import compute_miou + miou = compute_miou(seg_pred, masks) + + # Update running averages + total_loss += loss.item() + total_miou += miou + + # Update progress bar + pbar.set_postfix({ + 'Loss': f'{loss.item():.4f}', + 'mIoU': f'{miou:.3f}', + 'LR': f'{optimizer.param_groups[0]["lr"]:.2e}' + }) + + # Log batch metrics + if batch_idx % 20 == 0 and WANDB_AVAILABLE and not args.no_wandb: + log_dict = { + 'train/batch_loss': loss.item(), + 'train/batch_miou': miou, + 'train/lr': optimizer.param_groups[0]['lr'], + 'epoch': epoch + } + + # Add loss components if available + if isinstance(loss_dict, dict): + for key, value in loss_dict.items(): + if key != 'total': + log_dict[f'train/batch_{key}'] = value + + wandb.log(log_dict) + + avg_loss = total_loss / num_batches + avg_miou = total_miou / num_batches + + return avg_loss, avg_miou + + +def validate_epoch(model, val_loader, criterion, device, config): + """Validate for one epoch""" + model.eval() + + total_loss = 0.0 + total_miou = 0.0 + num_batches = len(val_loader) + + # Create metrics tracker for detailed evaluation + dataset_config = config['dataset'] + metrics_tracker = create_metrics_tracker( + num_classes=dataset_config['num_classes'], + class_names=dataset_config.get('class_names', None) + ) + + with torch.no_grad(): + pbar = tqdm(val_loader, desc='Validation', leave=False) + + for images, masks in pbar: + images = images.to(device, non_blocking=True) + masks = masks.to(device, non_blocking=True) + + # Forward pass + outputs = model(images) + seg_pred = outputs['segmentation'] + + # Calculate loss + if hasattr(criterion, 'forward') and 'boundary' in str(criterion.forward.__code__.co_varnames): + boundary_pred = outputs.get('boundary', None) + loss, loss_dict = criterion(seg_pred, masks, boundary_pred) + else: + loss, loss_dict = criterion(seg_pred, masks) + + # Calculate metrics + from core.metrics import compute_miou + miou = compute_miou(seg_pred, masks) + + # Update metrics tracker + metrics_tracker.update(seg_pred, masks) + + # Update running averages + total_loss += loss.item() + total_miou += miou + + # Update progress bar + pbar.set_postfix({ + 'Loss': f'{loss.item():.4f}', + 'mIoU': f'{miou:.3f}' + }) + + avg_loss = total_loss / num_batches + avg_miou = total_miou / num_batches + + # Get detailed metrics + detailed_metrics = metrics_tracker.get_summary() + + return avg_loss, avg_miou, detailed_metrics + + +def save_checkpoint(model, optimizer, scheduler, epoch, best_miou, config, args, filename=None): + """Save training checkpoint""" + output_config = config['output'] + output_dir = args.output_dir or output_config['base_dir'] + + if filename is None: + filename = f'checkpoint_epoch_{epoch}.pth' + + filepath = os.path.join(output_dir, filename) + + # Save hyperparameters from config + hyperparameters = { + 'tier': args.tier, + 'feature_dim': config.get('model', {}).get('segmentation_head', {}).get('feature_dim', 256), + 'dropout_rate': config.get('model', {}).get('segmentation_head', {}).get('dropout_rate', 0.1), + 'num_classes': config['dataset']['num_classes'], + 'image_size': config['dataset']['image_size'] + } + + torch.save({ + 'epoch': epoch, + 'model_state_dict': model.state_dict(), + 'optimizer_state_dict': optimizer.state_dict(), + 'scheduler_state_dict': scheduler.state_dict(), + 'best_miou': best_miou, + 'config': config, + 'args': vars(args), + 'hyperparameters': hyperparameters, + 'model_info': model.get_model_info() + }, filepath) + + print(f"💾 Saved checkpoint: {filepath}") + + +def setup_wandb(config: dict, args): + """Setup Weights & Biases logging""" + if args.no_wandb or not WANDB_AVAILABLE: + return + + logging_config = config.get('logging', {}) + wandb_config = logging_config.get('wandb', {}) + + # Determine project name + project_name = ( + args.wandb_project or + wandb_config.get('project') or + f'dfine-{args.tier}-segmentation' + ) + + # Determine run name + run_name = ( + args.wandb_run_name or + f'{args.tier}_seg_{config["training"]["epochs"]}ep' + ) + + # Flatten config for wandb + wandb_config_dict = {} + + def flatten_dict(d, parent_key='', sep='/'): + if d is None: # Handle None values + return {} + + items = [] + for k, v in d.items(): + new_key = f"{parent_key}{sep}{k}" if parent_key else k + if v is None: + items.append((new_key, None)) + elif isinstance(v, dict): + items.extend(flatten_dict(v, new_key, sep=sep).items()) + else: + items.append((new_key, v)) + return dict(items) + + wandb_config_dict.update(flatten_dict(config)) + wandb_config_dict.update(vars(args)) + + wandb.init( + project=project_name, + name=run_name, + config=wandb_config_dict, + tags=wandb_config.get('tags', []) + [args.tier] + ) + + print(f"📊 Initialized wandb: project={project_name}, run={run_name}") + + +def handle_backbone_unfreezing(model, optimizer, epoch, config): + """Handle backbone unfreezing based on configuration""" + training_config = config['training'] + + if not training_config.get('freeze_detection', True): + return # Detection is not frozen + + unfreeze_epoch = training_config.get('unfreeze_epoch', None) + if unfreeze_epoch and epoch == unfreeze_epoch: + print("🔓 Unfreezing detection backbone") + model.unfreeze_detection() + + # Add backbone parameters to optimizer with lower learning rate + backbone_params = [p for p in model.dfine_model.backbone.parameters() if p.requires_grad] + if backbone_params: + backbone_lr_factor = training_config.get('backbone_lr_factor', 0.1) + optimizer.add_param_group({ + 'params': backbone_params, + 'lr': training_config['learning_rate'] * backbone_lr_factor + }) + print(f" Added {len(backbone_params)} backbone parameters with LR factor {backbone_lr_factor}") + + +def main(): + """Main training function""" + args = parse_args() + + # Load tier configuration + try: + if args.config_file: + config = load_tier_config(args.tier, args.config_dir, config_file=args.config_file) + else: + config = load_tier_config(args.tier, args.config_dir) + except FileNotFoundError as e: + print(f"❌ Error: {e}") + print(f"Available configs in {args.config_dir}: {os.listdir(args.config_dir)}") + sys.exit(1) + + # Override config with command line arguments + if args.dataset: + config['dataset']['root_dir'] = args.dataset + if args.num_workers: + config.setdefault('system', {})['num_workers'] = args.num_workers + + # Print configuration summary + print("=" * 80) + print(f"🚀 UNIFIED DFINE SEGMENTATION TRAINING - {args.tier.upper()} TIER") + print("=" * 80) + print(f"📋 Configuration:") + print(f" Tier: {args.tier}") + print(f" Dataset: {config['dataset']['root_dir']}") + print(f" Image size: {config['dataset']['image_size']}") + print(f" Batch size: {config['training']['batch_size']}") + print(f" Epochs: {config['training']['epochs']}") + print(f" Learning rate: {config['training']['learning_rate']}") + print(f" Multi-scale: {config['training'].get('multi_scale_training', False)}") + + # Print tier-specific features + if args.tier == 'lightweight': + print(f"🚀 Lightweight features: Speed optimized, mobile-ready") + elif args.tier == 'standard': + print(f"⚖️ Standard features: Balanced performance and accuracy") + elif args.tier == 'advanced': + print(f"🎯 Advanced features: Maximum accuracy, attention mechanisms") + + print("=" * 80) + + if args.dry_run: + print("🧪 Dry run completed - configuration loaded successfully") + return + + # Setup device + device = setup_device(args.device) + + # Create output directory + output_config = config['output'] + output_dir = args.output_dir or output_config['base_dir'] + os.makedirs(output_dir, exist_ok=True) + args.output_dir = output_dir + + # Initialize wandb + setup_wandb(config, args) + + # Create model + model = create_model(config, args.tier, device) + + # Create datasets + train_loader, val_loader = create_datasets(config, args.tier, args.dataset) + + # Create loss function + loss_config = config.get('loss', {}) + loss_type = loss_config.get('type', args.tier) + + criterion = create_loss( + tier=loss_type, + num_classes=config['dataset']['num_classes'], + ignore_index=loss_config.get('ignore_index', 255) + ) + + # Create optimizer and scheduler + optimizer, scheduler = create_optimizer_and_scheduler(model, config) + + # Resume training if specified + start_epoch = 0 + best_miou = 0.0 + + if args.resume: + checkpoint = torch.load(args.resume, map_location=device) + model.load_state_dict(checkpoint['model_state_dict']) + if args.finetune: + print(f"🔧 Fine-tuning from: {args.resume} (weights only)") + else: + optimizer.load_state_dict(checkpoint['optimizer_state_dict']) + scheduler.load_state_dict(checkpoint['scheduler_state_dict']) + start_epoch = checkpoint['epoch'] + 1 + best_miou = checkpoint.get('best_miou', 0.0) + print(f"📚 Resumed from epoch {start_epoch}, best mIoU: {best_miou:.4f}") + + # Training loop + epochs = config['training']['epochs'] + save_every = output_config.get('save_every', 10) + + print(f"\n🚀 Starting training for {epochs - start_epoch} epochs...") + print(f"📁 Output directory: {output_dir}") + + for epoch in range(start_epoch, epochs): + print(f"\n📅 Epoch {epoch+1}/{epochs}") + + # Handle backbone unfreezing + handle_backbone_unfreezing(model, optimizer, epoch, config) + + # Train + train_loss, train_miou = train_epoch( + model, train_loader, criterion, optimizer, device, epoch, config, args + ) + + # Validate + val_loss, val_miou, detailed_metrics = validate_epoch( + model, val_loader, criterion, device, config + ) + + # Update scheduler + scheduler.step() + + # Print epoch results + print(f"📊 Train - Loss: {train_loss:.4f}, mIoU: {train_miou:.3f}") + print(f"📊 Val - Loss: {val_loss:.4f}, mIoU: {val_miou:.3f}") + + # Print tier-specific performance info + targets = config.get('targets', {}) + if 'accuracy' in targets: + target_miou = targets['accuracy'].get('miou', 0.0) + if target_miou > 0: + progress = (val_miou / target_miou) * 100 + print(f"🎯 Target progress: {progress:.1f}% of {target_miou:.3f} mIoU target") + + # Log to wandb + if WANDB_AVAILABLE and not args.no_wandb: + log_dict = { + 'epoch': epoch, + 'train/loss': train_loss, + 'train/miou': train_miou, + 'val/loss': val_loss, + 'val/miou': val_miou, + 'lr': optimizer.param_groups[0]['lr'] + } + + # Add detailed metrics + for key, value in detailed_metrics.items(): + log_dict[f'val/{key}'] = value + + wandb.log(log_dict) + + # Save best model + if val_miou > best_miou: + best_miou = val_miou + save_checkpoint(model, optimizer, scheduler, epoch, best_miou, config, args, 'best_model.pth') + print(f"🏆 New best model! mIoU: {best_miou:.4f}") + + # Save regular checkpoint + if (epoch + 1) % save_every == 0: + save_checkpoint(model, optimizer, scheduler, epoch, best_miou, config, args) + + print(f"\n🎉 Training completed!") + print(f"🏆 Best validation mIoU: {best_miou:.4f}") + print(f"📁 Models saved in: {output_dir}") + + # Print tier-specific final summary + targets = config.get('targets', {}) + if 'accuracy' in targets: + target_miou = targets['accuracy'].get('miou', 0.0) + if target_miou > 0: + if best_miou >= target_miou: + print(f"✅ Target achieved! {best_miou:.4f} >= {target_miou:.4f}") + else: + print(f"📊 Target progress: {(best_miou/target_miou)*100:.1f}% of {target_miou:.4f} target") + + if WANDB_AVAILABLE and not args.no_wandb: + wandb.finish() + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/src/nn/arch/dfine_segmentation.py b/src/nn/arch/dfine_segmentation.py new file mode 100644 index 00000000..dd92e3a2 --- /dev/null +++ b/src/nn/arch/dfine_segmentation.py @@ -0,0 +1,92 @@ +""" +D-FINE with Body Parts Segmentation +""" + +import torch +import torch.nn as nn +from ...core import register + +__all__ = ["DFineSegmentation"] + + +@register() +class DFineSegmentation(nn.Module): + """ + D-FINE with added body parts segmentation head + """ + __inject__ = [ + "backbone", + "neck", + "head", + "seg_head" + ] + + def __init__( + self, + backbone: nn.Module, + neck: nn.Module, + head: nn.Module, + seg_head: nn.Module, + freeze_detection=True + ): + super().__init__() + self.backbone = backbone + self.neck = neck + self.head = head + self.seg_head = seg_head + self.freeze_detection = freeze_detection + + # Freeze detection components if needed + if freeze_detection: + self._freeze_detection() + + def _freeze_detection(self): + """Freeze all detection-related parameters""" + frozen_params = 0 + total_params = 0 + + for name, param in self.named_parameters(): + total_params += param.numel() + if 'seg_head' not in name: + param.requires_grad = False + frozen_params += param.numel() + + trainable_params = total_params - frozen_params + + print(f"🔒 Frozen detection components:") + print(f" 📊 Total parameters: {total_params:,}") + print(f" ❄️ Frozen parameters: {frozen_params:,} ({frozen_params/total_params*100:.1f}%)") + print(f" 🎯 Trainable parameters: {trainable_params:,} ({trainable_params/total_params*100:.1f}%)") + + def forward(self, x, targets=None): + # Get backbone features + backbone_features = self.backbone(x) + + outputs = {} + + # Detection branch (frozen during training) + if not self.training or not self.freeze_detection: + neck_features = self.neck(backbone_features) + det_outputs = self.head(neck_features) + outputs.update(det_outputs) # pred_logits, pred_boxes + + # Segmentation branch (trainable) + seg_logits = self.seg_head(backbone_features) + + # Upsample to input resolution + seg_logits = F.interpolate( + seg_logits, size=x.shape[-2:], + mode='bilinear', align_corners=False + ) + + outputs['segmentation'] = seg_logits + + return outputs + + def deploy(self): + """Prepare model for deployment""" + self.eval() + for m in self.modules(): + if m is not self and hasattr(m, 'deploy'): + m.deploy() + return self \ No newline at end of file diff --git a/src/nn/segmentation/__init__.py b/src/nn/segmentation/__init__.py new file mode 100644 index 00000000..664fc2ac --- /dev/null +++ b/src/nn/segmentation/__init__.py @@ -0,0 +1 @@ +from .segmentation_head import BodyPartsSegmentationHead \ No newline at end of file diff --git a/src/nn/segmentation/segmentation_head.py b/src/nn/segmentation/segmentation_head.py new file mode 100644 index 00000000..8a1c12b8 --- /dev/null +++ b/src/nn/segmentation/segmentation_head.py @@ -0,0 +1,116 @@ +""" +Body Parts Segmentation Head for D-FINE +Designed for HGNetv2 backbone features +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from ..backbone.common import ConvNormLayer, get_activation +from ...core import register + +__all__ = ["BodyPartsSegmentationHead"] + + +class FPN(nn.Module): + """Feature Pyramid Network for multi-scale feature fusion""" + + def __init__(self, in_channels_list, out_channels=256): + super().__init__() + + # Lateral connections + self.lateral_convs = nn.ModuleList([ + nn.Conv2d(in_ch, out_channels, 1) + for in_ch in in_channels_list + ]) + + # Output convolutions + self.fpn_convs = nn.ModuleList([ + nn.Conv2d(out_channels, out_channels, 3, padding=1) + for _ in in_channels_list + ]) + + def forward(self, features): + """ + Args: + features: List of features [P2, P3, P4, P5] from HGNetv2 + """ + # Build laterals + laterals = [conv(feat) for conv, feat in zip(self.lateral_convs, features)] + + # Build top-down pathway + for i in range(len(laterals) - 1, 0, -1): + # Upsample higher level feature + upsampled = F.interpolate( + laterals[i], size=laterals[i-1].shape[-2:], + mode='bilinear', align_corners=False + ) + laterals[i-1] = laterals[i-1] + upsampled + + # Apply final convs + outputs = [conv(lateral) for conv, lateral in zip(self.fpn_convs, laterals)] + + return outputs + + +@register() +class BodyPartsSegmentationHead(nn.Module): + """ + Body Parts Segmentation Head for D-FINE + 6 body parts: head, torso, arms, hands, legs, feet + """ + + def __init__( + self, + in_channels_list, # HGNetv2 output channels [128, 256, 512, 1024] for example + num_classes=7, # 6 body parts + background + feature_dim=256, + dropout_rate=0.1 + ): + super().__init__() + self.num_classes = num_classes + self.in_channels_list = in_channels_list + + print(f"🏗️ Creating segmentation head with channels: {in_channels_list}") + + # Feature Pyramid Network + self.fpn = FPN(in_channels_list, feature_dim) + + # Segmentation decoder + self.decoder = nn.Sequential( + ConvNormLayer(feature_dim, feature_dim, 3, 1, act='relu'), + ConvNormLayer(feature_dim, feature_dim//2, 3, 1, act='relu'), + nn.Dropout2d(dropout_rate) if dropout_rate > 0 else nn.Identity(), + nn.Conv2d(feature_dim//2, num_classes, 1) + ) + + self._init_weights() + + def _init_weights(self): + """Initialize weights""" + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, nn.BatchNorm2d): + nn.init.ones_(m.weight) + nn.init.zeros_(m.bias) + + def forward(self, features): + """ + Args: + features: List of feature maps from HGNetv2 backbone + Returns: + Segmentation logits + """ + # Use FPN to fuse multi-scale features + fpn_features = self.fpn(features) + + # Use the highest resolution feature (P2) for final prediction + final_feature = fpn_features[0] + + # Generate segmentation prediction + seg_logits = self.decoder(final_feature) + + return seg_logits \ No newline at end of file diff --git a/src/zoo/dfine/__init__.py b/src/zoo/dfine/__init__.py index f61af2a4..b314ab01 100644 --- a/src/zoo/dfine/__init__.py +++ b/src/zoo/dfine/__init__.py @@ -9,3 +9,5 @@ from .hybrid_encoder import HybridEncoder from .matcher import HungarianMatcher from .postprocessor import DFINEPostProcessor + +from . import detrpose # noqa: F401 diff --git a/src/zoo/dfine/detrpose/__init__.py b/src/zoo/dfine/detrpose/__init__.py new file mode 100644 index 00000000..de64c8b0 --- /dev/null +++ b/src/zoo/dfine/detrpose/__init__.py @@ -0,0 +1,10 @@ +""" +DETRPose-style pose transformer components vendored into this repo. + +We keep this in a separate package so we can: +- integrate a DETRPose/GroupPose-style decoder without breaking existing DFINETransformer +- switch via YAML config / training script flags +""" + +from .transformer import DETRPoseTransformer # noqa: F401 + diff --git a/src/zoo/dfine/detrpose/criterion.py b/src/zoo/dfine/detrpose/criterion.py new file mode 100644 index 00000000..934d94f3 --- /dev/null +++ b/src/zoo/dfine/detrpose/criterion.py @@ -0,0 +1,229 @@ +""" +Criterion for DETRPose-style outputs (pose-only). + +Adapted from: + DETRPose/src/models/detrpose/criterion.py + +This criterion is designed to work with our pose_estimation_berna dataset format: +- targets[i]["keypoints"]: [N,K,3] (x_px, y_px, v) +- targets[i]["boxes"]: [N,4] cxcywh normalized (for area proxy) +""" + +from __future__ import annotations + +from typing import Dict, List + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .matcher import HungarianMatcherDETRPose, _area_from_targets_px + + +def sigmoid_focal_loss(inputs, targets, num_boxes, alpha: float = 0.25, gamma: float = 2.0): + prob = inputs.sigmoid() + ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none") + p_t = prob * targets + (1 - prob) * (1 - targets) + loss = ce_loss * ((1 - p_t) ** gamma) + if alpha >= 0: + alpha_t = alpha * targets + (1 - alpha) * (1 - targets) + loss = alpha_t * loss + return loss.mean(1).sum() / num_boxes + + +class DETRPoseCriterion(nn.Module): + def __init__( + self, + num_classes: int, + matcher: HungarianMatcherDETRPose, + weight_dict: Dict[str, float], + focal_alpha: float = 0.25, + gamma: float = 2.0, + num_keypoints: int = 17, + ): + super().__init__() + self.num_classes = int(num_classes) + self.matcher = matcher + self.weight_dict = dict(weight_dict) + self.focal_alpha = float(focal_alpha) + self.gamma = float(gamma) + self.num_keypoints = int(num_keypoints) + + def _get_src_permutation_idx(self, indices): + batch_idx = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)]) + src_idx = torch.cat([src for (src, _) in indices]) + return batch_idx, src_idx + + def _num_boxes(self, targets: List[dict], device): + n = sum(len(t["labels"]) for t in targets) + return torch.as_tensor([max(1, n)], dtype=torch.float32, device=device) + + def loss_labels(self, outputs, targets, indices, num_boxes): + src_logits = outputs["pred_logits"] + idx = self._get_src_permutation_idx(indices) + target_classes_o = torch.cat([t["labels"][J] for t, (_, J) in zip(targets, indices)], dim=0) + target_classes = torch.full( + src_logits.shape[:2], self.num_classes, dtype=torch.int64, device=src_logits.device + ) + target_classes[idx] = target_classes_o + target_onehot = F.one_hot(target_classes, num_classes=self.num_classes + 1)[..., :-1].to(src_logits.dtype) + loss_ce = sigmoid_focal_loss( + src_logits, target_onehot, num_boxes, alpha=self.focal_alpha, gamma=self.gamma + ) * src_logits.shape[1] + return {"loss_ce": loss_ce} + + def loss_keypoints(self, outputs, targets, indices, num_boxes): + idx = self._get_src_permutation_idx(indices) + src_kpt = outputs["pred_keypoints"][idx] # [M,2K] in [0,1] + if src_kpt.numel() == 0: + z = outputs["pred_logits"].sum() * 0.0 + return {"loss_keypoints": z, "loss_oks": z} + + tgt_kpt = torch.cat([t["keypoints"][J] for t, (_, J) in zip(targets, indices)], dim=0).to( + src_kpt.device + ) # [M,K,3] px + + # normalize gt to [0,1] + sizes = [] + for t in targets: + if "size" in t: + h, w = t["size"].tolist() + else: + w, h = t["orig_size"].tolist() + sizes.append([w, h]) + sizes = torch.tensor(sizes, device=src_kpt.device, dtype=src_kpt.dtype) + batch_idx = idx[0] + wh = sizes[batch_idx].clamp(min=1.0) + Z_gt = (tgt_kpt[..., :2] / wh[:, None, :]).reshape(tgt_kpt.shape[0], -1) + V_gt = (tgt_kpt[..., 2] > 0).to(src_kpt.dtype) # [M,K] + + # area proxy in pixel^2 for OKS denom + # (OKS is computed in matcher; here we keep a simple 1-oks-like term via weighted l1) + pose_loss = F.l1_loss(src_kpt, Z_gt, reduction="none") + pose_loss = pose_loss * V_gt.repeat_interleave(2, dim=1) + loss_keypoints = pose_loss.sum() / num_boxes + + # Optional extra term: encourage high OKS by penalizing large errors on visible joints (already captured in L1). + loss_oks = (pose_loss.sum(dim=1) / (V_gt.repeat_interleave(2, dim=1).sum(dim=1).clamp(min=1.0))).mean() + return {"loss_keypoints": loss_keypoints, "loss_oks": loss_oks} + + def loss_vfl(self, outputs, targets, indices, num_boxes): + """ + VariFocal-like: target quality is OKS for matched pairs. + """ + src_logits = outputs["pred_logits"] + idx = self._get_src_permutation_idx(indices) + + # compute matched OKS as quality target + src_kpt = outputs["pred_keypoints"][idx] # [M,2K] + tgt_kpt = torch.cat([t["keypoints"][J] for t, (_, J) in zip(targets, indices)], dim=0).to( + src_kpt.device + ) + sizes = [] + for t in targets: + if "size" in t: + h, w = t["size"].tolist() + else: + w, h = t["orig_size"].tolist() + sizes.append([w, h]) + sizes = torch.tensor(sizes, device=src_kpt.device, dtype=src_kpt.dtype) + wh = sizes[idx[0]].clamp(min=1.0) + Z_gt = (tgt_kpt[..., :2] / wh[:, None, :]).reshape(tgt_kpt.shape[0], -1) + V_gt = (tgt_kpt[..., 2] > 0).to(src_kpt.dtype) # [M,K] + + # area in px^2 (use all targets, then pick matched) + areas_all = _area_from_targets_px(targets, device=src_kpt.device, dtype=src_kpt.dtype) + # build mapping from concatenated targets -> matched rows + # indices pairs already align per batch; simplest: recompute per matched by gathering target boxes + tgt_boxes = torch.cat([t["boxes"][J] for t, (_, J) in zip(targets, indices)], dim=0).to(src_kpt.device) + # compute area per matched in px^2 + areas_m = [] + for b, (_, J) in enumerate(indices): + if J.numel() == 0: + continue + t = targets[b] + boxes = t["boxes"][J].to(device=src_kpt.device, dtype=src_kpt.dtype) + if "size" in t: + h, w = t["size"].tolist() + else: + w, h = t["orig_size"].tolist() + bw = (boxes[:, 2] * float(w)).clamp(min=1.0) + bh = (boxes[:, 3] * float(h)).clamp(min=1.0) + areas_m.append(bw * bh) + area = torch.cat(areas_m, dim=0) if areas_m else src_kpt.new_ones((0,)) + + # OKS quality target (in pixel space) + K = self.num_keypoints + sigmas = src_kpt.new_tensor( + [0.026, 0.025, 0.025, 0.035, 0.035, 0.079, 0.079, 0.072, 0.072, 0.062, 0.062, 0.107, 0.107, 0.087, 0.087, 0.089, 0.089] + )[:K] + vars_ = (sigmas * 2) ** 2 + pred_xy = src_kpt.reshape(-1, K, 2) * wh[:, None, :] # px + gt_xy = Z_gt.reshape(-1, K, 2) * wh[:, None, :] + d2 = ((pred_xy - gt_xy) ** 2).sum(-1) + oks = torch.exp(-d2 / (area[:, None] * vars_[None, :] * 2 + 1e-6)) * V_gt + oks = oks.sum(-1) / V_gt.sum(-1).clamp(min=1.0) + oks = oks.detach().clamp(0.0, 1.0) + + target_classes_o = torch.cat([t["labels"][J] for t, (_, J) in zip(targets, indices)], dim=0) + target_classes = torch.full( + src_logits.shape[:2], self.num_classes, dtype=torch.int64, device=src_logits.device + ) + target_classes[idx] = target_classes_o + target = F.one_hot(target_classes, num_classes=self.num_classes + 1)[..., :-1] + + target_score_o = torch.zeros_like(target_classes, dtype=src_logits.dtype) + target_score_o[idx] = oks.to(target_score_o.dtype) + target_score = target_score_o.unsqueeze(-1) * target + + pred_score = torch.sigmoid(src_logits).detach() + weight = self.focal_alpha * pred_score.pow(self.gamma) * (1 - target) + target_score + loss = F.binary_cross_entropy_with_logits(src_logits, target_score, weight=weight, reduction="none") + loss = loss.mean(1).sum() * src_logits.shape[1] / num_boxes + return {"loss_vfl": loss} + + def forward(self, outputs: dict, targets: List[dict]): + # match + indices = self.matcher({"pred_logits": outputs["pred_logits"], "pred_keypoints": outputs["pred_keypoints"]}, targets)[ + "indices" + ] + num_boxes = float(self._num_boxes(targets, device=outputs["pred_logits"].device).item()) + + # base losses + losses = {} + losses.update(self.loss_labels(outputs, targets, indices, num_boxes)) + losses.update(self.loss_vfl(outputs, targets, indices, num_boxes)) + losses.update(self.loss_keypoints(outputs, targets, indices, num_boxes)) + + # aux losses (decoder layers) + if "aux_outputs" in outputs and isinstance(outputs["aux_outputs"], list): + for i, aux in enumerate(outputs["aux_outputs"]): + aux_idx = self.matcher( + {"pred_logits": aux["pred_logits"], "pred_keypoints": aux["pred_keypoints"]}, targets + )["indices"] + l = {} + l.update(self.loss_labels(aux, targets, aux_idx, num_boxes)) + l.update(self.loss_vfl(aux, targets, aux_idx, num_boxes)) + l.update(self.loss_keypoints(aux, targets, aux_idx, num_boxes)) + for k, v in l.items(): + losses[f"{k}_{i}"] = v + + # apply weights (return weighted values only; train.py sums dict values) + weighted = {} + for k, v in losses.items(): + # Aux losses are named like: loss_ce_0, loss_keypoints_3, ... + # Map them back to their base key in weight_dict (loss_ce, loss_keypoints, ...). + base_k = k + if k.startswith("loss_"): + parts = k.rsplit("_", 1) + if len(parts) == 2 and parts[1].isdigit(): + base_k = parts[0] + # Only include losses that are explicitly weighted (matches DETRPose config style). + if base_k not in self.weight_dict and k not in self.weight_dict: + continue + w = self.weight_dict.get(base_k, self.weight_dict.get(k, 0.0)) + if float(w) == 0.0: + continue + weighted[k] = v * float(w) + return weighted + diff --git a/src/zoo/dfine/detrpose/dn_component.py b/src/zoo/dfine/detrpose/dn_component.py new file mode 100644 index 00000000..ac5b36a3 --- /dev/null +++ b/src/zoo/dfine/detrpose/dn_component.py @@ -0,0 +1,236 @@ +""" +Pose contrastive denoising component vendored from DETRPose. + +Source: /DETRPose/src/models/detrpose/dn_component.py + +NOTE: This file expects targets to provide: +- labels: [N] +- boxes: [N,4] (cxcywh normalized in our pipeline) +- keypoints: [N,K,3] (x_px, y_px, v) in resized image pixels +- size: [2] (H,W) resized +We adapt DETRPose logic to derive area if not provided. +""" + +import numpy as np +import torch +import torch.nn.functional as F + +from .utils import inverse_sigmoid + + +def get_sigmas(num_keypoints, device): + if num_keypoints == 17: + sigmas = np.array( + [ + 0.26, + 0.25, + 0.25, + 0.35, + 0.35, + 0.79, + 0.79, + 0.72, + 0.72, + 0.62, + 0.62, + 1.07, + 1.07, + 0.87, + 0.87, + 0.89, + 0.89, + ], + dtype=np.float32, + ) / 10.0 + elif num_keypoints == 14: + sigmas = ( + np.array([0.79, 0.79, 0.72, 0.72, 0.62, 0.62, 1.07, 1.07, 0.87, 0.87, 0.89, 0.89, 0.79, 0.79]) + / 10.0 + ) + elif num_keypoints == 3: + sigmas = np.array([1.07, 1.07, 0.67], dtype=np.float32) / 10.0 + else: + raise ValueError(f"Unsupported keypoints number {num_keypoints}") + # prepend sigma for the center of the human + sigmas = np.concatenate([[0.1], sigmas]) + sigmas = torch.tensor(sigmas, device=device, dtype=torch.float32) + return sigmas[None, :, None] + + +def _area_from_boxes_cxcywh_norm(targets): + # area in pixel^2, consistent with DETRPose's use in OKS denom + areas = [] + for t in targets: + boxes = t["boxes"] + if "size" in t: + h, w = t["size"].tolist() + else: + w, h = t["orig_size"].tolist() + wh = boxes.new_tensor([w, h, w, h]).view(1, 4).clamp(min=1.0) + bw = (boxes[:, 2] * wh[:, 0]).clamp(min=1.0) + bh = (boxes[:, 3] * wh[:, 1]).clamp(min=1.0) + areas.append((bw * bh).to(boxes.dtype)) + return torch.cat(areas, dim=0) + + +def prepare_for_cdn( + dn_args, + training, + num_queries, + num_classes, + num_keypoints, + hidden_dim, + label_enc, + pose_enc, + device, +): + """ + Return: + input_query_label: [B, pad, (K+1), hidden_dim] + input_query_pose: [B, pad, (K+1), 2] (unsigmoid coords) + attn_mask: [pad+Q, pad+Q] + dn_meta: dict + """ + if training: + targets, dn_number, label_noise_ratio = dn_args + dn_number = dn_number * 2 # pos+neg + known = [torch.ones_like(t["labels"]) for t in targets] + batch_size = len(known) + known_num = [int(k.sum().item()) for k in known] + if int(max(known_num)) == 0: + return None, None, None, None + + dn_number = dn_number // (int(max(known_num) * 2)) + dn_number = 1 if dn_number == 0 else dn_number + + unmask_bbox = unmask_label = torch.cat(known) + + labels = torch.cat([t["labels"] for t in targets]) + batch_idx = torch.cat([torch.full_like(t["labels"].long(), i) for i, t in enumerate(targets)]) + + known_indice = torch.nonzero(unmask_label + unmask_bbox).view(-1) + known_indice = known_indice.repeat(2 * dn_number, 1).view(-1) + + known_labels = labels.repeat(2 * dn_number, 1).view(-1) + known_labels_expaned = known_labels.clone() + + known_bid = batch_idx.repeat(2 * dn_number, 1).view(-1) + + if label_noise_ratio > 0: + p = torch.rand_like(known_labels_expaned.float()) + chosen_indice = torch.nonzero(p < (label_noise_ratio * 0.5)).view(-1) + new_label = torch.randint_like(chosen_indice, 0, num_classes) + known_labels_expaned.scatter_(0, chosen_indice, new_label) + + # keypoint noise + boxes = torch.cat([t["boxes"] for t in targets]) # cxcywh normalized + xy = boxes[:, :2] # center in [0,1] + keypoints = torch.cat([t["keypoints"] for t in targets]) # [N,K,3] px + + # derive areas in pixel^2 (preferred for OKS denom) + if "area" in targets[0]: + areas = torch.cat([t["area"] for t in targets]).to(device=device) + else: + areas = _area_from_boxes_cxcywh_norm(targets).to(device=device) + + poses = keypoints[..., :2] # px + # normalize to [0,1] in resized image + sizes = [] + for t in targets: + if "size" in t: + h, w = t["size"].tolist() + else: + w, h = t["orig_size"].tolist() + sizes.append([w, h]) + sizes = torch.tensor(sizes, device=device, dtype=poses.dtype) + # expand per-instance + rep = torch.cat([torch.full((len(t["labels"]),), i, device=device, dtype=torch.long) for i, t in enumerate(targets)]) + wh = sizes[rep].clamp(min=1.0) # [N,2] + poses = (poses / wh[:, None, :]).clamp(0.0, 1.0) + poses = poses.reshape(poses.shape[0], -1) # [N, K*2] + poses = torch.cat([xy, poses], dim=1) # prepend center: [N, 2 + K*2] + + non_viz = (keypoints[..., 2] <= 0).reshape(keypoints.shape[0], -1) + non_viz = torch.cat((torch.zeros_like(non_viz[:, 0:1]).bool(), non_viz), dim=1) + + vars_ = (2 * get_sigmas(num_keypoints, device)) ** 2 + + known_poses = poses.repeat(2 * dn_number, 1).reshape(-1, num_keypoints + 1, 2) + known_areas = areas.repeat(2 * dn_number)[..., None, None] + known_non_viz = non_viz.repeat(2 * dn_number, 1) + + single_pad = int(max(known_num)) + pad_size = int(single_pad * 2 * dn_number) + positive_idx = torch.arange(len(poses), device=device).long().unsqueeze(0).repeat(dn_number, 1) + positive_idx += (torch.arange(dn_number, device=device).long() * len(poses) * 2).unsqueeze(1) + positive_idx = positive_idx.flatten() + negative_idx = positive_idx + len(poses) + + eps = np.finfo("float32").eps + rand_vector = torch.rand_like(known_poses) + rand_vector = F.normalize(rand_vector, -1) + rand_alpha = torch.zeros_like(known_poses[..., :1]).uniform_(-np.log(1), -np.log(0.5)) + rand_alpha[negative_idx] = rand_alpha[negative_idx].uniform_(-np.log(0.5), -np.log(0.1)) + rand_alpha *= 2 * (known_areas + eps) * vars_ + # normalize by max(H,W) in pixels (approx using resized size) + img_dim = targets[0]["size"].tolist() if "size" in targets[0] else targets[0]["orig_size"].tolist()[::-1] + rand_alpha = torch.sqrt(rand_alpha) / float(max(img_dim)) + rand_alpha[known_non_viz] = 0.0 + + known_poses_expand = (known_poses + rand_alpha * rand_vector).clamp(0.0, 1.0) + + m = known_labels_expaned.long().to(device) + input_label_embed = label_enc(m) + input_label_pose_embed = pose_enc.weight[None].repeat(known_poses_expand.size(0), 1, 1) + input_label_embed = torch.cat([input_label_embed.unsqueeze(1), input_label_pose_embed], dim=1).flatten(1) + + input_pose_embed = inverse_sigmoid(known_poses_expand) + + padding_label = torch.zeros(pad_size, hidden_dim * (num_keypoints + 1), device=device) + padding_pose = torch.zeros(pad_size, num_keypoints + 1, device=device) + + input_query_label = padding_label.repeat(batch_size, 1, 1) + input_query_pose = padding_pose[..., None].repeat(batch_size, 1, 1, 2) + + map_known_indice = torch.tensor([], device=device) + if len(known_num): + map_known_indice = torch.cat([torch.arange(num, device=device) for num in known_num]) + map_known_indice = torch.cat([map_known_indice + single_pad * i for i in range(2 * dn_number)]).long() + if len(known_bid): + input_query_label[(known_bid.long(), map_known_indice)] = input_label_embed + input_query_pose[(known_bid.long(), map_known_indice)] = input_pose_embed + + tgt_size = pad_size + num_queries + attn_mask = torch.ones(tgt_size, tgt_size, device=device) < 0 + attn_mask[pad_size:, :pad_size] = True + for i in range(dn_number): + if i == 0: + attn_mask[single_pad * 2 * i : single_pad * 2 * (i + 1), single_pad * 2 * (i + 1) : pad_size] = True + if i == dn_number - 1: + attn_mask[single_pad * 2 * i : single_pad * 2 * (i + 1), : single_pad * i * 2] = True + else: + attn_mask[single_pad * 2 * i : single_pad * 2 * (i + 1), single_pad * 2 * (i + 1) : pad_size] = True + attn_mask[single_pad * 2 * i : single_pad * 2 * (i + 1), : single_pad * 2 * i] = True + + dn_meta = {"pad_size": pad_size, "num_dn_group": dn_number} + else: + input_query_label = None + input_query_pose = None + attn_mask = None + dn_meta = None + + return input_query_label.unflatten(-1, (-1, hidden_dim)), input_query_pose, attn_mask, dn_meta + + +def dn_post_process(outputs_class, outputs_keypoints, dn_meta, aux_loss, _set_aux_loss): + if dn_meta and dn_meta["pad_size"] > 0: + output_known_class = outputs_class[:, :, : dn_meta["pad_size"], :] + output_known_keypoints = outputs_keypoints[:, :, : dn_meta["pad_size"], :] + outputs_class = outputs_class[:, :, dn_meta["pad_size"] :, :] + outputs_keypoints = outputs_keypoints[:, :, dn_meta["pad_size"] :, :] + out = {"pred_logits": output_known_class[-1], "pred_keypoints": output_known_keypoints[-1]} + if aux_loss: + out["aux_outputs"] = _set_aux_loss(output_known_class, output_known_keypoints) + dn_meta["output_known_lbs_keypoints"] = out + return outputs_class, outputs_keypoints + diff --git a/src/zoo/dfine/detrpose/matcher.py b/src/zoo/dfine/detrpose/matcher.py new file mode 100644 index 00000000..ba29ba57 --- /dev/null +++ b/src/zoo/dfine/detrpose/matcher.py @@ -0,0 +1,171 @@ +""" +Hungarian matcher for DETRPose-style outputs (pose-only). + +Adapted from: + DETRPose/src/models/detrpose/matcher.py + +Differences vs our DFINE matcher: +- no pred_boxes needed +- matching uses class cost + keypoint L1 (visible only) + OKS cost +""" + +from __future__ import annotations + +import numpy as np +import torch +import torch.nn.functional as F +from scipy.optimize import linear_sum_assignment +from torch import nn + + +def _coco_sigmas(num_keypoints: int, device, dtype): + if num_keypoints == 17: + sigmas = np.array( + [ + 0.26, + 0.25, + 0.25, + 0.35, + 0.35, + 0.79, + 0.79, + 0.72, + 0.72, + 0.62, + 0.62, + 1.07, + 1.07, + 0.87, + 0.87, + 0.89, + 0.89, + ], + dtype=np.float32, + ) / 10.0 + else: + raise ValueError(f"Unsupported num_keypoints={num_keypoints} for COCO sigmas") + return torch.tensor(sigmas, device=device, dtype=dtype) + + +def _area_from_targets_px(targets: list[dict], device, dtype) -> torch.Tensor: + """Area per instance in pixel^2, using target boxes (cxcywh normalized).""" + areas = [] + for t in targets: + boxes = t["boxes"].to(device=device, dtype=dtype) + if "size" in t: + h, w = t["size"].tolist() + else: + w, h = t["orig_size"].tolist() + w = max(float(w), 1.0) + h = max(float(h), 1.0) + bw = (boxes[:, 2] * w).clamp(min=1.0) + bh = (boxes[:, 3] * h).clamp(min=1.0) + areas.append((bw * bh).to(dtype)) + return torch.cat(areas, dim=0) if areas else torch.zeros((0,), device=device, dtype=dtype) + + +class HungarianMatcherDETRPose(nn.Module): + def __init__( + self, + cost_class: float = 1.0, + focal_alpha: float = 0.25, + cost_keypoints: float = 1.0, + cost_oks: float = 0.01, + num_keypoints: int = 17, + ): + super().__init__() + self.cost_class = float(cost_class) + self.cost_keypoints = float(cost_keypoints) + self.cost_oks = float(cost_oks) + self.focal_alpha = float(focal_alpha) + self.num_keypoints = int(num_keypoints) + + @torch.no_grad() + def forward(self, outputs: dict, targets: list[dict]): + """ + outputs: + - pred_logits: [B,Q,C] logits + - pred_keypoints: [B,Q,2K] normalized (x,y) in [0,1] (resized image space) + targets: + - labels: [N] + - keypoints: [N,K,3] (x_px,y_px,v) + - boxes: [N,4] cxcywh normalized (for area only) + - size/orig_size + """ + bs, num_queries = outputs["pred_logits"].shape[:2] + out_prob = outputs["pred_logits"].flatten(0, 1).sigmoid() # [B*Q,C] + out_kpt = outputs["pred_keypoints"].flatten(0, 1) # [B*Q,2K] + + tgt_ids = torch.cat([v["labels"] for v in targets], dim=0) + tgt_kpt = torch.cat([v["keypoints"] for v in targets], dim=0).to(out_kpt.device) # [sumN,K,3] + # normalize gt to [0,1] (resized image space) + sizes = [] + for t in targets: + if "size" in t: + h, w = t["size"].tolist() + else: + w, h = t["orig_size"].tolist() + sizes.append([w, h]) + sizes = torch.tensor(sizes, device=out_kpt.device, dtype=out_kpt.dtype) + batch_idx = torch.cat([torch.full((len(t["labels"]),), i, device=out_kpt.device, dtype=torch.long) for i, t in enumerate(targets)], dim=0) + wh = sizes[batch_idx].clamp(min=1.0) # [sumN,2] + Z_gt = (tgt_kpt[..., :2] / wh[:, None, :]).reshape(tgt_kpt.shape[0], -1) # [sumN,2K] + V_gt = (tgt_kpt[..., 2] > 0).to(out_kpt.dtype) # [sumN,K] + + tgt_area = _area_from_targets_px(targets, device=out_kpt.device, dtype=out_kpt.dtype) # [sumN] + + # class cost (focal-style) + alpha = self.focal_alpha + gamma = 2.0 + neg_cost_class = (1 - alpha) * (out_prob**gamma) * (-(1 - out_prob + 1e-8).log()) + pos_cost_class = alpha * ((1 - out_prob) ** gamma) * (-(out_prob + 1e-8).log()) + cost_class = pos_cost_class[:, tgt_ids] - neg_cost_class[:, tgt_ids] # [B*Q, sumN] + + # keypoint L1 visible-only + cost_keypoints = torch.abs(out_kpt[:, None, :] - Z_gt[None, :, :]) # [B*Q,sumN,2K] + vis_rep = V_gt.repeat_interleave(2, dim=1) # [sumN,2K] + cost_keypoints = (cost_keypoints * vis_rep[None, :, :]).sum(-1) # [B*Q,sumN] + + # OKS cost (1-oks) -- guarded against NaN/Inf and degenerate areas + sigmas = _coco_sigmas(self.num_keypoints, device=out_kpt.device, dtype=out_kpt.dtype) + variances = (sigmas * 2) ** 2 # [K] + kpt_preds = out_kpt.reshape(-1, self.num_keypoints, 2) + kpt_gts = Z_gt.reshape(-1, self.num_keypoints, 2) + + # squared distance per keypoint, normalized in resized image space + d2 = (kpt_preds[:, None, :, 0] - kpt_gts[None, :, :, 0]) ** 2 + (kpt_preds[:, None, :, 1] - kpt_gts[None, :, :, 1]) ** 2 + + # Use sqrt(area) as a scale proxy; ensure areas are positive and finite + safe_tgt_area = tgt_area.clone() + safe_tgt_area[~torch.isfinite(safe_tgt_area)] = 1.0 + safe_tgt_area = safe_tgt_area.clamp(min=1.0) + scale = torch.sqrt(safe_tgt_area) # [sumN] + + # convert normalized dist to pixel dist by multiplying by scale^2 + d2 = d2 * (scale[None, :, None] ** 2) + + # Replace any non-finite d2 with large numbers to force low OKS + if not torch.isfinite(d2).all(): + d2 = torch.where(torch.isfinite(d2), d2, torch.full_like(d2, 1e6)) + + denom = (safe_tgt_area[None, :, None] * variances[None, None, :] * 2.0) + # Prevent extremely small denominators that cause huge exponent values + denom = denom.clamp(min=1e-4) + oks = torch.exp(-d2 / denom) * V_gt[None, :, :] + oks = oks.sum(dim=-1) / (V_gt.sum(dim=-1).clamp(min=1.0)[None, :]) + + # clamp oks into reasonable range and compute cost + oks = oks.clamp(min=0.0, max=1.0) + cost_oks = (1.0 - oks).clamp(min=0.0, max=1.0) + + C = self.cost_class * cost_class + self.cost_keypoints * cost_keypoints + self.cost_oks * cost_oks + C = C.view(bs, num_queries, -1).cpu() + + # Replace any remaining NaN/Inf with large finite values so Hungarian never crashes + C = torch.nan_to_num(C, nan=100.0, posinf=100.0, neginf=-100.0) + + sizes_split = [len(v["boxes"]) for v in targets] + indices = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes_split, -1))] + indices = [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices] + return {"indices": indices} + diff --git a/src/zoo/dfine/detrpose/ms_deform_attn.py b/src/zoo/dfine/detrpose/ms_deform_attn.py new file mode 100644 index 00000000..c5056d7c --- /dev/null +++ b/src/zoo/dfine/detrpose/ms_deform_attn.py @@ -0,0 +1,139 @@ +""" +Multi-Scale Deformable Attention (PyTorch implementation). + +Vendored from DETRPose: + /DETRPose/src/models/detrpose/ms_deform_attn.py +""" + +import math + +import torch +from torch import nn +import torch.nn.functional as F +from torch.nn.init import constant_ + + +def _spatial_shapes_to_list(value_spatial_shapes): + """Normalize spatial shapes to a Python list of (H, W) ints. + + Using a Python list here avoids tensor iteration/unflatten patterns that are + fragile during ONNX/TensorRT export for this module. + """ + if isinstance(value_spatial_shapes, torch.Tensor): + # Value is expected static for export (fixed image size). + return [(int(h), int(w)) for h, w in value_spatial_shapes.detach().cpu().tolist()] + return [(int(h), int(w)) for h, w in value_spatial_shapes] + + +def ms_deform_attn_core_pytorch(value, value_spatial_shapes, sampling_locations, attention_weights): + # value: list[L] of tensors shaped [N_*M_, D_, H_*W_] (later unflattened) + _, D_, _ = value[0].shape + N_, Lq_, M_, L_, P_, _ = sampling_locations.shape + + sampling_grids = 2 * sampling_locations - 1 + sampling_grids = sampling_grids.transpose(1, 2).flatten(0, 1) + + sampling_value_list = [] + spatial_shapes_list = _spatial_shapes_to_list(value_spatial_shapes) + for lid_, (H_, W_) in enumerate(spatial_shapes_list): + # Avoid Tensor.unflatten with symbolic shape values in export paths. + value_l_ = value[lid_].reshape(value[lid_].shape[0], D_, H_, W_) + sampling_grid_l_ = sampling_grids[:, :, lid_] + sampling_value_l_ = F.grid_sample( + value_l_, sampling_grid_l_, mode="bilinear", padding_mode="zeros", align_corners=False + ) + sampling_value_list.append(sampling_value_l_) + + attention_weights = attention_weights.transpose(1, 2).reshape(N_ * M_, 1, Lq_, L_ * P_) + output = (torch.concat(sampling_value_list, dim=-1) * attention_weights).sum(-1).view(N_, M_ * D_, Lq_) + return output.transpose(1, 2) + + +class MSDeformAttn(nn.Module): + def __init__(self, d_model=256, n_levels=4, n_heads=8, n_points=4, use_4D_normalizer=False): + super().__init__() + if d_model % n_heads != 0: + raise ValueError(f"d_model must be divisible by n_heads, got {d_model} and {n_heads}") + + self.d_model = d_model + self.n_levels = n_levels + self.n_heads = n_heads + self.n_points = n_points + + self.sampling_offsets = nn.Linear(d_model, n_heads * n_levels * n_points * 2) + self.attention_weights = nn.Linear(d_model, n_heads * n_levels * n_points) + self.use_4D_normalizer = use_4D_normalizer + + self._reset_parameters() + + def _reset_parameters(self): + constant_(self.sampling_offsets.weight.data, 0.0) + thetas = torch.arange(self.n_heads, dtype=torch.float32) * (2.0 * math.pi / self.n_heads) + grid_init = torch.stack([thetas.cos(), thetas.sin()], -1) + grid_init = ( + (grid_init / grid_init.abs().max(-1, keepdim=True)[0]) + .view(self.n_heads, 1, 1, 2) + .repeat(1, self.n_levels, self.n_points, 1) + ) + for i in range(self.n_points): + grid_init[:, :, i, :] *= i % 4 + 1 + with torch.no_grad(): + self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1)) + if self.n_points % 4 != 0: + constant_(self.sampling_offsets.bias, 0.0) + constant_(self.attention_weights.weight.data, 0.0) + constant_(self.attention_weights.bias.data, 0.0) + + def forward(self, query, reference_points, value, input_spatial_shapes): + N, Len_q, _ = query.shape + + sampling_offsets = self.sampling_offsets(query).view( + N, Len_q, self.n_heads, self.n_levels, self.n_points, 2 + ) + attention_weights = self.attention_weights(query).view( + N, Len_q, self.n_heads, self.n_levels * self.n_points + ) + attention_weights = F.softmax(attention_weights, -1).view( + N, Len_q, self.n_heads, self.n_levels, self.n_points + ) + + # reference_points: [N, Len_q, n_levels, ..., 2/4] -> [N, Len_q*n_levels, ..., 2/4] + reference_points = torch.transpose(reference_points, 2, 3).flatten(1, 2) + + if reference_points.shape[-1] == 2: + # Keep tensor type/device consistent for export/runtime. + if isinstance(input_spatial_shapes, torch.Tensor): + offset_normalizer = input_spatial_shapes.to(device=query.device, dtype=query.dtype) + else: + offset_normalizer = torch.as_tensor(input_spatial_shapes, device=query.device, dtype=query.dtype) + offset_normalizer = offset_normalizer.flip([1]).reshape(1, 1, 1, self.n_levels, 1, 2) + sampling_locations = reference_points[:, :, None, :, None, :] + sampling_offsets / offset_normalizer + elif reference_points.shape[-1] == 4: + if self.use_4D_normalizer: + if not isinstance(input_spatial_shapes, torch.Tensor): + input_spatial_shapes = torch.as_tensor( + input_spatial_shapes, device=query.device, dtype=query.dtype + ) + offset_normalizer = torch.stack( + [input_spatial_shapes[..., 1], input_spatial_shapes[..., 0]], -1 + ) + sampling_locations = ( + reference_points[:, :, None, :, None, :2] + + sampling_offsets + / offset_normalizer[None, None, None, :, None, :] + * reference_points[:, :, None, :, None, 2:] + * 0.5 + ) + else: + sampling_locations = ( + reference_points[:, :, None, :, None, :2] + + sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5 + ) + else: + raise ValueError( + f"Last dim of reference_points must be 2 or 4, got {reference_points.shape[-1]}" + ) + + output = ms_deform_attn_core_pytorch(value, input_spatial_shapes, sampling_locations, attention_weights) + return output + diff --git a/src/zoo/dfine/detrpose/transformer.py b/src/zoo/dfine/detrpose/transformer.py new file mode 100644 index 00000000..48e72d1c --- /dev/null +++ b/src/zoo/dfine/detrpose/transformer.py @@ -0,0 +1,615 @@ +""" +DETRPose-style transformer (decoder) vendored into this repo. + +This implements: +- (K+1) tokens per query: 1 instance token + K keypoint tokens +- within-instance + across-instance self-attention (GroupPose-style) +- deformable cross-attention +- Pose-LQE using feature sampling at predicted keypoints +- pose denoising (cdn) producing positive/negative queries based on OKS noise + +Source inspiration: + /DETRPose/src/models/detrpose/transformer.py + +Integration notes for this repo: +- We register DETRPoseTransformer so it can be used as DFINE.decoder in YAMLConfig. +- This module is pose-centric (no boxes). It outputs: + out["pred_logits"]: [B, Q, num_classes] + out["pred_keypoints"]: [B, Q, K*2] normalized (x,y) in [0,1] in resized image space. +""" + +import copy +import math +from typing import Optional, List + +import torch +import torch.nn.functional as F +from torch import nn, Tensor + +from ....core import register + +from .ms_deform_attn import MSDeformAttn +from .dn_component import prepare_for_cdn, dn_post_process +from .utils import inverse_sigmoid, MLP, _get_activation_fn + + +def _get_clones(module, N, layer_share=False): + if layer_share: + return nn.ModuleList([module for _ in range(N)]) + return nn.ModuleList([copy.deepcopy(module) for _ in range(N)]) + + +def weighting_function(reg_max, up, reg_scale, deploy=False): + if deploy: + upper_bound1 = (abs(up[0]) * abs(reg_scale)).item() + upper_bound2 = (abs(up[0]) * abs(reg_scale) * 2).item() + step = (upper_bound1 + 1) ** (2 / (reg_max - 2)) + left_values = [-((step) ** i) + 1 for i in range(reg_max // 2 - 1, 0, -1)] + right_values = [(step) ** i - 1 for i in range(1, reg_max // 2)] + values = [-upper_bound2] + left_values + [torch.zeros_like(up[0][None])] + right_values + [upper_bound2] + return torch.tensor([values], dtype=up.dtype, device=up.device) + upper_bound1 = abs(up[0]) * abs(reg_scale) + upper_bound2 = abs(up[0]) * abs(reg_scale) * 2 + step = (upper_bound1 + 1) ** (2 / (reg_max - 2)) + left_values = [-((step) ** i) + 1 for i in range(reg_max // 2 - 1, 0, -1)] + right_values = [(step) ** i - 1 for i in range(1, reg_max // 2)] + values = [-upper_bound2] + left_values + [torch.zeros_like(up[0][None])] + right_values + [upper_bound2] + return torch.cat(values, 0) + + +def distance2pose(points, distance, reg_scale): + reg_scale = abs(reg_scale) + x1 = points[..., 0] + distance[..., 0] / reg_scale + y1 = points[..., 1] + distance[..., 1] / reg_scale + return torch.stack([x1, y1], -1) + + +class Gate(nn.Module): + def __init__(self, d_model): + super().__init__() + self.gate = nn.Linear(2 * d_model, 2 * d_model) + bias = float(-math.log((1 - 0.5) / 0.5)) + nn.init.constant_(self.gate.bias, bias) + nn.init.constant_(self.gate.weight, 0) + self.norm = nn.LayerNorm(d_model) + + def forward(self, x1, x2): + gate_input = torch.cat([x1, x2], dim=-1) + gates = torch.sigmoid(self.gate(gate_input)) + gate1, gate2 = gates.chunk(2, dim=-1) + return self.norm(gate1 * x1 + gate2 * x2) + + +class Integral(nn.Module): + def __init__(self, reg_max=32): + super().__init__() + self.reg_max = reg_max + + def forward(self, x, project): + shape = x.shape + x = F.softmax(x.reshape(-1, self.reg_max + 1), dim=1) + x = F.linear(x, project.to(x.device)).reshape(-1, 4) + return x.reshape(list(shape[:-1]) + [-1]) + + +class LQE(nn.Module): + def __init__(self, topk, hidden_dim, num_layers, num_body_points): + super().__init__() + self.k = topk + self.reg_conf = MLP(num_body_points * (topk + 1), hidden_dim, 1, num_layers) + nn.init.constant_(self.reg_conf.layers[-1].weight.data, 0) + nn.init.constant_(self.reg_conf.layers[-1].bias.data, 0) + self.num_body_points = num_body_points + + def forward(self, scores, pred_poses, feat): + B, L = pred_poses.shape[:2] + pred_poses = pred_poses.reshape(B, L, self.num_body_points, 2) + sampling_values = ( + F.grid_sample(feat, 2 * pred_poses - 1, mode="bilinear", padding_mode="zeros", align_corners=False) + .permute(0, 2, 3, 1) + ) + prob_topk = sampling_values.topk(self.k, dim=-1)[0] + stat = torch.cat([prob_topk, prob_topk.mean(dim=-1, keepdim=True)], dim=-1) + quality_score = self.reg_conf(stat.reshape(B, L, -1)) + return scores + quality_score + + +class DeformableTransformerDecoderLayer(nn.Module): + def __init__( + self, + d_model=256, + d_ffn=1024, + dropout=0.1, + activation="relu", + n_levels=4, + n_heads=8, + n_points=4, + ): + super().__init__() + # within-instance self-attention + self.within_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True) + self.within_dropout = nn.Dropout(dropout) + self.within_norm = nn.LayerNorm(d_model) + # across-instance self-attention + self.across_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True) + self.across_dropout = nn.Dropout(dropout) + self.across_norm = nn.LayerNorm(d_model) + # deformable cross-attention + self.cross_attn = MSDeformAttn(d_model, n_levels, n_heads, n_points) + self.dropout1 = nn.Dropout(dropout) + # gate + self.gateway = Gate(d_model) + # FFN + self.linear1 = nn.Linear(d_model, d_ffn) + self.activation = _get_activation_fn(activation, d_model=d_ffn, batch_dim=1) + self.dropout2 = nn.Dropout(dropout) + self.linear2 = nn.Linear(d_ffn, d_model) + self.dropout3 = nn.Dropout(dropout) + self.norm2 = nn.LayerNorm(d_model) + + self._reset_parameters() + + def _reset_parameters(self): + nn.init.xavier_uniform_(self.linear1.weight) + nn.init.xavier_uniform_(self.linear2.weight) + + @staticmethod + def with_pos_embed(tensor, pos): + if pos is not None: + np_ = pos.shape[2] + # Avoid in-place modification which breaks autograd when `tensor` requires_grad. + # Create a new tensor for the addition to keep the computational graph intact. + out = tensor.clone() + out[:, :, -np_:] = out[:, :, -np_:] + pos + return out + return tensor + + def forward_FFN(self, tgt): + tgt2 = self.linear2(self.dropout2(self.activation(self.linear1(tgt)))) + tgt = tgt + self.dropout3(tgt2) + tgt = self.norm2(tgt.clamp(min=-65504, max=65504)) + return tgt + + def forward( + self, + tgt_pose: Optional[Tensor], + tgt_pose_query_pos: Optional[Tensor], + tgt_pose_reference_points: Optional[Tensor], + attn_mask: Optional[Tensor] = None, + memory: Optional[Tensor] = None, + memory_spatial_shapes: Optional[Tensor] = None, + ): + bs, nq, num_kpt, d_model = tgt_pose.shape + + # within-instance self-attention + q = k = self.with_pos_embed(tgt_pose, tgt_pose_query_pos).flatten(0, 1) + tgt2 = self.within_attn(q, k, tgt_pose.flatten(0, 1))[0].reshape(bs, nq, num_kpt, d_model) + tgt_pose = tgt_pose + self.within_dropout(tgt2) + tgt_pose = self.within_norm(tgt_pose) + + # across-instance self-attention + tgt_pose = tgt_pose.transpose(1, 2).flatten(0, 1) # bs*num_kpt, nq, d + tgt2_pose = self.across_attn(tgt_pose, tgt_pose, tgt_pose, attn_mask=attn_mask)[0].reshape( + bs * num_kpt, nq, d_model + ) + tgt_pose = tgt_pose + self.across_dropout(tgt2_pose) + tgt_pose = ( + self.across_norm(tgt_pose).reshape(bs, num_kpt, nq, d_model).transpose(1, 2) + ) # bs,nq,num_kpt,d + + # deformable cross-attention + tgt2_pose = self.cross_attn( + self.with_pos_embed(tgt_pose, tgt_pose_query_pos).flatten(1, 2), + tgt_pose_reference_points, + memory, + memory_spatial_shapes, + ).reshape(bs, nq, num_kpt, d_model) + tgt_pose = self.gateway(tgt_pose, self.dropout1(tgt2_pose)) + tgt_pose = self.forward_FFN(tgt_pose) + return tgt_pose + + +class TransformerDecoder(nn.Module): + def __init__(self, decoder_layer, num_layers, return_intermediate=False, hidden_dim=256, num_body_points=17): + super().__init__() + self.layers = _get_clones(decoder_layer, num_layers, layer_share=False) if num_layers > 0 else [] + self.hidden_dim = hidden_dim + self.num_layers = num_layers + self.num_body_points = num_body_points + self.return_intermediate = return_intermediate + self.class_embed = None + self.pose_embed = None + self.half_pose_ref_point_head = MLP(hidden_dim, hidden_dim, hidden_dim, 2) + self.eval_idx = num_layers - 1 + + dim_t = torch.arange(hidden_dim // 2, dtype=torch.float32) + dim_t = 10000 ** (2 * (dim_t // 2) / (hidden_dim // 2)) + self.register_buffer("dim_t", dim_t) + self.scale = 2 * math.pi + + def sine_embedding(self, pos_tensor): + x_embed = pos_tensor[..., 0:1] * self.scale + y_embed = pos_tensor[..., 1:2] * self.scale + pos_x = x_embed / self.dim_t + pos_y = y_embed / self.dim_t + pos_x = torch.stack((pos_x[..., 0::2].sin(), pos_x[..., 1::2].cos()), dim=4).flatten(3) + pos_y = torch.stack((pos_y[..., 0::2].sin(), pos_y[..., 1::2].cos()), dim=4).flatten(3) + if pos_tensor.size(-1) == 2: + return torch.cat((pos_y, pos_x), dim=3) + raise ValueError(f"Unknown pos_tensor last dim {pos_tensor.size(-1)}") + + def forward( + self, + tgt, + memory, + refpoints_sigmoid, + pre_pose_head, + pose_head, + class_head, + lqe_head, + feat_lqe, + integral, + up, + reg_scale, + reg_max, + project, + attn_mask=None, + spatial_shapes: Optional[Tensor] = None, + ): + output = tgt + refpoint_pose = refpoints_sigmoid + output_pose_detach = pred_corners_undetach = 0 + + dec_out_poses = [] + dec_out_logits = [] + dec_out_refs = [] + dec_out_pred_corners = [] + + for layer_id, layer in enumerate(self.layers): + refpoint_pose_input = refpoint_pose[:, :, None] + refpoint_only_pose = refpoint_pose[:, :, 1:] + pose_query_sine_embed = self.sine_embedding(refpoint_only_pose) + pose_query_pos = self.half_pose_ref_point_head(pose_query_sine_embed) + + output = layer( + tgt_pose=output, + tgt_pose_query_pos=pose_query_pos, + tgt_pose_reference_points=refpoint_pose_input, + attn_mask=attn_mask, + memory=memory, + memory_spatial_shapes=spatial_shapes, + ) + + output_pose = output[:, :, 1:] + output_instance = output[:, :, 0] + + if layer_id == 0: + pre_poses = torch.sigmoid(pre_pose_head(output_pose) + inverse_sigmoid(refpoint_only_pose)) + pre_scores = class_head[0](output_instance) + ref_pose_initial = pre_poses.detach() + + pred_corners = pose_head[layer_id](output_pose + output_pose_detach) + pred_corners_undetach + refpoint_pose_without_center = distance2pose(ref_pose_initial, integral(pred_corners, project), reg_scale) + + refpoint_center_pose = torch.mean(refpoint_pose_without_center, dim=2, keepdim=True) + refpoint_pose = torch.cat([refpoint_center_pose, refpoint_pose_without_center], dim=2) + + if self.training or layer_id == self.eval_idx: + score = class_head[layer_id](output_instance) + logit = lqe_head[layer_id](score, refpoint_pose_without_center, feat_lqe) + dec_out_logits.append(logit) + dec_out_poses.append(refpoint_pose_without_center) + dec_out_pred_corners.append(pred_corners) + dec_out_refs.append(ref_pose_initial) + if not self.training: + break + + pred_corners_undetach = pred_corners + if self.training: + refpoint_pose = refpoint_pose.detach() + output_pose_detach = output_pose.detach() + + return ( + torch.stack(dec_out_poses), + torch.stack(dec_out_logits), + torch.stack(dec_out_pred_corners), + torch.stack(dec_out_refs), + pre_poses, + pre_scores, + ) + + +@register() +class DETRPoseTransformer(nn.Module): + """ + Drop-in decoder for DFINE model: expects encoder features list (already 256-d). + """ + + __share__ = ["num_classes", "eval_spatial_size"] + + def __init__( + self, + num_classes=80, + hidden_dim=256, + num_queries=300, + num_decoder_layers=6, + dim_feedforward=1024, + dropout=0.0, + activation="relu", + num_feature_levels=3, + dec_n_points=4, + nhead=8, + aux_loss=True, + num_body_points=17, + feat_strides=(8, 16, 32), + eval_spatial_size=None, + reg_max=32, + reg_scale=4.0, + dn_number=20, + dn_label_noise_ratio=0.5, + ): + super().__init__() + self.num_feature_levels = int(num_feature_levels) + self.num_decoder_layers = int(num_decoder_layers) + self.num_queries = int(num_queries) + self.num_classes = int(num_classes) + self.aux_loss = bool(aux_loss) + self.num_body_points = int(num_body_points) + + decoder_layer = DeformableTransformerDecoderLayer( + hidden_dim, dim_feedforward, dropout, activation, self.num_feature_levels, nhead, dec_n_points + ) + self.decoder = TransformerDecoder( + decoder_layer, + self.num_decoder_layers, + return_intermediate=True, + hidden_dim=hidden_dim, + num_body_points=self.num_body_points, + ) + + # shared priors + self.keypoint_embedding = nn.Embedding(self.num_body_points, hidden_dim) + self.instance_embedding = nn.Embedding(1, hidden_dim) + + self.label_enc = nn.Embedding(self.num_classes + 1, hidden_dim) + self.pose_enc = nn.Embedding(self.num_body_points, hidden_dim) + + # class head + pose heads + _class_embed = nn.Linear(hidden_dim, self.num_classes) + prior_prob = 0.01 + bias_value = -math.log((1 - prior_prob) / prior_prob) + _class_embed.bias.data = torch.ones(self.num_classes) * bias_value + + _pre_point_embed = MLP(hidden_dim, hidden_dim, 2, 3) + nn.init.constant_(_pre_point_embed.layers[-1].weight.data, 0) + nn.init.constant_(_pre_point_embed.layers[-1].bias.data, 0) + + _point_embed = MLP(hidden_dim, hidden_dim, 2 * (reg_max + 1), 3) + nn.init.constant_(_point_embed.layers[-1].weight.data, 0) + nn.init.constant_(_point_embed.layers[-1].bias.data, 0) + + _lqe_embed = LQE(4, hidden_dim, 2, self.num_body_points) + + self.class_embed = nn.ModuleList([copy.deepcopy(_class_embed) for _ in range(self.num_decoder_layers)]) + self.pose_embed = nn.ModuleList([copy.deepcopy(_point_embed) for _ in range(self.num_decoder_layers)]) + self.lqe_embed = nn.ModuleList([copy.deepcopy(_lqe_embed) for _ in range(self.num_decoder_layers)]) + self.pre_pose_embed = _pre_point_embed + + self.integral = Integral(reg_max) + self.up = nn.Parameter(torch.tensor([1 / 2]), requires_grad=False) + self.reg_max = int(reg_max) + self.reg_scale = nn.Parameter(torch.tensor([reg_scale]), requires_grad=False) + + # two-stage encoder heads + self.enc_output = nn.Linear(hidden_dim, hidden_dim) + self.enc_output_norm = nn.LayerNorm(hidden_dim) + self.enc_out_class_embed = copy.deepcopy(_class_embed) + self.enc_pose_embed = MLP(hidden_dim, 2 * hidden_dim, 2 * self.num_body_points, 4) + nn.init.constant_(self.enc_pose_embed.layers[-1].weight.data, 0) + nn.init.constant_(self.enc_pose_embed.layers[-1].bias.data, 0) + + self.feat_strides = list(feat_strides) if feat_strides is not None else [8, 16, 32] + self.eval_spatial_size = eval_spatial_size + if self.eval_spatial_size: + anchors, valid_mask = self._generate_anchors() + self.register_buffer("anchors", anchors) + self.register_buffer("valid_mask", valid_mask) + + self.dn_number = int(dn_number) + self.dn_label_noise_ratio = float(dn_label_noise_ratio) + + self._reset_parameters() + + def _reset_parameters(self): + for p in self.parameters(): + if p.dim() > 1: + nn.init.xavier_uniform_(p) + for m in self.modules(): + if isinstance(m, MSDeformAttn): + m._reset_parameters() + + def _get_encoder_input(self, feats: List[torch.Tensor]): + feat_flatten = [] + spatial_shapes = [] + split_sizes = [] + for feat in feats: + _, _, h, w = feat.shape + feat_flatten.append(feat.flatten(2).permute(0, 2, 1)) + spatial_shapes.append([h, w]) + split_sizes.append(h * w) + feat_flatten = torch.concat(feat_flatten, 1) + # Keep spatial shapes as Python list for export stability in DETRPose ms-deform-attn. + return feat_flatten, spatial_shapes, split_sizes + + def _generate_anchors(self, spatial_shapes=None, device="cpu"): + if spatial_shapes is None: + spatial_shapes = [] + eval_h, eval_w = self.eval_spatial_size + for s in self.feat_strides: + spatial_shapes.append([int(eval_h / s), int(eval_w / s)]) + anchors = [] + for (H_, W_) in spatial_shapes: + grid_y, grid_x = torch.meshgrid( + torch.linspace(0, H_ - 1, H_, dtype=torch.float32, device=device), + torch.linspace(0, W_ - 1, W_, dtype=torch.float32, device=device), + indexing="ij", + ) + grid = torch.stack([grid_x, grid_y], -1) + grid = (grid.unsqueeze(0).expand(1, -1, -1, -1) + 0.5) / torch.tensor( + [W_, H_], dtype=torch.float32, device=device + ) + anchors.append(grid.view(1, -1, 2)) + anchors = torch.cat(anchors, 1) + valid_mask = ((anchors > 0.01) & (anchors < 0.99)).all(-1, keepdim=True) + anchors = torch.log(anchors / (1 - anchors)) + return anchors, ~valid_mask + + def convert_to_deploy(self): + self.project = weighting_function(self.reg_max, self.up, self.reg_scale, deploy=True) + self.lqe_embed = nn.ModuleList([nn.Identity()] * (self.num_decoder_layers - 1) + [self.lqe_embed[-1]]) + + def forward(self, feats, targets=None): + memory, spatial_shapes, split_sizes = self._get_encoder_input(feats) + + if self.training: + output_proposals, valid_mask = self._generate_anchors(spatial_shapes, memory.device) + output_memory = memory.masked_fill(valid_mask, float(0)) + output_proposals = output_proposals.repeat(memory.size(0), 1, 1) + else: + output_proposals = self.anchors.repeat(memory.size(0), 1, 1) + output_memory = memory.masked_fill(self.valid_mask, float(0)) + + output_memory = self.enc_output_norm(self.enc_output(output_memory)) + topk = self.num_queries + enc_outputs_class_unselected = self.enc_out_class_embed(output_memory) + topk_idx = torch.topk(enc_outputs_class_unselected.max(-1)[0], topk, dim=1)[1] + + topk_memory = output_memory.gather(dim=1, index=topk_idx.unsqueeze(-1).repeat(1, 1, output_memory.shape[-1])) + topk_anchors = output_proposals.gather(dim=1, index=topk_idx.unsqueeze(-1).repeat(1, 1, 2)) + + bs, nq = topk_memory.shape[:2] + delta_unsig_keypoint = self.enc_pose_embed(topk_memory).reshape(bs, nq, self.num_body_points, 2) + enc_outputs_pose_coord = torch.sigmoid(delta_unsig_keypoint + topk_anchors.unsqueeze(-2)) + enc_outputs_center_coord = torch.mean(enc_outputs_pose_coord, dim=2, keepdim=True) + enc_outputs_pose_coord = torch.cat([enc_outputs_center_coord, enc_outputs_pose_coord], dim=2) + refpoint_pose_sigmoid = enc_outputs_pose_coord.detach() + + tgt = topk_memory.detach().unsqueeze(-2) + tgt_pose = self.keypoint_embedding.weight[None, None].expand(bs, topk, -1, -1) + tgt + tgt_global = self.instance_embedding.weight[None, None].expand(bs, topk, -1, -1) + tgt_pose = torch.cat([tgt_global, tgt_pose], dim=2) + + # Denoising (pose CDN) + if self.training and targets is not None: + input_query_label, input_query_pose, attn_mask, dn_meta = prepare_for_cdn( + dn_args=(targets, self.dn_number, self.dn_label_noise_ratio), + training=True, + num_queries=self.num_queries, + num_classes=self.num_classes, + num_keypoints=self.num_body_points, + hidden_dim=self.class_embed[0].in_features, + label_enc=self.label_enc, + pose_enc=self.pose_enc, + device=feats[0].device, + ) + tgt_pose = torch.cat([input_query_label, tgt_pose], dim=1) + refpoint_pose_sigmoid = torch.cat([input_query_pose.sigmoid(), refpoint_pose_sigmoid], dim=1) + else: + attn_mask = None + dn_meta = None + + # preprocess memory for deformable attention (same as DETRPose) + value = memory.unflatten(2, (self.decoder.layers[0].cross_attn.n_heads, -1)) + value = value.permute(0, 2, 3, 1).flatten(0, 1).split(split_sizes, dim=-1) + + if not hasattr(self, "project"): + project = weighting_function(self.reg_max, self.up, self.reg_scale) + else: + project = self.project + + out_poses, out_logits, out_corners, out_refs, out_pre_poses, out_pre_scores = self.decoder( + tgt=tgt_pose, + memory=value, + refpoints_sigmoid=refpoint_pose_sigmoid, + spatial_shapes=spatial_shapes, + attn_mask=attn_mask, + pre_pose_head=self.pre_pose_embed, + pose_head=self.pose_embed, + class_head=self.class_embed, + lqe_head=self.lqe_embed, + feat_lqe=feats[0], + up=self.up, + reg_max=self.reg_max, + reg_scale=self.reg_scale, + integral=self.integral, + project=project, + ) + + # flatten keypoints to [B, Q, K*2] (no center) + out_poses = out_poses.flatten(-2) # [L,B,Q,K*2] + + if self.training and dn_meta is not None: + out_pre_poses = out_pre_poses.flatten(-2) + dn_out_poses, out_poses = torch.split(out_poses, [dn_meta["pad_size"], self.num_queries], dim=2) + dn_out_logits, out_logits = torch.split(out_logits, [dn_meta["pad_size"], self.num_queries], dim=2) + dn_out_corners, out_corners = torch.split(out_corners, [dn_meta["pad_size"], self.num_queries], dim=2) + dn_out_refs, out_refs = torch.split(out_refs, [dn_meta["pad_size"], self.num_queries], dim=2) + dn_out_pre_poses, out_pre_poses = torch.split(out_pre_poses, [dn_meta["pad_size"], self.num_queries], dim=1) + dn_out_pre_scores, out_pre_scores = torch.split(out_pre_scores, [dn_meta["pad_size"], self.num_queries], dim=1) + + out = {"pred_logits": out_logits[-1], "pred_keypoints": out_poses[-1]} + + if self.training and self.aux_loss: + out.update( + { + "pred_corners": out_corners[-1], + "ref_points": out_refs[-1], + "up": self.up, + "reg_scale": self.reg_scale, + "reg_max": self.reg_max, + } + ) + out["aux_outputs"] = self._set_aux_loss2( + out_logits[:-1], + out_poses[:-1], + out_corners[:-1], + out_refs[:-1], + out_corners[-1], + out_logits[-1], + ) + out["aux_pre_outputs"] = {"pred_logits": out_pre_scores, "pred_keypoints": out_pre_poses} + if dn_meta is not None: + out["dn_aux_outputs"] = self._set_aux_loss2( + dn_out_logits, + dn_out_poses, + dn_out_corners, + dn_out_refs, + dn_out_corners[-1], + dn_out_logits[-1], + ) + out["dn_aux_pre_outputs"] = {"pred_logits": dn_out_pre_scores, "pred_keypoints": dn_out_pre_poses} + out["dn_meta"] = dn_meta + return out + + @torch.jit.unused + def _set_aux_loss2( + self, + outputs_class, + outputs_keypoints, + outputs_corners, + outputs_ref, + teacher_corners=None, + teacher_logits=None, + ): + return [ + { + "pred_logits": a, + "pred_keypoints": b, + "pred_corners": c, + "ref_points": d, + "teacher_corners": teacher_corners, + "teacher_logits": teacher_logits, + } + for a, b, c, d in zip(outputs_class, outputs_keypoints, outputs_corners, outputs_ref) + ] + diff --git a/src/zoo/dfine/detrpose/utils.py b/src/zoo/dfine/detrpose/utils.py new file mode 100644 index 00000000..3d51c1c8 --- /dev/null +++ b/src/zoo/dfine/detrpose/utils.py @@ -0,0 +1,80 @@ +""" +Small utilities vendored from DETRPose. + +Source: /DETRPose/src/models/detrpose/utils.py +""" + +import math + +import torch +import torch.nn.functional as F +from torch import nn, Tensor + + +def gen_encoder_output_proposals(memory: Tensor, spatial_shapes: Tensor): + """ + Args: + memory: [bs, sum(hw), d_model] + spatial_shapes: [n_levels, 2] (H,W) + Returns: + output_memory: [bs, sum(hw), d_model] + output_proposals: [bs, sum(hw), 2] (unsigmoid xy anchors) + """ + N_, S_, C_ = memory.shape + proposals = [] + for (H_, W_) in spatial_shapes: + grid_y, grid_x = torch.meshgrid( + torch.linspace(0, H_ - 1, int(H_), dtype=torch.float32, device=memory.device), + torch.linspace(0, W_ - 1, int(W_), dtype=torch.float32, device=memory.device), + indexing="ij", + ) + grid = torch.stack([grid_x, grid_y], -1) # H,W,2 + grid = (grid.unsqueeze(0).expand(N_, -1, -1, -1) + 0.5) / torch.tensor( + [W_, H_], dtype=torch.float32, device=memory.device + ) + proposal = grid.view(N_, -1, 2) + proposals.append(proposal) + output_proposals = torch.cat(proposals, 1) # [bs, sum(hw), 2] in [0,1] + output_proposals_valid = ((output_proposals > 0.01) & (output_proposals < 0.99)).all(-1, keepdim=True) + output_proposals = torch.log(output_proposals / (1 - output_proposals)) # unsigmoid + output_proposals = output_proposals.masked_fill(~output_proposals_valid, float("inf")) + + output_memory = memory.masked_fill(~output_proposals_valid, float(0)) + return output_memory, output_proposals + + +class MLP(nn.Module): + """Very simple multi-layer perceptron (FFN).""" + + def __init__(self, input_dim, hidden_dim, output_dim, num_layers): + super().__init__() + self.num_layers = num_layers + h = [hidden_dim] * (num_layers - 1) + self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])) + + def forward(self, x): + for i, layer in enumerate(self.layers): + x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x) + return x + + +def _get_activation_fn(activation, d_model=256, batch_dim=0): + if activation == "relu": + return F.relu + if activation == "gelu": + return F.gelu + if activation == "glu": + return F.glu + if activation == "prelu": + return nn.PReLU() + if activation == "selu": + return F.selu + raise RuntimeError(f"activation should be relu/gelu, not {activation}.") + + +def inverse_sigmoid(x, eps=1e-3): + x = x.clamp(min=0, max=1) + x1 = x.clamp(min=eps) + x2 = (1 - x).clamp(min=eps) + return torch.log(x1 / x2) + diff --git a/test/test.py b/test/test.py new file mode 100644 index 00000000..b72e5401 --- /dev/null +++ b/test/test.py @@ -0,0 +1,101 @@ +import argparse +import sys +import time + +try: + import serial + from serial.tools import list_ports +except ModuleNotFoundError: + print( + "Missing dependency: pyserial\n" + "Install it in your active environment with:\n" + " python3 -m pip install pyserial", + file=sys.stderr, + ) + raise SystemExit(1) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Read raw UART data from a serial port and print it." + ) + parser.add_argument( + "port", + nargs="?", + default="/dev/ttyUSB0", + help="Serial port, for example /dev/ttyUSB0 or COM3", + ) + parser.add_argument( + "--baudrate", + type=int, + default=38400, + help="Serial baudrate (default: 38400)", + ) + parser.add_argument( + "--timeout", + type=float, + default=1.0, + help="Read timeout in seconds (default: 1.0)", + ) + parser.add_argument( + "--list-ports", + action="store_true", + help="List detected serial ports and exit", + ) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + + if args.list_ports: + ports = list(list_ports.comports()) + if not ports: + print("No serial ports detected.") + return 1 + + print("Detected serial ports:") + for port in ports: + print(f" {port.device} - {port.description}") + return 0 + + try: + ser = serial.Serial( + port=args.port, + baudrate=args.baudrate, + timeout=args.timeout, + ) + except serial.SerialException as exc: + print(f"Failed to open serial port {args.port}: {exc}", file=sys.stderr) + return 1 + + print(f"Listening on {args.port} at {args.baudrate} baud. Press Ctrl+C to stop.") + + try: + while True: + waiting = ser.in_waiting + chunk = ser.read(waiting or 1) + + if not chunk: + continue + + timestamp = time.strftime("%H:%M:%S") + print(f"[{timestamp}] Raw bytes: {chunk!r}") + print(f"[{timestamp}] Hex : {chunk.hex(' ')}") + + try: + decoded = chunk.decode("utf-8") + print(f"[{timestamp}] Text : {decoded.rstrip()}") + except UnicodeDecodeError: + decoded = chunk.decode("latin-1", errors="replace") + print(f"[{timestamp}] Text* : {decoded.rstrip()}") + except KeyboardInterrupt: + print("\nStopped by user.") + finally: + ser.close() + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/benchmark/video_latency_compare.py b/tools/benchmark/video_latency_compare.py new file mode 100644 index 00000000..bc4c668d --- /dev/null +++ b/tools/benchmark/video_latency_compare.py @@ -0,0 +1,382 @@ +#!/usr/bin/env python3 +"""Compare inference latency for PyTorch, ONNX Runtime and TensorRT on the same video frames.""" + +import argparse +import os +import statistics +import sys +import time +from dataclasses import dataclass +from typing import Callable, Dict, List, Optional, Tuple + +import cv2 +import numpy as np +import torch +import torchvision.transforms as T +from PIL import Image + +REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")) +if REPO not in sys.path: + sys.path.insert(0, REPO) + +from src.core import YAMLConfig +from tools.inference.trt_inf import TRTInference +from tools.model_surgery.shared_arch import SegmentationHead, SharedBackboneDualDecoder, load_any_state + +try: + import onnxruntime as ort +except Exception: + ort = None + + +@dataclass +class BackendResult: + name: str + frames: int + mean_ms: float + p50_ms: float + p95_ms: float + + +def resize_with_aspect_ratio(image: Image.Image, size: int) -> Tuple[Image.Image, int, int]: + original_width, original_height = image.size + ratio = min(size / original_width, size / original_height) + new_width = int(original_width * ratio) + new_height = int(original_height * ratio) + image = image.resize((new_width, new_height), Image.BILINEAR) + + padded = Image.new("RGB", (size, size)) + pad_w = (size - new_width) // 2 + pad_h = (size - new_height) // 2 + padded.paste(image, (pad_w, pad_h)) + return padded, new_width, new_height + + +def load_video_inputs(video_path: str, input_size: int, max_frames: int) -> Tuple[List[np.ndarray], List[np.ndarray]]: + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + raise RuntimeError(f"Could not open video: {video_path}") + + to_tensor = T.ToTensor() + images_np: List[np.ndarray] = [] + orig_sizes_np: List[np.ndarray] = [] + + while True: + if max_frames > 0 and len(images_np) >= max_frames: + break + + ok, frame = cap.read() + if not ok: + break + + frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frame_pil = Image.fromarray(frame_rgb) + resized, new_w, new_h = resize_with_aspect_ratio(frame_pil, input_size) + x = to_tensor(resized).unsqueeze(0).numpy().astype(np.float32) + sz = np.array([[new_h, new_w]], dtype=np.float32) + + images_np.append(x) + orig_sizes_np.append(sz) + + cap.release() + if not images_np: + raise RuntimeError("No frames decoded from video") + return images_np, orig_sizes_np + + +def measure_backend( + name: str, + run_one: Callable[[np.ndarray, np.ndarray], None], + images_np: List[np.ndarray], + orig_sizes_np: List[np.ndarray], + warmup: int, +) -> BackendResult: + if warmup > 0: + for i in range(min(warmup, len(images_np))): + run_one(images_np[i], orig_sizes_np[i]) + + times_ms: List[float] = [] + for x_np, sz_np in zip(images_np, orig_sizes_np): + t0 = time.perf_counter() + run_one(x_np, sz_np) + dt = (time.perf_counter() - t0) * 1000.0 + times_ms.append(dt) + + return BackendResult( + name=name, + frames=len(times_ms), + mean_ms=float(statistics.fmean(times_ms)), + p50_ms=float(np.percentile(np.array(times_ms), 50)), + p95_ms=float(np.percentile(np.array(times_ms), 95)), + ) + + +def build_torch_runner(config: str, resume: str, device: str): + cfg = YAMLConfig(config, resume=resume) + if "HGNetv2" in cfg.yaml_cfg: + cfg.yaml_cfg["HGNetv2"]["pretrained"] = False + + checkpoint = torch.load(resume, map_location="cpu") + if "ema" in checkpoint: + state = checkpoint["ema"]["module"] + else: + state = checkpoint["model"] + + incompatible = cfg.model.load_state_dict(state, strict=False) + if incompatible.missing_keys or incompatible.unexpected_keys: + print("[torch] WARNING: non-strict checkpoint load") + + model = cfg.model.deploy().to(device).eval() + + use_cuda = str(device).startswith("cuda") and torch.cuda.is_available() + + @torch.inference_mode() + def run_one(x_np: np.ndarray, sz_np: np.ndarray): + x = torch.from_numpy(x_np).to(device) + # torch deploy model ignores orig size; keep signature consistent with others + if use_cuda: + torch.cuda.synchronize() + _ = model(x) + if use_cuda: + torch.cuda.synchronize() + + return run_one + + +def _disable_hgnet_pretrained(cfg: YAMLConfig): + if hasattr(cfg, "yaml_cfg") and isinstance(cfg.yaml_cfg, dict) and "HGNetv2" in cfg.yaml_cfg: + node = cfg.yaml_cfg["HGNetv2"] + if isinstance(node, dict): + node["pretrained"] = False + + +def build_torch_multimodel_runner( + det_config: str, + pose_config: str, + merged_ckpt: str, + device: str, + input_size: int = 640, + seg_num_classes: int = 7, + seg_feature_dim: int = 384, + seg_dropout: float = 0.1, +): + det_cfg = YAMLConfig(det_config) + pose_cfg = YAMLConfig(pose_config) + _disable_hgnet_pretrained(det_cfg) + _disable_hgnet_pretrained(pose_cfg) + + det_model = det_cfg.model + pose_model = pose_cfg.model + + det_model.eval() + with torch.no_grad(): + dummy = torch.randn(1, 3, int(input_size), int(input_size)) + feats = det_model.backbone(dummy) + in_channels = [int(f.shape[1]) for f in feats] + + ckpt_state = load_any_state(merged_ckpt) + fpn_key = "seg_head.fpn.lateral_convs.0.weight" + feature_dim = int(ckpt_state[fpn_key].shape[0]) if fpn_key in ckpt_state else int(seg_feature_dim) + seg_head = SegmentationHead(in_channels, int(seg_num_classes), int(feature_dim), float(seg_dropout)) + + model = SharedBackboneDualDecoder( + backbone=det_model.backbone, + encoder=det_model.encoder, + det_decoder=det_model.decoder, + pose_decoder=pose_model.decoder, + seg_head=seg_head, + ) + missing, unexpected = model.load_state_dict(ckpt_state, strict=False) + if missing or unexpected: + print( + f"[torch-multimodel] WARNING non-strict load: " + f"missing={len(missing)} unexpected={len(unexpected)}" + ) + + model = model.to(device).eval() + use_cuda = str(device).startswith("cuda") and torch.cuda.is_available() + + @torch.inference_mode() + def run_one(x_np: np.ndarray, sz_np: np.ndarray): + x = torch.from_numpy(x_np).to(device) + if use_cuda: + torch.cuda.synchronize() + _ = model(x) + if use_cuda: + torch.cuda.synchronize() + + return run_one + + +def build_onnx_runner(onnx_path: str, use_gpu: bool): + if ort is None: + raise RuntimeError("onnxruntime is not installed") + + available = list(ort.get_available_providers()) + providers = ["CPUExecutionProvider"] + if use_gpu: + if "CUDAExecutionProvider" in available: + providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] + else: + print( + "[onnx] WARNING: CUDAExecutionProvider is not available in this environment. " + f"Falling back to CPU. available={available}" + ) + + providers = [p for p in providers if p in available] + if not providers: + raise RuntimeError(f"No usable ONNX Runtime providers found. available={available}") + + try: + sess = ort.InferenceSession(onnx_path, providers=providers) + except Exception as e: + msg = str(e) + if "GatherElementsAxis1Plugin" in msg: + raise RuntimeError( + "ONNX model includes TensorRT-only custom op 'GatherElementsAxis1Plugin'. " + "Use an ORT-compatible ONNX (before TRT plugin rewrite), e.g. " + "'outputs/phase2_run/singlepass_best_dynamo.onnx'." + ) from e + raise + print(f"[onnx] providers={sess.get_providers()}") + input_names = {i.name for i in sess.get_inputs()} + + def run_one(x_np: np.ndarray, sz_np: np.ndarray): + feed: Dict[str, np.ndarray] = {} + if "images" in input_names: + feed["images"] = x_np + if "orig_target_sizes" in input_names: + feed["orig_target_sizes"] = sz_np + if not feed: + # fallback for unnamed/custom single-input graphs + first_name = sess.get_inputs()[0].name + feed[first_name] = x_np + _ = sess.run(None, feed) + + return run_one + + +def build_trt_runner(engine_path: str, device: str, plugin_libs: List[str]): + trt_runner = TRTInference( + engine_path, + device=device, + max_batch_size=1, + plugin_libs=plugin_libs, + ) + + def run_one(x_np: np.ndarray, sz_np: np.ndarray): + blob: Dict[str, torch.Tensor] = {} + x_t = torch.from_numpy(x_np).to(device) + sz_t = torch.from_numpy(sz_np).to(device) + + if "images" in trt_runner.input_names: + blob["images"] = x_t + if "orig_target_sizes" in trt_runner.input_names: + blob["orig_target_sizes"] = sz_t + + # Fallback if engine uses unnamed/custom input names. + for n in trt_runner.input_names: + if n not in blob: + blob[n] = x_t if len(blob) == 0 else sz_t + + _ = trt_runner(blob) + trt_runner.synchronize() + + return run_one + + +def print_table(results: List[BackendResult]): + print("\nLatency (samme video):") + print(f"{'Modell':<12} {'Frames':>8} {'Mean (ms)':>12} {'P50 (ms)':>10} {'P95 (ms)':>10}") + for r in results: + print(f"{r.name:<12} {r.frames:>8d} {r.mean_ms:>12.2f} {r.p50_ms:>10.2f} {r.p95_ms:>10.2f}") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--video", required=True, help="Path to input video") + ap.add_argument("--input-size", type=int, default=640) + ap.add_argument("--max-frames", type=int, default=300, help="0 = all frames") + ap.add_argument("--warmup", type=int, default=30) + + ap.add_argument("--torch-config", type=str, default=None) + ap.add_argument("--torch-resume", type=str, default=None) + ap.add_argument("--torch-device", type=str, default="cuda:0") + ap.add_argument("--torch-det-config", type=str, default=None, help="Detection config for multimodel torch benchmark") + ap.add_argument("--torch-pose-config", type=str, default=None, help="Pose config for multimodel torch benchmark") + ap.add_argument("--torch-merged-ckpt", type=str, default=None, help="Merged model-surgery checkpoint (.pth)") + ap.add_argument("--torch-seg-num-classes", type=int, default=7) + ap.add_argument("--torch-seg-feature-dim", type=int, default=384) + ap.add_argument("--torch-seg-dropout", type=float, default=0.1) + + ap.add_argument("--onnx", type=str, default=None) + ap.add_argument("--onnx-gpu", action="store_true") + + ap.add_argument("--trt", type=str, default=None) + ap.add_argument("--trt-device", type=str, default="cuda:0") + ap.add_argument("--plugin-lib", action="append", default=[]) + + args = ap.parse_args() + + max_frames = args.max_frames if args.max_frames > 0 else 10**9 + images_np, orig_sizes_np = load_video_inputs(args.video, args.input_size, max_frames) + print(f"Loaded {len(images_np)} frames from {args.video}") + + results: List[BackendResult] = [] + failures: List[str] = [] + + if args.torch_det_config and args.torch_pose_config and args.torch_merged_ckpt: + print("Running PyTorch benchmark (multimodel)...") + try: + torch_run = build_torch_multimodel_runner( + det_config=args.torch_det_config, + pose_config=args.torch_pose_config, + merged_ckpt=args.torch_merged_ckpt, + device=args.torch_device, + input_size=args.input_size, + seg_num_classes=args.torch_seg_num_classes, + seg_feature_dim=args.torch_seg_feature_dim, + seg_dropout=args.torch_seg_dropout, + ) + results.append(measure_backend("PyTorch", torch_run, images_np, orig_sizes_np, args.warmup)) + except Exception as e: + failures.append(f"PyTorch failed: {e}") + elif args.torch_config and args.torch_resume: + print("Running PyTorch benchmark...") + try: + torch_run = build_torch_runner(args.torch_config, args.torch_resume, args.torch_device) + results.append(measure_backend("PyTorch", torch_run, images_np, orig_sizes_np, args.warmup)) + except Exception as e: + failures.append(f"PyTorch failed: {e}") + + if args.onnx: + print("Running ONNX Runtime benchmark...") + try: + onnx_run = build_onnx_runner(args.onnx, args.onnx_gpu) + results.append(measure_backend("ONNX", onnx_run, images_np, orig_sizes_np, args.warmup)) + except Exception as e: + failures.append(f"ONNX failed: {e}") + + if args.trt: + print("Running TensorRT benchmark...") + try: + trt_run = build_trt_runner(args.trt, args.trt_device, args.plugin_lib) + results.append(measure_backend("TensorRT", trt_run, images_np, orig_sizes_np, args.warmup)) + except Exception as e: + failures.append(f"TensorRT failed: {e}") + + if not results: + raise RuntimeError( + "No successful benchmark runs. " + + (" ".join(failures) if failures else "Provide at least one of --torch-*, --onnx, --trt") + ) + + print_table(results) + if failures: + print("\nBackends that failed:") + for m in failures: + print(f"- {m}") + + +if __name__ == "__main__": + main() diff --git a/train.py b/train.py deleted file mode 100644 index d4aa4518..00000000 --- a/train.py +++ /dev/null @@ -1,113 +0,0 @@ -""" -D-FINE: Redefine Regression Task of DETRs as Fine-grained Distribution Refinement -Copyright (c) 2024 The D-FINE Authors. All Rights Reserved. ---------------------------------------------------------------------------------- -Modified from RT-DETR (https://github.com/lyuwenyu/RT-DETR) -Copyright (c) 2023 lyuwenyu. All Rights Reserved. -""" - -import os -import sys -import torch - -sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")) - -import argparse - -from src.core import YAMLConfig, yaml_utils -from src.misc import dist_utils -from src.solver import TASKS -from pprint import pprint - -debug = False - -if debug: - def custom_repr(self): - return f"{{Tensor:{tuple(self.shape)}}} {original_repr(self)}" - - original_repr = torch.Tensor.__repr__ - torch.Tensor.__repr__ = custom_repr - - -def safe_get_rank(): - if torch.distributed.is_available() and torch.distributed.is_initialized(): - return torch.distributed.get_rank() - else: - return 0 - - -def main(args) -> None: - """main""" - dist_utils.setup_distributed(args.print_rank, args.print_method, seed=args.seed) - - assert not all( - [args.tuning, args.resume] - ), "Only support from_scrach or resume or tuning at one time" - - update_dict = yaml_utils.parse_cli(args.update) - update_dict.update( - { - k: v - for k, v in args.__dict__.items() - if k - not in [ - "update", - ] - and v is not None - } - ) - - cfg = YAMLConfig(args.config, **update_dict) - - if args.resume or args.tuning: - if "HGNetv2" in cfg.yaml_cfg: - cfg.yaml_cfg["HGNetv2"]["pretrained"] = False - - if safe_get_rank() == 0: - print("cfg: ") - pprint(cfg.__dict__) - - solver = TASKS[cfg.yaml_cfg["task"]](cfg) - - if args.test_only: - solver.val() - else: - solver.fit() - - dist_utils.cleanup() - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - # priority 0 - parser.add_argument("-c", "--config", type=str, required=True) - parser.add_argument("-r", "--resume", type=str, help="resume from checkpoint") - parser.add_argument("-t", "--tuning", type=str, help="tuning from checkpoint") - parser.add_argument( - "-d", - "--device", - type=str, - help="device", - ) - parser.add_argument("--seed", type=int, help="exp reproducibility") - parser.add_argument("--use-amp", action="store_true", help="auto mixed precision training") - parser.add_argument("--output-dir", type=str, help="output directoy") - parser.add_argument("--summary-dir", type=str, help="tensorboard summry") - parser.add_argument( - "--test-only", - action="store_true", - default=False, - ) - - # priority 1 - parser.add_argument("-u", "--update", nargs="+", help="update yaml config") - - # env - parser.add_argument("--print-method", type=str, default="builtin", help="print method") - parser.add_argument("--print-rank", type=int, default=0, help="print rank id") - - parser.add_argument("--local-rank", type=int, help="local rank id") - args = parser.parse_args() - - main(args) diff --git a/webcam_client.py b/webcam_client.py new file mode 100644 index 00000000..9f0a253d --- /dev/null +++ b/webcam_client.py @@ -0,0 +1,500 @@ +#!/usr/bin/env python3 +""" +Real-time Webcam Client for TensorRT Segmentation +Runs on your local computer to capture webcam and display results. + +Usage: + python webcam_client.py --server-ip YOUR_SERVER_IP --server-port 65432 +""" + +import cv2 +import numpy as np +import socket +import pickle +import struct +import time +import argparse +import threading +from collections import deque +from dataclasses import dataclass +from typing import Optional, Tuple +import sys + +@dataclass +class FrameData: + """Container for frame data""" + frame: np.ndarray + timestamp: float + frame_id: int + +class WebcamClient: + """Real-time webcam client""" + + def __init__(self, server_ip: str, server_port: int, webcam_id: int = 0): + self.server_ip = server_ip + self.server_port = server_port + self.webcam_id = webcam_id + + # Network + self.socket = None + self.connected = False + + # Webcam + self.cap = None + + # Threading + self.running = False + self.send_thread = None + self.receive_thread = None + + # Frame management + self.frame_queue = deque(maxlen=5) # Smaller buffer to reduce latency + self.result_queue = deque(maxlen=5) # Smaller buffer to reduce latency + self.frame_counter = 0 + + # Performance tracking + self.fps_history = deque(maxlen=30) + self.latency_history = deque(maxlen=30) + + # Display settings + self.show_fps = True + self.show_latency = True + self.display_scale = 1.0 + + def connect_to_server(self) -> bool: + """Connect to the inference server""" + try: + print(f"🔗 Connecting to server {self.server_ip}:{self.server_port}...") + + self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + + # Set timeouts + self.socket.settimeout(10.0) + + self.socket.connect((self.server_ip, self.server_port)) + self.connected = True + + print("✅ Connected to server successfully!") + return True + + except Exception as e: + print(f"❌ Failed to connect to server: {e}") + if self.socket: + self.socket.close() + self.socket = None + return False + + def initialize_webcam(self) -> bool: + """Initialize webcam capture""" + try: + print(f"📷 Initializing webcam {self.webcam_id}...") + + # Try different backends for better compatibility + backends = [ + cv2.CAP_DSHOW, # DirectShow (Windows) + cv2.CAP_V4L2, # Video4Linux (Linux) + cv2.CAP_AVFOUNDATION, # AVFoundation (macOS) + cv2.CAP_ANY # Auto-detect + ] + + self.cap = None + for backend in backends: + try: + self.cap = cv2.VideoCapture(self.webcam_id, backend) + if self.cap.isOpened(): + print(f"✅ Using backend: {backend}") + break + else: + self.cap.release() + self.cap = None + except: + continue + + if self.cap is None or not self.cap.isOpened(): + # Try without specifying backend + self.cap = cv2.VideoCapture(self.webcam_id) + if not self.cap.isOpened(): + raise ValueError(f"Could not open webcam {self.webcam_id}") + + # Set webcam properties for better performance + self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640) + self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480) + self.cap.set(cv2.CAP_PROP_FPS, 30) + self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) # Reduce latency + + # Test capture + ret, test_frame = self.cap.read() + if not ret or test_frame is None: + raise ValueError("Failed to capture test frame") + + # Get actual properties + width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + fps = self.cap.get(cv2.CAP_PROP_FPS) + + print(f"✅ Webcam initialized: {width}x{height} @ {fps:.1f} FPS") + return True + + except Exception as e: + print(f"❌ Failed to initialize webcam: {e}") + if self.cap: + self.cap.release() + self.cap = None + return False + + def send_frame_worker(self): + """Worker thread for sending frames to server""" + print("🔄 Send worker thread started") + + while self.running and self.connected: + try: + if len(self.frame_queue) == 0: + time.sleep(0.001) # Short sleep if no frames + continue + + frame_data = self.frame_queue.popleft() + + # Serialize frame data + try: + serialized_data = pickle.dumps(frame_data, protocol=pickle.HIGHEST_PROTOCOL) + except Exception as e: + print(f"❌ Serialization error: {e}") + continue + + # Send data size first + data_size = len(serialized_data) + size_data = struct.pack("!I", data_size) + + try: + self.socket.sendall(size_data) + self.socket.sendall(serialized_data) + except socket.error as e: + print(f"❌ Socket send error: {e}") + self.connected = False + break + + except Exception as e: + if self.running: # Only print error if we're supposed to be running + print(f"❌ Send worker error: {e}") + self.connected = False + break + + print("🔄 Send worker thread stopped") + + def receive_result_worker(self): + """Worker thread for receiving processed frames from server""" + print("🔄 Receive worker thread started") + + while self.running and self.connected: + try: + # Receive data size + size_data = self._recv_all(4) + if not size_data: + print("❌ Failed to receive size data") + break + + data_size = struct.unpack("!I", size_data)[0] + + # Sanity check on data size + if data_size <= 0 or data_size > 100 * 1024 * 1024: # Max 100MB + print(f"❌ Invalid data size: {data_size}") + break + + # Receive actual data + serialized_data = self._recv_all(data_size) + if not serialized_data: + print("❌ Failed to receive frame data") + break + + # Deserialize result + try: + result_data = pickle.loads(serialized_data) + except Exception as e: + print(f"❌ Deserialization error: {e}") + continue + + # Add to result queue + if len(self.result_queue) >= self.result_queue.maxlen: + # Drop oldest if queue is full + self.result_queue.popleft() + + self.result_queue.append(result_data) + + except Exception as e: + if self.running: # Only print error if we're supposed to be running + print(f"❌ Receive worker error: {e}") + self.connected = False + break + + print("🔄 Receive worker thread stopped") + + def _recv_all(self, size: int) -> Optional[bytes]: + """Receive exactly 'size' bytes from socket""" + data = b'' + while len(data) < size: + try: + chunk = self.socket.recv(size - len(data)) + if not chunk: + return None + data += chunk + except socket.timeout: + print("⚠️ Socket timeout during receive") + return None + except socket.error as e: + print(f"❌ Socket receive error: {e}") + return None + except Exception as e: + print(f"❌ Unexpected receive error: {e}") + return None + return data + + def capture_and_display_loop(self): + """Main loop for capturing frames and displaying results""" + print("🎬 Starting capture and display loop...") + print("Press 'q' to quit, 's' to toggle stats, 'f' to toggle fullscreen") + + last_capture_time = time.perf_counter() + fullscreen = False + window_name = "Real-time Segmentation" + + # Create window + cv2.namedWindow(window_name, cv2.WINDOW_NORMAL) + + frame_skip_counter = 0 + + while self.running: + try: + # Capture frame + ret, frame = self.cap.read() + if not ret or frame is None: + print("❌ Failed to capture frame") + time.sleep(0.1) # Wait a bit before retrying + continue + + current_time = time.perf_counter() + + # Skip frames to reduce queue buildup (send every 2nd frame) + frame_skip_counter += 1 + if frame_skip_counter % 2 == 0 and self.connected: + # Add frame to send queue (skip if queue is full to maintain real-time) + if len(self.frame_queue) < self.frame_queue.maxlen: + frame_data = FrameData( + frame=frame.copy(), + timestamp=current_time, + frame_id=self.frame_counter + ) + self.frame_queue.append(frame_data) + self.frame_counter += 1 + + # Calculate capture FPS + capture_fps = 1.0 / (current_time - last_capture_time) if current_time > last_capture_time else 0 + self.fps_history.append(capture_fps) + last_capture_time = current_time + + # Display frame (either processed or original) + display_frame = frame.copy() + + # Try to get processed result + if len(self.result_queue) > 0: + result_data = self.result_queue.popleft() + + # Calculate latency + if 'client_timestamp' in result_data: + latency = current_time - result_data['client_timestamp'] + self.latency_history.append(latency * 1000) # Convert to ms + + # Use processed frame + if 'processed_frame' in result_data and result_data['processed_frame'] is not None: + display_frame = result_data['processed_frame'] + + # Add performance overlay + if self.show_fps or self.show_latency: + self._add_performance_overlay(display_frame) + + # Scale frame if needed + if self.display_scale != 1.0: + new_width = int(display_frame.shape[1] * self.display_scale) + new_height = int(display_frame.shape[0] * self.display_scale) + display_frame = cv2.resize(display_frame, (new_width, new_height)) + + # Display frame + cv2.imshow(window_name, display_frame) + + # Handle key presses + key = cv2.waitKey(1) & 0xFF + if key == ord('q'): + print("👋 Quit requested by user") + break + elif key == ord('s'): + self.show_fps = not self.show_fps + self.show_latency = not self.show_latency + print(f"📊 Stats display: {'ON' if self.show_fps else 'OFF'}") + elif key == ord('f'): + fullscreen = not fullscreen + if fullscreen: + cv2.setWindowProperty(window_name, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN) + else: + cv2.setWindowProperty(window_name, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_NORMAL) + elif key == ord('+') or key == ord('='): + self.display_scale = min(2.0, self.display_scale + 0.1) + print(f"🔍 Display scale: {self.display_scale:.1f}x") + elif key == ord('-'): + self.display_scale = max(0.5, self.display_scale - 0.1) + print(f"🔍 Display scale: {self.display_scale:.1f}x") + elif key == ord('r'): # Reconnect + if not self.connected: + print("🔄 Attempting to reconnect...") + if self.connect_to_server(): + self._restart_worker_threads() + + except KeyboardInterrupt: + print("\n👋 Interrupted by user") + break + except Exception as e: + print(f"❌ Display loop error: {e}") + import traceback + traceback.print_exc() + time.sleep(0.1) # Brief pause before continuing + + cv2.destroyAllWindows() + + def _restart_worker_threads(self): + """Restart worker threads after reconnection""" + if self.send_thread and self.send_thread.is_alive(): + return # Already running + + self.send_thread = threading.Thread(target=self.send_frame_worker, daemon=True) + self.receive_thread = threading.Thread(target=self.receive_result_worker, daemon=True) + + self.send_thread.start() + self.receive_thread.start() + + def _add_performance_overlay(self, frame: np.ndarray): + """Add performance information overlay to frame""" + overlay_y = 30 + + if self.show_fps and len(self.fps_history) > 0: + avg_fps = sum(self.fps_history) / len(self.fps_history) + fps_text = f"Capture FPS: {avg_fps:.1f}" + cv2.putText(frame, fps_text, (10, overlay_y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) + overlay_y += 30 + + if self.show_latency and len(self.latency_history) > 0: + avg_latency = sum(self.latency_history) / len(self.latency_history) + latency_text = f"Latency: {avg_latency:.1f}ms" + cv2.putText(frame, latency_text, (10, overlay_y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2) + overlay_y += 30 + + # Connection status + status_text = f"Server: {'Connected' if self.connected else 'Disconnected'}" + status_color = (0, 255, 0) if self.connected else (0, 0, 255) + cv2.putText(frame, status_text, (10, overlay_y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, status_color, 2) + overlay_y += 30 + + # Queue status + send_queue_size = len(self.frame_queue) + result_queue_size = len(self.result_queue) + queue_text = f"Queues: S:{send_queue_size} R:{result_queue_size}" + cv2.putText(frame, queue_text, (10, overlay_y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2) + + def start(self): + """Start the client""" + print("🚀 Starting webcam client...") + + # Initialize webcam first + if not self.initialize_webcam(): + print("❌ Cannot start without webcam") + return False + + # Connect to server + if not self.connect_to_server(): + print("⚠️ Starting without server connection (you can press 'r' to reconnect)") + + # Start worker threads + self.running = True + + if self.connected: + self.send_thread = threading.Thread(target=self.send_frame_worker, daemon=True) + self.receive_thread = threading.Thread(target=self.receive_result_worker, daemon=True) + + self.send_thread.start() + self.receive_thread.start() + + # Start main loop + try: + self.capture_and_display_loop() + finally: + self.stop() + + return True + + def stop(self): + """Stop the client""" + print("🛑 Stopping webcam client...") + + self.running = False + + # Wait for threads to finish + if self.send_thread and self.send_thread.is_alive(): + self.send_thread.join(timeout=2.0) + + if self.receive_thread and self.receive_thread.is_alive(): + self.receive_thread.join(timeout=2.0) + + # Close connections + if self.socket: + try: + self.socket.close() + except: + pass + self.socket = None + + if self.cap: + self.cap.release() + self.cap = None + + self.connected = False + print("✅ Client stopped") + +def parse_args(): + """Parse command line arguments""" + parser = argparse.ArgumentParser(description='Real-time Webcam Client for TensorRT Segmentation') + parser.add_argument('--server-ip', type=str, required=True, + help='IP address of the inference server') + parser.add_argument('--server-port', type=int, default=65432, + help='Port of the inference server (default: 65432)') + parser.add_argument('--webcam-id', type=int, default=0, + help='Webcam device ID (default: 0)') + return parser.parse_args() + +def main(): + """Main function""" + args = parse_args() + + print("=" * 60) + print("📷 Real-time Webcam Segmentation Client") + print("=" * 60) + print(f"🖥️ Server: {args.server_ip}:{args.server_port}") + print(f"📷 Webcam: {args.webcam_id}") + print("=" * 60) + + client = WebcamClient(args.server_ip, args.server_port, args.webcam_id) + + try: + success = client.start() + if not success: + print("❌ Failed to start client") + sys.exit(1) + except KeyboardInterrupt: + print("\n👋 Interrupted by user") + except Exception as e: + print(f"❌ Client error: {e}") + import traceback + traceback.print_exc() + sys.exit(1) + + print("👋 Client finished") + +if __name__ == "__main__": + main()