oneformer_coco_swin_large

Performs semantic, instance, and panoptic segmentation in a single unified model architecture using task-conditioned prompting. The model uses a Swin Transformer backbone with a unified segmentation head that accepts a task token (semantic/instance/panoptic) as input conditioning, enabling dynamic task selection at inference time without model switching. This eliminates the need for separate task-specific models while maintaining competitive performance across all three segmentation paradigms through a shared feature extraction and decoding pathway.

Extracts multi-scale hierarchical image features using a Swin Transformer backbone with shifted window attention mechanisms. The backbone operates in 4 stages (C1-C4) producing feature maps at 4×, 8×, 16×, and 32× downsampling ratios. Shifted window attention reduces computational complexity from O(n²) to O(n log n) by partitioning feature maps into local windows and shifting window positions between layers, enabling efficient processing of high-resolution images while maintaining global receptive fields through cross-window connections.

Decodes multi-scale backbone features into segmentation predictions using a cross-attention based decoder that progressively fuses features from all 4 backbone stages. The decoder uses learnable query embeddings that attend to backbone features at each scale through cross-attention mechanisms, enabling selective feature aggregation and adaptive weighting of information from different scales. This approach avoids simple concatenation by learning task-aware feature combinations that emphasize relevant scales for each prediction location.

Generates task-specific segmentation predictions (semantic/instance/panoptic) from decoded features using a task-conditioned prediction head that dynamically routes computation based on the input task token. The head uses separate prediction branches for semantic segmentation (per-pixel class logits) and instance segmentation (mask logits + class predictions), with task conditioning controlling which branches are active and how features are processed. For panoptic segmentation, both branches execute and their outputs are combined through learned fusion weights that depend on the task token.

Provides pre-trained weights optimized for COCO dataset segmentation with a 133-class vocabulary covering 80 thing classes (objects) and 53 stuff classes (background regions). The model was trained on COCO 2017 train split (118K images) using multi-task learning across semantic, instance, and panoptic segmentation objectives. Pre-training uses a combination of cross-entropy loss for semantic predictions and dice loss for instance masks, with class-balanced sampling to handle long-tail class distributions in COCO.

Supports mixed-precision inference (FP16/BF16) to reduce memory consumption and latency while maintaining accuracy. The model can run in FP32 (full precision) for maximum accuracy or FP16 (half precision) for 2× memory reduction and 1.5-2× speedup on NVIDIA GPUs with Tensor Cores. BF16 precision is supported on newer hardware (A100, H100) for better numerical stability than FP16. Automatic mixed precision (AMP) can be enabled to selectively cast operations to lower precision while keeping numerically sensitive operations in FP32.

Processes multiple images in a single batch with support for variable input resolutions through dynamic padding and batching strategies. Images are padded to a common size within each batch (typically the maximum resolution in the batch) to enable efficient GPU computation. The model supports arbitrary input resolutions from 256×256 to 2048×2048, automatically adjusting internal computation to handle different aspect ratios and sizes. Post-processing includes resolution-aware upsampling to restore predictions to original image dimensions.

Refines instance segmentation predictions through post-processing that includes non-maximum suppression (NMS), mask refinement, and boundary smoothing. The post-processor takes raw mask logits and class predictions from the model and applies learned refinement operations including morphological operations (dilation/erosion) to clean up small artifacts, boundary smoothing using Gaussian filtering, and instance-level filtering to remove low-confidence predictions. NMS is applied in mask space rather than box space, enabling more accurate instance separation for overlapping objects.

Integrates with HuggingFace Model Hub for one-line model loading and inference through the transformers library. The model is registered with model ID 'shi-labs/oneformer_coco_swin_large' and can be loaded using AutoModel.from_pretrained() with automatic weight downloading and caching. The integration includes model card documentation, inference examples, and compatibility with HuggingFace's inference API for serverless deployment. Model weights are versioned and cached locally to avoid repeated downloads.

Provides pre-computed benchmark results on COCO 2017 validation set using standard evaluation metrics including mIoU (mean Intersection-over-Union) for semantic segmentation, AP (Average Precision) for instance segmentation, and PQ (Panoptic Quality) for panoptic segmentation. Results are computed using official COCO evaluation scripts with IoU thresholds at 0.5:0.95 (standard COCO metric). The model achieves 45.1 PQ on COCO panoptic segmentation, competitive with state-of-the-art methods while maintaining unified architecture.