Spaces:
Sleeping
Sleeping
File size: 2,879 Bytes
fb60bec | 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 73 74 75 76 77 78 79 | import streamlit as st
import cv2
from ultralytics import YOLO
import numpy as np
from PIL import Image
# Initialize the YOLO model
model_path = 'yolov11x1.1-trained.pt' # Ensure this model file is in the same directory
model = YOLO(model_path)
# Temporary fix: add placeholder names for missing classes
expected_classes = 13 # Set this to the correct number of classes
for i in range(expected_classes):
if i not in model.names:
model.names[i] = f"class_{i}"
def annotate_image(input_image_path, output_image_path, confidence=0.25):
"""Loads an image, runs YOLO model to detect skin issues with specified confidence, and saves annotated image."""
# Load the image
img = cv2.imread(input_image_path)
if img is None:
raise ValueError(f"Image at path {input_image_path} could not be loaded.")
# Run YOLO model inference with the specified confidence threshold
results = model.predict(img, conf=confidence)
# Get the annotated image from results
annotated_img = results[0].plot()
# Save the annotated image to the specified output path
cv2.imwrite(output_image_path, annotated_img)
print(f"Annotated image saved at {output_image_path} with confidence threshold {confidence}")
# Streamlit UI
st.title("Skin Issue Detection with YOLO")
st.write("Upload an image to detect and annotate skin issues with a confidence threshold.")
# Image uploader
uploaded_file = st.file_uploader("Choose an image...", type=['jpg', 'jpeg', 'png'])
# Confidence slider
confidence_threshold = st.slider("Confidence Threshold", min_value=0.0, max_value=1.0, value=0.25)
if uploaded_file is not None:
# Save the uploaded file locally as 'test1.jpeg'
input_image_path = 'test1.jpeg'
with open(input_image_path, "wb") as f:
f.write(uploaded_file.getbuffer())
output_image_path = 'annotated_test1.jpeg'
# Annotate image using your existing code function
try:
annotate_image(input_image_path, output_image_path, confidence=confidence_threshold)
# Display the original and annotated images side by side
col1, col2 = st.columns(2)
with col1:
st.subheader("Original Image")
st.image(uploaded_file, use_column_width=True)
with col2:
st.subheader("Annotated Image")
annotated_img = Image.open(output_image_path)
st.image(annotated_img, use_column_width=True)
# Provide a download link for the annotated image
with open(output_image_path, "rb") as file:
btn = st.download_button(
label="Download Annotated Image",
data=file,
file_name="annotated_test1.jpeg",
mime="image/jpeg"
)
except Exception as e:
st.error(f"An error occurred: {str(e)}")
|