import gradio as gr
import torch
import torchvision.transforms as T
from torchvision.models.detection import maskrcnn_resnet50_fpn
from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
from google_drive_downloader import GoogleDriveDownloader as gdd

# Download and load the RAG model and tokenizer
gdd.download_file_from_google_drive(file_id='your_model_file_id', dest_path='./model.pt')
gdd.download_file_from_google_drive(file_id='your_tokenizer_file_id', dest_path='./tokenizer')

tokenizer = RagTokenizer.from_pretrained('./tokenizer')
retriever = RagRetriever.from_pretrained('./model.pt')
model = RagSequenceForGeneration.from_pretrained('./model.pt')

# Load the Mask R-CNN model
model_rcnn = maskrcnn_resnet50_fpn(pretrained=True)
model_rcnn.eval()

# Define the class labels for COCO dataset
class_labels = [
    '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign',
    'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
    'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A',
    'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
    'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
    'bottle', 'N/A', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
    'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
    'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table',
    'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard',
    'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book',
    'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]

# Define the image-to-text object segmentation function
def image_to_text_segmentation(image):
    # Convert the image to the expected format (RGB and tensor)
    image = T.ToTensor()(image)
    image = image.unsqueeze(0)

    # Run the image through the Mask R-CNN model
    with torch.no_grad():
        predictions = model_rcnn(image)

    # Extract the bounding boxes, labels, and masks from the predictions
    boxes = predictions[0]['boxes'].tolist()
    labels = [class_labels[i] for i in predictions[0]['labels'].tolist()]
    masks = predictions[0]['masks'].squeeze().detach().cpu().numpy()

    # Generate the segmented text for each object
    segmented_text = []
    for i in range(len(boxes)):
        mask = masks[i]
        object_text = ""
        for j in range(mask.shape[0]):
            for k in range(mask.shape[1]):
                if mask[j][k]:
                    object_text += labels[i] + " "
        segmented_text.append(object_text.strip())

    return segmented_text

# Define the Gradio interface
input_image = gr.inputs.Image(label="Input Image")
input_text = gr.inputs.Textbox(label="Question")
output_text = gr.outputs.Textbox(label="Generated Text")

title = "RAG Text Generation and Object Segmentation"
description = "Generate text based on the given question using RAG model and perform object segmentation on the input image."

gr.Interface(
    fn=generate_text,
    inputs=input_text,
    outputs=output_text,
    title=title,
    description=description,
    examples=[
        ["What is the capital of France?"],
        ["Who is the president of the United States?"],
    ]
).launch()

gr.Interface(
    fn=image_to_text_segmentation,
    inputs=input_image,
    outputs=output_text,
    title="Image-to-Text Object Segmentation",
    description="Segment objects in the image and generate corresponding text.",
).launch()