sergioska's picture
text size
bcd0b17
Raw
History Blame Contribute Delete
1.87 kB
import torch
import streamlit as st
import numpy as np
from PIL import Image, ImageDraw
from transformers import pipeline
from tempfile import NamedTemporaryFile
imagepipe = pipeline("image-classification", model="flatmoon102/fruits_and_vegetables_image_classification")
detector = pipeline(model="google/owlvit-base-patch32", task="zero-shot-object-detection")
uploaded_image_file = st.file_uploader("Choose an image file")
if uploaded_image_file is not None:
with NamedTemporaryFile() as temp:
temp.write(uploaded_image_file.getvalue())
temp.seek(0)
result = imagepipe(temp.name)
st.write(result)
st.title('Upload an image file to detection')
uploaded_image_zero_file = st.file_uploader("Choose an image file (zero)")
texts = st.text_input('tags')
if uploaded_image_zero_file is not None:
image = Image.open(uploaded_image_zero_file)
#temp.write(uploaded_image_file.getvalue())
#temp.seek(0)
outputImage = np.array(image)
predictions = detector(
image,
candidate_labels=['eggs', 'apple', 'pear']
)
st.image(outputImage)
if st.button('apply tag'):
tags = [['eggs', 'apple', 'pear']]
#inputs = processor(text=tags, images=image, return_tensors="pt")
#outputs = model(**inputs)
#target_sizes = torch.Tensor([image.size[::-1]])
#results = processor.post_process_object_detection(outputs=outputs, threshold=0.1, target_sizes=target_sizes)
draw = ImageDraw.Draw(image)
for prediction in predictions:
box = prediction["box"]
label = prediction["label"]
score = prediction["score"]
xmin, ymin, xmax, ymax = box.values()
draw.rectangle((xmin, ymin, xmax, ymax), outline="red", width=3)
draw.text((xmin, ymin), f"{label}: {round(score,5)}", fill="white")
st.image(image)
st.write(predictions)