sergioska's picture
text size
bcd0b17
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)