Spaces:
Sleeping
Sleeping
| 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() |