Update utils.py

utils.py

@@ -1,42 +1,134 @@
 import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models import swin_t
 from torchvision import transforms
 from PIL import Image
-from timm import create_model
+import time
 
-# Function to load a Swin Transformer model with specified number of output classes
-def load_model(model_path, num_classes):
-    model = create_model("swin_tiny_patch4_window7_224", pretrained=False, num_classes=num_classes)
-    model.load_state_dict(torch.load(model_path, map_location='cpu'))
-    model.eval()
-    return model
+# Define MMIM model structure
+class MMIM(nn.Module):
+    def __init__(self, num_classes):
+        super(MMIM, self).__init__()
+        self.backbone = swin_t(weights=None)  # Load without pretrained weights
+        self.backbone.head = nn.Identity()
+        self.classifier = nn.Sequential(
+            nn.Linear(768, 512),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(512, num_classes)
+        )
+
+    def forward(self, x):
+        x = self.backbone(x)
+        return self.classifier(x)
+
+# Load all models and compute class offsets
+def load_all_models():
+    model_defs = [
+        ("MMIM_best1.pth", 9),   # Model 1
+        ("MMIM_best3.pth", 4),   # Model 2
+        ("MMIM_best2.pth", 12)   # Model 3
+    ]
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+    models = []
+    offsets = []
+    total_classes = 0
+
+    for path, num_classes in model_defs:
+        model = MMIM(num_classes)
+        state_dict = torch.load(path, map_location=device)
+        model.load_state_dict(state_dict)
+        model.to(device)
+        model.eval()
+        models.append(model)
+        offsets.append(total_classes)
+        total_classes += num_classes
+
+    idx_to_class = {i: f"class{i}" for i in range(total_classes)}
+    return models, offsets, idx_to_class
+
+# Predict image using confidence-based model router
+def predict_image(image, models, offsets, idx_to_class, threshold=0.9):
+    start = time.time()
 
-# Function to preprocess an uploaded image Get the predicted class + confidence from each.(resize, normalize, add batch dimension)
-def preprocess_image(image):
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
     transform = transforms.Compose([
-        transforms.Resize((224, 224)),
+        transforms.Resize((128, 128)),  # Faster inference
         transforms.ToTensor(),
-        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
+        transforms.Normalize([0.5]*3, [0.5]*3)
     ])
-    image = Image.open(image).convert("RGB")
-    return transform(image).unsqueeze(0)  # Shape: (1, 3, 224, 224)
+    image_tensor = transform(i
+runtime error
+Exit code: 1. Reason: .5.4.mlp.0.bias", "backbone.features.5.4.mlp.3.weight", "backbone.features.5.4.mlp.3.bias", "backbone.features.5.5.norm1.weight", "backbone.features.5.5.norm1.bias", "backbone.features.5.5.attn.relative_position_bias_table", "backbone.features.5.5.attn.relative_position_index", "backbone.features.5.5.attn.qkv.weight", "backbone.features.5.5.attn.qkv.bias", "backbone.features.5.5.attn.proj.weight", "backbone.features.5.5.attn.proj.bias", "backbone.features.5.5.norm2.weight", "backbone.features.5.5.norm2.bias", "backbone.features.5.5.mlp.0.weight", "backbone.features.5.5.mlp.0.bias", "backbone.features.5.5.mlp.3.weight", "backbone.features.5.5.mlp.3.bias", "backbone.features.6.reduction.weight", "backbone.features.6.norm.weight", "backbone.features.6.norm.bias", "backbone.features.7.0.norm1.weight", "backbone.features.7.0.norm1.bias", "backbone.features.7.0.attn.relative_position_bias_table", "backbone.features.7.0.attn.relative_position_index", "backbone.features.7.0.attn.qkv.weight", "backbone.features.7.0.attn.qkv.bias", "backbone.features.7.0.attn.proj.weight", "backbone.features.7.0.attn.proj.bias", "backbone.features.7.0.norm2.weight", "backbone.features.7.0.norm2.bias", "backbone.features.7.0.mlp.0.weight", "backbone.features.7.0.mlp.0.bias", "backbone.features.7.0.mlp.3.weight", "backbone.features.7.0.mlp.3.bias", "backbone.features.7.1.norm1.weight", "backbone.features.7.1.norm1.bias", "backbone.features.7.1.attn.relative_position_bias_table", "backbone.features.7.1.attn.relative_position_index", "backbone.features.7.1.attn.qkv.weight", "backbone.features.7.1.attn.qkv.bias", "backbone.features.7.1.attn.proj.weight", "backbone.features.7.1.attn.proj.bias", "backbone.features.7.1.norm2.weight", "backbone.features.7.1.norm2.bias", "backbone.features.7.1.mlp.0.weight", "backbone.features.7.1.mlp.0.bias", "backbone.features.7.1.mlp.3.weight", "backbone.features.7.1.mlp.3.bias", "backbone.norm.weight", "backbone.norm.bias", "classifier.0.weight", "classifier.0.bias", "classifier.3.weight", "classifier.3.bias". 
+
+Container logs:
+
+===== Application Startup at 2025-06-26 06:03:02 =====
+
+2025-06-26 08:04:04.986 WARNING streamlit.runtime.scriptrunner_utils.script_run_context: Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
+2025-06-26 08:04:04.987 WARNING streamlit.runtime.scriptrunner_utils.script_run_context: Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
+2025-06-26 08:04:04.999 
+  Warning: to view this Streamlit app on a browser, run it with the following
+  command:
+
+    streamlit run app.py [ARGUMENTS]
+2025-06-26 08:04:04.999 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
+2025-06-26 08:04:04.999 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
+2025-06-26 08:04:04.999 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
+2025-06-26 08:04:04.999 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
+2025-06-26 08:04:04.999 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
+2025-06-26 08:04:05.000 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
+2025-06-26 08:04:05.000 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
+2025-06-26 08:04:05.000 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
+2025-06-26 08:04:05.000 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
+2025-06-26 08:04:05.548 Thread 'Thread-1': missing ScriptRunContext! This warning can be ignored when running in bare mode.
+2025-06-26 08:04:05.548 Thread 'Thread-1': missing ScriptRunContext! This warning can be ignored when running in bare mode.
+2025-06-26 08:04:05.548 Thread 'Thread-1': missing ScriptRunContext! This warning can be ignored when running in bare mode.
+2025-06-26 08:04:06.139 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
+2025-06-26 08:04:06.139 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
+2025-06-26 08:04:06.139 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
+Traceback (most recent call last):
+  File "/home/user/app/app.py", line 35, in <module>
+    model_defs = load_all_models()
+  File "/usr/local/lib/python3.10/site-packages/streamlit/runtime/caching/cache_utils.py", line 219, in __call__
+    return self._get_or_create_cached_value(args, kwargs, spinner_message)
+  File "/usr/local/lib/python3.10/site-packages/streamlit/runtime/caching/cache_utils.py", line 261, in _get_or_create_cached_value
+    return self._handle_cache_miss(cache, value_key, func_args, func_kwargs)
+  File "/usr/local/lib/python3.10/site-packages/streamlit/runtime/caching/cache_utils.py", line 320, in _handle_cache_miss
+    computed_value = self._info.func(*func_args, **func_kwargs)
+  File "/home/user/app/app.py", line 32, in load_all_models
+    model_def["model"] = load_model(model_def["path"], len(model_def["class_names"]))
+  File "/home/user/app/utils.py", line 9, in load_model
+    model.load_state_dict(torch.load(model_path, map_location='cpu'))
+  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2593, in load_state_dict
+    raise RuntimeError(
+RuntimeError: Error(s) in loading state_dict for SwinTransformer:
+	Missing key(s) in state_dict: "patch_embed.proj.weight", "patch_embed.proj.bias", "patch_embed.norm.weight", "patch_embed.norm.bias", "layers.0.blocks.0.norm1.weight", "layers.0.blocks.0.norm1.bias", "layers.0.blocks.0.attn.relative_position_bias_table", "layers.0.blocks.0.attn.qkv.weight", "layers.0.blocks.0.attn.qkv.bias", "layers.0.blocks.0.attn.proj.weight", "layers.0.blocks.0.attn.proj.bias", "layers.0.blocks.0.norm2.weight", "layers.0.blocks.0.norm2.bias", "layers.0.blocks.0.mlp.fc1.weight", "layers.0.blocks.0.mlp.fc1.bias", "layers.0.blocks.0.mlp.fc2.weight", "layers.0.blocks.0.mlp.fc2.bias", "layers.0.blocks.1.norm1.weight", "layers.0.blocks.1.norm1.bias", "layers.0.blocks.1.attn.relative_position_bias_table", "layers.0.blocks.1.attn.qkv.weight", "layers.0.blocks.1.attn.qkv.bias", "layers.0.blocks.1.attn.proj.weight", "layers.0.blocks.1.attn.proj.bias", "layers.0.blocks.1.norm2.weight", "layers.0.blocks.1.norm2.bias", "layers.0.blocks.1.mlp.fc1.weight", "layers.0.blocks.1.mlp.fc1.bias", "layers.0.blocks.1.mlp.fc2.weight", "layers.0.blocks.1.mlp.fc2.bias", "layers.1.downsample.norm.weight", "layers.1.downsample.norm.bias", "layers.1.downsample.reduction.weight", "layers.1.blocks.0.norm1.weight", "layers.1.blocks.0.norm1.bias", "layers.1.blocks.0.attn.relative_position_bias_table", "layers.1.blocks.0.attn.qkv.weight", "layers.1.blocks.0.attn.qkv.bias", "layers.1.blocks.0.attn.proj.weight", "layers.1.blocks.0.attn.proj.bias", "layers.1.blocks.0.norm2.weight", "layers.1.blocks.0.norm2.bias", "layers.1.blocks.0.mlp.fc1.weight", "layers.1.blocks.0.mlp.fc1.bias", "layers.1.blocks.0.mlp.fc2.weight", "layers.1.blocks.0.mlp.fc2.bias", "layers.1.blocks.1.norm1.weight", "layers.1.blocks.1.norm1.bias", "layers.1.blocks.1.attn.relative_position_bias_table", "layers.1.blocks.1.attn.qkv.weight", "layers.1.blocks.1.attn.qkv.bias", "layers.1.blocks.1.attn.proj.weight", "layers.1.blocks.1.attn.proj.bias", "layers.1.blocks.1.norm2.weight", "layers.1.blocks.1.norm2.bias", "layers.1.blocks.1.mlp.fc1.weight", "layers.1.blocks.1.mlp.fc1.bias", "layers.1.blocks.1.mlp.fc2.weight", "layers.1.blocks.1.mlp.fc2.bias", "layers.2.downsample.norm.weight", "layers.2.downsample.norm.bias", "layers.2.downsample.reduction.weight", "layers.2.blocks.0.norm1.weight", "layers.2.blocks.0.norm1.bias", "layers.2.blocks.0.attn.relative_position_bias_table", "layers.2.blocks.0.attn.qkv.weight", "layers.2.blocks.0.attn.qkv.bias", "layers.2.blocks.0.attn.proj.weight", "layers.2.blocks.0.attn.proj.bias", "layers.2.blocks.0.norm2.weight", "layers.2.blocks.0.norm2.bias", "layers.2.blocks.0.mlp.fc1.weight", "layers.2.blocks.0.mlp.fc1.bias", "layers.2.blocks.0.mlp.fc2.weight", "layers.2.blocks.0.mlp.fc2.bias", "layers.2.blocks.1.norm1.weight", "layers.2.blocks.1.norm1.bias", "layers.2.blocks.1.attn.relative_position_bias_table", "layers.2.blocks.1.attn.qkv.weight", "layers.2.blocks.1.attn.qkv.bias", "layers.2.blocks.1.attn.proj.weight", "layers.2.blocks.1.attn.proj.bias", "layers.2.blocks.1.norm2.weight", "layers.2.blocks.1.norm2.bias", "layers.2.blocks.1.mlp.fc1.weight", "layers.2.blocks.1.mlp.fc1.bias", "layers.2.blocks.1.mlp.fc2.weight", "layers.2.blocks.1.mlp.fc2.bias", "layers.2.blocks.2.norm1.weight", "layers.2.blocks.2.norm1.bias", "layers.2.blocks.2.attn.relative_position_bias_table", "layers.2.blocks.2.attn.qkv.weight", "layers.2.blocks.2.attn.qkv.bias", "layers.2.blocks.2.attn.proj.weight", "layers.2.blocks.2.attn.proj.bias", "layers.2.blocks.2.norm2.weight", "layers.2.blocks.2.norm2.bias", "layers.2.blocks.2.mlp.fc1.weight", "layers.2.blocks.2.mlp.fc1.bias", "layers.2.blocks.2.mlp.fc2.weight", "layers.2.blocks.2.mlp.fc2.bias", "layers.2.blocks.3.norm1.weight", "layers.2.blocks.3.norm1.bias", "layers.2.blocks.3.attn.relative_position_bias_table", "layers.2.blocks.3.attn.qkv.weight", "layers.2.blocks.3.attn.qkv.bias", "layers.2.blocks.3.attn.proj.weight", "layers.2.blocks.3.attn.proj.bias", "layers.2.blocks.3.norm2.weight", "layers.2.blocks.3.norm2.bias", "layers.2.blocks.3.mlp.fc1.weight", "layers.2.blocks.3.mlp.fc1.bias", "layers.2.blocks.3.mlp.fc2.weight", "layers.2.blocks.3.mlp.fc2.bias", "layers.2.blocks.4.norm1.weight", "layers.2.blocks.4.norm1.bias", "layers.2.blocks.4.attn.relative_position_bias_table", "layers.2.blocks.4.attn.qkv.weight", "layers.2.blocks.4.attn.qkv.bias", "layers.2.blocks.4.attn.proj.weight", "layers.2.blocks.4.attn.proj.bias", "layers.2.blocks.4.norm2.weight", "layers.2.blocks.4.norm2.bias", "layers.2.blocks.4.mlp.fc1.weight", "layers.2.blocks.4.mlp.fc1.bias", "layers.2.blocks.4.mlp.fc2.weight", "layers.2.blocks.4.mlp.fc2.bias", "layers.2.blocks.5.norm1.weight", "layers.2.blocks.5.norm1.bias", "layers.2.blocks.5.attn.relative_position_bias_table", "layers.2.blocks.5.attn.qkv.weight", "layers.2.blocks.5.attn.qkv.bias", "layers.2.blocks.5.attn.proj.weight", "layers.2.blocks.5.attn.proj.bias", "layers.2.blocks.5.norm2.weight", "layers.2.blocks.5.norm2.bias", "layers.2.blocks.5.mlp.fc1.weight", "layers.2.blocks.5.mlp.fc1.bias", "layers.2.blocks.5.mlp.fc2.weight", "layers.2.blocks.5.mlp.fc2.bias", "layers.3.downsample.norm.weight", "layers.3.downsample.norm.bias", "layers.3.downsample.reduction.weight", "layers.3.blocks.0.norm1.weight", "layers.3.blocks.0.norm1.bias", "layers.3.blocks.0.attn.relative_position_bias_table", "layers.3.blocks.0.attn.qkv.weight", "layers.3.blocks.0.attn.qkv.bias", "layers.3.blocks.0.attn.proj.weight", "layers.3.blocks.0.attn.proj.bias", "layers.3.blocks.0.norm2.weight", "layers.3.blocks.0.norm2.bias", "layers.3.blocks.0.mlp.fc1.weight", "layers.3.blocks.0.mlp.fc1.bias", "layers.3.blocks.0.mlp.fc2.weight", "layers.3.blocks.0.mlp.fc2.bias", "layers.3.blocks.1.norm1.weight", "layers.3.blocks.1.norm1.bias", "layers.3.blocks.1.attn.relative_position_bias_table", "layers.3.blocks.1.attn.qkv.weight", "layers.3.blocks.1.attn.qkv.bias", "layers.3.blocks.1.attn.proj.weight", "layers.3.blocks.1.attn.proj.bias", "layers.3.blocks.1.norm2.weight", "layers.3.blocks.1.norm2.bias", "layers.3.blocks.1.mlp.fc1.weight", "layers.3.blocks.1.mlp.fc1.bias", "layers.3.blocks.1.mlp.fc2.weight", "layers.3.blocks.1.mlp.fc2.bias", "norm.weight", "norm.bias", "head.fc.weight", "head.fc.bias". 
+	Unexpected key(s) in state_dict: "backbone.features.0.0.weight", "backbone.features.0.0.bias", "backbone.features.0.2.weight", "backbone.features.0.2.bias", "backbone.features.1.0.norm1.weight", "backbone.features.1.0.norm1.bias", "backbone.features.1.0.attn.relative_position_bias_table", "backbone.features.1.0.attn.relative_position_index", "backbone.features.1.0.attn.qkv.weight", "backbone.features.1.0.attn.qkv.bias", "backbone.features.1.0.attn.proj.weight", "backbone.features.1.0.attn.proj.bias", "backbone.features.1.0.norm2.weight", "backbone.features.1.0.norm2.bias", "backbone.features.1.0.mlp.0.weight", "backbone.features.1.0.mlp.0.bias", "backbone.features.1.0.mlp.3.weight", "backbone.features.1.0.mlp.3.bias", "backbone.features.1.1.norm1.weight", "backbone.features.1.1.norm1.bias", "backbone.features.1.1.attn.relative_position_bias_table", "backbone.features.1.1.attn.relative_position_index", "backbone.features.1.1.attn.qkv.weight", "backbone.features.1.1.attn.qkv.bias", "backbone.features.1.1.attn.proj.weight", "backbone.features.1.1.attn.proj.bias", "backbone.features.1.1.norm2.weight", "backbone.features.1.1.norm2.bias", "backbone.features.1.1.mlp.0.weight", "backbone.features.1.1.mlp.0.bias", "backbone.features.1.1.mlp.3.weight", "backbone.features.1.1.mlp.3.bias", "backbone.features.2.reduction.weight", "backbone.features.2.norm.weight", "backbone.features.2.norm.bias", "backbone.features.3.0.norm1.weight", "backbone.features.3.0.norm1.bias", "backbone.features.3.0.attn.relative_position_bias_table", "backbone.features.3.0.attn.relative_position_index", "backbone.features.3.0.attn.qkv.weight", "backbone.features.3.0.attn.qkv.bias", "backbone.features.3.0.attn.proj.weight", "backbone.features.3.0.attn.proj.bias", "backbone.features.3.0.norm2.weight", "backbone.features.3.0.norm2.bias", "backbone.features.3.0.mlp.0.weight", "backbone.features.3.0.mlp.0.bias", "backbone.features.3.0.mlp.3.weight", "backbone.features.3.0.mlp.3.bias", "backbone.features.3.1.norm1.weight", "backbone.features.3.1.norm1.bias", "backbone.features.3.1.attn.relative_position_bias_table", "backbone.features.3.1.attn.relative_position_index", "backbone.features.3.1.attn.qkv.weight", "backbone.features.3.1.attn.qkv.bias", "backbone.features.3.1.attn.proj.weight", "backbone.features.3.1.attn.proj.bias", "backbone.features.3.1.norm2.weight", "backbone.features.3.1.norm2.bias", "backbone.features.3.1.mlp.0.weight", "backbone.features.3.1.mlp.0.bias", "backbone.features.3.1.mlp.3.weight", "backbone.features.3.1.mlp.3.bias", "backbone.features.4.reduction.weight", "backbone.features.4.norm.weight", "backbone.features.4.norm.bias", "backbone.features.5.0.norm1.weight", "backbone.features.5.0.norm1.bias", "backbone.features.5.0.attn.relative_position_bias_table", "backbone.features.5.0.attn.relative_position_index", "backbone.features.5.0.attn.qkv.weight", "backbone.features.5.0.attn.qkv.bias", "backbone.features.5.0.attn.proj.weight", "backbone.features.5.0.attn.proj.bias", "backbone.features.5.0.norm2.weight", "backbone.features.5.0.norm2.bias", "backbone.features.5.0.mlp.0.weight", "backbone.features.5.0.mlp.0.bias", "backbone.features.5.0.mlp.3.weight", "backbone.features.5.0.mlp.3.bias", "backbone.features.5.1.norm1.weight", "backbone.features.5.1.norm1.bias", "backbone.features.5.1.attn.relative_position_bias_table", "backbone.features.5.1.attn.relative_position_index", "backbone.features.5.1.attn.qkv.weight", "backbone.features.5.1.attn.qkv.bias", "backbone.features.5.1.attn.proj.weight", "backbone.features.5.1.attn.proj.bias", "backbone.features.5.1.norm2.weight", "backbone.features.5.1.norm2.bias", "backbone.features.5.1.mlp.0.weight", "backbone.features.5.1.mlp.0.bias", "backbone.features.5.1.mlp.3.weight", "backbone.features.5.1.mlp.3.bias", "backbone.features.5.2.norm1.weight", "backbone.features.5.2.norm1.bias", "backbone.features.5.2.attn.relative_position_bias_table", "backbone.features.5.2.attn.relative_position_index", "backbone.features.5.2.attn.qkv.weight", "backbone.features.5.2.attn.qkv.bias", "backbone.features.5.2.attn.proj.weight", "backbone.features.5.2.attn.proj.bias", "backbone.features.5.2.norm2.weight", "backbone.features.5.2.norm2.bias", "backbone.features.5.2.mlp.0.weight", "backbone.features.5.2.mlp.0.bias", "backbone.features.5.2.mlp.3.weight", "backbone.features.5.2.mlp.3.bias", "backbone.features.5.3.norm1.weight", "backbone.features.5.3.norm1.bias", "backbone.features.5.3.attn.relative_position_bias_table", "backbone.features.5.3.attn.relative_position_index", "backbone.features.5.3.attn.qkv.weight", "backbone.features.5.3.attn.qkv.bias", "backbone.features.5.3.attn.proj.weight", "backbone.features.5.3.attn.proj.bias", "backbone.features.5.3.norm2.weight", "backbone.features.5.3.norm2.bias", "backbone.features.5.3.mlp.0.weight", "backbone.features.5.3.mlp.0.bias", "backbone.features.5.3.mlp.3.weight", "backbone.features.5.3.mlp.3.bias", "backbone.features.5.4.norm1.weight", "backbone.features.5.4.norm1.bias", "backbone.features.5.4.attn.relative_position_bias_table", "backbone.features.5.4.attn.relative_position_index", "backbone.features.5.4.attn.qkv.weight", "backbone.features.5.4.attn.qkv.bias", "backbone.features.5.4.attn.proj.weight", "backbone.features.5.4.attn.proj.bias", "backbone.features.5.4.norm2.weight", "backbone.features.5.4.norm2.bias", "backbone.features.5.4.mlp.0.weight", "backbone.features.5.4.mlp.0.bias", "backbone.features.5.4.mlp.3.weight", "backbone.features.5.4.mlp.3.bias", "backbone.features.5.5.norm1.weight", "backbone.features.5.5.norm1.bias", "backbone.features.5.5.attn.relative_position_bias_table", "backbone.features.5.5.attn.relative_position_index", "backbone.features.5.5.attn.qkv.weight", "backbone.features.5.5.attn.qkv.bias", "backbone.features.5.5.attn.proj.weight", "backbone.features.5.5.attn.proj.bias", "backbone.features.5.5.norm2.weight", "backbone.features.5.5.norm2.bias", "backbone.features.5.5.mlp.0.weight", "backbone.features.5.5.mlp.0.bias", "backbone.features.5.5.mlp.3.weight", "backbone.features.5.5.mlp.3.bias", "backbone.features.6.reduction.weight", "backbone.features.6.norm.weight", "backbone.features.6.norm.bias", "backbone.features.7.0.norm1.weight", "backbone.features.7.0.norm1.bias", "backbone.features.7.0.attn.relative_position_bias_table", "backbone.features.7.0.attn.relative_position_index", "backbone.features.7.0.attn.qkv.weight", "backbone.features.7.0.attn.qkv.bias", "backbone.features.7.0.attn.proj.weight", "backbone.features.7.0.attn.proj.bias", "backbone.features.7.0.norm2.weight", "backbone.features.7.0.norm2.bias", "backbone.features.7.0.mlp.0.weight", "backbone.features.7.0.mlp.0.bias", "backbone.features.7.0.mlp.3.weight", "backbone.features.7.0.mlp.3.bias", "backbone.features.7.1.norm1.weight", "backbone.features.7.1.norm1.bias", "backbone.features.7.1.attn.relative_position_bias_table", "backbone.features.7.1.attn.relative_position_index", "backbone.features.7.1.attn.qkv.weight", "backbone.features.7.1.attn.qkv.bias", "backbone.features.7.1.attn.proj.weight", "backbone.features.7.1.attn.proj.bias", "backbone.features.7.1.norm2.weight", "backbone.features.7.1.norm2.bias", "backbone.features.7.1.mlp.0.weight", "backbone.features.7.1.mlp.0.bias", "backbone.features.7.1.mlp.3.weight", "backbone.features.7.1.mlp.3.bias", "backbone.norm.weight", "backbone.norm.bias", "classifier.0.weight", "classifier.0.bias", "classifier.3.weight", "classifier.3.bias". 
+
+mage).unsqueeze(0).to(device)
 
-# Function to run image through all models and return the most confident class
-def predict_final_class(image_tensor, model_defs):
-    best_class = None
-    best_conf = -1
+    temperatures = [1.2, 0.9, 1.0]  # Adjust if needed
 
-    for model_def in model_defs:
-        model = model_def["model"]
-        class_names = model_def["class_names"]
+    max_score = float('-inf')
+    final_pred = -1
+    confidence = 0.0
 
+    for model, offset, temp in zip(models, offsets, temperatures):
         with torch.no_grad():
-            outputs = model(image_tensor)
-            probs = torch.nn.functional.softmax(outputs, dim=1)
-            top_idx = torch.argmax(probs, dim=1).item()
-            conf = probs[0][top_idx].item()
+            logits = model(image_tensor) / temp
+            probs = F.softmax(logits, dim=1).squeeze(0)
+            score, pred = torch.max(probs, dim=0)
+            if score.item() > max_score:
+                max_score = score.item()
+                final_pred = pred.item() + offset
+                confidence = score.item()
 
-        if conf > best_conf:
-            best_conf = conf
-            best_class = class_names[top_idx]
+    print(f"Inference took: {time.time() - start:.2f} seconds")  # Log time
 
-    return best_class  # ✅ Only return the class name
+    if confidence >= threshold:
+        return {idx_to_class[final_pred]: round(confidence, 4)}
+    else:
+        return {"Low confidence": round(confidence, 4)}