Spaces:
Sleeping
Sleeping
| import os | |
| import cv2 | |
| import glob | |
| import torch | |
| import streamlit as st | |
| import numpy as np | |
| import base64 | |
| from random import randrange | |
| from PIL import Image, ImageDraw | |
| from transformers import pipeline | |
| from tempfile import NamedTemporaryFile | |
| from st_clickable_images import clickable_images | |
| def clearEphimeral(): | |
| for filename in glob.glob('/tmp/cropped_image_*'): | |
| os.remove(filename) | |
| st.set_page_config(layout="wide") | |
| clearEphimeral() | |
| if 'wip' not in st.session_state: | |
| st.session_state['wip'] = False | |
| st.session_state['image'] = None | |
| detector = pipeline("object-detection", model="sergioska/yolo-finetuned-flyers") | |
| col1, col2, col3 = st.columns([0.5, 0.3, 0.2]) | |
| gibsOk = False | |
| totalGibs = 0 | |
| with col1: | |
| clearEphimeral() | |
| st.header('flyer') | |
| uploaded_image_zero_file = st.file_uploader("Choose an image file (zero)") | |
| if not st.session_state['wip']: | |
| imageLocation = st.empty() | |
| else: | |
| imageLocation = st.image(st.session_state['image']) | |
| if uploaded_image_zero_file is not None: | |
| image = Image.open(uploaded_image_zero_file) | |
| outputImage = np.array(image) | |
| predictions = detector( | |
| image, | |
| candidate_labels=['product'] | |
| ) | |
| imageLocation.image(outputImage) | |
| if st.button('detect') or st.session_state['wip']: | |
| draw = ImageDraw.Draw(image) | |
| arrImg = np.array(image) | |
| imageCv = cv2.cvtColor(arrImg, cv2.COLOR_RGB2BGR) | |
| i = 0 | |
| 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=5) | |
| draw.text((xmin, ymin), f"{label}: {round(score,5)}", fill="white") | |
| crop_img = imageCv[ymin:ymax, xmin:xmax] | |
| if crop_img.size > 0: | |
| cv2.imwrite("/tmp/cropped_image_{}.jpg".format(i), crop_img) | |
| i += 1 | |
| totalGibs = i | |
| gibsOk = True | |
| imageLocation.image(image) | |
| st.session_state['wip'] = True | |
| st.session_state['image'] = image | |
| st.write(predictions) | |
| with col2: | |
| st.header('gibs') | |
| images = [] | |
| imagesForOpen = [] | |
| for cropped in glob.glob('/tmp/cropped_image_*'): | |
| with open(cropped, "rb") as image: | |
| encoded = base64.b64encode(image.read()).decode() | |
| images.append(f"data:image/jpeg;base64,{encoded}") | |
| imagesForOpen.append(cropped) | |
| clicked = clickable_images( | |
| images, | |
| titles=[f"Image #{str(i)}" for i in range(totalGibs)], | |
| div_style={"display": "flex", "justify-content": "center", "flex-wrap": "wrap"}, | |
| img_style={"margin": "5px", "height": "150px"}, | |
| key="cropped_images" | |
| ) | |
| with col3: | |
| st.header('details') | |
| colors = ["red", "yellow", "purple", "blue", "green"] | |
| if gibsOk: | |
| if clicked == -1: | |
| clicked = 0 | |
| detectorForContents = pipeline("object-detection", model="sergioska/yolo-finetuned-flyers-content") | |
| imageForDetails = Image.open(imagesForOpen[clicked]) | |
| outputImageForDetails = np.array(imageForDetails) | |
| predictionsForDetails = detectorForContents( | |
| imageForDetails, | |
| candidate_labels=['title', 'price'] | |
| ) | |
| drawForDetails = ImageDraw.Draw(imageForDetails) | |
| c = 0 | |
| for predictionsForDetail in predictionsForDetails: | |
| box = predictionsForDetail["box"] | |
| label = predictionsForDetail["label"] | |
| score = predictionsForDetail["score"] | |
| xmin, ymin, xmax, ymax = box.values() | |
| drawForDetails.rectangle((xmin, ymin, xmax, ymax), outline=colors[c], width=5) | |
| drawForDetails.text((xmin, ymin), f"{label}: {round(score,5)}", fill="white") | |
| c += 1 | |
| st.image(imageForDetails) | |
| st.write(predictionsForDetails) | |
| st.markdown(f"Image #{clicked} clicked" if clicked > -1 else "No image clicked") | |