Spaces:
Sleeping
Sleeping
| 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)}") | |