# 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())