Update model_loader.py

model_loader.py

@@ -2,7 +2,6 @@ import torch
 import torch.nn as nn
 from torchvision import models, transforms
 from PIL import Image
-import os
 
 def build_alexnet(num_classes=2):
     model = models.alexnet(pretrained=False)
@@ -10,26 +9,20 @@ def build_alexnet(num_classes=2):
     model.classifier[6] = nn.Linear(in_features, num_classes)
     return model
 
-def load_alexnet_model(model_path):
-    checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
-    model = build_alexnet(len(checkpoint['classes']))
-    model.load_state_dict(checkpoint['model_state'])
+def load_alexnet_model(model_path, device=None):
+    # Load weights on CPU first (safer with CUDA init)
+    checkpoint = torch.load(model_path, map_location="cpu")
+    model = build_alexnet(len(checkpoint["classes"]))
+    model.load_state_dict(checkpoint["model_state"])
+    if device is not None:
+        model.to(device)
     model.eval()
-    return model, checkpoint['classes']
+    return model, checkpoint["classes"]
 
-DEFAULT_IMG_SIZE = int(os.getenv("IMG_SIZE", "224"))  # changed default to match training (224x224)
-
-def preprocess_image(image: Image.Image, normalize: bool = True) -> torch.Tensor:
-    transform_list = [
-        transforms.Resize((DEFAULT_IMG_SIZE, DEFAULT_IMG_SIZE), Image.Resampling.LANCZOS),
+def preprocess_image(image: Image.Image) -> torch.Tensor:
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
         transforms.ToTensor(),
-    ]
-
-    # default to ImageNet stats used in training; can override with IMG_MEAN / IMG_STD env vars
-    IMAGENET_MEAN = list(map(float, os.getenv("IMG_MEAN", "0.485,0.456,0.406").split(",")))
-    IMAGENET_STD  = list(map(float, os.getenv("IMG_STD",  "0.229,0.224,0.225").split(",")))
-
-    if normalize:
-        transform_list.append(transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD))
-    transform = transforms.Compose(transform_list)
-    return transform(image).unsqueeze(0)
+        transforms.Normalize([0.4914,0.4822,0.4465], [0.2470,0.2435,0.2616]), # CIFAR MEAN and STD
+    ])
+    return transform(image).unsqueeze(0)