import torch
from ultralytics import YOLO
import cv2
import numpy as np
import os
import threading
import gradio as gr
from telegram import Update
from telegram.ext import Updater, CommandHandler, MessageHandler, Filters
from PIL import Image

# Load YOLO model
model = YOLO("yolov8n.pt")  # Make sure you have this model

# Telegram Bot Token
TOKEN = "7686883158:AAHiw5ad5P-eZVjymLuwJ1o2sX6Zdpen_v4"

# Function for Telegram bot
async def start(update: Update, context: CallbackContext) -> None:
    await update.message.reply_text("Send me an image, and I'll detect objects in it!")

async def detect_objects(update: Update, context: CallbackContext) -> None:
    photo = await update.message.photo[-1].get_file()
    await photo.download_to_drive("input.jpg")

    # Run YOLO detection
    results = model("input.jpg")

    # Save output image
    for result in results:
        result.save("output.jpg")

    # Send result back
    await update.message.reply_photo(photo=open("output.jpg", "rb"))

# Run Telegram bot in a separate thread
def run_telegram_bot():
    app = Application.builder().token(TOKEN).build()
    app.add_handler(CommandHandler("start", start))
    app.add_handler(MessageHandler(filters.PHOTO, detect_objects))

    print("Telegram bot is running...")
    app.run_polling()

# Gradio function
def process_image(img):
    img_path = "gradio_input.jpg"
    img.save(img_path)

    # Run YOLO model
    results = model(img_path)

    # Save the output
    output_path = "gradio_output.jpg"
    for result in results:
        result.save(output_path)

    return output_path

# Launch Gradio web app
def launch_gradio():
    interface = gr.Interface(fn=process_image, inputs="image", outputs="image", title="YOLO Object Detection")
    interface.launch()

# Run both Telegram bot and Gradio
if __name__ == "__main__":
    threading.Thread(target=run_telegram_bot, daemon=True).start()
    launch_gradio()