Update app.py

app.py

@@ -5,27 +5,27 @@ from torchvision import models, transforms
 from PIL import Image
 
 # 1. SETUP MODEL ARCHITECTURE
-# Based on your logs, you used ResNet50. 
-# If you actually used ResNet18, change 'resnet50' to 'resnet18' below.
+# (Matches the ResNet50 logs you provided)
 model = models.resnet50(weights=None) 
 
 # 2. MATCH THE FINAL LAYER
-# ResNet50 has 2048 input features in the final layer. 
-# (If you used ResNet18, this number would be 512).
 num_ftrs = model.fc.in_features 
 model.fc = nn.Linear(num_ftrs, 2)
 
 # 3. LOAD WEIGHTS
-# Replace 'fire_detection_resnet18.pth' with the EXACT name of the file you uploaded
+# Ensure this matches the EXACT filename you uploaded to the Files tab
 model_path = "fire_detection_resnet50.pth" 
 
-# Load weights for CPU (since HF Spaces Free Tier uses CPU)
-state_dict = torch.load(model_path, map_location=torch.device('cpu'))
-model.load_state_dict(state_dict)
+try:
+    state_dict = torch.load(model_path, map_location=torch.device('cpu'))
+    model.load_state_dict(state_dict)
+    print("Model weights loaded successfully.")
+except Exception as e:
+    print(f"Error loading model weights: {e}")
+
 model.eval()
 
 # 4. DEFINE PREPROCESSING
-# This must match what you used during training
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
@@ -33,32 +33,35 @@ transform = transforms.Compose([
 ])
 
 # 5. PREDICTION FUNCTION
-labels = ['Non-Fire', 'Fire'] # 0 is Non-Fire, 1 is Fire
+labels = ['Non-Fire', 'Fire']
 
 def predict(image):
     if image is None:
         return None
     
-    # Preprocess
-    image = image.convert('RGB')
-    image_tensor = transform(image).unsqueeze(0) # Add batch dimension
-    
-    # Inference
-    with torch.no_grad():
-        outputs = model(image_tensor)
-        probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
-    
-    # Return dictionary for Gradio Label
-    return {labels[i]: float(probabilities[i]) for i in range(len(labels))}
+    try:
+        # Preprocess
+        image = image.convert('RGB')
+        image_tensor = transform(image).unsqueeze(0) 
+        
+        # Inference
+        with torch.no_grad():
+            outputs = model(image_tensor)
+            probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
+        
+        # Return dictionary for Gradio Label
+        return {labels[i]: float(probabilities[i]) for i in range(len(labels))}
+    except Exception as e:
+        return {f"Error: {str(e)}": 0.0}
 
 # 6. LAUNCH GRADIO UI
+# Removed 'examples' to prevent the crash
 interface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil", label="Upload Image"),
     outputs=gr.Label(num_top_classes=2, label="Prediction"),
-    title="Fire Detection System tements",
-    description="Upload an image to detect if fire is present. (Model: ResNet50)",
-    examples=["fire.jpg", "forest.jpg"] # Optional: Upload these images to your space for users to click
+    title="Fire Detection System",
+    description="Upload an image to detect if fire is present."
 )
 
 if __name__ == "__main__":