Update app.py

app.py

@@ -1,65 +1,73 @@
-from PIL import Image
-import requests
+import ruamel_yaml as yaml
+import numpy as np
+import random
+
 import torch
-from torchvision import transforms
-from torchvision.transforms.functional import InterpolationMode
+import torchvision.transforms as transforms
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+from PIL import Image
+from models.tag2text import tag2text_caption
 
+import gradio as gr
 
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
+image_size = 384
 
 
-import gradio as gr
+normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                std=[0.229, 0.224, 0.225])
+transform = transforms.Compose([transforms.Resize((image_size, image_size)),transforms.ToTensor(),normalize])
 
-# from models.blip import blip_decoder
-from transformers import BlipProcessor, BlipForConditionalGeneration
 
-model_id = "Salesforce/blip-image-captioning-base"
+#######Swin Version
+pretrained = '/home/notebook/code/personal/S9049611/BLIP/output/blip_tagtotext_14m/blip_tagtotext_encoderdiv_tar_random_swin/caption_coco_finetune_tagparse_tagfinetune_threshold075_bceloss_tagsingle_5e6_epoch19_negative_1_05_pos_1_10/checkpoint_05.pth'
 
-model = BlipForConditionalGeneration.from_pretrained(model_id)
+config_file = 'configs/tag2text_caption.yaml'
+config = yaml.load(open(config_file, 'r'), Loader=yaml.Loader)
 
 
-image_size = 384
-transform = transforms.Compose([
-    transforms.Resize((image_size,image_size),interpolation=InterpolationMode.BICUBIC),
-    transforms.ToTensor(),
-    transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
-    ]) 
+model = tag2text_caption(pretrained=pretrained, image_size=image_size, vit=config['vit'], 
+                    vit_grad_ckpt=config['vit_grad_ckpt'], vit_ckpt_layer=config['vit_ckpt_layer'],
+                    prompt=config['prompt'],config=config,threshold = 0.75 )
 
-# model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_large_caption.pth'
-    
-# model = blip_decoder(pretrained=model_url, image_size=384, vit='large')
 model.eval()
 model = model.to(device)
 
 
-# from models.blip_vqa import blip_vqa
-
-# image_size_vq = 480
-# transform_vq = transforms.Compose([
-#     transforms.Resize((image_size_vq,image_size_vq),interpolation=InterpolationMode.BICUBIC),
-#     transforms.ToTensor(),
-#     transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
-#     ]) 
-
-# model_url_vq = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model*_vqa.pth'
-    
-# model_vq = blip_vqa(pretrained=model_url_vq, image_size=480, vit='base')
-# model_vq.eval()
-# model_vq = model_vq.to(device)
-
-
-
-def inference(raw_image, model_n, question, strategy):
+def inference(raw_image, model_n, input_tag, strategy):
     if model_n == 'Image Captioning':
+        raw_image = raw_image.resize((image_size, image_size))
         image = transform(raw_image).unsqueeze(0).to(device)   
+        model.threshold = 0.7
+        if input_tag == '' or input_tag == 'none' or input_tag == 'None':
+            input_tag_list = None
+        else:
+            input_tag_list = []
+            input_tag_list.append(input_tag.replace(',',' | '))
         with torch.no_grad():
-          if strategy == "Beam search":
-              caption = model.generate(image, sample=False, num_beams=3, max_length=20, min_length=5)
-          else:
-              caption = model.generate(image, sample=True, top_p=0.9, max_length=20, min_length=5)
-          return 'caption: '+caption[0]
+            if strategy == "Beam search":
+                
+
+                caption, tag_predict = model.generate(image,tag_input = input_tag_list, return_tag_predict = True)
+                if input_tag_list == None:
+                    tag_1 = tag_predict
+                    tag_2 = ['none']
+                else:
+                    _, tag_1 = model.generate(image,tag_input = None, return_tag_predict = True)
+                    tag_2 = tag_predict
+
+            else:
+
+                caption,tag_predict = model.generate(image,  tag_input = input_tag_list,sample=True, top_p=0.9, max_length=20, min_length=5, return_tag_predict = True)
+                if input_tag_list == None:
+                    tag_1 = tag_predict
+                    tag_2 = ['none']
+                else:
+                    _, tag_1 = model.generate(image,tag_input = None, return_tag_predict = True)
+                    tag_2 = tag_predict
+            return tag_1[0],tag_2[0],caption[0]
+            
 
     else:   
         image_vq = transform_vq(raw_image).unsqueeze(0).to(device)  
@@ -67,16 +75,26 @@ def inference(raw_image, model_n, question, strategy):
             answer = model_vq(image_vq, question, train=False, inference='generate') 
         return  'answer: '+answer[0]
     
-# inputs = [gr.inputs.Image(type='pil'),gr.inputs.Radio(choices=['Image Captioning',"Visual Question Answering"], type="value", default="Image Captioning", label="Task"),gr.inputs.Textbox(lines=2, label="Question"),gr.inputs.Radio(choices=['Beam search','Nucleus sampling'], type="value", default="Nucleus sampling", label="Caption Decoding Strategy")]
-inputs = [gr.inputs.Image(type='pil'),gr.inputs.Radio(choices=['Image Captioning'], type="value", default="Image Captioning", label="Task"),gr.inputs.Textbox(lines=2, label="Question"),gr.inputs.Radio(choices=['Beam search','Nucleus sampling'], type="value", default="Nucleus sampling", label="Caption Decoding Strategy")]
+inputs = [gr.inputs.Image(type='pil'),gr.inputs.Radio(choices=['Image Captioning'], type="value", default="Image Captioning", label="Task"),gr.inputs.Textbox(lines=2, label="User Identified Tags (Optional, Enter with commas)"),gr.inputs.Radio(choices=['Beam search','Nucleus sampling'], type="value", default="Beam search", label="Caption Decoding Strategy")]
+
+outputs = [gr.outputs.Textbox(label="Model Identified Tags"),gr.outputs.Textbox(label="User Identified Tags"), gr.outputs.Textbox(label="Image Caption") ]
+
+title = "Tag2Text"
+
+description = "Gradio demo for Tag2Text: Guiding Language-Image Model via Image Tagging (Fudan University, OPPO Research Institute, International Digital Economy Academy)."
 
-outputs = gr.outputs.Textbox(label="Output")
+article = "<p style='text-align: center'><a href='' target='_blank'>Tag2Text: Guiding Language-Image Model via Image Tagging</a> | <a href='' target='_blank'>Github Repo</a></p>"
 
-title = "BLIP"
+demo = gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, examples=[['images/COCO_val2014_000000551338.jpg',"Image Captioning","none","Beam search"], 
+                                                                                                                ['images/COCO_val2014_000000551338.jpg',"Image Captioning","fence, sky","Beam search"],
+                                                                                                                # ['images/COCO_val2014_000000551338.jpg',"Image Captioning","grass","Beam search"],
+                                                                                                                 ['images/COCO_val2014_000000483108.jpg',"Image Captioning","none","Beam search"],
+                                                                                                                 ['images/COCO_val2014_000000483108.jpg',"Image Captioning","electric cable","Beam search"],
+                                                                                                                  # ['images/COCO_val2014_000000483108.jpg',"Image Captioning","sky, train","Beam search"],
+                                                                                                                 ['images/COCO_val2014_000000483108.jpg',"Image Captioning","track, train","Beam search"] ,    
+                                                                                                                 ['images/COCO_val2014_000000483108.jpg',"Image Captioning","grass","Beam search"]     
+                                                                                                                ])
 
-description = "Gradio demo for BLIP: Bootstrapping Language-Image Pre-training for Unified Vision-Language Understanding and Generation (Salesforce Research). To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
 
-article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2201.12086' target='_blank'>BLIP: Bootstrapping Language-Image Pre-training for Unified Vision-Language Understanding and Generation</a> | <a href='https://github.com/salesforce/BLIP' target='_blank'>Github Repo</a></p>"
 
 
-gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, examples=[['starrynight.jpeg',"Image Captioning","None","Nucleus sampling"]]).launch(enable_queue=True)