Update app.py

app.py

@@ -1,57 +1,74 @@
+import os
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
-import torchvision.transforms as transforms
+from torch.utils.data import Dataset, DataLoader, random_split
+from torchvision import transforms
 from PIL import Image
 import gradio as gr
 
-# ------------------- Model Definition -------------------
+# -------- CONFIG --------
+data_dir = "D:/Dataset/face_age"
+checkpoint_path = "D:/Dataset/age_prediction_model2.pth"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+# -------- SIMPLE CNN MODEL --------
 class SimpleCNN(nn.Module):
-    def __init__(self, num_classes=1):
+    def __init__(self):
         super(SimpleCNN, self).__init__()
-        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
-        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
-        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
-        self.pool = nn.MaxPool2d(2, 2)
-        self.fc1 = nn.Linear(128 * 28 * 28, 512)
-        self.fc2 = nn.Linear(512, num_classes)
+        self.features = nn.Sequential(
+            nn.Conv2d(3, 32, kernel_size=3, padding=1), nn.ReLU(),
+            nn.MaxPool2d(2),  # 64x64
+            nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.ReLU(),
+            nn.MaxPool2d(2),  # 32x32
+            nn.Conv2d(64, 128, kernel_size=3, padding=1), nn.ReLU(),
+            nn.MaxPool2d(2),  # 16x16
+        )
+        self.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(128 * 16 * 16, 256), nn.ReLU(),
+            nn.Linear(256, 1)  # Output: age (regression)
+        )
 
     def forward(self, x):
-        x = self.pool(F.relu(self.conv1(x)))
-        x = self.pool(F.relu(self.conv2(x)))
-        x = self.pool(F.relu(self.conv3(x)))
-        x = x.view(-1, 128 * 28 * 28)
-        x = F.relu(self.fc1(x))
-        x = self.fc2(x)
+        x = self.features(x)
+        x = self.classifier(x)
         return x
 
-# ------------------- Load Model -------------------
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = SimpleCNN()
-model.load_state_dict(torch.load("age_prediction_model1.pth", map_location=device))
-model.to(device)
-model.eval()
+# -------- LOAD MODEL --------
+model = SimpleCNN().to(device)
 
-# ------------------- Transform -------------------
+# Check if checkpoint exists before loading
+if os.path.exists(checkpoint_path):
+    model.load_state_dict(torch.load(checkpoint_path))
+    model.eval()  # Set the model to evaluation mode
+    print(f"Model loaded from {checkpoint_path}")
+else:
+    print(f"Error: Checkpoint file not found at {checkpoint_path}. Please check the path.")
+
+# -------- PREPROCESSING --------
 transform = transforms.Compose([
-    transforms.Resize((224, 224)),
+    transforms.Resize((128, 128)),
     transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406],
-                         [0.229, 0.224, 0.225])
+    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
 ])
 
-# ------------------- Prediction Function -------------------
-def predict(image):
+# -------- PREDICTION FUNCTION --------
+def predict_age(image):
     image = transform(image).unsqueeze(0).to(device)
     with torch.no_grad():
-        output = model(image).squeeze().item()
-    return f"Predicted Age: {round(output, 2)} years"
-
-# ------------------- Gradio Interface -------------------
-iface = gr.Interface(fn=predict, 
-                     inputs=gr.Image(type="pil"), 
-                     outputs="text",
-                     title="Face Age Prediction",
-                     description="Upload a face image and get a predicted age")
+        output = model(image)
+        age = output.item()  # Convert to a single scalar
+    return f"Predicted Age: {age:.2f}"
+
+# -------- GRADIO INTERFACE --------
+iface = gr.Interface(
+    fn=predict_age,
+    inputs=gr.inputs.Image(shape=(128, 128), image_mode='RGB', source='upload'),
+    outputs="text",
+    title="Age Prediction Model",
+    description="Upload an image to predict the age.",
+    live=True
+)
 
 iface.launch()