prithivMLmods commited on
Commit
974c881
·
verified ·
1 Parent(s): f347918

Create gender_classification.py

Browse files
Files changed (1) hide show
  1. gender_classification.py +38 -0
gender_classification.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import spaces
3
+ from transformers import AutoImageProcessor
4
+ from transformers import SiglipForImageClassification
5
+ from PIL import Image
6
+ import torch
7
+
8
+ # Load model and processor
9
+ model_name = "prithivMLmods/Gender-Classifier-Mini"
10
+ model = SiglipForImageClassification.from_pretrained(model_name)
11
+ processor = AutoImageProcessor.from_pretrained(model_name)
12
+
13
+ @spaces.GPU
14
+ def gender_classification(image):
15
+ """Predicts gender category for an image."""
16
+ image = Image.fromarray(image).convert("RGB")
17
+ inputs = processor(images=image, return_tensors="pt")
18
+
19
+ with torch.no_grad():
20
+ outputs = model(**inputs)
21
+ logits = outputs.logits
22
+ probs = torch.nn.functional.softmax(logits, dim=1).squeeze().tolist()
23
+
24
+ labels = {"0": "Female ♀", "1": "Male ♂"}
25
+ predictions = {labels[str(i)]: round(probs[i], 3) for i in range(len(probs))}
26
+
27
+ return predictions
28
+
29
+ if __name__ == "__main__":
30
+ import gradio as gr
31
+ iface = gr.Interface(
32
+ fn=gender_classification,
33
+ inputs=gr.Image(type="numpy"),
34
+ outputs=gr.Label(label="Prediction Scores"),
35
+ title="Gender Classification",
36
+ description="Upload an image to classify its gender."
37
+ )
38
+ iface.launch()