tribunal-kit 9.2.3 → 9.2.4

package/skills/huggingface-vision-trainer/references/finetune_sam2_trainer.md

@@ -0,0 +1,254 @@
+# Fine-tuning SAM2 with HF Trainer
+
+Fine-tune SAM2.1 on a small part of the MicroMat dataset for image matting,
+using the Hugging Face Trainer with a custom loss function.
+
+```python
+!pip install -q transformers datasets monai trackio
+```
+
+## Load and explore the dataset
+
+```python
+from datasets import load_dataset
+
+dataset = load_dataset("merve/MicroMat-mini", split="train")
+dataset
+```
+
+```python
+dataset = dataset.train_test_split(test_size=0.1)
+train_ds = dataset["train"]
+val_ds = dataset["test"]
+```
+
+```python
+import json
+
+train_ds[0]
+```
+
+```python
+json.loads(train_ds["prompt"][0])["bbox"]
+```
+
+## Visualize a sample
+
+```python
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def show_mask(mask, ax, bbox):
+    color = np.array([0.12, 0.56, 1.0, 0.6])
+    mask = np.array(mask)
+    h, w = mask.shape
+    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, 4)
+    ax.imshow(mask_image)
+    x0, y0, x1, y1 = bbox
+    ax.add_patch(
+        plt.Rectangle(
+            (x0, y0), x1 - x0, y1 - y0, fill=False, edgecolor="lime", linewidth=2
+        )
+    )
+
+
+example = train_ds[0]
+image = np.array(example["image"])
+ground_truth_mask = np.array(example["mask"])
+
+fig, ax = plt.subplots()
+ax.imshow(image)
+show_mask(ground_truth_mask, ax, json.loads(example["prompt"])["bbox"])
+ax.set_title("Ground truth mask")
+ax.set_axis_off()
+plt.show()
+```
+
+## Build the dataset and collator
+
+`SAMDataset` wraps each sample into the format expected by the SAM2 processor.
+Ground-truth masks are stored under the key `"labels"` so the Trainer
+automatically pops them before calling `model.forward()`.
+
+```python
+from torch.utils.data import Dataset
+import torch
+import torch.nn.functional as F
+
+
+class SAMDataset(Dataset):
+    def __init__(self, dataset, processor):
+        self.dataset = dataset
+        self.processor = processor
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        item = self.dataset[idx]
+        image = item["image"]
+        prompt = json.loads(item["prompt"])["bbox"]
+        inputs = self.processor(image, input_boxes=[[prompt]], return_tensors="pt")
+        inputs["labels"] = (np.array(item["mask"]) > 0).astype(np.float32)
+        inputs["original_image_size"] = torch.tensor(image.size[::-1])
+        return inputs
+
+
+def collate_fn(batch):
+    pixel_values = torch.cat([item["pixel_values"] for item in batch], dim=0)
+    original_sizes = torch.stack([item["original_sizes"] for item in batch])
+    input_boxes = torch.cat([item["input_boxes"] for item in batch], dim=0)
+    labels = torch.cat(
+        [
+            F.interpolate(
+                torch.as_tensor(x["labels"]).unsqueeze(0).unsqueeze(0).float(),
+                size=(256, 256),
+                mode="nearest",
+            )
+            for x in batch
+        ],
+        dim=0,
+    ).long()
+
+    return {
+        "pixel_values": pixel_values,
+        "original_sizes": original_sizes,
+        "input_boxes": input_boxes,
+        "labels": labels,
+        "original_image_size": torch.stack(
+            [item["original_image_size"] for item in batch]
+        ),
+        "multimask_output": False,
+    }
+```
+
+```python
+from transformers import Sam2Processor
+
+processor = Sam2Processor.from_pretrained("facebook/sam2.1-hiera-small")
+
+train_dataset = SAMDataset(dataset=train_ds, processor=processor)
+val_dataset = SAMDataset(dataset=val_ds, processor=processor)
+```
+
+## Load model and freeze encoder layers
+
+```python
+from transformers import Sam2Model
+
+model = Sam2Model.from_pretrained("facebook/sam2.1-hiera-small")
+
+for name, param in model.named_parameters():
+    if name.startswith("vision_encoder") or name.startswith("prompt_encoder"):
+        param.requires_grad_(False)
+```
+
+## Inference before training
+
+```python
+item = val_ds[1]
+img = item["image"]
+bbox = json.loads(item["prompt"])["bbox"]
+inputs = processor(images=img, input_boxes=[[bbox]], return_tensors="pt").to(
+    model.device
+)
+
+with torch.no_grad():
+    outputs = model(**inputs)
+
+masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])[0]
+preds = masks.squeeze(0)
+mask = (preds[0] > 0).cpu().numpy()
+
+overlay = np.asarray(img, dtype=np.uint8).copy()
+overlay[mask] = 0.55 * overlay[mask] + 0.45 * np.array([0, 255, 0], dtype=np.float32)
+
+plt.imshow(overlay)
+plt.title("Before training")
+plt.axis("off")
+plt.show()
+```
+
+## Define custom loss
+
+SAM2 does not compute loss in its `forward()`, so we provide a
+`compute_loss_func` to the Trainer. The Trainer pops `"labels"` from the
+batch before calling `model(**inputs)`, then passes `(outputs, labels)` to
+this function.
+
+```python
+import monai
+from transformers import Trainer, TrainingArguments
+import trackio
+
+seg_loss = monai.losses.DiceCELoss(sigmoid=True, squared_pred=True, reduction="mean")
+
+
+def compute_loss(outputs, labels, num_items_in_batch=None):
+    predicted_masks = outputs.pred_masks.squeeze(1)
+    return seg_loss(predicted_masks, labels.float())
+```
+
+## Train with Trainer
+
+Key settings:
+- `remove_unused_columns=False`: the Trainer must keep `input_boxes`,
+  `original_sizes`, etc. that are not in the model's `forward()` signature.
+- `compute_loss_func`: our custom DiceCE loss.
+- `report_to="trackio"`: logs the training loss to trackio.
+
+```python
+training_args = TrainingArguments(
+    output_dir="sam2-finetuned",
+    num_train_epochs=30,
+    per_device_train_batch_size=4,
+    learning_rate=1e-5,
+    weight_decay=0,
+    logging_steps=1,
+    save_strategy="epoch",
+    save_total_limit=2,
+    remove_unused_columns=False,
+    dataloader_pin_memory=False,
+    report_to="trackio",
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    data_collator=collate_fn,
+    compute_loss_func=compute_loss,
+)
+
+trainer.train()
+```
+
+## Inference after training
+
+```python
+item = val_ds[1]
+img = item["image"]
+bbox = json.loads(item["prompt"])["bbox"]
+
+inputs = processor(images=img, input_boxes=[[bbox]], return_tensors="pt").to(
+    model.device
+)
+
+with torch.no_grad():
+    outputs = model(**inputs)
+
+preds = processor.post_process_masks(
+    outputs.pred_masks.cpu(), inputs["original_sizes"]
+)[0]
+preds = preds.squeeze(0)
+mask = (preds[0] > 0).cpu().numpy()
+
+overlay = np.asarray(img, dtype=np.uint8).copy()
+overlay[mask] = 0.55 * overlay[mask] + 0.45 * np.array([0, 255, 0], dtype=np.float32)
+
+plt.imshow(overlay)
+plt.title("After training")
+plt.axis("off")
+plt.show()
+```