update app.py

app.py

@@ -6,7 +6,7 @@ from PIL import Image
 
 from open_flamingo.train.distributed import init_distributed_device, world_info_from_env
 import string
-
+import cv2
 
 
 import gradio as gr
@@ -44,14 +44,13 @@ def generate(
     idx,
     image,
     text,
-    tsvfile,
     vis_embed_size=256,
     rank=0,
     world_size=1,
 ):
     if image is None:
         raise gr.Error("Please upload an image.")
-    flamingo.eval().cuda()
+    flamingo.eval()
     loc_token_ids = []
     for i in range(1000):
         loc_token_ids.append(int(tokenizer(f"<loc_{i}>", add_special_tokens=False)["input_ids"][-1]))
@@ -70,7 +69,12 @@ def generate(
     height = image.height
     image = image.resize((224, 224))
     batch_images = image_processor(image).unsqueeze(0).unsqueeze(1).unsqueeze(0)
-    prompt = [f"<|#image#|>{tokenizer.pad_token*vis_embed_size}<|#endofimage#|><|#obj#|>{text.rstrip('.')}"]
+    if idx ==1:
+        prompt = [f"<|#image#|>{tokenizer.pad_token*vis_embed_size}<|#endofimage#|><|#obj#|>{text.rstrip('.')}"]
+        bad_words_ids = bad_words_ids
+    else:
+        prompt = [f"<|#image#|>{tokenizer.pad_token * vis_embed_size}<|#endofimage#|>{text.rstrip('.')}"]
+        bad_words_ids = None
     encodings = tokenizer(
         prompt,
         padding="longest",
@@ -85,13 +89,13 @@ def generate(
     image_nums = [1] * len(input_ids)
     outputs = get_outputs(
         model=flamingo,
-        batch_images=batch_images.cuda(),
-        attention_mask=attention_mask.cuda(),
+        batch_images=batch_images,
+        attention_mask=attention_mask,
         max_generation_length=5,
         min_generation_length=4,
         num_beams=1,
         length_penalty=1.0,
-        input_ids=input_ids.cuda(),
+        input_ids=input_ids,
         bad_words_ids=bad_words_ids,
         image_start_index_list=image_start_index_list,
         image_nums=image_nums,
@@ -106,12 +110,23 @@ def generate(
     #     tqdm.write(f"output: {tokenizer.batch_decode(outputs)}")
     #     tqdm.write(f"prompt: {prompt}")
 
+    if len(box) == 4:
+        img = cv2.cvtColor(np.array(image_ori), cv2.COLOR_RGB2BGR)
+        out = cv2.rectangle(img, (int(box[0] * width / 1000), int(box[1] * height / 1000)),
+                            (int(box[2] * width / 1000), int(box[3] * height / 1000)), color=(255, 0, 255), thickness=2)
+        out = cv2.cvtColor(out, cv2.COLOR_BGR2RGB)
+        out_image = Image.fromarray(out)
+    # else:
+    #     tqdm.write(f"output: {tokenizer.batch_decode(outputs)}")
+    #     tqdm.write(f"prompt: {prompt}")
+
     gen_text = tokenizer.batch_decode(outputs)
-    return (
-        f"Output:{gen_text}"
-        if idx != 2
-        else f"Question: {text.strip()} Answer: {gen_text}"
-    )
+    if idx == 1:
+        return f"Output:{gen_text}", out_image
+    elif idx == 2:
+        return (f"Question: {text.strip()} Answer: {gen_text}")
+    else:
+        return (f"Output:{gen_text}")
 
 
 with gr.Blocks() as demo: