ajcdp commited on
Commit
1b94d1b
1 Parent(s): 4212c74

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +40 -0
app.py ADDED
@@ -0,0 +1,40 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import numpy as np
3
+ from PIL import Image
4
+ from transformers import pipeline
5
+
6
+ # Create the pipeline object
7
+ pipe = pipeline("zero-shot-image-classification", model="openai/clip-vit-base-patch32")
8
+
9
+ # Define the function that will be used by the interface
10
+ def zero_shot_classification(image, labels_text):
11
+ # Convert image to a PIL image object
12
+ pil_image = Image.fromarray(np.uint8(image)).convert("RGB")
13
+
14
+ # Split the labels text into a list of labels
15
+ labels = labels_text.split(",")
16
+
17
+ # Use the pipeline to classify the image with the given labels
18
+ res = pipe(
19
+ images=pil_image,
20
+ candidate_labels=labels,
21
+ hypothesis_template= "This is a photo of a {}"
22
+ )
23
+
24
+ # Return a dictionary mapping labels to scores
25
+ return {dic["label"]: dic["score"] for dic in res}
26
+
27
+ # Create the interface
28
+ iface = gr.Interface(
29
+ zero_shot_classification,
30
+ ["image", "text"],
31
+ "label",
32
+ examples=[
33
+ ["dog.jpg", "dog,cat,horse,zebra"],
34
+ ],
35
+ description="Add a picture and a list of labels separated by commas",
36
+ title="Zero-shot Image Classification"
37
+ )
38
+
39
+ # Launch the interface
40
+ iface.launch()