Spaces:
Sleeping
Sleeping
| # Import necessary libraries | |
| import os | |
| import cv2 | |
| import numpy as np | |
| from ultralytics import YOLO | |
| import gradio as gr | |
| import sqlite3 | |
| from datetime import datetime | |
| import traceback | |
| print("Libraries imported successfully") | |
| # Initialize YOLOv8 model | |
| try: | |
| model = YOLO('yolov8m.pt') | |
| print("YOLOv8 model loaded successfully") | |
| except Exception as e: | |
| print(f"Error loading YOLO model: {str(e)}") | |
| print(traceback.format_exc()) | |
| # Set up SQLite database | |
| def setup_database(): | |
| try: | |
| conn = sqlite3.connect('detections.db', check_same_thread=False) | |
| cursor = conn.cursor() | |
| cursor.execute(''' | |
| CREATE TABLE IF NOT EXISTS detections | |
| (id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| timestamp TEXT, | |
| object TEXT, | |
| confidence REAL) | |
| ''') | |
| conn.commit() | |
| print("Database set up successfully") | |
| return conn | |
| except Exception as e: | |
| print(f"Error setting up database: {str(e)}") | |
| print(traceback.format_exc()) | |
| return None | |
| db_conn = setup_database() | |
| def detect_objects(image): | |
| try: | |
| if image is None: | |
| return None, "Error: No image provided" | |
| print(f"Received image shape: {image.shape}") | |
| # Ensure image is in RGB format | |
| if len(image.shape) == 2: # Grayscale | |
| image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB) | |
| elif image.shape[2] == 4: # RGBA | |
| image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB) | |
| # Perform object detection | |
| results = model(image) | |
| print(f"Detection completed. Results: {results}") | |
| # Process results | |
| detected_objects = [] | |
| for r in results: | |
| boxes = r.boxes | |
| for box in boxes: | |
| x1, y1, x2, y2 = box.xyxy[0] | |
| conf = box.conf[0] | |
| cls = int(box.cls[0]) | |
| name = model.names[cls] | |
| # Draw bounding box | |
| cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2) | |
| cv2.putText(image, f'{name} {conf:.2f}', (int(x1), int(y1) - 10), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2) | |
| # Save to database | |
| timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S") | |
| with sqlite3.connect('detections.db', check_same_thread=False) as conn: | |
| cursor = conn.cursor() | |
| cursor.execute("INSERT INTO detections (timestamp, object, confidence) VALUES (?, ?, ?)", | |
| (timestamp, name, float(conf))) | |
| conn.commit() | |
| detected_objects.append(f"{name}: {conf:.2f}") | |
| if not detected_objects: | |
| print("No objects detected in the image") | |
| return image, "No objects detected" | |
| else: | |
| print(f"Detected objects: {detected_objects}") | |
| return image, "\n".join(detected_objects) | |
| except Exception as e: | |
| print(f"Error in detect_objects: {str(e)}") | |
| print(traceback.format_exc()) | |
| return None, f"Error: {str(e)}" | |
| # Set up Gradio interface | |
| try: | |
| iface = gr.Interface( | |
| fn=detect_objects, | |
| inputs=gr.Image(type="numpy"), | |
| outputs=[gr.Image(type="numpy"), gr.Textbox(label="Detected Objects")], | |
| title="YOLOv8 Object Detection", | |
| description="Upload an image to detect objects using YOLOv8." | |
| ) | |
| print("Gradio interface set up successfully") | |
| except Exception as e: | |
| print(f"Error setting up Gradio interface: {str(e)}") | |
| print(traceback.format_exc()) | |
| # Launch the interface# 啟動介面 | |
| try: | |
| iface.launch(debug=True) | |
| print("Gradio interface launched successfully") | |
| except Exception as e: | |
| print(f"Error launching Gradio interface: {str(e)}") | |
| print(traceback.format_exc()) | |