Update model_loader.py

model_loader.py

@@ -2,6 +2,9 @@ import torch
 import torch.nn as nn
 from torchvision import models, transforms
 from PIL import Image
+import os
+
+DEFAULT_IMG_SIZE = int(os.getenv("IMG_SIZE", "32"))
 
 def build_alexnet(num_classes=2):
     model = models.alexnet(pretrained=False)
@@ -9,20 +12,26 @@ def build_alexnet(num_classes=2):
     model.classifier[6] = nn.Linear(in_features, num_classes)
     return model
 
-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)
+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'])
     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) -> torch.Tensor:
-    transform = transforms.Compose([
-        transforms.Resize((224, 224)),
+def preprocess_image(image: Image.Image, normalize: bool = True) -> torch.Tensor:
+    transform_list = [
+        transforms.Resize((DEFAULT_IMG_SIZE, DEFAULT_IMG_SIZE), Image.Resampling.LANCZOS),
         transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-    ])
-    return transform(image).unsqueeze(0)
+    ]
+
+    # 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)