|
import gradio as gr |
|
import os |
|
from pathlib import Path |
|
from ultralytics import YOLO |
|
import cv2 |
|
import logging |
|
import numpy as np |
|
|
|
|
|
def setup_logging(): |
|
logging.basicConfig(level=logging.INFO, |
|
format='%(asctime)s - %(levelname)s - %(message)s') |
|
|
|
|
|
def process_image(model_path, image): |
|
try: |
|
|
|
model = YOLO(model_path) |
|
logging.info(f'Loaded model from: {model_path}') |
|
|
|
|
|
logging.info(f'Processing file') |
|
|
|
results = model(image) |
|
|
|
for result in results: |
|
|
|
result_img = result.plot() |
|
|
|
logging.info("Image processing completed.") |
|
return result_img |
|
except Exception as e: |
|
logging.error(f'Error occurred: {e}') |
|
return None |
|
|
|
|
|
def yolo_detection(image): |
|
model_path = 'model.pt' |
|
result_image = process_image(model_path, image) |
|
if result_image is not None: |
|
return cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB) |
|
else: |
|
return None |
|
|
|
|
|
with gr.Blocks() as demo: |
|
gr.Markdown("# 👁️ tiny YOLOv8 Open/closed eye detection\nUpload an image and see the detection results.") |
|
|
|
image_input = gr.Image(label="Input Image") |
|
image_output = gr.Image(label="Detected Objects") |
|
|
|
detect_button = gr.Button("Detect Objects") |
|
|
|
detect_button.click(fn=yolo_detection, inputs=image_input, outputs=image_output) |
|
|
|
demo.launch(share=False) |
|
|