Spaces:
Runtime error
Runtime error
File size: 2,621 Bytes
05a0149 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
import streamlit as st
from PIL import Image
import numpy as np
from ultralytics import YOLO
import inference
from langchain.chat_models import AzureChatOpenAI
import os
from langchain.schema import HumanMessage
import json
yolo_model = YOLO('best.pt')
roboflow_model = inference.get_model("web-icon-classification/1")
chat4 = AzureChatOpenAI(
openai_api_base=os.environ['BASE_URL'],
openai_api_version="2024-02-15-preview",
deployment_name="gpt-4",
openai_api_key=os.environ["OPENAI_API_KEY"],
openai_api_type="azure",
temperature=0,
request_timeout=30,
max_retries=3
)
def initiate_prompt(icon_name):
prompt = '''Given the name of an app icon, return a list of alternative names that represent similar functionality in the context of a web or mobile app.
User Input: "Settings"
Expected Output: Generate a list of alternative names that convey the same or similar functionality as "Settings" in the context of web or mobile apps.
Model Response: {
"alternatives": ["Preferences", "Options", "Controls", "Configuration", "Setup"]
}
User Input:''' + icon_name +'\n'+ " Model Response:"
return prompt
st.title("App/Web Icon Classification Comparison")
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
image = Image.open(uploaded_file)
st.image(image, caption='Uploaded Image', use_column_width=True)
if st.button("Classify Image"):
with st.spinner('Classifying...'):
try:
prediction = yolo_model(image)
class_id_1 = prediction[0].names[prediction[0].probs.top1]
classes_1 = json.loads(chat4.predict_messages(messages=[HumanMessage(content=initiate_prompt(class_id_1))]).content)['alternatives']
classes_1.insert(0, class_id_1)
except:
classes_1 = "None"
try:
prediction = roboflow_model.infer(image)
class_id_2 = prediction[0].predicted_classes[0]
classes_2 = json.loads(chat4.predict_messages(messages=[HumanMessage(content=initiate_prompt(class_id_2))]).content)['alternatives']
classes_2.insert(0, class_id_2)
except:
classes_2 = "None"
col1, col2 = st.columns(2)
with col1:
st.subheader("Yolov8-x Prediction")
st.write(f"Predicted Class: {classes_1}")
with col2:
st.subheader("ViT Prediction")
st.write(f"Predicted Class: {classes_2}")
|