MoinulwithAI commited on
Commit
4abc82f
·
verified ·
1 Parent(s): 15eff2e

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +77 -0
app.py ADDED
@@ -0,0 +1,77 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ from torchvision import transforms
4
+ from PIL import Image
5
+ import gradio as gr
6
+
7
+ # -------- MODEL DEFINITION --------
8
+ class ImprovedCNN(nn.Module):
9
+ def __init__(self):
10
+ super(ImprovedCNN, self).__init__()
11
+ self.features = nn.Sequential(
12
+ nn.Conv2d(3, 32, kernel_size=3, padding=1),
13
+ nn.BatchNorm2d(32),
14
+ nn.ReLU(),
15
+ nn.MaxPool2d(2),
16
+
17
+ nn.Conv2d(32, 64, kernel_size=3, padding=1),
18
+ nn.BatchNorm2d(64),
19
+ nn.ReLU(),
20
+ nn.MaxPool2d(2),
21
+
22
+ nn.Conv2d(64, 128, kernel_size=3, padding=1),
23
+ nn.BatchNorm2d(128),
24
+ nn.ReLU(),
25
+ nn.MaxPool2d(2),
26
+ )
27
+ self.classifier = nn.Sequential(
28
+ nn.Flatten(),
29
+ nn.Linear(128 * 16 * 16, 512),
30
+ nn.ReLU(),
31
+ nn.Dropout(0.5),
32
+ nn.Linear(512, 1)
33
+ )
34
+
35
+ def forward(self, x):
36
+ x = self.features(x)
37
+ x = self.classifier(x)
38
+ return x
39
+
40
+ # -------- LOAD MODEL --------
41
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
42
+ print(f"Using device: {device}")
43
+ model = ImprovedCNN().to(device)
44
+
45
+ model_path = "age_prediction_model3.pth"
46
+ model.load_state_dict(torch.load(model_path, map_location=device))
47
+ model.eval()
48
+ print(f"✅ Model loaded from {model_path}")
49
+
50
+ # -------- IMAGE PREPROCESSING --------
51
+ transform = transforms.Compose([
52
+ transforms.Resize((128, 128)),
53
+ transforms.ToTensor(),
54
+ transforms.Normalize(mean=[0.485, 0.456, 0.406],
55
+ std=[0.229, 0.224, 0.225])
56
+ ])
57
+
58
+ # -------- PREDICTION FUNCTION --------
59
+ def predict_age(image: Image.Image) -> float:
60
+ image_tensor = transform(image).unsqueeze(0).to(device)
61
+ with torch.no_grad():
62
+ output = model(image_tensor)
63
+ age = output.item()
64
+ return round(age, 2)
65
+
66
+ # -------- GRADIO UI --------
67
+ demo = gr.Interface(
68
+ fn=predict_age,
69
+ inputs=gr.Image(type="pil", image_mode="RGB", label="Upload Face Image"),
70
+ outputs=gr.Number(label="Predicted Age"),
71
+ title="Face Age Prediction",
72
+ description="Upload a face image to predict age using a CNN model."
73
+ )
74
+
75
+ # -------- LAUNCH --------
76
+ if __name__ == "__main__":
77
+ demo.launch()