update cap

app.py

@@ -248,7 +248,7 @@ def gradio_ask(user_message, chatbot, chat_state,radio):
 
 
 def gradio_answer(chatbot, chat_state, img_list,  radio, text,num_beams, temperature):
-    image == None
+    image = None
     llm_message,image = \
     chat.answer(conv=chat_state, img_list=img_list, max_new_tokens=300, num_beams=1, temperature=temperature,
                 max_length=2000,radio = radio,text_input = text)

multimodal/open_flamingo/chat/conversation.py

@@ -19,6 +19,7 @@ import gradio as gr
 from huggingface_hub import hf_hub_download, login
 
 from open_flamingo.src.factory import create_model_and_transforms
+from open_flamingo.eval.task.caption import captioner
 
 class SeparatorStyle(Enum):
     """Different separator style."""
@@ -403,56 +404,59 @@ class Chat:
         image_start_index_list = [[x] for x in image_start_index_list]
         image_nums = [1] * len(input_ids)
         added_bbox_list = []
-        with torch.inference_mode():
-            text_outputs = self.model.generate(
-                batch_images,
-                input_ids,
-                attention_mask=attention_mask,
-                max_new_tokens=20,
-                # min_new_tokens=8,
-                num_beams=1,
-                # length_penalty=0,
-                image_start_index_list=image_start_index_list,
-                image_nums=image_nums,
-                added_bbox_list=added_bbox_list if len(added_bbox_list) != 0 else None,
-            )
-        # and torch.cuda.amp.autocast(dtype=torch.float16)
-        with torch.no_grad():
-            outputs = self.model(
-                vision_x=batch_images,
-                lang_x=input_ids,
-                attention_mask=attention_mask,
-                image_nums=image_nums,
-                image_start_index_list=image_start_index_list,
-                added_bbox_list=None,
-                add_box=False,
-            )
-        boxes = outputs["boxes"]
-        scores = outputs["scores"]
-        if len(scores) > 0:
-            box = boxes[scores.argmax()] / 224
-        print(f"{box}")
-        out_image = None
-
-        if len(boxes)>0:
-            width, height = image_ori.size
-            open_cv_image = np.array(image_ori)
-            # Convert RGB to BGR
-            open_cv_image = open_cv_image[:, :, ::-1].copy()
-            box = box * [width, height, width, height]
-            # for box in boxes:
-            open_cv_image = cv2.rectangle(open_cv_image, box[:2].astype(int), box[2:].astype(int), (255, 0, 0), 2)
-            out_image = Image.fromarray(cv2.cvtColor(open_cv_image, cv2.COLOR_BGR2RGB))
-            
-            
-        # output_token = outputs[0, input_ids.shape[1]:]
-        # output_text = tokenizer.decode(output_token, skip_special_tokens=True).strip()
-        # conv[-1]["value"] = output_text
-        # # conv.messages[-1][1] = output_text
-        # print(
-        #     f"### Assistant: {tokenizer.decode(outputs[0, input_ids.shape[1]:], skip_special_tokens=True).strip()}")
-        output_text = self.tokenizer.decode(text_outputs[0])
-        output_text = re.findall(r'Assistant:(.+)', output_text)[-1]
+        if radio in ["Cap"]:
+            output_text, out_image = captioner(self.model,self.tokenizer,image_ori,batch_images,input_ids,attention_mask,image_start_index_list,image_nums,added_bbox_list)
+        else:
+            with torch.inference_mode():
+                text_outputs = self.model.generate(
+                    batch_images,
+                    input_ids,
+                    attention_mask=attention_mask,
+                    max_new_tokens=20,
+                    # min_new_tokens=8,
+                    num_beams=1,
+                    # length_penalty=0,
+                    image_start_index_list=image_start_index_list,
+                    image_nums=image_nums,
+                    added_bbox_list=added_bbox_list if len(added_bbox_list) != 0 else None,
+                )
+            # and torch.cuda.amp.autocast(dtype=torch.float16)
+            with torch.no_grad():
+                outputs = self.model(
+                    vision_x=batch_images,
+                    lang_x=input_ids,
+                    attention_mask=attention_mask,
+                    image_nums=image_nums,
+                    image_start_index_list=image_start_index_list,
+                    added_bbox_list=None,
+                    add_box=False,
+                )
+            boxes = outputs["boxes"]
+            scores = outputs["scores"]
+            if len(scores) > 0:
+                box = boxes[scores.argmax()] / 224
+            print(f"{box}")
+            out_image = None
+
+            if len(boxes)>0:
+                width, height = image_ori.size
+                open_cv_image = np.array(image_ori)
+                # Convert RGB to BGR
+                open_cv_image = open_cv_image[:, :, ::-1].copy()
+                box = box * [width, height, width, height]
+                # for box in boxes:
+                open_cv_image = cv2.rectangle(open_cv_image, box[:2].astype(int), box[2:].astype(int), (255, 0, 0), 2)
+                out_image = Image.fromarray(cv2.cvtColor(open_cv_image, cv2.COLOR_BGR2RGB))
+
+
+            # output_token = outputs[0, input_ids.shape[1]:]
+            # output_text = tokenizer.decode(output_token, skip_special_tokens=True).strip()
+            # conv[-1]["value"] = output_text
+            # # conv.messages[-1][1] = output_text
+            # print(
+            #     f"### Assistant: {tokenizer.decode(outputs[0, input_ids.shape[1]:], skip_special_tokens=True).strip()}")
+            output_text = self.tokenizer.decode(text_outputs[0])
+            output_text = re.findall(r'Assistant:(.+)', output_text)[-1]
 
         return output_text, out_image
 

multimodal/open_flamingo/eval/task/caption.py

@@ -7,7 +7,7 @@ import json
 import time
 import os
 from transformers import LogitsProcessor, MinNewTokensLengthLogitsProcessor, ForcedEOSTokenLogitsProcessor
-
+from PIL import Image
 
 class VisualLogitsProcessor(LogitsProcessor):
     def __init__(self, tokenizer):
@@ -51,6 +51,136 @@ def prepare_batch_images(batch, image_processor):
     return batch_images
 
 
+def captioner(
+    model,tokenizer,image_ori,batch_images,input_ids,attention_mask,image_start_index_list,image_nums,added_bbox_list,debug=False):
+    """Evaluate a model on COCO dataset.
+    Returns:
+        float: CIDEr score
+
+    """
+    visual_logits_processor = VisualLogitsProcessor(tokenizer)
+    model.eval()
+    # model.eval().cuda()
+    lang_encoder_name = model.lang_encoder.__class__.__name__.lower()
+    media_token_id = tokenizer("<|#image#|>", add_special_tokens=False)["input_ids"][-1]
+    endofmedia_token_id = tokenizer("<|#endofimage#|>", add_special_tokens=False)["input_ids"][-1]
+    pad_token_id = tokenizer(tokenizer.pad_token, add_special_tokens=False)["input_ids"][-1]
+    bos_token_id = tokenizer(tokenizer.bos_token, add_special_tokens=False)["input_ids"][-1]
+    previsual_token_id = tokenizer("<|#previsual#|>", add_special_tokens=False)["input_ids"][-1]
+    visual_token_id = tokenizer("<|#visual#|>", add_special_tokens=False)["input_ids"][-1]
+    box_token = "<|#box#|>"
+    prebox_token = "<|#prebox#|>"
+    endofobject_token = "<|#endofobject#|>"
+    object_token = "<|#object#|>"
+    ori_prompt_length = len(input_ids[0])
+    have_prebox = False
+    while True:
+        batch_images = batch_images
+        input_ids = input_ids
+        attention_mask = attention_mask
+        image_start_index_list = image_start_index_list
+        image_nums = image_nums
+        if debug:
+            print("input--->",tokenizer.decode(input_ids[0]))
+        p1 = MinNewTokensLengthLogitsProcessor(
+            prompt_length_to_skip=input_ids.shape[-1],
+            min_new_tokens=5,
+            eos_token_id=bos_token_id,
+        )
+        with torch.inference_mode():
+            outputs = model.generate(
+                batch_images,
+                input_ids,
+                attention_mask=attention_mask,
+                max_new_tokens=20,
+                # min_new_tokens=8,
+                num_beams=1,
+                # length_penalty=0,
+                image_start_index_list=image_start_index_list,
+                image_nums=image_nums,
+                added_bbox_list=added_bbox_list if len(added_bbox_list) != 0 else None,
+                logits_processor_list=[p1, visual_logits_processor],
+            )
+        if debug:
+            print("outputs--->",tokenizer.decode(outputs[0]))
+        if outputs[0, -2] in [previsual_token_id, visual_token_id] and outputs[0, -1] == bos_token_id:
+            prompt = tokenizer.decode(outputs.clone()[0])
+            is_visual = (outputs[0, -2] == visual_token_id)
+            batch_text = tokenizer.batch_decode(outputs[:, :-1])
+            encodings = tokenizer(
+                batch_text,
+                padding="longest",
+                truncation=True,
+                return_tensors="pt",
+                max_length=2000,
+            )
+            input_ids = encodings["input_ids"]
+            attention_mask = encodings["attention_mask"]
+            image_start_index_list = ((input_ids == media_token_id).nonzero(as_tuple=True)[-1] + 1).tolist()
+            image_start_index_list = [[x] for x in image_start_index_list]
+            image_nums = [1] * len(input_ids)
+            if debug:
+                print("get the visual bbox--->",tokenizer.decode(input_ids[0]))
+            with torch.no_grad():
+                outputs = model(
+                    vision_x=batch_images,
+                    lang_x=input_ids,
+                    attention_mask=attention_mask,
+                    image_nums=image_nums,
+                    image_start_index_list=image_start_index_list,
+                    added_bbox_list=added_bbox_list if len(added_bbox_list) != 0 else None,
+                    add_box=added_bbox_list is not None and len(added_bbox_list) != 0,
+                )
+            boxes = outputs["boxes"]
+            scores = outputs["scores"]
+            # if not model.valid:
+            #     import pdb; pdb.set_trace()
+            if boxes is not None:
+                if is_visual:
+                    if have_prebox:
+                        added_bbox_list.pop()
+                        prompt = prompt.replace("<|#previsual#|><|#prebox#|><|#object#|>", "")
+                        have_prebox = False
+                        if debug:
+                            print("find previsual and remove it--->", prompt)
+                    first_box = boxes[scores.argmax()]
+                    added_bbox_list += [torch.tensor(first_box).unsqueeze(0) / 224]
+                    prompt = prompt[:-len(tokenizer.eos_token)]
+                    prompt += box_token + endofobject_token
+                    if debug:
+                        print("after inserting visual---->", prompt)
+                else:
+                    import numpy as np
+                    import cv2
+                    open_cv_image = np.array(image_ori)
+                    open_cv_image = open_cv_image[:, :, ::-1].copy()
+                    for i, pre_box in enumerate(boxes):
+                        open_cv_image = cv2.rectangle(open_cv_image, pre_box[:2].astype(int), pre_box[2:].astype(int), (0, 255, 0), i+1)
+                    out_image = Image.fromarray(cv2.cvtColor(open_cv_image, cv2.COLOR_BGR2RGB))
+                    # exit()
+                    pre_box = boxes[scores.argmax()]
+                    added_bbox_list += [torch.tensor(pre_box).unsqueeze(0).cuda() / 224]
+                    prompt = prompt[:-len(tokenizer.eos_token)]
+                    prompt += prebox_token + object_token
+                    have_prebox = True
+                    if debug:
+                        print("after inserting previsual---->", prompt)
+            else:
+                if debug:
+                    import pdb;pdb.set_trace()
+                prompt = tokenizer.decode(outputs[0, :-2].clone()[0])
+        else:
+            break
+    outputs = outputs[:, ori_prompt_length:]
+    outputs = postprocess_captioning_generation(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]).replace('"', "")
+    # new_predictions = [
+    #     postprocess_captioning_generation(out).replace('"', "")
+    #     for out in tokenizer.batch_decode(outputs, skip_special_tokens=True)
+    # ]
+        # import pdb; pdb.set_trace()
+    return outputs, out_image
+
+
 def evaluate_coco_flickr(
     model,
     tokenizer,
@@ -94,6 +224,7 @@ def evaluate_coco_flickr(
         if ii % world_size != rank:
             continue
         cnt += len(batch)
+        batch[0]["image"] = Image.open("/gpfs/u/home/LMCG/LMCGljnn/scratch/images/img3.jpg").resize((224, 224))
         batch_images = prepare_batch_images(
             batch=batch,
             image_processor=image_processor,
@@ -194,13 +325,14 @@ def evaluate_coco_flickr(
                         if debug:
                             print("after inserting visual---->", prompt)
                     else:
-                        # import numpy as np
-                        # import cv2
-                        # open_cv_image = np.array(batch[0]["image"])
-                        # open_cv_image = open_cv_image[:, :, ::-1].copy()
-                        # for pre_box in boxes:
-                        #     open_cv_image = cv2.rectangle(open_cv_image, pre_box[:2].astype(int), pre_box[2:].astype(int), (0, 255, 0), 2)
-                        # cv2.imwrite("Atest.png", open_cv_image)
+                        import numpy as np
+                        import cv2
+                        open_cv_image = np.array(batch[0]["image"])
+                        open_cv_image = open_cv_image[:, :, ::-1].copy()
+                        for i, pre_box in enumerate(boxes):
+                            open_cv_image = cv2.rectangle(open_cv_image, pre_box[:2].astype(int), pre_box[2:].astype(int), (0, 255, 0), i+1)
+                        cv2.imwrite("Atest.png", open_cv_image)
+                        exit()
                         pre_box = boxes[scores.argmax()]
                         added_bbox_list += [torch.tensor(pre_box).unsqueeze(0).cuda() / 224]
                         prompt = prompt[:-len(tokenizer.eos_token)]
@@ -225,6 +357,8 @@ def evaluate_coco_flickr(
             predictions[int(sample["image_id"])] = {
                 "caption": new_predictions[i],
             }
+        print(new_predictions)
+        exit()
     results_path = (
         f"flickrresults_{lang_encoder_name}_{rank}_{id}.json"
         if is_flickr