Spaces:
Sleeping
Sleeping
File size: 1,305 Bytes
87ecabf c120c7a 87ecabf c120c7a 87ecabf c120c7a |
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 |
import streamlit as st
from PIL import Image
import numpy as np
from ultralytics import YOLO # Make sure this import works in your Hugging Face environment
# Load the model
@st.cache_resource(allow_output_mutation=True)
def load_model():
model = YOLO("weights.pt") # Adjust path if needed
return model
model = load_model()
st.title("Circuit Sketch Recognition")
# File uploader allows user to add their own image
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
image = Image.open(uploaded_file).convert("RGB")
st.image(image, caption='Uploaded Image', use_column_width=True)
st.write("")
st.write("Detecting...")
# Perform inference
results = model.predict(uploaded_file)
r = results[0]
im_bgr = r.plot(conf=False, pil=True, font_size=48, line_width=3) # Returns a PIL image if pil=True
im_rgb = Image.fromarray(im_bgr[..., ::-1]) # Convert BGR to RGB
# Display the prediction
st.image(im_rgb, caption='Prediction', use_column_width=True)
# Optionally, display pre-computed example images
if st.checkbox('Show Example Results'):
st.image('example1.jpg', use_column_width=True, caption='Example 1')
st.image('example2.jpg', use_column_width=True, caption='Example 2')
|