Update app.py

app.py

@@ -52,15 +52,20 @@ def add_contour(img, mask, color=(1., 1., 1.)):
 @spaces.GPU(duration=120)
 def generate_masks(image, mask_list, mask_raw_list):
     """
-    Generate masks from user-drawn annotations on an image.
-    
+    Generates segmentation masks for selected regions in an image using SAM.
+
     Args:
-        image: Dictionary containing the image editor state with background and layers
-        mask_list: List of generated mask images with labels
-        mask_raw_list: List of raw numpy arrays of masks
-        
+        image (dict): A dictionary containing image data, typically from a Gradio ImageEditor,
+                      with 'background' (PIL Image) and 'layers' (list of PIL Image layers).
+        mask_list (list): A list to accumulate (mask_image, label) tuples for display in a gallery.
+        mask_raw_list (list): A list to accumulate raw NumPy mask arrays.
+
     Returns:
-        Tuple containing updated mask_list, image editor state, mask_list, and mask_raw_list
+        tuple: A tuple containing:
+               - mask_list (list): Updated list of mask images for display.
+               - image (dict): Updated image dictionary with layers cleared.
+               - mask_list (list): Redundant return of mask_list (for Gradio update).
+               - mask_raw_list (list): Updated list of raw mask arrays.
     """
     image['image'] = image['background'].convert('RGB')
     # del image['background'], image['composite']
@@ -89,15 +94,21 @@ def generate_masks(image, mask_list, mask_raw_list):
 @spaces.GPU(duration=120)
 def generate_masks_video(image, mask_list_video, mask_raw_list_video):
     """
-    Generate masks from user-drawn annotations on a video frame.
-    
+    Generates segmentation masks for selected regions in the first frame of a video using SAM.
+
     Args:
-        image: Dictionary containing the image editor state with background and layers
-        mask_list_video: List of generated mask images with labels for video
-        mask_raw_list_video: List of raw numpy arrays of masks for video
-        
+        image (dict): A dictionary containing image data (first frame of video),
+                      typically from a Gradio ImageEditor, with 'background' (PIL Image)
+                      and 'layers' (list of PIL Image layers).
+        mask_list_video (list): A list to accumulate (mask_image, label) tuples for display.
+        mask_raw_list_video (list): A list to accumulate raw NumPy mask arrays for video processing.
+
     Returns:
-        Tuple containing updated mask_list_video, image editor state, mask_list_video, and mask_raw_list_video
+        tuple: A tuple containing:
+               - mask_list_video (list): Updated list of mask images for display.
+               - image (dict): Updated image dictionary with layers cleared.
+               - mask_list_video (list): Redundant return of mask_list_video (for Gradio update).
+               - mask_raw_list_video (list): Updated list of raw mask arrays.
     """
     image['image'] = image['background'].convert('RGB')
     # del image['background'], image['composite']
@@ -127,16 +138,19 @@ def generate_masks_video(image, mask_list_video, mask_raw_list_video):
 @spaces.GPU(duration=120)
 def describe(image, mode, query, masks):
     """
-    Generate descriptions or answer questions about regions in an image.
-    
+    Describes an image based on selected regions or answers a question about them.
+
     Args:
-        image: Dictionary containing the image editor state
-        mode: Either "Caption" or "QA" mode
-        query: Question to ask about the image (used in QA mode)
-        masks: List of mask arrays for the regions
-        
-    Returns:
-        Generator yielding image with contours, generated text, and updated image state
+        image (dict): A dictionary containing image data, typically from a Gradio ImageEditor,
+                      with 'background' (PIL Image) and 'layers' (list of PIL Image layers).
+        mode (str): The operational mode, either "Caption" (to describe a selected region)
+                    or "QA" (to answer a question about one or more regions).
+        query (str): The question to ask in "QA" mode. Ignored in "Caption" mode.
+        masks (list): A list of raw NumPy mask arrays representing previously generated masks.
+
+    Yields:
+        tuple: An image with contours and the generated text description/answer,
+               or updates for Gradio components during streaming.
     """
     # Create an image object from the uploaded image
     # print(image.keys())
@@ -229,13 +243,16 @@ def describe(image, mode, query, masks):
   
 def load_first_frame(video_path):
     """
-    Load and return the first frame of a video.
-    
+    Loads the first frame of a given video file.
+
     Args:
-        video_path: Path to the video file
-        
+        video_path (str): The file path to the video.
+
     Returns:
-        PIL Image of the first frame
+        PIL.Image.Image: The first frame of the video as a PIL Image.
+
+    Raises:
+        gr.Error: If the video file cannot be read.
     """
     cap = cv2.VideoCapture(video_path)
     ret, frame = cap.read()
@@ -249,18 +266,23 @@ def load_first_frame(video_path):
 @spaces.GPU(duration=120)
 def describe_video(video_path, mode, query, annotated_frame, masks, mask_list_video):
     """
-    Generate descriptions or answer questions about regions in a video.
-    
+    Describes a video based on selected regions in its first frame or answers a question about them.
+
     Args:
-        video_path: Path to the video file
-        mode: Either "Caption" or "QA" mode
-        query: Question to ask about the video (used in QA mode)
-        annotated_frame: Dictionary containing the annotated first frame
-        masks: List of mask arrays for the regions
-        mask_list_video: List of mask images with labels
-        
-    Returns:
-        Generator yielding frame image, generated text, and updated mask lists
+        video_path (str): The file path to the video.
+        mode (str): The operational mode, either "Caption" (to describe a selected region)
+                    or "QA" (to answer a question about one or more regions).
+        query (str): The question to ask in "QA" mode. Ignored in "Caption" mode.
+        annotated_frame (dict): A dictionary containing the first frame's image data
+                                from a Gradio ImageEditor, with 'background' (PIL Image)
+                                and 'layers' (list of PIL Image layers).
+        masks (list): A list of raw NumPy mask arrays representing previously generated masks
+                      for objects in the video.
+        mask_list_video (list): A list to accumulate (mask_image, label) tuples for display.
+
+    Yields:
+        tuple: The annotated first frame, the generated text description/answer,
+               and updated mask lists for Gradio components during streaming.
     """
     # Create a temporary directory to save extracted video frames
     cap = cv2.VideoCapture(video_path)
@@ -351,4 +373,286 @@ def describe_video(video_path, mode, query, annotated_frame, masks, mask_list_vi
         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, mask_list_video
+    yield frame_img, text, mask_list_video, mask_list_video
+
+    for token in get_model_output(
+        video_tensor,
+        query,
+        model=model,
+        tokenizer=tokenizer,
+        masks=masks,
+        mask_ids=mask_ids,
+        modal='video',
+        streaming=True,
+    ):
+        text += token
+        yield gr.update(), text, gr.update(), gr.update()
+
+
+@spaces.GPU(duration=120)
+def apply_sam(image, input_points):
+    """
+    Applies the Segment Anything Model (SAM) to an image based on input points
+    to generate a segmentation mask.
+
+    Args:
+        image (PIL.Image.Image): The input image.
+        input_points (list): A list of lists, where each inner list contains
+                             [x, y] coordinates representing points used for segmentation.
+
+    Returns:
+        numpy.ndarray: The selected binary segmentation mask as a NumPy array (H, W).
+    """
+    inputs = sam_processor(image, input_points=input_points, return_tensors="pt").to(device)
+
+    with torch.no_grad():
+        outputs = sam_model(**inputs)
+
+    masks = sam_processor.image_processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu())[0][0]
+    scores = outputs.iou_scores[0, 0]
+
+    mask_selection_index = scores.argmax()
+
+    mask_np = masks[mask_selection_index].numpy()
+
+    return mask_np
+
+
+def clear_masks():
+    """
+    Clears the stored lists of masks and raw masks.
+
+    Returns:
+        tuple: Three empty lists, intended to reset Gradio components
+               displaying masks.
+    """
+    return [], [], []
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="VideoRefer gradio demo")
+    parser.add_argument("--model-path", type=str, default="DAMO-NLP-SG/VideoRefer-VideoLLaMA3-7B", help="Path to the model checkpoint")
+    parser.add_argument("--prompt-mode", type=str, default="focal_prompt", help="Prompt mode")
+    parser.add_argument("--conv-mode", type=str, default="v1", help="Conversation mode")
+    parser.add_argument("--temperature", type=float, default=0.2, help="Sampling temperature")
+    parser.add_argument("--top_p", type=float, default=0.5, help="Top-p for sampling")
+
+    args_cli = parser.parse_args()
+
+    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>
+                <h5 style="margin: 0;">Feel free to click on anything that grabs your interest!</h5>
+                <h5 style="margin: 0;">If this demo please you, please give us a star ⭐ on Github or 💖 on this space.</h5>
+            </div>
+            </div>
+            <div style="display: flex; justify-content: left; margin-top: 10px;">
+            <a href="https://arxiv.org/pdf/2501.00599"><img src="https://img.shields.io/badge/Arxiv-2501.00599-ECA8A7" style="margin-right: 5px;"></a>
+            <a href="https://github.com/DAMO-NLP-SG/VideoRefer"><img src='https://img.shields.io/badge/Github-VideoRefer-F7C97E' style="margin-right: 5px;"></a>
+            <a href="https://github.com/DAMO-NLP-SG/VideoLLaMA3"><img src='https://img.shields.io/badge/Github-VideoLLaMA3-9DC3E6' style="margin-right: 5px;"></a>
+            </div>
+            """)
+
+        with gr.Row():
+            with gr.Column():
+                gr.HTML(HEADER)
+    
+
+        image_tips = """
+                ### 💡 Tips:
+
+                🧸 Upload an image, and you can use the drawing tool✍️ to highlight the areas you're interested in.
+            
+                🔖 For single-object caption mode, simply select the area and click the 'Generate Caption' button to receive a caption for the object.
+                
+                🔔 In QA mode, you can generate multiple masks by clicking the 'Generate Mask' button multiple times. Afterward, use the corresponding object id to ask questions.
+                
+                📌 Click the button 'Clear Masks' to clear the current generated masks.
+                
+                """
+        
+        video_tips = """
+                ### 💡 Tips:
+                ⚠️ For video mode, we only support masking on the first frame in this demo.
+
+                🧸 Upload an video, and you can use the drawing tool✍️ to highlight the areas you're interested in the first frame.
+            
+                🔖 For single-object caption mode, simply select the area and click the 'Generate Caption' button to receive a caption for the object.
+                
+                🔔 In QA mode, you can generate multiple masks by clicking the 'Generate Mask' button multiple times. Afterward, use the corresponding object id to ask questions.
+                
+                📌 Click the button 'Clear Masks' to clear the current generated masks.
+                
+                """
+  
+
+        with gr.TabItem("Image"):
+            with gr.Row():
+                with gr.Column():
+                    image_input = gr.ImageEditor(
+                        label="Image",
+                        type="pil", 
+                        sources=['upload'], 
+                        brush=gr.Brush(colors=["#ED7D31"], color_mode="fixed", default_size=10),
+                        eraser=True,
+                        layers=False,
+                        transforms=[],
+                        height=300,
+                    )
+                    generate_mask_btn = gr.Button("1️⃣ Generate Mask", visible=False, variant="primary")
+                    mode = gr.Radio(label="Mode", choices=["Caption", "QA"], value="Caption")
+                    query = gr.Textbox(label="Question", value="What is the relationship between <region0> and <region1>?", interactive=True, visible=False)
+                    
+                    submit_btn = gr.Button("Generate Caption", variant="primary")
+                    submit_btn1 = gr.Button("2️⃣ Generate Answer", variant="primary", visible=False)
+                    gr.Examples([f"./demo/images/{i+1}.jpg" for i in range(8)], inputs=image_input, label="Examples")
+    
+                with gr.Column():
+                    mask_output = gr.Gallery(label="Referred Masks", object_fit='scale-down', visible=False)
+                    output_image = gr.Image(label="Image with Mask", visible=True, height=400)
+                    description = gr.Textbox(label="Output", visible=True)
+                    
+                    clear_masks_btn = gr.Button("Clear Masks", variant="secondary", visible=False)
+            gr.Markdown(image_tips)
+
+        with gr.TabItem("Video"):
+            with gr.Row():
+                with gr.Column():
+                    video_input = gr.Video(label="Video")
+                    # load_btn = gr.Button("🖼️ Load First Frame", variant="secondary")
+                    first_frame = gr.ImageEditor(
+                        label="Annotate First Frame",
+                        type="pil", 
+                        sources=['upload'], 
+                        brush=gr.Brush(colors=["#ED7D31"], color_mode="fixed", default_size=10),
+                        eraser=True,
+                        layers=False,
+                        transforms=[],
+                        height=300,
+                    )
+                    generate_mask_btn_video = gr.Button("1️⃣ Generate Mask", visible=False, variant="primary")
+                    gr.Examples([f"./demo/videos/{i+1}.mp4" for i in range(4)], inputs=video_input, label="Examples")
+
+                with gr.Column():
+                    mode_video = gr.Radio(label="Mode", choices=["Caption", "QA"], value="Caption")
+                    mask_output_video = gr.Gallery(label="Referred Masks", object_fit='scale-down')
+
+                    query_video = gr.Textbox(label="Question", value="What is the relationship between <object0> and <object1>?", interactive=True, visible=False)
+
+                    submit_btn_video = gr.Button("Generate Caption", variant="primary")
+                    submit_btn_video1 = gr.Button("2️⃣ Generate Answer", variant="primary", visible=False)
+                    description_video = gr.Textbox(label="Output", visible=True)
+                    
+                    clear_masks_btn_video = gr.Button("Clear Masks", variant="secondary")
+
+            gr.Markdown(video_tips)
+
+        
+        def toggle_query_and_generate_button(mode):
+            """
+            Toggles the visibility of query-related Gradio components based on the selected mode.
+            Also clears mask states.
+
+            Args:
+                mode (str): The selected mode ("Caption" or "QA").
+
+            Returns:
+                tuple: A series of gr.update() calls and empty lists to update Gradio components.
+            """
+            query_visible = mode == "QA"
+            caption_visible = mode == "Caption"
+            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, mask_list, mask_raw_list, mask_list_video, mask_raw_list_video])
+        
+        def toggle_query_and_generate_button_video(mode):
+            """
+            Toggles the visibility of query-related Gradio components for video mode
+            based on the selected mode. Also clears mask states.
+
+            Args:
+                mode (str): The selected mode ("Caption" or "QA").
+
+            Returns:
+                tuple: A series of gr.update() calls and empty lists to update Gradio components.
+            """
+            query_visible = mode == "QA"
+            caption_visible = mode == "Caption"
+            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, mask_list, mask_raw_list, mask_list_video, mask_raw_list_video])
+
+        submit_btn.click(
+            fn=describe,
+            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, mask_raw_list],
+            outputs=[output_image, description, image_input],
+            api_name="describe"
+        )
+
+        generate_mask_btn.click(
+            fn=generate_masks,
+            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, 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, mask_list, mask_raw_list]
+        )
+
+        clear_masks_btn_video.click(
+            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, 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, mask_raw_list_video, mask_list_video],
+            outputs=[first_frame, description_video, mask_output_video, mask_list_video],
+            api_name="describe_video"
+        )
+
+
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    sam_model = SamModel.from_pretrained("facebook/sam-vit-huge").to(device)
+    sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
+
+    disable_torch_init()
+
+
+    model, processor, tokenizer = model_init(args_cli.model_path)
+    
+
+    demo.launch(mcp_server=True)