Update app.py

app.py

@@ -31,7 +31,8 @@ for package, version in packages.items():
         print(f"Installing {package}...")
         install(f"{package}=={version}")
 
-# Import all required libraries
+
+
 import numpy as np
 import torch
 import torch.nn as nn
@@ -40,7 +41,7 @@ import torchvision.transforms as transforms
 from PIL import Image
 import gradio as gr
 
-# Define the model
+# Define the model exactly as in training
 class ModifiedLargeNet(nn.Module):
     def __init__(self):
         super(ModifiedLargeNet, self).__init__()
@@ -59,26 +60,47 @@ class ModifiedLargeNet(nn.Module):
         x = self.fc2(x)
         return x
 
-# Load the trained model
-model = ModifiedLargeNet()
-model.load_state_dict(torch.load("modified_large_net.pt", map_location=torch.device("cpu")))
-model.eval()
+# Load the trained model with error handling
+try:
+    model = ModifiedLargeNet()
+    state_dict = torch.load("modified_large_net.pt", map_location=torch.device("cpu"))
+    model.load_state_dict(state_dict)
+    print("Model loaded successfully")
+    model.eval()
+except Exception as e:
+    print(f"Error loading model: {str(e)}")
+    traceback.print_exc()
 
-# Define image transformation pipeline
+# Define image transformation pipeline to match training
 transform = transforms.Compose([
     transforms.Resize((128, 128)),
     transforms.ToTensor(),
-    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+    # Using standard normalization as in training
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
 def process_image(image):
     if image is None:
         return None
     
-    # Convert to RGB if necessary
-    if image.mode != 'RGB':
-        image = image.convert('RGB')
-    return image
+    try:
+        # Convert numpy array to PIL Image
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+        
+        # Convert to RGB if necessary
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+            
+        # Print debug information
+        print(f"Processed image size: {image.size}")
+        print(f"Processed image mode: {image.mode}")
+        
+        return image
+    except Exception as e:
+        print(f"Error in process_image: {str(e)}")
+        traceback.print_exc()
+        return None
 
 def predict(image):
     if image is None:
@@ -86,30 +108,37 @@ def predict(image):
     
     try:
         # Process the image
-        processed_image = process_image(Image.fromarray(image))
+        processed_image = process_image(image)
         if processed_image is None:
             return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
         
         # Transform for model
         tensor_image = transform(processed_image).unsqueeze(0)
+        print(f"Input tensor shape: {tensor_image.shape}")
         
         # Make prediction
         with torch.no_grad():
             outputs = model(tensor_image)
+            print(f"Raw outputs: {outputs}")
+            
             probabilities = F.softmax(outputs, dim=1)[0].cpu().numpy()
+            print(f"Probabilities: {probabilities}")
             
         # Return results
         classes = ["Rope", "Hammer", "Other"]
-        return {cls: float(prob) for cls, prob in zip(classes, probabilities)}
+        results = {cls: float(prob) for cls, prob in zip(classes, probabilities)}
+        print(f"Final results: {results}")
+        return results
     
     except Exception as e:
         print(f"Prediction error: {str(e)}")
+        traceback.print_exc()
         return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
 
 # Gradio interface
 interface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="numpy"),
+    inputs=gr.Image(),  # Accept any image format
     outputs=gr.Label(num_top_classes=3),
     title="Mechanical Tools Classifier",
     description="Upload an image of a tool to classify it as 'Rope', 'Hammer', or 'Other'.",