1pip install torch transformers pillow matplotlib numpy opencv-python albumentations scipy scikit-learn
2Basic Inference
3python
4import torch
5import numpy as np
6from PIL import Image
7import matplotlib.pyplot as plt
8
9# Import the model architecture (same as training)
10from model import DINOv3Encoder, ShallowStem, UNetDecoder, PolypSegmentationModel
11
12# Load model
13model = PolypSegmentationModel.from_pretrained(
14 "your-username/dinov3-polyp-seg",
15 device="cuda" if torch.cuda.is_available() else "cpu"
16)
17
18# Preprocess image
19def preprocess_image(image_path, target_size=(256, 256)):
20 image = Image.open(image_path).convert('RGB')
21 image = image.resize(target_size, Image.Resampling.BILINEAR)
22
23 # Convert to numpy and normalize
24 image_array = np.array(image).astype(np.float32) / 255.0
25 mean = np.array([0.485, 0.456, 0.406]).reshape(1, 1, 3)
26 std = np.array([0.229, 0.224, 0.225]).reshape(1, 1, 3)
27 image_array = (image_array - mean) / std
28
29 # Convert to tensor [B, C, H, W]
30 image_tensor = torch.from_numpy(image_array).permute(2, 0, 1).unsqueeze(0)
31 return image_tensor, image
32
33# Run inference
34image_tensor, original_image = preprocess_image("colonoscopy_image.jpg")
35
36with torch.no_grad():
37 prediction = model(image_tensor)
38 mask = torch.sigmoid(prediction)
39 binary_mask = (mask > 0.5).float()
40 mask_np = binary_mask.squeeze().cpu().numpy()
41
42# Visualize
43fig, axes = plt.subplots(1, 3, figsize=(15, 5))
44axes[0].imshow(original_image)
45axes[0].set_title("Input Image")
46axes[1].imshow(mask_np, cmap='gray')
47axes[1].set_title("Polyp Segmentation")
48axes[2].imshow(original_image)
49axes[2].imshow(mask_np, cmap='Reds', alpha=0.5)
50axes[2].set_title("Overlay")
51plt.show()
52Advanced Usage with Metrics
53python
54from scipy.ndimage import morphology
55
56def compute_hd95(pred, target):
57 """Compute Hausdorff Distance 95th percentile"""
58 if pred.sum() == 0 or target.sum() == 0:
59 return float('inf')
60
61 pred_border = pred - morphology.binary_erosion(pred)
62 target_border = target - morphology.binary_erosion(target)
63
64 pred_coords = np.argwhere(pred_border > 0)
65 target_coords = np.argwhere(target_border > 0)
66
67 distances = []
68 for p in pred_coords:
69 dist = np.min(np.sqrt(np.sum((target_coords - p) ** 2, axis=1)))
70 distances.append(dist)
71
72 return np.percentile(distances, 95)
73
74# Batch inference
75dataloader = DataLoader(dataset, batch_size=16, shuffle=False)
76
77all_metrics = {'dice': [], 'iou': [], 'hd95': []}
78for images, masks in dataloader:
79 with torch.no_grad():
80 predictions = model(images)
81
82 # Calculate metrics for each image
83 for pred, mask in zip(predictions, masks):
84 pred_binary = (torch.sigmoid(pred) > 0.5).float()
85
86 # Dice
87 intersection = (pred_binary * mask).sum()
88 dice = (2. * intersection) / (pred_binary.sum() + mask.sum() + 1e-6)
89
90 # IoU
91 union = pred_binary.sum() + mask.sum() - intersection
92 iou = intersection / (union + 1e-6)
93
94 # HD95
95 hd95 = compute_hd95(pred_binary.numpy().squeeze(), mask.numpy().squeeze())
96
97 all_metrics['dice'].append(dice.item())
98 all_metrics['iou'].append(iou.item())
99 all_metrics['hd95'].append(hd95)
100
101print(f"Average Dice: {np.mean(all_metrics['dice']):.4f} ± {np.std(all_metrics['dice']):.4f}")
102print(f"Average IoU: {np.mean(all_metrics['iou']):.4f} ± {np.std(all_metrics['iou']):.4f}")
103print(f"Average HD95: {np.mean(all_metrics['hd95']):.2f} ± {np.std(all_metrics['hd95']):.2f}")
104Model Limitations
105Input size: Fixed to 256×256 pixels (resize your images accordingly)
106
107Domain: Trained only on colonoscopy images from Kvasir-SEG
108
109Polyp types: May not generalize to all polyp morphologies
110
111Image quality: Best performance with standard white-light colonoscopy images
112
113## Dataset
114Trained on the Kvasir-SEG dataset, which contains 1000 polyp images with corresponding ground truth masks from colonoscopy procedures.
115
116## License
117This model is released under the MIT License.
118
119## Citation
120If you use this model in your research, please cite:
121
122bibtex
123@software{dinov3_polyp_seg,
124 author = {Amirreza Mehrzadian},
125 title = {DINOv3 Polyp Segmentation with U-Net Decoder},
126 year = {2024},
127 url = {https://huggingface.co/uncleMehrzad/dinov3-polyp-seg}
128}
129## Acknowledgments
130DINOv3 team for the powerful vision backbone
131
132Kvasir-SEG dataset providers for the polyp segmentation data
133
134HuggingFace for model hosting infrastructure
135
136
137
138
139
140```python
141class PolypSegmentationModel(nn.Module):
142 """Complete model wrapper matching training architecture"""
143
144 def __init__(self, encoder, stem, decoder):
145 super().__init__()
146 self.encoder = encoder
147 self.stem = stem
148 self.decoder = decoder
149
150 def forward(self, x):
151 vit_features = self.encoder(x)
152 skip_features = self.stem(x)
153 return self.decoder(vit_features, skip_features)
154
155 @classmethod
156 def from_pretrained(cls, model_path, config, device="cpu"):
157 """Load the complete model from checkpoint"""
158 checkpoint = torch.load(model_path, map_location=device)
159
160 # Initialize components
161 encoder = DINOv3Encoder(
162 model_name=config.model_name,
163 local_path=config.local_model_path,
164 freeze=True,
165 layers=config.multi_scale_layers
166 )
167
168 stem = ShallowStem(in_channels=3, base_channels=64)
169
170 decoder = UNetDecoder(
171 vit_channels=encoder.out_channels,
172 stem_channels=[512, 256, 128],
173 num_classes=1
174 )
175
176 # Load weights
177 decoder.load_state_dict(checkpoint['decoder_state_dict'])
178 stem.load_state_dict(checkpoint['stem_state_dict'])
179
180 model = cls(encoder, stem, decoder)
181 model.to(device)
182 model.eval()
183
184 return model
