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)