lib/firefly_rgb.py

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, PretrainedConfig
from transformers.modeling_outputs import SemanticSegmenterOutput

# ==============================================================================
# SHARED UTILITIES & LOSS FUNCTIONS
# ==============================================================================

def dice_loss(pred, target, skip_classes=None, eps=1e-6):
    N, C = pred.shape
    if N == 0:
        return pred.sum() * 0.0

    p = torch.softmax(pred, dim=1)
    t = torch.nn.functional.one_hot(target, C).float()

    if skip_classes is None:
        skip_classes = []

    present = (t.sum(dim=0) > 0)
    for sc in skip_classes:
        if 0 <= sc < C:
            present[sc] = False

    if present.sum() == 0:
        return pred.new_tensor(0.)

    p_sel = p[:, present]
    t_sel = t[:, present]

    inter = (p_sel * t_sel).sum(dim=0)
    union = p_sel.sum(dim=0) + t_sel.sum(dim=0)
    dice = (2 * inter + eps) / (union + eps)

    return 1 - dice.mean()


class DiceLoss(nn.Module):
    def __init__(self, skip_classes=None, eps=1e-6):
        super(DiceLoss, self).__init__()
        self.skip_classes = skip_classes
        self.eps = eps

    def forward(self, pred, target):
        return dice_loss(pred, target, skip_classes=self.skip_classes, eps=self.eps)


def pad_to_multiple(x, multiple):
    _, _, h, w = x.shape
    pad_h = (multiple - h % multiple) % multiple
    pad_w = (multiple - w % multiple) % multiple
    padding = (0, pad_w, 0, pad_h)
    x_padded = F.pad(x, padding, mode='reflect')
    return x_padded, padding


def crop_to_shape(x, target_h, target_w):
    return x[:, :, :target_h, :target_w]

# ==============================================================================
# BASE CONFIGURATION & MODEL
# ==============================================================================

class FireflyBaseConfig(PretrainedConfig):
    """Base Configuration for all Firefly models."""
    def __init__(
        self,
        num_labels: int = 2,
        image_size: int = 640,
        embedding_dim: int = 256,
        backbone_embed_dim: int = 768,
        patch_size: int = 16,
        repo_dir: str = None,
        model_name: str = "dinov3_vitb16",
        weights_path: str = None,
        semantic_loss_ignore_index: int = 255,
        dropout_ratio: float = 0.1,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.num_labels = num_labels
        self.image_size = image_size
        self.embedding_dim = embedding_dim
        self.backbone_embed_dim = backbone_embed_dim
        self.patch_size = patch_size
        self.repo_dir = repo_dir
        self.model_name = model_name
        self.weights_path = weights_path
        self.semantic_loss_ignore_index = semantic_loss_ignore_index
        self.dropout_ratio = dropout_ratio


class FireflyBaseModel(PreTrainedModel):
    """
    Base Model for Firefly. Handles backbone initialization,
    weight loading, and parameter freezing.
    """
    def __init__(self, config: FireflyBaseConfig):
        super().__init__(config)
        self.config = config
        self.num_labels = config.num_labels
        self.patch_size = config.patch_size
        
        if config.repo_dir and config.model_name:
            print(f"Loading backbone: {config.model_name} from {config.repo_dir}")
            self.backbone = torch.hub.load(
                config.repo_dir,
                config.model_name,
                source="local",
                weights=config.weights_path,
                pretrained=False
            )
        else:
            raise ValueError("repo_dir and model_name must be specified in the config.")
        
        if config.weights_path and os.path.exists(config.weights_path):
            print(f"Loading fine-tuned backbone weights from: {config.weights_path}")
            checkpoint = torch.load(config.weights_path, map_location='cpu')
            
            if isinstance(checkpoint, dict):
                if "student_model" in checkpoint:
                    state_dict = checkpoint["student_model"]
                elif "model" in checkpoint:
                    state_dict = checkpoint["model"]
                else:
                    state_dict = checkpoint
            else:
                state_dict = checkpoint

            new_state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
            msg = self.backbone.load_state_dict(new_state_dict, strict=False)
            print(f"Weight loading result: {msg}")
        else:
            print(f"Warning: Weights path '{config.weights_path}' not found or not provided. Backbone initialized randomly/default.")
        
        self.backbone.eval()
        self._setup_trainable_params()

    def _setup_trainable_params(self):
        """Freeze backbone parameters initially."""
        for param in self.backbone.parameters():
            param.requires_grad = False
            
    def print_trainable_params(self):
        total = sum(p.numel() for p in self.parameters())
        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
        print(f"Total Parameters: {total / 1e6:.2f}M")
        print(f"Trainable Parameters: {trainable / 1e6:.2f}M")
        print(f"Trainable Ratio: {100 * trainable / total:.2f}%")

# ==============================================================================
# SEGMENTATION HEAD & MODEL
# ==============================================================================

class MLPHead(nn.Module):
    def __init__(self, in_channels=768, hidden_dim=256, num_classes=12, dropout_prob=0.1):
        super().__init__()
        self.model = nn.Sequential(
            nn.Conv2d(in_channels, hidden_dim, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Dropout2d(p=dropout_prob),
            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Dropout2d(p=dropout_prob),
            nn.Conv2d(hidden_dim, num_classes, kernel_size=1)
        )
    def forward(self, x):
        return self.model(x)

class FireflyConfigRGB(FireflyBaseConfig):
    model_type = "firefly-rgb"
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

class FireflyForSemanticSegmentationRGB(FireflyBaseModel):
    config_class = FireflyConfigRGB

    def __init__(self, config: FireflyConfigRGB):
        super().__init__(config)
        self.head = MLPHead(
            in_channels=config.backbone_embed_dim,
            hidden_dim=config.embedding_dim,
            num_classes=config.num_labels,
            dropout_prob=config.dropout_ratio
        )
        self.loss_fn = DiceLoss()

    def _extract_features(self, pixel_values: torch.Tensor, h_pad: int, w_pad: int):
        features_raw = self.backbone.get_intermediate_layers(pixel_values, n=1)
        feat = features_raw[0]
        if isinstance(feat, tuple): feat = feat[0]

        if feat.ndim == 3:
            B, N, Dim = feat.shape
            H_grid, W_grid = h_pad // self.config.patch_size, w_pad // self.config.patch_size
            num_spatial_tokens = H_grid * W_grid
            if N > num_spatial_tokens:
                feat = feat[:, -num_spatial_tokens:, :]
            feat = feat.permute(0, 2, 1).view(B, Dim, H_grid, W_grid)
        return feat

    def forward(self, pixel_values: torch.Tensor, labels: torch.Tensor = None, **kwargs):
        original_h, original_w = pixel_values.shape[-2:]
        x_padded, padding = pad_to_multiple(pixel_values, self.config.patch_size)
        pad_h, pad_w = x_padded.shape[-2], x_padded.shape[-1]
        
        features = self._extract_features(x_padded, pad_h, pad_w)
        out_upscaled = F.interpolate(self.head(features), size=(pad_h, pad_w), mode='bilinear', align_corners=False)
        logits = crop_to_shape(out_upscaled, original_h, original_w)
        
        loss = None
        if labels is not None:
            valid_mask = labels != self.config.semantic_loss_ignore_index
            logits_masked = logits.permute(0, 2, 3, 1)[valid_mask]
            loss = self.loss_fn(logits_masked, labels[valid_mask])
        
        return SemanticSegmenterOutput(loss=loss, logits=logits)