chanelcolgate commited on
Commit
b4031e4
·
1 Parent(s): 7cbea20

Add application file

Browse files
Files changed (2) hide show
  1. app.py +25 -0
  2. requirements.txt +2 -0
app.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import gradio as gr
3
+ from transformers import AutoModelForImageClassification, AutoFeatureExtractor
4
+
5
+ model_id = f"chanelcolgate/vit-base-patch16-224-finetunned-flower"
6
+ labels = ["daisy", "dandelion", "roses", "sunflowers", "tulips"]
7
+
8
+
9
+ def classify_image(image):
10
+ model = AutoModelForImageClassification.from_pretrained(model_id)
11
+ feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)
12
+ inp = feature_extractor(image, return_tensors="pt")
13
+ outp = model(**inp)
14
+ pred = torch.nn.functional.softmax(outp.logits, dim=-1)
15
+ preds = pred[0].cpu().detach().numpy()
16
+ confidence = {label: float(preds[i]) for i, label in enumerate(labels)}
17
+ return confidence
18
+
19
+
20
+ interface = gr.Interface(
21
+ fn=classify_image,
22
+ inputs="image",
23
+ examples=["flower-1.jpeg", "flower-2.jpeg"],
24
+ outputs="label",
25
+ ).launch()
requirements.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ torch
2
+ transformers