Update model_loader.py

model_loader.py

@@ -1,16 +1,104 @@
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
+import timm
 
-class InceptionViT_Model(nn.Module):
-    def __init__(self, model_path):
-        super(InceptionViT_Model, self).__init__()
-        # Load the pretrained model and ensure it loads on CPU
-        self.model = torch.load(model_path, map_location=torch.device('cpu'))
-        self.model.eval()
+# --- PART 1: The Stem Class ---
+class InceptionStem(nn.Module):
+    def __init__(self, in_channels=3, out_channels=64):
+        super().__init__()
+        self.branch1 = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels//2, kernel_size=1, stride=1, padding=0),
+            nn.ReLU(inplace=True)
+        )
+        self.branch3 = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels//2, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(inplace=True)
+        )
+        self.pool_branch = nn.Sequential(
+            nn.MaxPool2d(3, stride=1, padding=1),
+            nn.Conv2d(in_channels, out_channels//2, kernel_size=1),
+            nn.ReLU(inplace=True)
+        )
+        self.project = nn.Conv2d(out_channels + out_channels//2, out_channels, kernel_size=1)
+        self.bn = nn.BatchNorm2d(out_channels)
 
     def forward(self, x):
-        return self.model(x)
+        b1 = self.branch1(x)
+        b3 = self.branch3(x)
+        p = self.pool_branch(x)
+        cat = torch.cat([b1, b3, p], dim=1)
+        out = F.relu(self.bn(self.project(cat)))
+        return out
 
+# --- PART 2: The Main Model Class ---
+class InceptionViT(nn.Module):
+    def __init__(self, vit_model_name='vit_base_patch16_224', pretrained=False, stem_out_channels=64, num_classes=2, dropout=0.3):
+        # NOTE: We set pretrained=False here because we are loading YOUR weights,
+        # so we don't need to download the ImageNet weights again.
+        super().__init__()
+        self.stem = InceptionStem(in_channels=3, out_channels=stem_out_channels)
+
+        self.vit = timm.create_model(vit_model_name, pretrained=False, num_classes=0, global_pool='')
+        self.embed_dim = self.vit.num_features
+        self.stem_pool = nn.AdaptiveAvgPool2d(1)
+        self.stem_proj = nn.Linear(stem_out_channels, self.embed_dim)
+
+        self.classifier = nn.Sequential(
+            nn.LayerNorm(self.embed_dim),
+            nn.Dropout(dropout),
+            nn.Linear(self.embed_dim, num_classes)
+        )
+
+    def forward(self, x):
+        stem_feat_map = self.stem(x)
+        
+        stem_vec = self.stem_pool(stem_feat_map).view(x.size(0), -1)
+        stem_emb = self.stem_proj(stem_vec)
+
+        B = x.size(0)
+        x_vit = self.vit.patch_embed(x)
+        cls_tokens = self.vit.cls_token.expand(B, -1, -1)
+        x_vit = torch.cat((cls_tokens, x_vit), dim=1)
+        x_vit = x_vit + self.vit.pos_embed[:, : x_vit.size(1), :].to(x.device)
+        x_vit = self.vit.pos_drop(x_vit)
+
+        for blk in self.vit.blocks:
+            x_vit = blk(x_vit)
+        x_vit = self.vit.norm(x_vit)
+
+        x_vit[:, 0, :] = x_vit[:, 0, :] + stem_emb
+
+        cls_emb = x_vit[:, 0, :]
+        out = self.classifier(cls_emb)
+        return out
+
+# --- PART 3: The Loader Function ---
 def load_model(model_path='InceptionViT_best_model.pth'):
-    model = InceptionViT_Model(model_path)
+    print(f"Loading InceptionViT from {model_path}...")
+    
+    # 1. Initialize the empty model structure
+    # CRITICAL: We must match the arguments you used in training!
+    # Looking at your code: num_classes=?, dropout=0.4
+    # Since I don't know your exact num_classes, I will inspect the weights to find it dynamically.
+    
+    # Load weights first to check dimensions
+    state_dict = torch.load(model_path, map_location='cpu')
+    
+    # Auto-detect number of classes from the last layer weight
+    # The key is usually "classifier.2.weight" based on your Sequential block
+    if "classifier.2.weight" in state_dict:
+        num_classes = state_dict["classifier.2.weight"].shape[0]
+        print(f"Auto-detected {num_classes} classes.")
+    else:
+        num_classes = 2 # Default fallback
+        print("Could not auto-detect classes, defaulting to 2.")
+
+    # Create the model
+    model = InceptionViT(num_classes=num_classes, dropout=0.4)
+    
+    # Load the weights
+    model.load_state_dict(state_dict)
+    model.eval()
+    
     return model