Create app.py

app.py

@@ -0,0 +1,90 @@
+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()