Spaces:
Sleeping
Sleeping
File size: 6,190 Bytes
c04bc97 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | import torch
import torch.nn as nn
import torchvision.models as models
# =====================================================================
# 1. Task 1: Classification Models
# =====================================================================
class BaselineCNN(nn.Module):
"""
A lightweight, custom 3-stage CNN to establish baseline classification performance.
Architecture: 3 Conv stages -> Global Average Pooling -> Linear head
"""
def __init__(self, num_classes: int = 9):
super().__init__()
self.features = nn.Sequential(
# Stage 1: Conv -> BN -> ReLU -> Pool
nn.Conv2d(3, 16, kernel_size=3, padding=1),
nn.BatchNorm2d(16),
nn.ReLU(inplace=True),
nn.MaxPool2d(2, 2),
# Stage 2: Conv -> BN -> ReLU -> Pool
nn.Conv2d(16, 32, kernel_size=3, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(inplace=True),
nn.MaxPool2d(2, 2),
# Stage 3: Conv -> BN -> ReLU -> Pool
nn.Conv2d(32, 64, kernel_size=3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.MaxPool2d(2, 2),
)
self.gap = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(64, num_classes)
def forward(self, x):
x = self.features(x)
x = self.gap(x)
x = torch.flatten(x, 1)
x = self.fc(x)
return x
def get_resnet18_model(num_classes: int = 9, pretrained: bool = True):
"""
Loads a ResNet-18 model, replaces the classification head, and returns it.
If pretrained is True, loads weights from ImageNet.
"""
if pretrained:
# Modern torchvision API for loading pretrained weights
from torchvision.models import ResNet18_Weights
weights = ResNet18_Weights.DEFAULT
model = models.resnet18(weights=weights)
else:
model = models.resnet18(weights=None)
# Replace the classification head
in_features = model.fc.in_features
model.fc = nn.Linear(in_features, num_classes)
return model
# =====================================================================
# 2. Task 2: Segmentation Model (Lightweight U-Net)
# =====================================================================
class DoubleConv(nn.Module):
"""(conv -> BN -> ReLU) * 2"""
def __init__(self, in_channels: int, out_channels: int):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True)
)
def forward(self, x):
return self.conv(x)
class UNet(nn.Module):
"""
A lightweight, configurable U-Net implementation with skip connections.
"""
def __init__(self, in_channels: int = 3, out_channels: int = 1, init_features: int = 32):
super().__init__()
# Encoder (Downsampling path)
self.enc1 = DoubleConv(in_channels, init_features)
self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
self.enc2 = DoubleConv(init_features, init_features * 2)
self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
self.enc3 = DoubleConv(init_features * 2, init_features * 4)
self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
# Bottleneck
self.bottleneck = DoubleConv(init_features * 4, init_features * 8)
# Decoder (Upsampling path)
self.up3 = nn.ConvTranspose2d(init_features * 8, init_features * 4, kernel_size=2, stride=2)
self.dec3 = DoubleConv(init_features * 8, init_features * 4)
self.up2 = nn.ConvTranspose2d(init_features * 4, init_features * 2, kernel_size=2, stride=2)
self.dec2 = DoubleConv(init_features * 4, init_features * 2)
self.up1 = nn.ConvTranspose2d(init_features * 2, init_features, kernel_size=2, stride=2)
self.dec1 = DoubleConv(init_features * 2, init_features)
# Final Convolution
self.conv = nn.Conv2d(init_features, out_channels, kernel_size=1)
def forward(self, x):
# Encoder
enc1 = self.enc1(x)
enc2 = self.enc2(self.pool1(enc1))
enc3 = self.enc3(self.pool2(enc2))
# Bottleneck
bottleneck = self.bottleneck(self.pool3(enc3))
# Decoder with skip connections
up3 = self.up3(bottleneck)
dec3 = self.dec3(torch.cat([up3, enc3], dim=1))
up2 = self.up2(dec3)
dec2 = self.dec2(torch.cat([up2, enc2], dim=1))
up1 = self.up1(dec2)
dec1 = self.dec1(torch.cat([up1, enc1], dim=1))
# Final outputs
return self.conv(dec1)
# =====================================================================
# 3. Task 3: Detection Model (Faster R-CNN)
# =====================================================================
def get_faster_rcnn_mobilenet_v3(num_classes: int = 4, pretrained: bool = True):
"""
Loads a Faster R-CNN model with a MobileNet-V3-Large FPN backbone.
Replacing the final predictor to support the requested number of classes (including background).
"""
from torchvision.models.detection import fasterrcnn_mobilenet_v3_large_fpn, FasterRCNN_MobileNet_V3_Large_FPN_Weights
if pretrained:
weights = FasterRCNN_MobileNet_V3_Large_FPN_Weights.DEFAULT
model = fasterrcnn_mobilenet_v3_large_fpn(weights=weights)
else:
model = fasterrcnn_mobilenet_v3_large_fpn(weights=None)
# Get number of input features for the classifier head
in_features = model.roi_heads.box_predictor.cls_score.in_features
# Replace the box predictor head
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
return model
|