prithivMLmods commited on
Commit
fa4245a
·
verified ·
1 Parent(s): df09cc1

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +45 -0
app.py ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoImageProcessor
3
+ from transformers import SiglipForImageClassification
4
+ from transformers.image_utils import load_image
5
+ from PIL import Image
6
+ import torch
7
+
8
+ # Load model and processor
9
+ model_name = "prithivMLmods/Age-Classification-SigLIP2"
10
+ model = SiglipForImageClassification.from_pretrained(model_name)
11
+ processor = AutoImageProcessor.from_pretrained(model_name)
12
+
13
+ def age_classification(image):
14
+ """Predicts the age group of a person from an image."""
15
+ image = Image.fromarray(image).convert("RGB")
16
+ inputs = processor(images=image, return_tensors="pt")
17
+
18
+ with torch.no_grad():
19
+ outputs = model(**inputs)
20
+ logits = outputs.logits
21
+ probs = torch.nn.functional.softmax(logits, dim=1).squeeze().tolist()
22
+
23
+ labels = {
24
+ "0": "Child 0-12",
25
+ "1": "Teenager 13-20",
26
+ "2": "Adult 21-44",
27
+ "3": "Middle Age 45-64",
28
+ "4": "Aged 65+"
29
+ }
30
+ predictions = {labels[str(i)]: round(probs[i], 3) for i in range(len(probs))}
31
+
32
+ return predictions
33
+
34
+ # Create Gradio interface
35
+ iface = gr.Interface(
36
+ fn=age_classification,
37
+ inputs=gr.Image(type="numpy"),
38
+ outputs=gr.Label(label="Prediction Scores"),
39
+ title="Age Group Classification",
40
+ description="Upload an image to predict the person's age group."
41
+ )
42
+
43
+ # Launch the app
44
+ if __name__ == "__main__":
45
+ iface.launch(ssr_mode=True)