Update app.py

app.py

@@ -1,14 +1,12 @@
 import os
 import torch
 import torch.nn as nn
-from torch.utils.data import Dataset, DataLoader, random_split
 from torchvision import transforms
 from PIL import Image
 import gradio as gr
 
 # -------- CONFIG --------
-data_dir = "D:/Dataset/face_age"
-checkpoint_path = "age_prediction_model2.pth"
+checkpoint_path = "age_prediction_model2.pth"  # Just the model file name for Hugging Face Spaces
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
@@ -38,9 +36,9 @@ class SimpleCNN(nn.Module):
 # -------- LOAD MODEL --------
 model = SimpleCNN().to(device)
 
-# Check if checkpoint exists before loading
+# Check if the checkpoint file exists and load
 if os.path.exists(checkpoint_path):
-    model.load_state_dict(torch.load(checkpoint_path))
+    model.load_state_dict(torch.load(checkpoint_path, map_location=device))  # Load to the correct device
     model.eval()  # Set the model to evaluation mode
     print(f"Model loaded from {checkpoint_path}")
 else:
@@ -61,12 +59,10 @@ def predict_age(image):
         age = output.item()  # Convert to a single scalar
     return f"Predicted Age: {age:.2f}"
 
-
-
-# Update the gr.Image initialization
+# -------- GRADIO INTERFACE --------
 iface = gr.Interface(
     fn=predict_age,
-    inputs=gr.Image(image_size=(128, 128), image_mode='RGB', source='upload'),
+    inputs=gr.Image(shape=(128, 128), image_mode='RGB', source='upload'),  # Updated input format
     outputs="text",
     title="Age Prediction Model",
     description="Upload an image to predict the age.",
@@ -74,5 +70,3 @@ iface = gr.Interface(
 )
 
 iface.launch()
-
-