Balajim57 commited on
Commit
d08e0d5
1 Parent(s): 6ea88d1

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +25 -0
app.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from PIL import Image
3
+ from transformers import pipeline
4
+
5
+ # Load your fine-tuned zero-shot model
6
+ classifier = pipeline("zero-shot-classification", model="Balajim57/zero-shot-vitb32")
7
+
8
+ def predict(image, prompt1, prompt2, prompt3):
9
+ # Perform zero-shot classification on the uploaded image with provided prompts
10
+ results = classifier(image, [prompt1, prompt2, prompt3])
11
+ return results["labels"]
12
+
13
+ # Streamlit UI components
14
+ st.title("Zero-Shot Image Classification")
15
+ uploaded_image = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
16
+ prompt1 = st.text_input("Prompt 1")
17
+ prompt2 = st.text_input("Prompt 2")
18
+ prompt3 = st.text_input("Prompt 3")
19
+ if uploaded_image:
20
+ image = Image.open(uploaded_image)
21
+ st.image(image, caption="Uploaded Image", use_column_width=True)
22
+ if st.button("Classify"):
23
+ results = predict(image, prompt1, prompt2, prompt3)
24
+ st.write("Classification Results:")
25
+ st.write(results)