Upload training/train_detection.py with huggingface_hub

training/train_detection.py

@@ -0,0 +1,491 @@
+#!/usr/bin/env python3
+"""
+OCULUS Detection Head Training
+
+Trains the detection (box) and point heads on COCO detection data.
+Uses the frozen vision encoders + trained projector, only trains the heads.
+"""
+
+import os
+import sys
+import json
+import time
+import random
+from pathlib import Path
+from dataclasses import dataclass
+from typing import List, Dict, Tuple, Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from PIL import Image
+
+OCULUS_ROOT = Path(__file__).parent
+
+# Add to path
+sys.path.insert(0, str(OCULUS_ROOT))
+
+from oculus_unified_model import OculusForConditionalGeneration, OculusConfig
+
+
+@dataclass
+class DetectionTrainingConfig:
+    """Training configuration."""
+    # Data
+    data_dir: str = "data/coco"
+    annotations_file: str = "annotations/instances_train2017.json"
+    images_subdir: str = "images"
+    
+    # Training
+    batch_size: int = 4
+    learning_rate: float = 1e-4
+    num_epochs: int = 3
+    warmup_steps: int = 100
+    max_samples: int = 3000  # Limit for faster training
+    
+    # Model
+    checkpoint_path: str = "checkpoints/oculus_coco/final"
+    
+    # Checkpointing
+    save_every: int = 200
+    checkpoint_dir: str = "checkpoints/oculus_detection"
+    
+    # Logging
+    log_every: int = 25
+
+
+class COCODetectionDataset(Dataset):
+    """COCO Detection dataset."""
+    
+    # COCO 80 class names
+    COCO_CLASSES = [
+        'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck',
+        'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench',
+        'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra',
+        'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
+        'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
+        'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
+        'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange',
+        'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
+        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse',
+        'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink',
+        'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
+        'toothbrush'
+    ]
+    
+    def __init__(self, data_dir: str, annotations_file: str, images_subdir: str, max_samples: int = None):
+        self.data_dir = Path(data_dir)
+        self.images_dir = self.data_dir / images_subdir
+        
+        # Load annotations
+        annotations_path = self.data_dir / annotations_file
+        print(f"  Loading annotations from {annotations_path}...")
+        
+        with open(annotations_path) as f:
+            coco_data = json.load(f)
+        
+        # Build category ID to index mapping
+        self.cat_id_to_idx = {}
+        for i, cat in enumerate(coco_data['categories']):
+            self.cat_id_to_idx[cat['id']] = i
+        
+        # Build image ID to annotations mapping
+        img_to_anns = {}
+        for ann in coco_data['annotations']:
+            img_id = ann['image_id']
+            if img_id not in img_to_anns:
+                img_to_anns[img_id] = []
+            img_to_anns[img_id].append(ann)
+        
+        # Build samples list
+        self.samples = []
+        for img_info in coco_data['images']:
+            img_id = img_info['id']
+            if img_id not in img_to_anns:
+                continue
+            
+            # Check if image exists
+            img_path = self.images_dir / img_info['file_name']
+            if not img_path.exists():
+                continue
+            
+            anns = img_to_anns[img_id]
+            
+            # Convert annotations to boxes
+            boxes = []
+            labels = []
+            for ann in anns:
+                if 'bbox' not in ann or ann.get('iscrowd', 0):
+                    continue
+                
+                # COCO bbox format: [x, y, width, height]
+                x, y, w, h = ann['bbox']
+                
+                # Convert to normalized [x1, y1, x2, y2]
+                x1 = x / img_info['width']
+                y1 = y / img_info['height']
+                x2 = (x + w) / img_info['width']
+                y2 = (y + h) / img_info['height']
+                
+                # Clamp to [0, 1]
+                x1, y1, x2, y2 = max(0, x1), max(0, y1), min(1, x2), min(1, y2)
+                
+                boxes.append([x1, y1, x2, y2])
+                labels.append(self.cat_id_to_idx[ann['category_id']])
+            
+            if boxes:
+                self.samples.append({
+                    'image_path': str(img_path),
+                    'boxes': boxes,
+                    'labels': labels,
+                    'width': img_info['width'],
+                    'height': img_info['height']
+                })
+            
+            if max_samples and len(self.samples) >= max_samples:
+                break
+        
+        print(f"  Loaded {len(self.samples):,} images with detections")
+    
+    def __len__(self):
+        return len(self.samples)
+    
+    def __getitem__(self, idx):
+        return self.samples[idx]
+
+
+class DetectionTrainer:
+    """Trainer for detection heads."""
+    
+    def __init__(self, config: DetectionTrainingConfig):
+        self.config = config
+        
+        print("\n" + "=" * 60)
+        print("🎯 OCULUS DETECTION TRAINER")
+        print("=" * 60)
+        
+        self._load_model()
+        self._load_dataset()
+        self._create_optimizer()
+        
+        self.checkpoint_dir = Path(config.checkpoint_dir)
+        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    
+    def _load_model(self):
+        """Load model with trained projector."""
+        print("\n[Loading Model]")
+        
+        checkpoint_path = OCULUS_ROOT / self.config.checkpoint_path
+        self.model = OculusForConditionalGeneration.from_pretrained(checkpoint_path)
+        
+        # Load vision encoders
+        self.model.vision_encoder.load_encoders()
+        
+        # Freeze vision encoder and projector
+        for param in self.model.vision_encoder.parameters():
+            param.requires_grad = False
+        for param in self.model.projector.parameters():
+            param.requires_grad = False
+        
+        # Make sure detection/point heads are trainable
+        for param in self.model.detection_head.parameters():
+            param.requires_grad = True
+        for param in self.model.point_head.parameters():
+            param.requires_grad = True
+        
+        # Count trainable params
+        trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
+        total = sum(p.numel() for p in self.model.parameters())
+        print(f"  ✓ Trainable: {trainable:,} / {total:,} parameters")
+    
+    def _load_dataset(self):
+        """Load COCO detection dataset."""
+        print("\n[Loading Dataset]")
+        self.dataset = COCODetectionDataset(
+            self.config.data_dir,
+            self.config.annotations_file,
+            self.config.images_subdir,
+            max_samples=self.config.max_samples
+        )
+    
+    def _create_optimizer(self):
+        """Create optimizer for detection heads only."""
+        print("\n[Optimizer]")
+        
+        # Only optimize detection heads
+        params = list(self.model.detection_head.parameters()) + \
+                 list(self.model.point_head.parameters())
+        
+        if self.model.vision_adapter is not None:
+            params += list(self.model.vision_adapter.parameters())
+        
+        self.optimizer = torch.optim.AdamW(params, lr=self.config.learning_rate, weight_decay=0.01)
+        print(f"  ✓ AdamW (lr={self.config.learning_rate})")
+    
+    def encode_image(self, image_path: str) -> torch.Tensor:
+        """Encode image to vision tokens."""
+        image = Image.open(image_path).convert('RGB')
+        
+        with torch.no_grad():
+            vision_tokens = self.model.encode_image(image)
+        
+        return vision_tokens
+    
+    def compute_detection_loss(
+        self,
+        vision_tokens: torch.Tensor,
+        target_boxes: List[List[float]],
+        target_labels: List[int]
+    ) -> Tuple[torch.Tensor, Dict]:
+        """Compute detection loss."""
+        
+        # Get predictions
+        cls_logits, box_preds = self.model.detection_head(vision_tokens)
+        
+        batch_size = vision_tokens.shape[0]
+        num_tokens = vision_tokens.shape[1]
+        
+        # For each ground truth box, assign it to the nearest predicted "slot"
+        total_cls_loss = 0
+        total_box_loss = 0
+        num_matches = 0
+        
+        target_boxes_t = torch.tensor(target_boxes, dtype=torch.float32)
+        target_labels_t = torch.tensor(target_labels, dtype=torch.long)
+        
+        for i in range(batch_size):
+            if len(target_boxes) == 0:
+                continue
+            
+            # Get predictions for this sample
+            pred_boxes = box_preds[i]  # [num_tokens, 4]
+            pred_cls = cls_logits[i]    # [num_tokens, num_classes]
+            
+            # For each GT box, find best matching prediction
+            for gt_idx, (gt_box, gt_label) in enumerate(zip(target_boxes, target_labels)):
+                gt_box_t = torch.tensor(gt_box, dtype=torch.float32)
+                
+                # Compute IoU with all predictions
+                ious = self._compute_iou(pred_boxes, gt_box_t.unsqueeze(0).expand(num_tokens, -1))
+                
+                # Find best match
+                best_idx = ious.argmax()
+                
+                # Classification loss for best match
+                cls_loss = F.cross_entropy(
+                    pred_cls[best_idx:best_idx+1],
+                    torch.tensor([gt_label], dtype=torch.long)
+                )
+                
+                # Box regression loss (L1)
+                box_loss = F.l1_loss(pred_boxes[best_idx], gt_box_t)
+                
+                total_cls_loss += cls_loss
+                total_box_loss += box_loss
+                num_matches += 1
+        
+        if num_matches > 0:
+            total_cls_loss /= num_matches
+            total_box_loss /= num_matches
+        
+        # Combined loss
+        total_loss = total_cls_loss + 5.0 * total_box_loss  # Weight box loss higher
+        
+        return total_loss, {
+            'cls_loss': float(total_cls_loss) if num_matches > 0 else 0,
+            'box_loss': float(total_box_loss) if num_matches > 0 else 0,
+            'num_matches': num_matches
+        }
+    
+    def _compute_iou(self, boxes1: torch.Tensor, boxes2: torch.Tensor) -> torch.Tensor:
+        """Compute IoU between two sets of boxes."""
+        # boxes format: [x1, y1, x2, y2]
+        x1 = torch.max(boxes1[:, 0], boxes2[:, 0])
+        y1 = torch.max(boxes1[:, 1], boxes2[:, 1])
+        x2 = torch.min(boxes1[:, 2], boxes2[:, 2])
+        y2 = torch.min(boxes1[:, 3], boxes2[:, 3])
+        
+        inter_area = torch.clamp(x2 - x1, min=0) * torch.clamp(y2 - y1, min=0)
+        
+        area1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])
+        area2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])
+        
+        union_area = area1 + area2 - inter_area + 1e-8
+        
+        return inter_area / union_area
+    
+    def train_step(self, sample: Dict) -> Tuple[float, Dict]:
+        """Single training step."""
+        
+        self.optimizer.zero_grad()
+        
+        try:
+            # Encode image (with gradients through adapter if needed)
+            image = Image.open(sample['image_path']).convert('RGB')
+            
+            # Get vision features from frozen encoders
+            with torch.no_grad():
+                vision_features = self.model.vision_encoder(image)
+            
+            # Check for dimension mismatch and create adapter
+            actual_dim = vision_features.shape[-1]
+            expected_dim = self.model.config.fused_vision_dim
+            
+            if actual_dim != expected_dim:
+                if self.model.vision_adapter is None:
+                    print(f"  [Adapter] Creating: {actual_dim} -> {expected_dim}")
+                    self.model.vision_adapter = nn.Linear(actual_dim, expected_dim)
+                    nn.init.xavier_uniform_(self.model.vision_adapter.weight)
+                    nn.init.zeros_(self.model.vision_adapter.bias)
+                    
+                    # Add adapter params to optimizer
+                    self.optimizer.add_param_group({
+                        'params': self.model.vision_adapter.parameters()
+                    })
+                
+                vision_features = self.model.vision_adapter(vision_features)
+            
+            # Project to tokens
+            vision_tokens = self.model.projector(vision_features)
+            
+            # Compute detection loss
+            loss, metrics = self.compute_detection_loss(
+                vision_tokens,
+                sample['boxes'],
+                sample['labels']
+            )
+            
+            if loss.requires_grad:
+                loss.backward()
+                self.optimizer.step()
+            
+            return float(loss), metrics
+            
+        except Exception as e:
+            print(f"  ⚠️ Error: {e}")
+            return 0.0, {}
+    
+    def save_checkpoint(self, step: int, loss: float):
+        """Save checkpoint."""
+        checkpoint_path = self.checkpoint_dir / f"step_{step:06d}"
+        checkpoint_path.mkdir(exist_ok=True)
+        
+        # Save detection heads
+        torch.save({
+            'detection': self.model.detection_head.state_dict(),
+            'point': self.model.point_head.state_dict(),
+            'adapter': self.model.vision_adapter.state_dict() if self.model.vision_adapter else None,
+        }, checkpoint_path / "heads.pth")
+        
+        # Save state
+        state = {'step': step, 'loss': loss}
+        with open(checkpoint_path / "state.json", "w") as f:
+            json.dump(state, f, indent=2)
+        
+        print(f"  💾 Checkpoint: {checkpoint_path}")
+    
+    def train(self):
+        """Main training loop."""
+        print("\n" + "=" * 60)
+        print("🚀 STARTING DETECTION TRAINING")
+        print("=" * 60)
+        print(f"  Dataset: {len(self.dataset):,} samples")
+        print(f"  Epochs: {self.config.num_epochs}")
+        print(f"  Learning rate: {self.config.learning_rate}")
+        
+        global_step = 0
+        best_loss = float('inf')
+        start_time = time.time()
+        
+        for epoch in range(self.config.num_epochs):
+            print(f"\n📚 Epoch {epoch + 1}/{self.config.num_epochs}")
+            print("-" * 40)
+            
+            # Shuffle
+            indices = list(range(len(self.dataset)))
+            random.shuffle(indices)
+            
+            epoch_loss = 0
+            epoch_box_loss = 0
+            epoch_cls_loss = 0
+            num_batches = 0
+            
+            for i, idx in enumerate(indices):
+                sample = self.dataset[idx]
+                
+                loss, metrics = self.train_step(sample)
+                
+                if loss == 0:
+                    continue
+                
+                epoch_loss += loss
+                epoch_box_loss += metrics.get('box_loss', 0)
+                epoch_cls_loss += metrics.get('cls_loss', 0)
+                num_batches += 1
+                global_step += 1
+                
+                # Logging
+                if global_step % self.config.log_every == 0:
+                    elapsed = time.time() - start_time
+                    avg_loss = epoch_loss / num_batches
+                    print(f"  Step {global_step:5d} | Loss: {loss:.4f} | "
+                          f"Avg: {avg_loss:.4f} | Box: {metrics.get('box_loss', 0):.4f} | "
+                          f"Cls: {metrics.get('cls_loss', 0):.4f} | {elapsed:.0f}s")
+                
+                # Checkpointing
+                if global_step % self.config.save_every == 0:
+                    self.save_checkpoint(global_step, loss)
+                    if loss < best_loss:
+                        best_loss = loss
+            
+            avg_epoch_loss = epoch_loss / max(num_batches, 1)
+            print(f"\n  ✓ Epoch {epoch + 1} | Avg loss: {avg_epoch_loss:.4f} | "
+                  f"Box: {epoch_box_loss/max(num_batches,1):.4f} | "
+                  f"Cls: {epoch_cls_loss/max(num_batches,1):.4f}")
+        
+        # Final save
+        print("\n" + "=" * 60)
+        print("💾 Saving Final Model")
+        print("=" * 60)
+        
+        final_path = self.checkpoint_dir / "final"
+        final_path.mkdir(exist_ok=True)
+        
+        # Save heads
+        torch.save({
+            'detection': self.model.detection_head.state_dict(),
+            'point': self.model.point_head.state_dict(),
+            'adapter': self.model.vision_adapter.state_dict() if self.model.vision_adapter else None,
+        }, final_path / "heads.pth")
+        
+        # Also copy over the projector
+        import shutil
+        src_projector = OCULUS_ROOT / self.config.checkpoint_path / "projector.npz"
+        src_config = OCULUS_ROOT / self.config.checkpoint_path / "config.json"
+        if src_projector.exists():
+            shutil.copy(src_projector, final_path / "projector.npz")
+        if src_config.exists():
+            shutil.copy(src_config, final_path / "config.json")
+        
+        print(f"✅ Training complete! Model: {final_path}")
+        return final_path
+
+
+def main():
+    config = DetectionTrainingConfig(
+        data_dir="data/coco",
+        max_samples=2000,  # Start smaller for faster iteration
+        num_epochs=2,
+        learning_rate=5e-4,
+        save_every=200,
+        log_every=25,
+    )
+    
+    trainer = DetectionTrainer(config)
+    trainer.train()
+
+
+if __name__ == "__main__":
+    main()