Upload training.py with huggingface_hub

training.py

@@ -0,0 +1,164 @@
+import io
+import os
+import zipfile
+import tempfile
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+from torch.utils.data import Dataset
+from PIL import Image
+
+
+class SAM3FineTuneDataset(Dataset):
+    """Dataset for fine-tuning SAM3 on accepted detections.
+
+    Each sample corresponds to one accepted detection. The processor is called
+    with the detection's bounding box as a box prompt, and the ground-truth
+    mask is returned as a binary tensor.
+    """
+
+    def __init__(self, images_dict, detections, processor):
+        """
+        Args:
+            images_dict: dict mapping filename -> PIL.Image
+            detections: list of detection dicts with keys:
+                image_path, box, mask (numpy H×W bool/uint8)
+            processor: Sam3Processor instance
+        """
+        self.images_dict = images_dict
+        self.detections = detections
+        self.processor = processor
+
+    def __len__(self):
+        return len(self.detections)
+
+    def __getitem__(self, idx):
+        det = self.detections[idx]
+        image = self.images_dict[det["image_path"]]
+        box = det["box"]  # [x1, y1, x2, y2]
+
+        inputs = self.processor(
+            images=image,
+            input_boxes=[[box]],
+            input_boxes_labels=[[1]],
+            return_tensors="pt",
+        )
+        # Keep batch dim from processor (batch-size-1 training)
+
+        mask_np = det["mask"].astype(np.float32)
+        mask_gt = torch.from_numpy(mask_np).unsqueeze(0)  # (1, H, W)
+
+        return inputs, mask_gt
+
+
+def freeze_encoder(model):
+    """Freeze vision encoder, keep mask decoder and prompt encoder trainable.
+
+    Returns (trainable_count, total_count).
+    """
+    for param in model.vision_encoder.parameters():
+        param.requires_grad = False
+
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total = sum(p.numel() for p in model.parameters())
+    return trainable, total
+
+
+def _dice_loss(pred, target):
+    """Compute Dice loss between predicted and target masks."""
+    pred_flat = pred.flatten(1)
+    target_flat = target.flatten(1)
+    intersection = (pred_flat * target_flat).sum(1)
+    return 1 - (2.0 * intersection + 1) / (pred_flat.sum(1) + target_flat.sum(1) + 1)
+
+
+def run_training(model, processor, dataset, epochs, learning_rate, progress_callback=None):
+    """Fine-tune SAM3 mask decoder + prompt encoder.
+
+    Args:
+        model: Sam3Model with encoder frozen
+        processor: Sam3Processor
+        dataset: SAM3FineTuneDataset
+        epochs: number of training epochs
+        learning_rate: AdamW learning rate
+        progress_callback: callable(epoch, step, total_steps, loss_val)
+
+    Returns dict with keys: model, loss_history (list of avg loss per epoch).
+    """
+    device = next(model.parameters()).device
+    optimizer = torch.optim.AdamW(
+        filter(lambda p: p.requires_grad, model.parameters()),
+        lr=learning_rate,
+    )
+
+    total_steps = len(dataset) * epochs
+    loss_history = []
+
+    model.train()
+    for epoch in range(epochs):
+        epoch_losses = []
+        for step_in_epoch in range(len(dataset)):
+            global_step = epoch * len(dataset) + step_in_epoch
+
+            inputs, mask_gt = dataset[step_in_epoch]
+            # Move inputs to device
+            inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
+            mask_gt = mask_gt.to(device)
+
+            outputs = model(**inputs)
+
+            # Get predicted masks — shape (batch, num_masks, H, W)
+            pred_masks = outputs.pred_masks
+            if pred_masks.dim() == 4:
+                pred_masks = pred_masks.squeeze(0)  # (num_masks, H, W)
+
+            # Resize prediction to match ground truth
+            if pred_masks.shape[-2:] != mask_gt.shape[-2:]:
+                pred_masks = F.interpolate(
+                    pred_masks.unsqueeze(0),
+                    size=mask_gt.shape[-2:],
+                    mode="bilinear",
+                    align_corners=False,
+                ).squeeze(0)
+
+            # Use first predicted mask
+            pred = pred_masks[0:1]  # (1, H, W)
+            pred_sigmoid = torch.sigmoid(pred)
+
+            bce = F.binary_cross_entropy_with_logits(pred, mask_gt)
+            dice = _dice_loss(pred_sigmoid, mask_gt).mean()
+            loss = bce + dice
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            loss_val = loss.item()
+            epoch_losses.append(loss_val)
+
+            if progress_callback:
+                progress_callback(epoch, global_step, total_steps, loss_val)
+
+        loss_history.append(sum(epoch_losses) / len(epoch_losses))
+
+    model.eval()
+    return {"model": model, "loss_history": loss_history}
+
+
+def get_model_zip_bytes(model, processor):
+    """Save fine-tuned model and processor to a zip file, return bytes."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        model_dir = os.path.join(tmpdir, "sam3_finetuned")
+        model.save_pretrained(model_dir)
+        processor.save_pretrained(model_dir)
+
+        buf = io.BytesIO()
+        with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
+            for root, _dirs, files in os.walk(model_dir):
+                for fname in files:
+                    filepath = os.path.join(root, fname)
+                    arcname = os.path.relpath(filepath, tmpdir)
+                    zf.write(filepath, arcname)
+        buf.seek(0)
+        return buf.getvalue()