# src/model.py

import torch
import torch.nn as nn
import torchvision.models as models

class ColonCancerModel(nn.Module):
    def __init__(self):
        super(ColonCancerModel, self).__init__()
        # Do not download ImageNet weights at runtime on the hosting server.
        self.mobilenet = models.mobilenet_v2(weights=None)
        # Freeze feature extractor (optional, can unfreeze later)
        for param in self.mobilenet.parameters():
            param.requires_grad = False
        
        # Replace classifier for binary output (1 unit; using BCEWithLogitsLoss later)
        in_features = self.mobilenet.classifier[1].in_features
        self.mobilenet.classifier = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(in_features, 1)
        )
    
    def forward(self, x):
        return self.mobilenet(x)

if __name__ == "__main__":
    model = ColonCancerModel()
    print(model)