import os
import gradio as gr

os.system("pip install -U gradio")
os.system("pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu102/torch1.9/index.html")
os.system("git clone https://github.com/facebookresearch/Detic.git --recurse-submodules")

# Importing necessary libraries
import numpy as np
import cv2
from PIL import Image
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg

# Configuring model and predictor
cfg = get_cfg()
cfg.merge_from_file("Detic/configs/Detic_LCOCOI21k_CLIP_SwinB_896b32_4x_ft4x_max-size.yaml")
cfg.MODEL.WEIGHTS = "https://dl.fbaipublicfiles.com/detic/Detic_LCOCOI21k_CLIP_SwinB_896b32_4x_ft4x_max-size.pth"
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
predictor = DefaultPredictor(cfg)

# Caption generator
from langchain.llms import OpenAIChat
session_token = os.environ.get("SessionToken")

def generate_caption(object_list_str, api_key, temperature):
    query = f"You are an intelligent image captioner. I will hand you the objects and their position, and you should give me a detailed description for the photo. In this photo we have the following objects\n{object_list_str}"
    llm = OpenAIChat(
        model_name="gpt-3.5-turbo", openai_api_key=api_key, temperature=temperature
    )

    try:
        caption = llm(query)
        caption = caption.strip()
    except:
        caption = "Sorry, something went wrong!"

    return caption

# Model Inference
def caption_image(img):
    im = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    outputs = predictor(im)["instances"]

    metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0])
    v = Visualizer(im[:, :, ::-1], metadata=metadata)
    out = v.draw_instance_predictions(outputs.to("cpu"))

    detected_objects = []
    object_list_str = []

    for i, prediction in enumerate(outputs):
        x0, y0, x1, y1 = prediction.pred_boxes.tensor[0].cpu().numpy()
        width = x1 - x0
        height = y1 - y0
        predicted_label = metadata.thing_classes[prediction.pred_classes[0]]
        detected_objects.append({
            "prediction": predicted_label,
            "x": int(x0),
            "y": int(y0),
            "w": int(width),
            "h": int(height)
        })
        object_list_str.append(f"{predicted_label} - X:({int(x0)} Y: {int(y0)} Width {int(width)} Height: {int(height)})")

    # GPT3 to generate caption
    api_key = session_token
    if api_key is not None:
        gpt_response = generate_caption(object_list_str, api_key, temperature=0.7)
    else:
        gpt_response = "Please paste your OpenAI key to use"

    return gpt_response

# Interface
image_input = gr.inputs.Image(shape=(896, 896))
caption_output = gr.outputs.Textbox()

gr.Interface(fn=caption_image, inputs=image_input, outputs=caption_output, title="Intelligent Image Captioning", description="Generate captions for an image with object detection.").launch()