Update app.py

app.py

@@ -1,10 +1,10 @@
 import spaces
 import gradio as gr
 import numpy as np
-import os
 import torch
 from transformers import SamModel, SamProcessor
 from PIL import Image
+import os
 import cv2
 import argparse
 import sys
@@ -25,11 +25,6 @@ color_rgbs = [
         (1.0, 0.0, 1.0),
     ]
 
-mask_list = []
-mask_raw_list = []
-mask_list_video = []
-mask_raw_list_video = []
-
 def extract_first_frame_from_video(video):
     cap = cv2.VideoCapture(video)
     success, frame = cap.read()
@@ -55,9 +50,7 @@ def add_contour(img, mask, color=(1., 1., 1.)):
     return img
 
 @spaces.GPU(duration=120)
-def generate_masks(image):
-    global mask_list
-    global mask_raw_list
+def generate_masks(image, mask_list, mask_raw_list):
     image['image'] = image['background'].convert('RGB')
     # del image['background'], image['composite']
     assert len(image['layers']) == 1, f"Expected 1 layer, got {len(image['layers'])}"
@@ -80,13 +73,10 @@ def generate_masks(image):
     # Return a list containing the mask image.
     image['layers'] = []
     image['composite'] = image['background']
-    return mask_list, image
-
+    return mask_list, image, mask_list, mask_raw_list
 
 @spaces.GPU(duration=120)
-def generate_masks_video(image):
-    global mask_list_video
-    global mask_raw_list_video
+def generate_masks_video(image, mask_list_video, mask_raw_list_video):
     image['image'] = image['background'].convert('RGB')
     # del image['background'], image['composite']
     assert len(image['layers']) == 1, f"Expected 1 layer, got {len(image['layers'])}"
@@ -109,7 +99,7 @@ def generate_masks_video(image):
     # Return a list containing the mask image.
     image['layers'] = []
     image['composite'] = image['background']
-    return mask_list_video, image
+    return mask_list_video, image, mask_list_video, mask_raw_list_video
 
 
 @spaces.GPU(duration=120)
@@ -152,13 +142,13 @@ def describe(image, mode, query, masks):
         img_with_contour_np = add_contour(img_np, mask_np, color=color_rgb)
         img_with_contour_pil = Image.fromarray((img_with_contour_np * 255.).astype(np.uint8))
     else:
-        masks = mask_raw_list
         img_with_contour_np = img_np.copy()
 
         mask_ids = []
         for i, mask_np in enumerate(masks):
-            img_with_contour_np = add_contour(img_with_contour_np, mask_np, color=color_rgbs[i])
-            img_with_contour_pil = Image.fromarray((img_with_contour_np * 255.).astype(np.uint8))
+            # img_with_contour_np = add_contour(img_with_contour_np, mask_np, color=color_rgbs[i])
+            # img_with_contour_pil = Image.fromarray((img_with_contour_np * 255.).astype(np.uint8))
+            img_with_contour_pil = Image.fromarray((img_with_contour_np* 255.).astype(np.uint8))
             mask_ids.append(0)
     
     masks = np.stack(masks, axis=0)
@@ -214,8 +204,7 @@ def load_first_frame(video_path):
     return image
 
 @spaces.GPU(duration=120)
-def describe_video(video_path, mode, query, annotated_frame, masks):
-    global mask_list_video
+def describe_video(video_path, mode, query, annotated_frame, masks, mask_list_video):
     # Create a temporary directory to save extracted video frames
     cap = cv2.VideoCapture(video_path)
 
@@ -267,7 +256,6 @@ def describe_video(video_path, mode, query, annotated_frame, masks):
 
 
     else:
-        masks = mask_raw_list_video
         img_with_contour_np = img_np.copy()
         
         mask_ids = []
@@ -306,7 +294,7 @@ def describe_video(video_path, mode, query, annotated_frame, masks):
         mask_image = Image.fromarray((mask_np[:,:,np.newaxis] * np.array(annotated_frame['image'])).astype(np.uint8))
         mask_list_video.append((mask_image, f"<object{len(mask_list_video)}>"))
     text = ""
-    yield frame_img, text, mask_list_video
+    yield frame_img, text, mask_list_video, mask_list_video
 
     for token in get_model_output(
         video_tensor,
@@ -319,7 +307,7 @@ def describe_video(video_path, mode, query, annotated_frame, masks):
         streaming=True,
     ):
         text += token
-        yield gr.update(), text, gr.update()
+        yield gr.update(), text, gr.update(), gr.update()
 
 
 @spaces.GPU(duration=120)
@@ -338,20 +326,9 @@ def apply_sam(image, input_points):
 
     return mask_np
 
-def clear_masks():
-    global mask_list
-    global mask_raw_list
-    mask_list = []
-    mask_raw_list = []
-    return []
-
 
-def clear_masks_video():
-    global mask_list_video
-    global mask_raw_list_video
-    mask_list_video = []
-    mask_raw_list_video = []
-    return []
+def clear_masks():
+    return [], [], []
 
 
 if __name__ == "__main__":
@@ -363,10 +340,15 @@ if __name__ == "__main__":
     parser.add_argument("--top_p", type=float, default=0.5, help="Top-p for sampling")
 
     args_cli = parser.parse_args()
-    print(args_cli.model_path)
 
     with gr.Blocks(theme=gr.themes.Soft(primary_hue="amber")) as demo:
 
+        mask_list = gr.State([])  
+        mask_raw_list = gr.State([])  
+        mask_list_video = gr.State([])  
+        mask_raw_list_video = gr.State([])  
+
+
         HEADER = ("""
             <div>
                 <h1>VideoRefer X VideoLLaMA3 Demo</h1>
@@ -479,75 +461,67 @@ if __name__ == "__main__":
         def toggle_query_and_generate_button(mode):
             query_visible = mode == "QA"
             caption_visible = mode == "Caption"
-            global mask_list
-            global mask_raw_list
-            mask_list = []
-            mask_raw_list = []
-            return gr.update(visible=query_visible), gr.update(visible=query_visible), gr.update(visible=query_visible), gr.update(visible=query_visible), gr.update(visible=query_visible), gr.update(visible=caption_visible), gr.update(visible=caption_visible), [], ""
+            return gr.update(visible=query_visible), gr.update(visible=query_visible), gr.update(visible=query_visible), gr.update(visible=query_visible), gr.update(visible=query_visible), gr.update(visible=caption_visible), gr.update(visible=caption_visible), [], "", [], [],[],[]
 
         video_input.change(load_first_frame, inputs=video_input, outputs=first_frame)
 
-        mode.change(toggle_query_and_generate_button, inputs=mode, outputs=[query, generate_mask_btn, clear_masks_btn, submit_btn1, mask_output, output_image, submit_btn, mask_output, description])
+        mode.change(toggle_query_and_generate_button, inputs=mode, outputs=[query, generate_mask_btn, clear_masks_btn, submit_btn1, mask_output, output_image, submit_btn, mask_output, description, mask_list, mask_raw_list, mask_list_video, mask_raw_list_video])
         
         def toggle_query_and_generate_button_video(mode):
             query_visible = mode == "QA"
             caption_visible = mode == "Caption"
-            global mask_list_video
-            global mask_raw_list_video
-            mask_list_video = []
-            mask_raw_list_video = []
-            return gr.update(visible=query_visible), gr.update(visible=query_visible), gr.update(visible=query_visible), gr.update(visible=caption_visible), []
+            return gr.update(visible=query_visible), gr.update(visible=query_visible), gr.update(visible=query_visible), gr.update(visible=caption_visible), [], [], [], [], []
 
 
-        mode_video.change(toggle_query_and_generate_button_video, inputs=mode_video, outputs=[query_video, generate_mask_btn_video, submit_btn_video1, submit_btn_video, mask_output_video])
+        mode_video.change(toggle_query_and_generate_button_video, inputs=mode_video, outputs=[query_video, generate_mask_btn_video, submit_btn_video1, submit_btn_video, mask_output_video, mask_list, mask_raw_list, mask_list_video, mask_raw_list_video])
 
         submit_btn.click(
             fn=describe,
-            inputs=[image_input, mode, query],
+            inputs=[image_input, mode, query, mask_raw_list],
             outputs=[output_image, description, image_input],
             api_name="describe"
         )
 
         submit_btn1.click(
             fn=describe,
-            inputs=[image_input, mode, query],
+            inputs=[image_input, mode, query, mask_raw_list],
             outputs=[output_image, description, image_input],
             api_name="describe"
         )
 
         generate_mask_btn.click(
             fn=generate_masks,
-            inputs=[image_input],
-            outputs=[mask_output, image_input]
+            inputs=[image_input, mask_list, mask_raw_list],
+            outputs=[mask_output, image_input, mask_list, mask_raw_list]
         )
 
         generate_mask_btn_video.click(
             fn=generate_masks_video,
-            inputs=[first_frame],
-            outputs=[mask_output_video, first_frame]
+            inputs=[first_frame, mask_list_video, mask_raw_list_video],
+            outputs=[mask_output_video, first_frame, mask_list_video, mask_raw_list_video]
         )
 
         clear_masks_btn.click(
             fn=clear_masks,
-            outputs=[mask_output]
+            outputs=[mask_output, mask_list, mask_raw_list]
         )
 
         clear_masks_btn_video.click(
-            fn=clear_masks_video,
-            outputs=[mask_output_video]
+            fn=clear_masks,
+            outputs=[mask_output_video, mask_list_video, mask_raw_list_video]
         )
 
         submit_btn_video.click(
             fn=describe_video,
-            inputs=[video_input, mode_video, query_video, first_frame],
-            outputs=[first_frame, description_video, mask_output_video],
+            inputs=[video_input, mode_video, query_video, first_frame, mask_raw_list_video, mask_list_video],
+            outputs=[first_frame, description_video, mask_output_video, mask_list_video],
             api_name="describe_video"
         )
 
         submit_btn_video1.click(
             fn=describe_video,
-            inputs=[video_input, mode_video, query_video, first_frame],
-            outputs=[first_frame, description_video, mask_output_video],
+            inputs=[video_input, mode_video, query_video, first_frame, mask_raw_list_video, mask_list_video],
+            outputs=[first_frame, description_video, mask_output_video, mask_list_video],
             api_name="describe_video"
         )