"modify model.py"

model.py

@@ -5,24 +5,31 @@ from PIL import Image
 import pickle
 
 def load_model():
+    """
+    Load the trained PyTorch model and class labels.
+    """
     # Load label mapping
     with open("model/label_maps.pkl", "rb") as f:
         label_maps = pickle.load(f)
-    
-    # Extract class names cleanly
-    if isinstance(next(iter(label_maps.values())), str):  # e.g., {0: "dry"}
+
+    # Extract class names based on label_maps format
+    if isinstance(next(iter(label_maps.values())), str):
+        # Example: {0: "dry", 1: "oily"}
         class_names = list(label_maps.values())
-    elif isinstance(next(iter(label_maps.values())), dict):  # e.g., {"dry": 0, "oily": 1}
+    elif isinstance(next(iter(label_maps.values())), dict):
+        # Example: {"dry": 0, "oily": 1}
         class_names = list(next(iter(label_maps.values())).keys())
     else:
         raise ValueError("Unexpected format in label_maps.pkl")
 
-    # Load full model
+    # Load the trained PyTorch model
     model = torch.load("model/best_skin_model_entire.pth", map_location="cpu")
     model.eval()
+
+    print(f"✅ Model loaded with {len(class_names)} classes: {class_names}")
     return model, class_names
 
-# Preprocessing pipeline
+# Image preprocessing pipeline (matches training setup)
 preprocess = transforms.Compose([
     transforms.Resize(256),
     transforms.CenterCrop(224),
@@ -32,16 +39,25 @@ preprocess = transforms.Compose([
 ])
 
 def predict(model, image: Image.Image, class_names):
-    img_tensor = preprocess(image).unsqueeze(0)
+    """
+    Perform prediction on a single image and return top class and probabilities.
+    """
+    img_tensor = preprocess(image).unsqueeze(0)  # Add batch dimension
+
     with torch.no_grad():
         outputs = model(img_tensor)
         probs = torch.nn.functional.softmax(outputs[0], dim=0)
-    
+
+    # Ensure number of classes matches model output size
+    num_classes = probs.shape[0]
+    class_names = class_names[:num_classes]
+
     predictions = [
         {"class": class_names[i], "confidence": round(float(probs[i]), 4)}
-        for i in range(len(class_names))
+        for i in range(num_classes)
     ]
     predictions.sort(key=lambda x: x["confidence"], reverse=True)
+
     return {
         "predictions": predictions,
         "top_class": predictions[0]["class"]