Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| from ultralyticsplus import YOLO, render_result | |
| from PIL import Image | |
| import os | |
| def yolov8_func(image, | |
| image_size, | |
| conf_thresold=0.4, | |
| iou_thresold=0.50): | |
| # Load the YOLOv8 model | |
| model_path = "best.pt" | |
| model = YOLO(model_path) | |
| # Make predictions | |
| result = model.predict(image, conf=conf_thresold, iou=iou_thresold, imgsz=image_size) | |
| # Access object detection results | |
| boxes = result[0].boxes | |
| num_boxes = len(boxes) | |
| # Categorize based on number of boxes (detections) and provide recommendations | |
| if num_boxes > 10: | |
| severity = "Worse" | |
| recommendation = "It is recommended to see a dermatologist and start stronger acne treatment." | |
| elif 5 <= num_boxes <= 10: | |
| severity = "Medium" | |
| recommendation = "You should follow a consistent skincare routine with proper cleansing and moisturizing." | |
| else: | |
| severity = "Good" | |
| recommendation = "Your skin looks good! Keep up with your current skincare routine." | |
| # Render the result (with bounding boxes/labels) | |
| render = render_result(model=model, image=image, result=result[0]) | |
| predicted_image_save_path = "predicted_image.jpg" | |
| render.save(predicted_image_save_path) | |
| return predicted_image_save_path, f"Acne condition: {severity}", recommendation | |
| # Create the Gradio | |
| with gr.Blocks() as yolo_app: | |
| gr.Markdown("# YOLOv8: An Object Detection for Acne") | |
| with gr.Row(): | |
| with gr.Column(scale=1): # Left side with input | |
| input_image = gr.Image(type="filepath", label="Input Image") | |
| image_size = gr.Slider(minimum=320, maximum=1280, step=32, value=640, label="Image Size") | |
| conf_thresh = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.15, label="Confidence Threshold") | |
| iou_thresh = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.2, label="IOU Threshold") | |
| submit_btn = gr.Button("Submit") | |
| with gr.Column(scale=1): # Right side with output | |
| output_image = gr.Image(type="filepath", label="Output Image") | |
| acne_condition = gr.Textbox(label="Acne Condition") | |
| recommendation = gr.Textbox(label="Recommendation") | |
| # Link the submit button to the function | |
| submit_btn.click(fn=yolov8_func, | |
| inputs=[input_image, image_size, conf_thresh, iou_thresh], | |
| outputs=[output_image, acne_condition, recommendation]) | |
| # Launch the app | |
| yolo_app.launch(debug=True) |