Initial commit

fbshipit-source-id: da6be2f26e3a1202f4bffde8cb980e2dcb851294

sam3/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+from .model_builder import build_sam3_image_model
+
+__version__ = "0.1.0"
+
+__all__ = ["build_sam3_image_model"]

sam3/agent/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved

sam3/agent/agent_core.py

@@ -0,0 +1,563 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import copy
+import json
+import os
+
+import cv2
+from PIL import Image
+
+from .client_llm import send_generate_request
+from .client_sam3 import call_sam_service
+from .viz import visualize
+
+
+def save_debug_messages(messages_list, debug, debug_folder_path, debug_jsonl_path):
+    """Save messages to debug jsonl file if debug is enabled"""
+    if debug and debug_jsonl_path:
+        # Ensure the debug directory exists before writing
+        os.makedirs(debug_folder_path, exist_ok=True)
+        with open(debug_jsonl_path, "w") as f:
+            for msg in messages_list:
+                f.write(json.dumps(msg, indent=4) + "\n")
+
+
+def cleanup_debug_files(debug, debug_folder_path, debug_jsonl_path):
+    """Clean up debug files when function successfully returns"""
+    if debug and debug_folder_path:
+        try:
+            if os.path.exists(debug_jsonl_path):
+                os.remove(debug_jsonl_path)
+            if os.path.exists(debug_folder_path):
+                os.rmdir(debug_folder_path)
+        except Exception as e:
+            print(f"Warning: Could not clean up debug files: {e}")
+
+
+def count_images(messages):
+    """Count the total number of images present in the messages history."""
+    total = 0
+    for message in messages:
+        # Check if message has content (should be a list)
+        if "content" in message and isinstance(message["content"], list):
+            # Iterate through each content item
+            for content_item in message["content"]:
+                # Check if content item is a dict with type "image"
+                if (
+                    isinstance(content_item, dict)
+                    and content_item.get("type") == "image"
+                ):
+                    total += 1
+    return total
+
+
+def _prune_messages_for_next_round(
+    messages_list,
+    used_text_prompts,
+    latest_sam3_text_prompt,
+    img_path,
+    initial_text_prompt,
+):
+    """Return a new messages list that contains only:
+    1) messages[:2] (with optional warning text added to the second message's content)
+    2) the latest assistant message (and everything after it) that contains a segment_phrase tool call
+    """
+    # There should not be more than 10 messages in the conversation history
+    assert len(messages_list) < 10
+
+    # Part 1: always keep the first two message JSONs
+    part1 = copy.deepcopy(messages_list[:2])
+
+    # Part 2: search backwards for the latest assistant message containing a segment_phrase tool call
+    part2_start_idx = None
+    for idx in range(len(messages_list) - 1, 1, -1):
+        msg = messages_list[idx]
+        # We only consider assistant messages with a "content" list
+        if msg.get("role") != "assistant" or "content" not in msg:
+            continue
+        # Look for any content element that is a text containing the segment_phrase tool call
+        for content in msg["content"]:
+            if (
+                isinstance(content, dict)
+                and content.get("type") == "text"
+                and "<tool>" in content.get("text", "")
+                and "segment_phrase" in content.get("text", "")
+            ):
+                part2_start_idx = idx
+                break
+        if part2_start_idx is not None:
+            break
+
+    part2 = messages_list[part2_start_idx:] if part2_start_idx is not None else []
+
+    # Part 3: decide whether to add warning text to the second message in part1
+    previously_used = (
+        [p for p in used_text_prompts if p != latest_sam3_text_prompt]
+        if latest_sam3_text_prompt
+        else list(used_text_prompts)
+    )
+    if part2 and len(previously_used) > 0:
+        warning_text = f'Note that we have previously called the segment_phrase tool with each "text_prompt" in this list: {list(previously_used)}, but none of the generated results were satisfactory. So make sure that you do not use any of these phrases as the "text_prompt" to call the segment_phrase tool again.'
+        # Replace the second message entirely to keep exactly 2 content items
+        part1[1] = {
+            "role": "user",
+            "content": [
+                {"type": "image", "image": img_path},
+                {
+                    "type": "text",
+                    "text": f"The above image is the raw input image. The initial user input query is: '{initial_text_prompt}'."
+                    + " "
+                    + warning_text,
+                },
+            ],
+        }
+        assert len(part1[1]["content"]) == 2
+
+    # Build the new messages list: part1 (with optional warning), then part2
+    new_messages = list(part1)
+    new_messages.extend(part2)
+    return new_messages
+
+
+def agent_inference(
+    img_path: str,
+    initial_text_prompt: str,
+    debug: bool = False,
+    send_generate_request=send_generate_request,
+    call_sam_service=call_sam_service,
+    max_generations: int = 100,
+    output_dir="../../sam3_agent_out",
+):
+    """
+    Given a text prompt and an image, this tool will perform all aspects of agentic problem solving,
+    while saving sam3 and MLLM outputs to their respective directories.
+
+    Args:
+        img_path: Path to the input image
+        initial_text_prompt: Initial text prompt from the user
+        debug: Whether to enable debug mode
+        max_generations: Maximum number of send_generate_request calls allowed (default: 100)
+    """
+    # setup dir
+    sam_output_dir = os.path.join(output_dir, "sam_out")
+    error_save_dir = os.path.join(output_dir, "none_out")
+    debug_save_dir = os.path.join(output_dir, "agent_debug_out")
+    os.makedirs(sam_output_dir, exist_ok=True)
+    os.makedirs(error_save_dir, exist_ok=True)
+    os.makedirs(debug_save_dir, exist_ok=True)
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    MLLM_SYSTEM_PROMPT_PATH = os.path.join(
+        current_dir, "system_prompts/system_prompt.txt"
+    )
+    ITERATIVE_CHECKING_SYSTEM_PROMPT_PATH = os.path.join(
+        current_dir, "system_prompts/system_prompt_iterative_checking.txt"
+    )
+    # init variables
+    PATH_TO_LATEST_OUTPUT_JSON = ""
+    LATEST_SAM3_TEXT_PROMPT = ""
+    USED_TEXT_PROMPTS = (
+        set()
+    )  # Track all previously used text prompts for segment_phrase
+    generation_count = 0  # Counter for number of send_generate_request calls
+
+    # debug setup
+    debug_folder_path = None
+    debug_jsonl_path = None
+    if debug:
+        debug_folder_path = os.path.join(
+            debug_save_dir, f"{img_path.rsplit('/', 1)[-1].rsplit('.', 1)[0]}"
+        )
+        debug_jsonl_path = os.path.join(debug_folder_path, "debug_history.json")
+        os.makedirs(debug_folder_path, exist_ok=True)
+
+    # The helper functions are now defined outside the agent_inference function
+    with open(MLLM_SYSTEM_PROMPT_PATH, "r") as f:
+        system_prompt = f.read().strip()
+    with open(ITERATIVE_CHECKING_SYSTEM_PROMPT_PATH, "r") as f:
+        iterative_checking_system_prompt = f.read().strip()
+
+    # Construct the initial message list
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "image": img_path},
+                {
+                    "type": "text",
+                    "text": f"The above image is the raw input image. The initial user input query is: '{initial_text_prompt}'.",
+                },
+            ],
+        },
+    ]
+    print(f"> Text prompt: {initial_text_prompt}")
+    print(f"> Image path: {img_path}")
+
+    print("\n\n")
+    print("-" * 30 + f" Round {str(generation_count + 1)}" + "-" * 30)
+    print("\n\n")
+    generated_text = send_generate_request(messages)
+    print(f"\n>>> MLLM Response [start]\n{generated_text}\n<<< MLLM Response [end]\n")
+    while generated_text is not None:
+        save_debug_messages(messages, debug, debug_folder_path, debug_jsonl_path)
+        assert (
+            "<tool>" in generated_text,
+            f"Generated text does not contain <tool> tag: {generated_text}",
+        )
+        generated_text = generated_text.split("</tool>", 1)[0] + "</tool>"
+        tool_call_json_str = (
+            generated_text.split("<tool>")[-1]
+            .split("</tool>")[0]
+            .strip()
+            .replace(r"}}}", r"}}")  # remove extra } if any
+        )
+        try:
+            tool_call = json.loads(tool_call_json_str)
+        except json.JSONDecodeError:
+            raise ValueError(f"Invalid JSON in tool call: {tool_call_json_str}")
+
+        if PATH_TO_LATEST_OUTPUT_JSON == "":
+            # The first tool call must be segment_phrase or report_no_mask
+            assert (
+                tool_call["name"] == "segment_phrase"
+                or tool_call["name"] == "report_no_mask"
+            )
+
+        if tool_call["name"] == "segment_phrase":
+            print("🔍 Calling segment_phrase tool...")
+            assert list(tool_call["parameters"].keys()) == ["text_prompt"]
+
+            # Check if this text_prompt has been used before
+            current_text_prompt = tool_call["parameters"]["text_prompt"]
+            if current_text_prompt in USED_TEXT_PROMPTS:
+                print(
+                    f"❌ Text prompt '{current_text_prompt}' has been used before. Requesting a different prompt."
+                )
+                duplicate_prompt_message = f"You have previously used '{current_text_prompt}' as your text_prompt to call the segment_phrase tool. You may not use it again. Please call the segment_phrase tool again with a different, perhaps more general, or more creative simple noun phrase prompt, while adhering to all the rules stated in the system prompt. You must also never use any of the following text_prompt(s): {str(list(USED_TEXT_PROMPTS))}."
+                messages.append(
+                    {
+                        "role": "assistant",
+                        "content": [{"type": "text", "text": generated_text}],
+                    }
+                )
+                messages.append(
+                    {
+                        "role": "user",
+                        "content": [{"type": "text", "text": duplicate_prompt_message}],
+                    }
+                )
+            else:
+                # Add the text_prompt to the set of used prompts
+                USED_TEXT_PROMPTS.add(current_text_prompt)
+                LATEST_SAM3_TEXT_PROMPT = current_text_prompt
+                PATH_TO_LATEST_OUTPUT_JSON = call_sam_service(
+                    image_path=img_path,
+                    text_prompt=current_text_prompt,
+                    output_folder_path=sam_output_dir,
+                )
+                sam3_outputs = json.load(open(PATH_TO_LATEST_OUTPUT_JSON, "r"))
+                sam3_output_image_path = sam3_outputs["output_image_path"]
+                num_masks = len(sam3_outputs["pred_boxes"])
+
+                messages.append(
+                    {
+                        "role": "assistant",
+                        "content": [{"type": "text", "text": generated_text}],
+                    }
+                )
+                if num_masks == 0:
+                    print("❌ No masks generated by SAM3, reporting no mask to Qwen.")
+                    sam3_output_text_message = f"The segment_phrase tool did not generate any masks for the text_prompt '{current_text_prompt}'. Now, please call the segment_phrase tool again with a different, perhaps more general, or more creative simple noun phrase text_prompt, while adhering to all the rules stated in the system prompt. Please be reminded that the original user query was '{initial_text_prompt}'."
+                    messages.append(
+                        {
+                            "role": "user",
+                            "content": [
+                                {"type": "text", "text": sam3_output_text_message}
+                            ],
+                        }
+                    )
+                else:
+                    sam3_output_text_message = rf"The segment_phrase tool generated {num_masks} available masks. All {num_masks} available masks are rendered in this image below, now you must analyze the {num_masks} available mask(s) carefully, compare them against the raw input image and the original user query, and determine your next action. Please be reminded that the original user query was '{initial_text_prompt}'."
+                    messages.append(
+                        {
+                            "role": "user",
+                            "content": [
+                                {"type": "text", "text": sam3_output_text_message},
+                                {"type": "image", "image": sam3_output_image_path},
+                            ],
+                        }
+                    )
+                print("\n\n>>> sam3_output_text_message:\n", sam3_output_text_message)
+
+        elif tool_call["name"] == "examine_each_mask":
+            print("🔍 Calling examine_each_mask tool...")
+            assert LATEST_SAM3_TEXT_PROMPT != ""
+
+            # Make sure that the last message is a image
+            assert (
+                messages[-1]["content"][1]["type"] == "image"
+            ), "Second content element should be an image"
+            messages.pop()  # Remove the last user message
+            # Add simplified replacement message
+            simplified_message = {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": "The segment_phrase tool generated several masks. Now you must analyze the mask(s) carefully, compare them against the raw input image and the original user query, and determine your next action.",
+                    }
+                ],
+            }
+            messages.append(simplified_message)
+
+            current_outputs = json.load(open(PATH_TO_LATEST_OUTPUT_JSON, "r"))
+            num_masks = len(current_outputs["pred_masks"])
+            masks_to_keep = []
+
+            # MLLM check the mask one by one
+            for i in range(num_masks):
+                print(f"🔍 Checking mask {i+1}/{num_masks}...")
+                image_w_mask_i, image_w_zoomed_in_mask_i = visualize(current_outputs, i)
+
+                image_w_zoomed_in_mask_i_path = os.path.join(
+                    sam_output_dir, rf"{LATEST_SAM3_TEXT_PROMPT}.png".replace("/", "_")
+                ).replace(".png", f"_zoom_in_mask_{i + 1}.png")
+                image_w_mask_i_path = os.path.join(
+                    sam_output_dir, rf"{LATEST_SAM3_TEXT_PROMPT}.png".replace("/", "_")
+                ).replace(".png", f"_selected_mask_{i + 1}.png")
+                image_w_zoomed_in_mask_i.save(image_w_zoomed_in_mask_i_path)
+                image_w_mask_i.save(image_w_mask_i_path)
+
+                iterative_checking_messages = [
+                    {"role": "system", "content": iterative_checking_system_prompt},
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": f"The raw input image: "},
+                            {"type": "image", "image": img_path},
+                            {
+                                "type": "text",
+                                "text": f"The initial user input query is: '{initial_text_prompt}'",
+                            },
+                            {
+                                "type": "text",
+                                "text": f"Image with the predicted segmentation mask rendered on it: ",
+                            },
+                            {"type": "image", "image": image_w_mask_i_path},
+                            {
+                                "type": "text",
+                                "text": f"Image with the zoomed-in mask: ",
+                            },
+                            {"type": "image", "image": image_w_zoomed_in_mask_i_path},
+                        ],
+                    },
+                ]
+                checking_generated_text = send_generate_request(
+                    iterative_checking_messages
+                )
+
+                # Process the generated text to determine if the mask should be kept or rejected
+                if checking_generated_text is None:
+                    raise ValueError(
+                        "Generated text is None, which is unexpected. Please check the Qwen server and the input parameters."
+                    )
+                print(f"Generated text for mask {i+1}: {checking_generated_text}")
+                verdict = (
+                    checking_generated_text.split("<verdict>")[-1]
+                    .split("</verdict>")[0]
+                    .strip()
+                )
+                if "Accept" in verdict:
+                    assert not "Reject" in verdict
+                    print(f"Mask {i+1} accepted, keeping it in the outputs.")
+                    masks_to_keep.append(i)
+                elif "Reject" in verdict:
+                    assert not "Accept" in verdict
+                    print(f"Mask {i+1} rejected, removing it from the outputs.")
+                else:
+                    raise ValueError(
+                        f"Unexpected verdict in generated text: {checking_generated_text}. Expected 'Accept' or 'Reject'."
+                    )
+
+            updated_outputs = {
+                "original_image_path": current_outputs["original_image_path"],
+                "orig_img_h": current_outputs["orig_img_h"],
+                "orig_img_w": current_outputs["orig_img_w"],
+                "pred_boxes": [current_outputs["pred_boxes"][i] for i in masks_to_keep],
+                "pred_scores": [
+                    current_outputs["pred_scores"][i] for i in masks_to_keep
+                ],
+                "pred_masks": [current_outputs["pred_masks"][i] for i in masks_to_keep],
+            }
+
+            image_w_check_masks = visualize(updated_outputs)
+            image_w_check_masks_path = os.path.join(
+                sam_output_dir, rf"{LATEST_SAM3_TEXT_PROMPT}.png"
+            ).replace(
+                ".png",
+                f"_selected_masks_{'-'.join(map(str, [i+1 for i in masks_to_keep]))}.png".replace(
+                    "/", "_"
+                ),
+            )
+            image_w_check_masks.save(image_w_check_masks_path)
+            # save the updated json outputs and append to message history
+            messages.append(
+                {
+                    "role": "assistant",
+                    "content": [{"type": "text", "text": generated_text}],
+                }
+            )
+            if len(masks_to_keep) == 0:
+                messages.append(
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "text",
+                                "text": f"The original user query was: '{initial_text_prompt}'. The examine_each_mask tool examined and rejected all of the masks generated by the segment_phrase tool. Now, please call the segment_phrase tool again with a different, perhaps more general, or more creative simple noun phrase text_prompt, while adhering to all the rules stated in the system prompt.",
+                            }
+                        ],
+                    }
+                )
+            else:
+                messages.append(
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "text",
+                                "text": f"The original user query was: '{initial_text_prompt}'. After calling the examine_each_mask tool on the available masks, the number of available masks is now {len(masks_to_keep)}. All {len(masks_to_keep)} available masks are rendered in this image below, now you must analyze the {len(masks_to_keep)} available mask(s) carefully, compare them against the raw input image and the original user query, and determine your next action.",
+                            },
+                            {"type": "image", "image": image_w_check_masks_path},
+                        ],
+                    }
+                )
+
+            # Create a new filename based on the original path to avoid filename length issues
+            base_path = PATH_TO_LATEST_OUTPUT_JSON
+            # Remove any existing "masks_" suffix to avoid duplication
+            if "masks_" in base_path:
+                base_path = base_path.split("masks_")[0] + ".json"
+            # Create new filename with current masks; use a clearer suffix when empty
+            if len(masks_to_keep) == 0:
+                PATH_TO_LATEST_OUTPUT_JSON = base_path.replace(
+                    ".json", "masks_none.json"
+                )
+            else:
+                PATH_TO_LATEST_OUTPUT_JSON = base_path.replace(
+                    ".json", f"masks_{'_'.join(map(str, masks_to_keep))}.json"
+                )
+            json.dump(updated_outputs, open(PATH_TO_LATEST_OUTPUT_JSON, "w"), indent=4)
+
+        elif tool_call["name"] == "select_masks_and_return":
+            print("🔍 Calling select_masks_and_return tool...")
+            current_outputs = json.load(open(PATH_TO_LATEST_OUTPUT_JSON, "r"))
+
+            assert list(tool_call["parameters"].keys()) == ["final_answer_masks"]
+            masks_to_keep = tool_call["parameters"]["final_answer_masks"]
+
+            # Keep only valid mask indices, remove duplicates, and preserve deterministic ascending order
+            available_masks = set(range(1, len(current_outputs["pred_masks"]) + 1))
+            masks_to_keep = sorted({i for i in masks_to_keep if i in available_masks})
+            # Change this to a update message telling the model to try again along with information about errors made.
+
+            final_outputs = {
+                "original_image_path": current_outputs["original_image_path"],
+                "orig_img_h": current_outputs["orig_img_h"],
+                "orig_img_w": current_outputs["orig_img_w"],
+                "pred_boxes": [
+                    current_outputs["pred_boxes"][i - 1] for i in masks_to_keep
+                ],
+                "pred_scores": [
+                    current_outputs["pred_scores"][i - 1] for i in masks_to_keep
+                ],
+                "pred_masks": [
+                    current_outputs["pred_masks"][i - 1] for i in masks_to_keep
+                ],
+            }
+
+            rendered_final_output = visualize(final_outputs)
+            messages.append(
+                {
+                    "role": "assistant",
+                    "content": [{"type": "text", "text": generated_text}],
+                }
+            )
+
+            # Clean up debug files before successful return
+            cleanup_debug_files(debug, debug_folder_path, debug_jsonl_path)
+            return messages, final_outputs, rendered_final_output
+
+        elif tool_call["name"] == "report_no_mask":
+            print("🔍 Calling report_no_mask tool...")
+            height, width = cv2.imread(img_path).shape[:2]
+            final_outputs = {
+                "original_image_path": img_path,
+                "orig_img_h": height,
+                "orig_img_w": width,
+                "pred_boxes": [],
+                "pred_scores": [],
+                "pred_masks": [],
+            }
+            rendered_final_output = Image.open(img_path)
+            messages.append(
+                {
+                    "role": "assistant",
+                    "content": [{"type": "text", "text": generated_text}],
+                }
+            )
+            return messages, final_outputs, rendered_final_output
+
+        else:
+            raise ValueError(f"Unknown tool call: {tool_call['name']}")
+
+        # sometimes the MLLM don't know when to stop, and generates multiple tool calls in one round, so we need to split the generated text by </tool> and only keep the first one
+
+        for message in messages:
+            if message["role"] == "assistant" and "content" in message:
+                for content in message["content"]:
+                    if (
+                        isinstance(content, dict)
+                        and content.get("type") == "text"
+                        and "text" in content
+                    ):
+                        content["text"] = (
+                            content["text"].split("</tool>", 1)[0] + "</tool>\n\n"
+                        )
+        # Prune the messages history before the next MLLM generation round according to the 3-part rules.
+        # This keeps history compact and ensures the model sees only the allowed parts.
+        messages = _prune_messages_for_next_round(
+            messages,
+            USED_TEXT_PROMPTS,
+            LATEST_SAM3_TEXT_PROMPT,
+            img_path,
+            initial_text_prompt,
+        )
+        # make sure there can never be more than 2 images in the context
+        assert count_images(messages) <= 2
+        generation_count += 1
+        if generation_count > max_generations:
+            raise ValueError(
+                f"Exceeded maximum number of allowed generation requests ({max_generations})"
+            )
+
+        print("\n\n")
+        print("-" * 30 + f" Round {str(generation_count + 1)}" + "-" * 30)
+        print("\n\n")
+        generated_text = send_generate_request(messages)
+        print(
+            f"\n>>> MLLM Response [start]\n{generated_text}\n<<< MLLM Response [end]\n"
+        )
+
+    print("\n\n>>> SAM 3 Agent execution ended.\n\n")
+
+    error_save_path = os.path.join(
+        error_save_dir,
+        f"{img_path.rsplit('/', 1)[-1].rsplit('.', 1)[0]}_error_history.json",
+    )
+    with open(error_save_path, "w") as f:
+        json.dump(messages, f, indent=4)
+    print("Saved messages history that caused error to:", error_save_path)
+    raise ValueError(
+        rf"Generated text is None, which is unexpected. Please check the Qwen server and the input parameters for image path: {img_path} and initial text prompt: {initial_text_prompt}."
+    )

sam3/agent/client_llm.py

@@ -0,0 +1,205 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import base64
+import os
+from typing import Any, Optional
+
+from openai import OpenAI
+
+
+def get_image_base64_and_mime(image_path):
+    """Convert image file to base64 string and get MIME type"""
+    try:
+        # Get MIME type based on file extension
+        ext = os.path.splitext(image_path)[1].lower()
+        mime_types = {
+            ".jpg": "image/jpeg",
+            ".jpeg": "image/jpeg",
+            ".png": "image/png",
+            ".gif": "image/gif",
+            ".webp": "image/webp",
+            ".bmp": "image/bmp",
+        }
+        mime_type = mime_types.get(ext, "image/jpeg")  # Default to JPEG
+
+        # Convert image to base64
+        with open(image_path, "rb") as image_file:
+            base64_data = base64.b64encode(image_file.read()).decode("utf-8")
+            return base64_data, mime_type
+    except Exception as e:
+        print(f"Error converting image to base64: {e}")
+        return None, None
+
+
+def send_generate_request(
+    messages,
+    server_url=None,
+    model="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
+    api_key=None,
+    max_tokens=4096,
+):
+    """
+    Sends a request to the OpenAI-compatible API endpoint using the OpenAI client library.
+
+    Args:
+        server_url (str): The base URL of the server, e.g. "http://127.0.0.1:8000"
+        messages (list): A list of message dicts, each containing role and content.
+        model (str): The model to use for generation (default: "llama-4")
+        max_tokens (int): Maximum number of tokens to generate (default: 4096)
+
+    Returns:
+        str: The generated response text from the server.
+    """
+    # Process messages to convert image paths to base64
+    processed_messages = []
+    for message in messages:
+        processed_message = message.copy()
+        if message["role"] == "user" and "content" in message:
+            processed_content = []
+            for c in message["content"]:
+                if isinstance(c, dict) and c.get("type") == "image":
+                    # Convert image path to base64 format
+                    image_path = c["image"]
+
+                    print("image_path", image_path)
+                    new_image_path = image_path.replace(
+                        "?", "%3F"
+                    )  # Escape ? in the path
+
+                    # Read the image file and convert to base64
+                    try:
+                        base64_image, mime_type = get_image_base64_and_mime(
+                            new_image_path
+                        )
+                        if base64_image is None:
+                            print(
+                                f"Warning: Could not convert image to base64: {new_image_path}"
+                            )
+                            continue
+
+                        # Create the proper image_url structure with base64 data
+                        processed_content.append(
+                            {
+                                "type": "image_url",
+                                "image_url": {
+                                    "url": f"data:{mime_type};base64,{base64_image}",
+                                    "detail": "high",
+                                },
+                            }
+                        )
+
+                    except FileNotFoundError:
+                        print(f"Warning: Image file not found: {new_image_path}")
+                        continue
+                    except Exception as e:
+                        print(f"Warning: Error processing image {new_image_path}: {e}")
+                        continue
+                else:
+                    processed_content.append(c)
+
+            processed_message["content"] = processed_content
+        processed_messages.append(processed_message)
+
+    # Create OpenAI client with custom base URL
+    client = OpenAI(api_key=api_key, base_url=server_url)
+
+    try:
+        print(f"🔍 Calling model {model}...")
+        response = client.chat.completions.create(
+            model=model,
+            messages=processed_messages,
+            max_completion_tokens=max_tokens,
+            n=1,
+        )
+        # print(f"Received response: {response.choices[0].message}")
+
+        # Extract the response content
+        if response.choices and len(response.choices) > 0:
+            return response.choices[0].message.content
+        else:
+            print(f"Unexpected response format: {response}")
+            return None
+
+    except Exception as e:
+        print(f"Request failed: {e}")
+        return None
+
+
+def send_direct_request(
+    llm: Any,
+    messages: list[dict[str, Any]],
+    sampling_params: Any,
+) -> Optional[str]:
+    """
+    Run inference on a vLLM model instance directly without using a server.
+
+    Args:
+        llm: Initialized vLLM LLM instance (passed from external initialization)
+        messages: List of message dicts with role and content (OpenAI format)
+        sampling_params: vLLM SamplingParams instance (initialized externally)
+
+    Returns:
+        str: Generated response text, or None if inference fails
+    """
+    try:
+        # Process messages to handle images (convert to base64 if needed)
+        processed_messages = []
+        for message in messages:
+            processed_message = message.copy()
+            if message["role"] == "user" and "content" in message:
+                processed_content = []
+                for c in message["content"]:
+                    if isinstance(c, dict) and c.get("type") == "image":
+                        # Convert image path to base64 format
+                        image_path = c["image"]
+                        new_image_path = image_path.replace("?", "%3F")
+
+                        try:
+                            base64_image, mime_type = get_image_base64_and_mime(
+                                new_image_path
+                            )
+                            if base64_image is None:
+                                print(
+                                    f"Warning: Could not convert image: {new_image_path}"
+                                )
+                                continue
+
+                            # vLLM expects image_url format
+                            processed_content.append(
+                                {
+                                    "type": "image_url",
+                                    "image_url": {
+                                        "url": f"data:{mime_type};base64,{base64_image}"
+                                    },
+                                }
+                            )
+                        except Exception as e:
+                            print(
+                                f"Warning: Error processing image {new_image_path}: {e}"
+                            )
+                            continue
+                    else:
+                        processed_content.append(c)
+
+                processed_message["content"] = processed_content
+            processed_messages.append(processed_message)
+
+        print("🔍 Running direct inference with vLLM...")
+
+        # Run inference using vLLM's chat interface
+        outputs = llm.chat(
+            messages=processed_messages,
+            sampling_params=sampling_params,
+        )
+
+        # Extract the generated text from the first output
+        if outputs and len(outputs) > 0:
+            generated_text = outputs[0].outputs[0].text
+            return generated_text
+        else:
+            print(f"Unexpected output format: {outputs}")
+            return None
+
+    except Exception as e:
+        print(f"Direct inference failed: {e}")
+        return None

sam3/agent/client_sam3.py

@@ -0,0 +1,138 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import json
+import os
+
+import torch
+from PIL import Image
+
+from sam3.model.box_ops import box_xyxy_to_xywh
+from sam3.train.masks_ops import rle_encode
+
+from .helpers.mask_overlap_removal import remove_overlapping_masks
+from .viz import visualize
+
+
+def sam3_inference(processor, image_path, text_prompt):
+    """Run SAM 3 image inference with text prompts and format the outputs"""
+    image = Image.open(image_path)
+    orig_img_w, orig_img_h = image.size
+
+    # model inference
+    inference_state = processor.set_image(image)
+    inference_state = processor.set_text_prompt(
+        state=inference_state, prompt=text_prompt
+    )
+
+    # format and assemble outputs
+    pred_boxes_xyxy = torch.stack(
+        [
+            inference_state["boxes"][:, 0] / orig_img_w,
+            inference_state["boxes"][:, 1] / orig_img_h,
+            inference_state["boxes"][:, 2] / orig_img_w,
+            inference_state["boxes"][:, 3] / orig_img_h,
+        ],
+        dim=-1,
+    )  # normalized in range [0, 1]
+    pred_boxes_xywh = box_xyxy_to_xywh(pred_boxes_xyxy).tolist()
+    pred_masks = rle_encode(inference_state["masks"].squeeze(1))
+    pred_masks = [m["counts"] for m in pred_masks]
+    outputs = {
+        "orig_img_h": orig_img_h,
+        "orig_img_w": orig_img_w,
+        "pred_boxes": pred_boxes_xywh,
+        "pred_masks": pred_masks,
+        "pred_scores": inference_state["scores"].tolist(),
+    }
+    return outputs
+
+
+def call_sam_service(
+    sam3_processor,
+    image_path: str,
+    text_prompt: str,
+    output_folder_path: str = "sam3_output",
+):
+    """
+    Loads an image, sends it with a text prompt to the service,
+    saves the results, and renders the visualization.
+    """
+    print(f"📞 Loading image '{image_path}' and sending with prompt '{text_prompt}'...")
+
+    text_prompt_for_save_path = (
+        text_prompt.replace("/", "_") if "/" in text_prompt else text_prompt
+    )
+
+    os.makedirs(
+        os.path.join(output_folder_path, image_path.replace("/", "-")), exist_ok=True
+    )
+    output_json_path = os.path.join(
+        output_folder_path,
+        image_path.replace("/", "-"),
+        rf"{text_prompt_for_save_path}.json",
+    )
+    output_image_path = os.path.join(
+        output_folder_path,
+        image_path.replace("/", "-"),
+        rf"{text_prompt_for_save_path}.png",
+    )
+
+    try:
+        # Send the image and text prompt as a multipart/form-data request
+        serialized_response = sam3_inference(sam3_processor, image_path, text_prompt)
+
+        # 1. Prepare the response dictionary
+        serialized_response = remove_overlapping_masks(serialized_response)
+        serialized_response = {
+            "original_image_path": image_path,
+            "output_image_path": output_image_path,
+            **serialized_response,
+        }
+
+        # 2. Reorder predictions by scores (highest to lowest) if scores are available
+        if "pred_scores" in serialized_response and serialized_response["pred_scores"]:
+            # Create indices sorted by scores in descending order
+            score_indices = sorted(
+                range(len(serialized_response["pred_scores"])),
+                key=lambda i: serialized_response["pred_scores"][i],
+                reverse=True,
+            )
+
+            # Reorder all three lists based on the sorted indices
+            serialized_response["pred_scores"] = [
+                serialized_response["pred_scores"][i] for i in score_indices
+            ]
+            serialized_response["pred_boxes"] = [
+                serialized_response["pred_boxes"][i] for i in score_indices
+            ]
+            serialized_response["pred_masks"] = [
+                serialized_response["pred_masks"][i] for i in score_indices
+            ]
+
+        # 3. Remove any invalid RLE masks that is too short (shorter than 5 characters)
+        valid_masks = []
+        valid_boxes = []
+        valid_scores = []
+        for i, rle in enumerate(serialized_response["pred_masks"]):
+            if len(rle) > 4:
+                valid_masks.append(rle)
+                valid_boxes.append(serialized_response["pred_boxes"][i])
+                valid_scores.append(serialized_response["pred_scores"][i])
+        serialized_response["pred_masks"] = valid_masks
+        serialized_response["pred_boxes"] = valid_boxes
+        serialized_response["pred_scores"] = valid_scores
+
+        with open(output_json_path, "w") as f:
+            json.dump(serialized_response, f, indent=4)
+        print(f"✅ Raw JSON response saved to '{output_json_path}'")
+
+        # 4. Render and save visualizations on the image and save it in the SAM3 output folder
+        print("🔍 Rendering visualizations on the image ...")
+        viz_image = visualize(serialized_response)
+        os.makedirs(os.path.dirname(output_image_path), exist_ok=True)
+        viz_image.save(output_image_path)
+        print("✅ Saved visualization at:", output_image_path)
+    except Exception as e:
+        print(f"❌ Error calling service: {e}")
+
+    return output_json_path

sam3/agent/helpers/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved

sam3/agent/helpers/boxes.py

@@ -0,0 +1,438 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import math
+from enum import IntEnum, unique
+from typing import List, Tuple, Union
+
+import numpy as np
+import torch
+from torch import device
+
+_RawBoxType = Union[List[float], Tuple[float, ...], torch.Tensor, np.ndarray]
+
+
+@unique
+class BoxMode(IntEnum):
+    """
+    Enum of different ways to represent a box.
+    """
+
+    XYXY_ABS = 0
+    """
+    (x0, y0, x1, y1) in absolute floating points coordinates.
+    The coordinates in range [0, width or height].
+    """
+    XYWH_ABS = 1
+    """
+    (x0, y0, w, h) in absolute floating points coordinates.
+    """
+    XYXY_REL = 2
+    """
+    Not yet supported!
+    (x0, y0, x1, y1) in range [0, 1]. They are relative to the size of the image.
+    """
+    XYWH_REL = 3
+    """
+    Not yet supported!
+    (x0, y0, w, h) in range [0, 1]. They are relative to the size of the image.
+    """
+    XYWHA_ABS = 4
+    """
+    (xc, yc, w, h, a) in absolute floating points coordinates.
+    (xc, yc) is the center of the rotated box, and the angle a is in degrees ccw.
+    """
+
+    @staticmethod
+    def convert(
+        box: _RawBoxType, from_mode: "BoxMode", to_mode: "BoxMode"
+    ) -> _RawBoxType:
+        """
+        Args:
+            box: can be a k-tuple, k-list or an Nxk array/tensor, where k = 4 or 5
+            from_mode, to_mode (BoxMode)
+
+        Returns:
+            The converted box of the same type.
+        """
+        if from_mode == to_mode:
+            return box
+
+        original_type = type(box)
+        is_numpy = isinstance(box, np.ndarray)
+        single_box = isinstance(box, (list, tuple))
+        if single_box:
+            assert len(box) == 4 or len(box) == 5, (
+                "BoxMode.convert takes either a k-tuple/list or an Nxk array/tensor,"
+                " where k == 4 or 5"
+            )
+            arr = torch.tensor(box)[None, :]
+        else:
+            # avoid modifying the input box
+            if is_numpy:
+                arr = torch.from_numpy(np.asarray(box)).clone()
+            else:
+                arr = box.clone()
+
+        assert to_mode not in [
+            BoxMode.XYXY_REL,
+            BoxMode.XYWH_REL,
+        ] and from_mode not in [
+            BoxMode.XYXY_REL,
+            BoxMode.XYWH_REL,
+        ], "Relative mode not yet supported!"
+
+        if from_mode == BoxMode.XYWHA_ABS and to_mode == BoxMode.XYXY_ABS:
+            assert (
+                arr.shape[-1] == 5
+            ), "The last dimension of input shape must be 5 for XYWHA format"
+            original_dtype = arr.dtype
+            arr = arr.double()
+
+            w = arr[:, 2]
+            h = arr[:, 3]
+            a = arr[:, 4]
+            c = torch.abs(torch.cos(a * math.pi / 180.0))
+            s = torch.abs(torch.sin(a * math.pi / 180.0))
+            # This basically computes the horizontal bounding rectangle of the rotated box
+            new_w = c * w + s * h
+            new_h = c * h + s * w
+
+            # convert center to top-left corner
+            arr[:, 0] -= new_w / 2.0
+            arr[:, 1] -= new_h / 2.0
+            # bottom-right corner
+            arr[:, 2] = arr[:, 0] + new_w
+            arr[:, 3] = arr[:, 1] + new_h
+
+            arr = arr[:, :4].to(dtype=original_dtype)
+        elif from_mode == BoxMode.XYWH_ABS and to_mode == BoxMode.XYWHA_ABS:
+            original_dtype = arr.dtype
+            arr = arr.double()
+            arr[:, 0] += arr[:, 2] / 2.0
+            arr[:, 1] += arr[:, 3] / 2.0
+            angles = torch.zeros((arr.shape[0], 1), dtype=arr.dtype)
+            arr = torch.cat((arr, angles), axis=1).to(dtype=original_dtype)
+        else:
+            if to_mode == BoxMode.XYXY_ABS and from_mode == BoxMode.XYWH_ABS:
+                arr[:, 2] += arr[:, 0]
+                arr[:, 3] += arr[:, 1]
+            elif from_mode == BoxMode.XYXY_ABS and to_mode == BoxMode.XYWH_ABS:
+                arr[:, 2] -= arr[:, 0]
+                arr[:, 3] -= arr[:, 1]
+            else:
+                raise NotImplementedError(
+                    "Conversion from BoxMode {} to {} is not supported yet".format(
+                        from_mode, to_mode
+                    )
+                )
+
+        if single_box:
+            return original_type(arr.flatten().tolist())
+        if is_numpy:
+            return arr.numpy()
+        else:
+            return arr
+
+
+class Boxes:
+    """
+    This structure stores a list of boxes as a Nx4 torch.Tensor.
+    It supports some common methods about boxes
+    (`area`, `clip`, `nonempty`, etc),
+    and also behaves like a Tensor
+    (support indexing, `to(device)`, `.device`, and iteration over all boxes)
+
+    Attributes:
+        tensor (torch.Tensor): float matrix of Nx4. Each row is (x1, y1, x2, y2).
+    """
+
+    def __init__(self, tensor: torch.Tensor):
+        """
+        Args:
+            tensor (Tensor[float]): a Nx4 matrix.  Each row is (x1, y1, x2, y2).
+        """
+        if not isinstance(tensor, torch.Tensor):
+            tensor = torch.as_tensor(
+                tensor, dtype=torch.float32, device=torch.device("cpu")
+            )
+        else:
+            tensor = tensor.to(torch.float32)
+        if tensor.numel() == 0:
+            # Use reshape, so we don't end up creating a new tensor that does not depend on
+            # the inputs (and consequently confuses jit)
+            tensor = tensor.reshape((-1, 4)).to(dtype=torch.float32)
+        assert tensor.dim() == 2 and tensor.size(-1) == 4, tensor.size()
+
+        self.tensor = tensor
+
+    def clone(self) -> "Boxes":
+        """
+        Clone the Boxes.
+
+        Returns:
+            Boxes
+        """
+        return Boxes(self.tensor.clone())
+
+    def to(self, device: torch.device):
+        # Boxes are assumed float32 and does not support to(dtype)
+        return Boxes(self.tensor.to(device=device))
+
+    def area(self) -> torch.Tensor:
+        """
+        Computes the area of all the boxes.
+
+        Returns:
+            torch.Tensor: a vector with areas of each box.
+        """
+        box = self.tensor
+        area = (box[:, 2] - box[:, 0]) * (box[:, 3] - box[:, 1])
+        return area
+
+    def clip(self, box_size: Tuple[int, int]) -> None:
+        """
+        Clip (in place) the boxes by limiting x coordinates to the range [0, width]
+        and y coordinates to the range [0, height].
+
+        Args:
+            box_size (height, width): The clipping box's size.
+        """
+        assert torch.isfinite(self.tensor).all(), "Box tensor contains infinite or NaN!"
+        h, w = box_size
+        x1 = self.tensor[:, 0].clamp(min=0, max=w)
+        y1 = self.tensor[:, 1].clamp(min=0, max=h)
+        x2 = self.tensor[:, 2].clamp(min=0, max=w)
+        y2 = self.tensor[:, 3].clamp(min=0, max=h)
+        self.tensor = torch.stack((x1, y1, x2, y2), dim=-1)
+
+    def nonempty(self, threshold: float = 0.0) -> torch.Tensor:
+        """
+        Find boxes that are non-empty.
+        A box is considered empty, if either of its side is no larger than threshold.
+
+        Returns:
+            Tensor:
+                a binary vector which represents whether each box is empty
+                (False) or non-empty (True).
+        """
+        box = self.tensor
+        widths = box[:, 2] - box[:, 0]
+        heights = box[:, 3] - box[:, 1]
+        keep = (widths > threshold) & (heights > threshold)
+        return keep
+
+    def __getitem__(self, item) -> "Boxes":
+        """
+        Args:
+            item: int, slice, or a BoolTensor
+
+        Returns:
+            Boxes: Create a new :class:`Boxes` by indexing.
+
+        The following usage are allowed:
+
+        1. `new_boxes = boxes[3]`: return a `Boxes` which contains only one box.
+        2. `new_boxes = boxes[2:10]`: return a slice of boxes.
+        3. `new_boxes = boxes[vector]`, where vector is a torch.BoolTensor
+           with `length = len(boxes)`. Nonzero elements in the vector will be selected.
+
+        Note that the returned Boxes might share storage with this Boxes,
+        subject to Pytorch's indexing semantics.
+        """
+        if isinstance(item, int):
+            return Boxes(self.tensor[item].view(1, -1))
+        b = self.tensor[item]
+        assert (
+            b.dim() == 2
+        ), "Indexing on Boxes with {} failed to return a matrix!".format(item)
+        return Boxes(b)
+
+    def __len__(self) -> int:
+        return self.tensor.shape[0]
+
+    def __repr__(self) -> str:
+        return "Boxes(" + str(self.tensor) + ")"
+
+    def inside_box(
+        self, box_size: Tuple[int, int], boundary_threshold: int = 0
+    ) -> torch.Tensor:
+        """
+        Args:
+            box_size (height, width): Size of the reference box.
+            boundary_threshold (int): Boxes that extend beyond the reference box
+                boundary by more than boundary_threshold are considered "outside".
+
+        Returns:
+            a binary vector, indicating whether each box is inside the reference box.
+        """
+        height, width = box_size
+        inds_inside = (
+            (self.tensor[..., 0] >= -boundary_threshold)
+            & (self.tensor[..., 1] >= -boundary_threshold)
+            & (self.tensor[..., 2] < width + boundary_threshold)
+            & (self.tensor[..., 3] < height + boundary_threshold)
+        )
+        return inds_inside
+
+    def get_centers(self) -> torch.Tensor:
+        """
+        Returns:
+            The box centers in a Nx2 array of (x, y).
+        """
+        return (self.tensor[:, :2] + self.tensor[:, 2:]) / 2
+
+    def scale(self, scale_x: float, scale_y: float) -> None:
+        """
+        Scale the box with horizontal and vertical scaling factors
+        """
+        self.tensor[:, 0::2] *= scale_x
+        self.tensor[:, 1::2] *= scale_y
+
+    @classmethod
+    def cat(cls, boxes_list: List["Boxes"]) -> "Boxes":
+        """
+        Concatenates a list of Boxes into a single Boxes
+
+        Arguments:
+            boxes_list (list[Boxes])
+
+        Returns:
+            Boxes: the concatenated Boxes
+        """
+        assert isinstance(boxes_list, (list, tuple))
+        if len(boxes_list) == 0:
+            return cls(torch.empty(0))
+        assert all([isinstance(box, Boxes) for box in boxes_list])
+
+        # use torch.cat (v.s. layers.cat) so the returned boxes never share storage with input
+        cat_boxes = cls(torch.cat([b.tensor for b in boxes_list], dim=0))
+        return cat_boxes
+
+    @property
+    def device(self) -> device:
+        return self.tensor.device
+
+    # type "Iterator[torch.Tensor]", yield, and iter() not supported by torchscript
+    # https://github.com/pytorch/pytorch/issues/18627
+    @torch.jit.unused
+    def __iter__(self):
+        """
+        Yield a box as a Tensor of shape (4,) at a time.
+        """
+        yield from self.tensor
+
+
+def pairwise_intersection(boxes1: Boxes, boxes2: Boxes) -> torch.Tensor:
+    """
+    Given two lists of boxes of size N and M,
+    compute the intersection area between __all__ N x M pairs of boxes.
+    The box order must be (xmin, ymin, xmax, ymax)
+
+    Args:
+        boxes1,boxes2 (Boxes): two `Boxes`. Contains N & M boxes, respectively.
+
+    Returns:
+        Tensor: intersection, sized [N,M].
+    """
+    boxes1, boxes2 = boxes1.tensor, boxes2.tensor
+    width_height = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) - torch.max(
+        boxes1[:, None, :2], boxes2[:, :2]
+    )  # [N,M,2]
+
+    width_height.clamp_(min=0)  # [N,M,2]
+    intersection = width_height.prod(dim=2)  # [N,M]
+    return intersection
+
+
+# implementation from https://github.com/kuangliu/torchcv/blob/master/torchcv/utils/box.py
+# with slight modifications
+def pairwise_iou(boxes1: Boxes, boxes2: Boxes) -> torch.Tensor:
+    """
+    Given two lists of boxes of size N and M, compute the IoU
+    (intersection over union) between **all** N x M pairs of boxes.
+    The box order must be (xmin, ymin, xmax, ymax).
+
+    Args:
+        boxes1,boxes2 (Boxes): two `Boxes`. Contains N & M boxes, respectively.
+
+    Returns:
+        Tensor: IoU, sized [N,M].
+    """
+    area1 = boxes1.area()  # [N]
+    area2 = boxes2.area()  # [M]
+    inter = pairwise_intersection(boxes1, boxes2)
+
+    # handle empty boxes
+    iou = torch.where(
+        inter > 0,
+        inter / (area1[:, None] + area2 - inter),
+        torch.zeros(1, dtype=inter.dtype, device=inter.device),
+    )
+    return iou
+
+
+def pairwise_ioa(boxes1: Boxes, boxes2: Boxes) -> torch.Tensor:
+    """
+    Similar to :func:`pariwise_iou` but compute the IoA (intersection over boxes2 area).
+
+    Args:
+        boxes1,boxes2 (Boxes): two `Boxes`. Contains N & M boxes, respectively.
+
+    Returns:
+        Tensor: IoA, sized [N,M].
+    """
+    area2 = boxes2.area()  # [M]
+    inter = pairwise_intersection(boxes1, boxes2)
+
+    # handle empty boxes
+    ioa = torch.where(
+        inter > 0, inter / area2, torch.zeros(1, dtype=inter.dtype, device=inter.device)
+    )
+    return ioa
+
+
+def pairwise_point_box_distance(points: torch.Tensor, boxes: Boxes):
+    """
+    Pairwise distance between N points and M boxes. The distance between a
+    point and a box is represented by the distance from the point to 4 edges
+    of the box. Distances are all positive when the point is inside the box.
+
+    Args:
+        points: Nx2 coordinates. Each row is (x, y)
+        boxes: M boxes
+
+    Returns:
+        Tensor: distances of size (N, M, 4). The 4 values are distances from
+            the point to the left, top, right, bottom of the box.
+    """
+    x, y = points.unsqueeze(dim=2).unbind(dim=1)  # (N, 1)
+    x0, y0, x1, y1 = boxes.tensor.unsqueeze(dim=0).unbind(dim=2)  # (1, M)
+    return torch.stack([x - x0, y - y0, x1 - x, y1 - y], dim=2)
+
+
+def matched_pairwise_iou(boxes1: Boxes, boxes2: Boxes) -> torch.Tensor:
+    """
+    Compute pairwise intersection over union (IOU) of two sets of matched
+    boxes that have the same number of boxes.
+    Similar to :func:`pairwise_iou`, but computes only diagonal elements of the matrix.
+
+    Args:
+        boxes1 (Boxes): bounding boxes, sized [N,4].
+        boxes2 (Boxes): same length as boxes1
+    Returns:
+        Tensor: iou, sized [N].
+    """
+    assert len(boxes1) == len(boxes2), (
+        "boxlists should have the same" "number of entries, got {}, {}".format(
+            len(boxes1), len(boxes2)
+        )
+    )
+    area1 = boxes1.area()  # [N]
+    area2 = boxes2.area()  # [N]
+    box1, box2 = boxes1.tensor, boxes2.tensor
+    lt = torch.max(box1[:, :2], box2[:, :2])  # [N,2]
+    rb = torch.min(box1[:, 2:], box2[:, 2:])  # [N,2]
+    wh = (rb - lt).clamp(min=0)  # [N,2]
+    inter = wh[:, 0] * wh[:, 1]  # [N]
+    iou = inter / (area1 + area2 - inter)  # [N]
+    return iou

sam3/agent/helpers/color_map.py

@@ -0,0 +1,150 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+"""
+An awesome colormap for really neat visualizations.
+Copied from Detectron, and removed gray colors.
+"""
+
+import random
+
+import numpy as np
+
+__all__ = ["colormap", "random_color", "random_colors"]
+
+
+# A list of 25 bright and sharp colors for segmentation masks,
+# generated from the edges of the sRGB color space for maximum intensity.
+_COLORS = (
+    np.array(
+        [
+            # The original 8 sharp colors
+            1.000,
+            1.000,
+            0.000,  # 1. Yellow
+            0.000,
+            1.000,
+            0.000,  # 2. Lime
+            0.000,
+            1.000,
+            1.000,  # 3. Cyan
+            1.000,
+            0.000,
+            1.000,  # 4. Magenta
+            1.000,
+            0.000,
+            0.000,  # 5. Red
+            1.000,
+            0.498,
+            0.000,  # 6. Orange
+            0.498,
+            1.000,
+            0.000,  # 7. Chartreuse
+            0.000,
+            1.000,
+            0.498,  # 8. Spring Green
+            1.000,
+            0.000,
+            0.498,  # 9. Rose
+            0.498,
+            0.000,
+            1.000,  # 10. Violet
+            0.753,
+            1.000,
+            0.000,  # 11. Electric Lime
+            1.000,
+            0.753,
+            0.000,  # 12. Vivid Orange
+            0.000,
+            1.000,
+            0.753,  # 13. Turquoise
+            0.753,
+            0.000,
+            1.000,  # 14. Bright Violet
+            1.000,
+            0.000,
+            0.753,  # 15. Bright Pink
+            1.000,
+            0.251,
+            0.000,  # 16. Fiery Orange
+            0.251,
+            1.000,
+            0.000,  # 17. Bright Chartreuse
+            0.000,
+            1.000,
+            0.251,  # 18. Malachite Green
+            0.251,
+            0.000,
+            1.000,  # 19. Deep Violet
+            1.000,
+            0.000,
+            0.251,  # 20. Hot Pink
+        ]
+    )
+    .astype(np.float32)
+    .reshape(-1, 3)
+)
+
+
+def colormap(rgb=False, maximum=255):
+    """
+    Args:
+        rgb (bool): whether to return RGB colors or BGR colors.
+        maximum (int): either 255 or 1
+
+    Returns:
+        ndarray: a float32 array of Nx3 colors, in range [0, 255] or [0, 1]
+    """
+    assert maximum in [255, 1], maximum
+    c = _COLORS * maximum
+    if not rgb:
+        c = c[:, ::-1]
+    return c
+
+
+def random_color(rgb=False, maximum=255):
+    """
+    Args:
+        rgb (bool): whether to return RGB colors or BGR colors.
+        maximum (int): either 255 or 1
+
+    Returns:
+        ndarray: a vector of 3 numbers
+    """
+    idx = np.random.randint(0, len(_COLORS))
+    ret = _COLORS[idx] * maximum
+    if not rgb:
+        ret = ret[::-1]
+    return ret
+
+
+def random_colors(N, rgb=False, maximum=255):
+    """
+    Args:
+        N (int): number of unique colors needed
+        rgb (bool): whether to return RGB colors or BGR colors.
+        maximum (int): either 255 or 1
+
+    Returns:
+        ndarray: a list of random_color
+    """
+    indices = random.sample(range(len(_COLORS)), N)
+    ret = [_COLORS[i] * maximum for i in indices]
+    if not rgb:
+        ret = [x[::-1] for x in ret]
+    return ret
+
+
+if __name__ == "__main__":
+    import cv2
+
+    size = 100
+    H, W = 10, 10
+    canvas = np.random.rand(H * size, W * size, 3).astype("float32")
+    for h in range(H):
+        for w in range(W):
+            idx = h * W + w
+            if idx >= len(_COLORS):
+                break
+            canvas[h * size : (h + 1) * size, w * size : (w + 1) * size] = _COLORS[idx]
+    cv2.imshow("a", canvas)
+    cv2.waitKey(0)

sam3/agent/helpers/keypoints.py

@@ -0,0 +1,244 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+from typing import Any, List, Tuple, Union
+
+import numpy as np
+import torch
+from torch.nn import functional as F
+
+
+class Keypoints:
+    """
+    Stores keypoint **annotation** data. GT Instances have a `gt_keypoints` property
+    containing the x,y location and visibility flag of each keypoint. This tensor has shape
+    (N, K, 3) where N is the number of instances and K is the number of keypoints per instance.
+
+    The visibility flag follows the COCO format and must be one of three integers:
+
+    * v=0: not labeled (in which case x=y=0)
+    * v=1: labeled but not visible
+    * v=2: labeled and visible
+    """
+
+    def __init__(self, keypoints: Union[torch.Tensor, np.ndarray, List[List[float]]]):
+        """
+        Arguments:
+            keypoints: A Tensor, numpy array, or list of the x, y, and visibility of each keypoint.
+                The shape should be (N, K, 3) where N is the number of
+                instances, and K is the number of keypoints per instance.
+        """
+        device = (
+            keypoints.device
+            if isinstance(keypoints, torch.Tensor)
+            else torch.device("cpu")
+        )
+        keypoints = torch.as_tensor(keypoints, dtype=torch.float32, device=device)
+        assert keypoints.dim() == 3 and keypoints.shape[2] == 3, keypoints.shape
+        self.tensor = keypoints
+
+    def __len__(self) -> int:
+        return self.tensor.size(0)
+
+    def to(self, *args: Any, **kwargs: Any) -> "Keypoints":
+        return type(self)(self.tensor.to(*args, **kwargs))
+
+    @property
+    def device(self) -> torch.device:
+        return self.tensor.device
+
+    def to_heatmap(self, boxes: torch.Tensor, heatmap_size: int) -> torch.Tensor:
+        """
+        Convert keypoint annotations to a heatmap of one-hot labels for training,
+        as described in :paper:`Mask R-CNN`.
+
+        Arguments:
+            boxes: Nx4 tensor, the boxes to draw the keypoints to
+
+        Returns:
+            heatmaps:
+                A tensor of shape (N, K), each element is integer spatial label
+                in the range [0, heatmap_size**2 - 1] for each keypoint in the input.
+            valid:
+                A tensor of shape (N, K) containing whether each keypoint is in the roi or not.
+        """
+        return _keypoints_to_heatmap(self.tensor, boxes, heatmap_size)
+
+    def __getitem__(self, item: Union[int, slice, torch.BoolTensor]) -> "Keypoints":
+        """
+        Create a new `Keypoints` by indexing on this `Keypoints`.
+
+        The following usage are allowed:
+
+        1. `new_kpts = kpts[3]`: return a `Keypoints` which contains only one instance.
+        2. `new_kpts = kpts[2:10]`: return a slice of key points.
+        3. `new_kpts = kpts[vector]`, where vector is a torch.ByteTensor
+           with `length = len(kpts)`. Nonzero elements in the vector will be selected.
+
+        Note that the returned Keypoints might share storage with this Keypoints,
+        subject to Pytorch's indexing semantics.
+        """
+        if isinstance(item, int):
+            return Keypoints([self.tensor[item]])
+        return Keypoints(self.tensor[item])
+
+    def __repr__(self) -> str:
+        s = self.__class__.__name__ + "("
+        s += "num_instances={})".format(len(self.tensor))
+        return s
+
+    @staticmethod
+    def cat(keypoints_list: List["Keypoints"]) -> "Keypoints":
+        """
+        Concatenates a list of Keypoints into a single Keypoints
+
+        Arguments:
+            keypoints_list (list[Keypoints])
+
+        Returns:
+            Keypoints: the concatenated Keypoints
+        """
+        assert isinstance(keypoints_list, (list, tuple))
+        assert len(keypoints_list) > 0
+        assert all(isinstance(keypoints, Keypoints) for keypoints in keypoints_list)
+
+        cat_kpts = type(keypoints_list[0])(
+            torch.cat([kpts.tensor for kpts in keypoints_list], dim=0)
+        )
+        return cat_kpts
+
+
+def _keypoints_to_heatmap(
+    keypoints: torch.Tensor, rois: torch.Tensor, heatmap_size: int
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Encode keypoint locations into a target heatmap for use in SoftmaxWithLoss across space.
+
+    Maps keypoints from the half-open interval [x1, x2) on continuous image coordinates to the
+    closed interval [0, heatmap_size - 1] on discrete image coordinates. We use the
+    continuous-discrete conversion from Heckbert 1990 ("What is the coordinate of a pixel?"):
+    d = floor(c) and c = d + 0.5, where d is a discrete coordinate and c is a continuous coordinate.
+
+    Arguments:
+        keypoints: tensor of keypoint locations in of shape (N, K, 3).
+        rois: Nx4 tensor of rois in xyxy format
+        heatmap_size: integer side length of square heatmap.
+
+    Returns:
+        heatmaps: A tensor of shape (N, K) containing an integer spatial label
+            in the range [0, heatmap_size**2 - 1] for each keypoint in the input.
+        valid: A tensor of shape (N, K) containing whether each keypoint is in
+            the roi or not.
+    """
+
+    if rois.numel() == 0:
+        return rois.new().long(), rois.new().long()
+    offset_x = rois[:, 0]
+    offset_y = rois[:, 1]
+    scale_x = heatmap_size / (rois[:, 2] - rois[:, 0])
+    scale_y = heatmap_size / (rois[:, 3] - rois[:, 1])
+
+    offset_x = offset_x[:, None]
+    offset_y = offset_y[:, None]
+    scale_x = scale_x[:, None]
+    scale_y = scale_y[:, None]
+
+    x = keypoints[..., 0]
+    y = keypoints[..., 1]
+
+    x_boundary_inds = x == rois[:, 2][:, None]
+    y_boundary_inds = y == rois[:, 3][:, None]
+
+    x = (x - offset_x) * scale_x
+    x = x.floor().long()
+    y = (y - offset_y) * scale_y
+    y = y.floor().long()
+
+    x[x_boundary_inds] = heatmap_size - 1
+    y[y_boundary_inds] = heatmap_size - 1
+
+    valid_loc = (x >= 0) & (y >= 0) & (x < heatmap_size) & (y < heatmap_size)
+    vis = keypoints[..., 2] > 0
+    valid = (valid_loc & vis).long()
+
+    lin_ind = y * heatmap_size + x
+    heatmaps = lin_ind * valid
+
+    return heatmaps, valid
+
+
+@torch.jit.script_if_tracing
+def heatmaps_to_keypoints(maps: torch.Tensor, rois: torch.Tensor) -> torch.Tensor:
+    """
+    Extract predicted keypoint locations from heatmaps.
+
+    Args:
+        maps (Tensor): (#ROIs, #keypoints, POOL_H, POOL_W). The predicted heatmap of logits for
+            each ROI and each keypoint.
+        rois (Tensor): (#ROIs, 4). The box of each ROI.
+
+    Returns:
+        Tensor of shape (#ROIs, #keypoints, 4) with the last dimension corresponding to
+        (x, y, logit, score) for each keypoint.
+
+    When converting discrete pixel indices in an NxN image to a continuous keypoint coordinate,
+    we maintain consistency with :meth:`Keypoints.to_heatmap` by using the conversion from
+    Heckbert 1990: c = d + 0.5, where d is a discrete coordinate and c is a continuous coordinate.
+    """
+
+    offset_x = rois[:, 0]
+    offset_y = rois[:, 1]
+
+    widths = (rois[:, 2] - rois[:, 0]).clamp(min=1)
+    heights = (rois[:, 3] - rois[:, 1]).clamp(min=1)
+    widths_ceil = widths.ceil()
+    heights_ceil = heights.ceil()
+
+    num_rois, num_keypoints = maps.shape[:2]
+    xy_preds = maps.new_zeros(rois.shape[0], num_keypoints, 4)
+
+    width_corrections = widths / widths_ceil
+    height_corrections = heights / heights_ceil
+
+    keypoints_idx = torch.arange(num_keypoints, device=maps.device)
+
+    for i in range(num_rois):
+        outsize = (int(heights_ceil[i]), int(widths_ceil[i]))
+        roi_map = F.interpolate(
+            maps[[i]], size=outsize, mode="bicubic", align_corners=False
+        )
+
+        # Although semantically equivalent, `reshape` is used instead of `squeeze` due
+        # to limitation during ONNX export of `squeeze` in scripting mode
+        roi_map = roi_map.reshape(roi_map.shape[1:])  # keypoints x H x W
+
+        # softmax over the spatial region
+        max_score, _ = roi_map.view(num_keypoints, -1).max(1)
+        max_score = max_score.view(num_keypoints, 1, 1)
+        tmp_full_resolution = (roi_map - max_score).exp_()
+        tmp_pool_resolution = (maps[i] - max_score).exp_()
+        # Produce scores over the region H x W, but normalize with POOL_H x POOL_W,
+        # so that the scores of objects of different absolute sizes will be more comparable
+        roi_map_scores = tmp_full_resolution / tmp_pool_resolution.sum(
+            (1, 2), keepdim=True
+        )
+
+        w = roi_map.shape[2]
+        pos = roi_map.view(num_keypoints, -1).argmax(1)
+
+        x_int = pos % w
+        y_int = (pos - x_int) // w
+
+        assert (
+            roi_map_scores[keypoints_idx, y_int, x_int]
+            == roi_map_scores.view(num_keypoints, -1).max(1)[0]
+        ).all()
+
+        x = (x_int.float() + 0.5) * width_corrections[i]
+        y = (y_int.float() + 0.5) * height_corrections[i]
+
+        xy_preds[i, :, 0] = x + offset_x[i]
+        xy_preds[i, :, 1] = y + offset_y[i]
+        xy_preds[i, :, 2] = roi_map[keypoints_idx, y_int, x_int]
+        xy_preds[i, :, 3] = roi_map_scores[keypoints_idx, y_int, x_int]
+
+    return xy_preds

sam3/agent/helpers/mask_overlap_removal.py

@@ -0,0 +1,128 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+from typing import Dict, List
+
+import numpy as np
+import torch
+
+try:
+    from pycocotools import mask as mask_utils
+except Exception:
+    mask_utils = None
+
+
+def mask_intersection(
+    masks1: torch.Tensor, masks2: torch.Tensor, block_size: int = 16
+) -> torch.Tensor:
+    assert masks1.shape[1:] == masks2.shape[1:]
+    assert masks1.dtype == torch.bool and masks2.dtype == torch.bool
+    N, M = masks1.shape[0], masks2.shape[0]
+    out = torch.zeros(N, M, device=masks1.device, dtype=torch.long)
+    for i in range(0, N, block_size):
+        for j in range(0, M, block_size):
+            a = masks1[i : i + block_size]
+            b = masks2[j : j + block_size]
+            inter = (a[:, None] & b[None, :]).flatten(-2).sum(-1)
+            out[i : i + block_size, j : j + block_size] = inter
+    return out
+
+
+def mask_iom(masks1: torch.Tensor, masks2: torch.Tensor) -> torch.Tensor:
+    assert masks1.shape[1:] == masks2.shape[1:]
+    assert masks1.dtype == torch.bool and masks2.dtype == torch.bool
+    inter = mask_intersection(masks1, masks2)
+    area1 = masks1.flatten(-2).sum(-1)  # (N,)
+    area2 = masks2.flatten(-2).sum(-1)  # (M,)
+    min_area = torch.min(area1[:, None], area2[None, :]).clamp_min(1)
+    return inter.float() / (min_area.float() + 1e-8)
+
+
+def _decode_single_mask(mask_repr, h: int, w: int) -> np.ndarray:
+    if isinstance(mask_repr, (list, tuple, np.ndarray)):
+        arr = np.array(mask_repr)
+        if arr.ndim != 2:
+            raise ValueError("Mask array must be 2D (H, W).")
+        return (arr > 0).astype(np.uint8)
+
+    if mask_utils is None:
+        raise ImportError(
+            "pycocotools is required to decode RLE mask strings. pip install pycocotools"
+        )
+
+    if not isinstance(mask_repr, (str, bytes)):
+        raise ValueError("Unsupported mask representation type for RLE decode.")
+
+    rle = {
+        "counts": mask_repr if isinstance(mask_repr, (str, bytes)) else str(mask_repr),
+        "size": [h, w],
+    }
+    decoded = mask_utils.decode(rle)
+    if decoded.ndim == 3:
+        decoded = decoded[:, :, 0]
+    return (decoded > 0).astype(np.uint8)
+
+
+def _decode_masks_to_torch_bool(pred_masks: List, h: int, w: int) -> torch.Tensor:
+    bin_masks = [_decode_single_mask(m, h, w) for m in pred_masks]
+    masks_np = np.stack(bin_masks, axis=0).astype(np.uint8)  # (N, H, W)
+    return torch.from_numpy(masks_np > 0)
+
+
+def remove_overlapping_masks(sample: Dict, iom_thresh: float = 0.3) -> Dict:
+    """
+    Greedy keep: sort by score desc; keep a mask if IoM to all kept masks <= threshold.
+    If pred_masks has length 0 or 1, returns sample unchanged (no extra keys).
+    """
+    # Basic presence checks
+    if "pred_masks" not in sample or not isinstance(sample["pred_masks"], list):
+        return sample  # nothing to do / preserve as-is
+
+    pred_masks = sample["pred_masks"]
+    N = len(pred_masks)
+
+    # --- Early exit: 0 or 1 mask -> do NOT modify the JSON at all ---
+    if N <= 1:
+        return sample
+
+    # From here on we have at least 2 masks
+    h = int(sample["orig_img_h"])
+    w = int(sample["orig_img_w"])
+    pred_scores = sample.get("pred_scores", [1.0] * N)  # fallback if scores missing
+    pred_boxes = sample.get("pred_boxes", None)
+
+    assert N == len(pred_scores), "pred_masks and pred_scores must have same length"
+    if pred_boxes is not None:
+        assert N == len(pred_boxes), "pred_masks and pred_boxes must have same length"
+
+    masks_bool = _decode_masks_to_torch_bool(pred_masks, h, w)  # (N, H, W)
+
+    order = sorted(range(N), key=lambda i: float(pred_scores[i]), reverse=True)
+    kept_idx: List[int] = []
+    kept_masks: List[torch.Tensor] = []
+
+    for i in order:
+        cand = masks_bool[i].unsqueeze(0)  # (1, H, W)
+        if len(kept_masks) == 0:
+            kept_idx.append(i)
+            kept_masks.append(masks_bool[i])
+            continue
+
+        kept_stack = torch.stack(kept_masks, dim=0)  # (K, H, W)
+        iom_vals = mask_iom(cand, kept_stack).squeeze(0)  # (K,)
+        if torch.any(iom_vals > iom_thresh):
+            continue  # overlaps too much with a higher-scored kept mask
+        kept_idx.append(i)
+        kept_masks.append(masks_bool[i])
+
+    kept_idx_sorted = sorted(kept_idx)
+
+    # Build filtered JSON (this *does* modify fields; only for N>=2 case)
+    out = dict(sample)
+    out["pred_masks"] = [pred_masks[i] for i in kept_idx_sorted]
+    out["pred_scores"] = [pred_scores[i] for i in kept_idx_sorted]
+    if pred_boxes is not None:
+        out["pred_boxes"] = [pred_boxes[i] for i in kept_idx_sorted]
+    out["kept_indices"] = kept_idx_sorted
+    out["removed_indices"] = [i for i in range(N) if i not in set(kept_idx_sorted)]
+    out["iom_threshold"] = float(iom_thresh)
+    return out

sam3/agent/helpers/masks.py

@@ -0,0 +1,560 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import copy
+import itertools
+from typing import Any, Iterator, List, Union
+
+import numpy as np
+import pycocotools.mask as mask_util
+import torch
+from torch import device
+
+from .boxes import Boxes
+from .memory import retry_if_cuda_oom
+
+from .roi_align import ROIAlign
+
+
+def polygon_area(x, y):
+    # Using the shoelace formula
+    # https://stackoverflow.com/questions/24467972/calculate-area-of-polygon-given-x-y-coordinates
+    return 0.5 * np.abs(np.dot(x, np.roll(y, 1)) - np.dot(y, np.roll(x, 1)))
+
+
+def polygons_to_bitmask(
+    polygons: List[np.ndarray], height: int, width: int
+) -> np.ndarray:
+    """
+    Args:
+        polygons (list[ndarray]): each array has shape (Nx2,)
+        height, width (int)
+
+    Returns:
+        ndarray: a bool mask of shape (height, width)
+    """
+    if len(polygons) == 0:
+        # COCOAPI does not support empty polygons
+        return np.zeros((height, width)).astype(bool)
+    rles = mask_util.frPyObjects(polygons, height, width)
+    rle = mask_util.merge(rles)
+    return mask_util.decode(rle).astype(bool)
+
+
+def rasterize_polygons_within_box(
+    polygons: List[np.ndarray], box: np.ndarray, mask_size: int
+) -> torch.Tensor:
+    """
+    Rasterize the polygons into a mask image and
+    crop the mask content in the given box.
+    The cropped mask is resized to (mask_size, mask_size).
+
+    This function is used when generating training targets for mask head in Mask R-CNN.
+    Given original ground-truth masks for an image, new ground-truth mask
+    training targets in the size of `mask_size x mask_size`
+    must be provided for each predicted box. This function will be called to
+    produce such targets.
+
+    Args:
+        polygons (list[ndarray[float]]): a list of polygons, which represents an instance.
+        box: 4-element numpy array
+        mask_size (int):
+
+    Returns:
+        Tensor: BoolTensor of shape (mask_size, mask_size)
+    """
+    # 1. Shift the polygons w.r.t the boxes
+    w, h = box[2] - box[0], box[3] - box[1]
+
+    polygons = copy.deepcopy(polygons)
+    for p in polygons:
+        p[0::2] = p[0::2] - box[0]
+        p[1::2] = p[1::2] - box[1]
+
+    # 2. Rescale the polygons to the new box size
+    # max() to avoid division by small number
+    ratio_h = mask_size / max(h, 0.1)
+    ratio_w = mask_size / max(w, 0.1)
+
+    if ratio_h == ratio_w:
+        for p in polygons:
+            p *= ratio_h
+    else:
+        for p in polygons:
+            p[0::2] *= ratio_w
+            p[1::2] *= ratio_h
+
+    # 3. Rasterize the polygons with coco api
+    mask = polygons_to_bitmask(polygons, mask_size, mask_size)
+    mask = torch.from_numpy(mask)
+    return mask
+
+
+class BitMasks:
+    """
+    This class stores the segmentation masks for all objects in one image, in
+    the form of bitmaps.
+
+    Attributes:
+        tensor: bool Tensor of N,H,W, representing N instances in the image.
+    """
+
+    def __init__(self, tensor: Union[torch.Tensor, np.ndarray]):
+        """
+        Args:
+            tensor: bool Tensor of N,H,W, representing N instances in the image.
+        """
+        if isinstance(tensor, torch.Tensor):
+            tensor = tensor.to(torch.bool)
+        else:
+            tensor = torch.as_tensor(
+                tensor, dtype=torch.bool, device=torch.device("cpu")
+            )
+        assert tensor.dim() == 3, tensor.size()
+        self.image_size = tensor.shape[1:]
+        self.tensor = tensor
+
+    @torch.jit.unused
+    def to(self, *args: Any, **kwargs: Any) -> "BitMasks":
+        return BitMasks(self.tensor.to(*args, **kwargs))
+
+    @property
+    def device(self) -> torch.device:
+        return self.tensor.device
+
+    @torch.jit.unused
+    def __getitem__(self, item: Union[int, slice, torch.BoolTensor]) -> "BitMasks":
+        """
+        Returns:
+            BitMasks: Create a new :class:`BitMasks` by indexing.
+
+        The following usage are allowed:
+
+        1. `new_masks = masks[3]`: return a `BitMasks` which contains only one mask.
+        2. `new_masks = masks[2:10]`: return a slice of masks.
+        3. `new_masks = masks[vector]`, where vector is a torch.BoolTensor
+           with `length = len(masks)`. Nonzero elements in the vector will be selected.
+
+        Note that the returned object might share storage with this object,
+        subject to Pytorch's indexing semantics.
+        """
+        if isinstance(item, int):
+            return BitMasks(self.tensor[item].unsqueeze(0))
+        m = self.tensor[item]
+        assert (
+            m.dim() == 3
+        ), "Indexing on BitMasks with {} returns a tensor with shape {}!".format(
+            item, m.shape
+        )
+        return BitMasks(m)
+
+    @torch.jit.unused
+    def __iter__(self) -> torch.Tensor:
+        yield from self.tensor
+
+    @torch.jit.unused
+    def __repr__(self) -> str:
+        s = self.__class__.__name__ + "("
+        s += "num_instances={})".format(len(self.tensor))
+        return s
+
+    def __len__(self) -> int:
+        return self.tensor.shape[0]
+
+    def nonempty(self) -> torch.Tensor:
+        """
+        Find masks that are non-empty.
+
+        Returns:
+            Tensor: a BoolTensor which represents
+                whether each mask is empty (False) or non-empty (True).
+        """
+        return self.tensor.flatten(1).any(dim=1)
+
+    @staticmethod
+    def from_polygon_masks(
+        polygon_masks: Union["PolygonMasks", List[List[np.ndarray]]],
+        height: int,
+        width: int,
+    ) -> "BitMasks":
+        """
+        Args:
+            polygon_masks (list[list[ndarray]] or PolygonMasks)
+            height, width (int)
+        """
+        if isinstance(polygon_masks, PolygonMasks):
+            polygon_masks = polygon_masks.polygons
+        masks = [polygons_to_bitmask(p, height, width) for p in polygon_masks]
+        if len(masks):
+            return BitMasks(torch.stack([torch.from_numpy(x) for x in masks]))
+        else:
+            return BitMasks(torch.empty(0, height, width, dtype=torch.bool))
+
+    @staticmethod
+    def from_roi_masks(roi_masks: "ROIMasks", height: int, width: int) -> "BitMasks":
+        """
+        Args:
+            roi_masks:
+            height, width (int):
+        """
+        return roi_masks.to_bitmasks(height, width)
+
+    def crop_and_resize(self, boxes: torch.Tensor, mask_size: int) -> torch.Tensor:
+        """
+        Crop each bitmask by the given box, and resize results to (mask_size, mask_size).
+        This can be used to prepare training targets for Mask R-CNN.
+        It has less reconstruction error compared to rasterization with polygons.
+        However we observe no difference in accuracy,
+        but BitMasks requires more memory to store all the masks.
+
+        Args:
+            boxes (Tensor): Nx4 tensor storing the boxes for each mask
+            mask_size (int): the size of the rasterized mask.
+
+        Returns:
+            Tensor:
+                A bool tensor of shape (N, mask_size, mask_size), where
+                N is the number of predicted boxes for this image.
+        """
+        assert len(boxes) == len(self), "{} != {}".format(len(boxes), len(self))
+        device = self.tensor.device
+
+        batch_inds = torch.arange(len(boxes), device=device).to(dtype=boxes.dtype)[
+            :, None
+        ]
+        rois = torch.cat([batch_inds, boxes], dim=1)  # Nx5
+
+        bit_masks = self.tensor.to(dtype=torch.float32)
+        rois = rois.to(device=device)
+        output = (
+            ROIAlign((mask_size, mask_size), 1.0, 0, aligned=True)
+            .forward(bit_masks[:, None, :, :], rois)
+            .squeeze(1)
+        )
+        output = output >= 0.5
+        return output
+
+    def get_bounding_boxes(self) -> Boxes:
+        """
+        Returns:
+            Boxes: tight bounding boxes around bitmasks.
+            If a mask is empty, it's bounding box will be all zero.
+        """
+        boxes = torch.zeros(self.tensor.shape[0], 4, dtype=torch.float32)
+        x_any = torch.any(self.tensor, dim=1)
+        y_any = torch.any(self.tensor, dim=2)
+        for idx in range(self.tensor.shape[0]):
+            x = torch.where(x_any[idx, :])[0]
+            y = torch.where(y_any[idx, :])[0]
+            if len(x) > 0 and len(y) > 0:
+                boxes[idx, :] = torch.as_tensor(
+                    [x[0], y[0], x[-1] + 1, y[-1] + 1], dtype=torch.float32
+                )
+        return Boxes(boxes)
+
+    @staticmethod
+    def cat(bitmasks_list: List["BitMasks"]) -> "BitMasks":
+        """
+        Concatenates a list of BitMasks into a single BitMasks
+
+        Arguments:
+            bitmasks_list (list[BitMasks])
+
+        Returns:
+            BitMasks: the concatenated BitMasks
+        """
+        assert isinstance(bitmasks_list, (list, tuple))
+        assert len(bitmasks_list) > 0
+        assert all(isinstance(bitmask, BitMasks) for bitmask in bitmasks_list)
+
+        cat_bitmasks = type(bitmasks_list[0])(
+            torch.cat([bm.tensor for bm in bitmasks_list], dim=0)
+        )
+        return cat_bitmasks
+
+
+class PolygonMasks:
+    """
+    This class stores the segmentation masks for all objects in one image, in the form of polygons.
+
+    Attributes:
+        polygons: list[list[ndarray]]. Each ndarray is a float64 vector representing a polygon.
+    """
+
+    def __init__(self, polygons: List[List[Union[torch.Tensor, np.ndarray]]]):
+        """
+        Arguments:
+            polygons (list[list[np.ndarray]]): The first
+                level of the list correspond to individual instances,
+                the second level to all the polygons that compose the
+                instance, and the third level to the polygon coordinates.
+                The third level array should have the format of
+                [x0, y0, x1, y1, ..., xn, yn] (n >= 3).
+        """
+        if not isinstance(polygons, list):
+            raise ValueError(
+                "Cannot create PolygonMasks: Expect a list of list of polygons per image. "
+                "Got '{}' instead.".format(type(polygons))
+            )
+
+        def _make_array(t: Union[torch.Tensor, np.ndarray]) -> np.ndarray:
+            # Use float64 for higher precision, because why not?
+            # Always put polygons on CPU (self.to is a no-op) since they
+            # are supposed to be small tensors.
+            # May need to change this assumption if GPU placement becomes useful
+            if isinstance(t, torch.Tensor):
+                t = t.cpu().numpy()
+            return np.asarray(t).astype("float64")
+
+        def process_polygons(
+            polygons_per_instance: List[Union[torch.Tensor, np.ndarray]],
+        ) -> List[np.ndarray]:
+            if not isinstance(polygons_per_instance, list):
+                raise ValueError(
+                    "Cannot create polygons: Expect a list of polygons per instance. "
+                    "Got '{}' instead.".format(type(polygons_per_instance))
+                )
+            # transform each polygon to a numpy array
+            polygons_per_instance = [_make_array(p) for p in polygons_per_instance]
+            for polygon in polygons_per_instance:
+                if len(polygon) % 2 != 0 or len(polygon) < 6:
+                    raise ValueError(
+                        f"Cannot create a polygon from {len(polygon)} coordinates."
+                    )
+            return polygons_per_instance
+
+        self.polygons: List[List[np.ndarray]] = [
+            process_polygons(polygons_per_instance)
+            for polygons_per_instance in polygons
+        ]
+
+    def to(self, *args: Any, **kwargs: Any) -> "PolygonMasks":
+        return self
+
+    @property
+    def device(self) -> torch.device:
+        return torch.device("cpu")
+
+    def get_bounding_boxes(self) -> Boxes:
+        """
+        Returns:
+            Boxes: tight bounding boxes around polygon masks.
+        """
+        boxes = torch.zeros(len(self.polygons), 4, dtype=torch.float32)
+        for idx, polygons_per_instance in enumerate(self.polygons):
+            minxy = torch.as_tensor([float("inf"), float("inf")], dtype=torch.float32)
+            maxxy = torch.zeros(2, dtype=torch.float32)
+            for polygon in polygons_per_instance:
+                coords = torch.from_numpy(polygon).view(-1, 2).to(dtype=torch.float32)
+                minxy = torch.min(minxy, torch.min(coords, dim=0).values)
+                maxxy = torch.max(maxxy, torch.max(coords, dim=0).values)
+            boxes[idx, :2] = minxy
+            boxes[idx, 2:] = maxxy
+        return Boxes(boxes)
+
+    def nonempty(self) -> torch.Tensor:
+        """
+        Find masks that are non-empty.
+
+        Returns:
+            Tensor:
+                a BoolTensor which represents whether each mask is empty (False) or not (True).
+        """
+        keep = [1 if len(polygon) > 0 else 0 for polygon in self.polygons]
+        return torch.from_numpy(np.asarray(keep, dtype=bool))
+
+    def __getitem__(
+        self, item: Union[int, slice, List[int], torch.BoolTensor]
+    ) -> "PolygonMasks":
+        """
+        Support indexing over the instances and return a `PolygonMasks` object.
+        `item` can be:
+
+        1. An integer. It will return an object with only one instance.
+        2. A slice. It will return an object with the selected instances.
+        3. A list[int]. It will return an object with the selected instances,
+           correpsonding to the indices in the list.
+        4. A vector mask of type BoolTensor, whose length is num_instances.
+           It will return an object with the instances whose mask is nonzero.
+        """
+        if isinstance(item, int):
+            selected_polygons = [self.polygons[item]]
+        elif isinstance(item, slice):
+            selected_polygons = self.polygons[item]
+        elif isinstance(item, list):
+            selected_polygons = [self.polygons[i] for i in item]
+        elif isinstance(item, torch.Tensor):
+            # Polygons is a list, so we have to move the indices back to CPU.
+            if item.dtype == torch.bool:
+                assert item.dim() == 1, item.shape
+                item = item.nonzero().squeeze(1).cpu().numpy().tolist()
+            elif item.dtype in [torch.int32, torch.int64]:
+                item = item.cpu().numpy().tolist()
+            else:
+                raise ValueError(
+                    "Unsupported tensor dtype={} for indexing!".format(item.dtype)
+                )
+            selected_polygons = [self.polygons[i] for i in item]
+        return PolygonMasks(selected_polygons)
+
+    def __iter__(self) -> Iterator[List[np.ndarray]]:
+        """
+        Yields:
+            list[ndarray]: the polygons for one instance.
+            Each Tensor is a float64 vector representing a polygon.
+        """
+        return iter(self.polygons)
+
+    def __repr__(self) -> str:
+        s = self.__class__.__name__ + "("
+        s += "num_instances={})".format(len(self.polygons))
+        return s
+
+    def __len__(self) -> int:
+        return len(self.polygons)
+
+    def crop_and_resize(self, boxes: torch.Tensor, mask_size: int) -> torch.Tensor:
+        """
+        Crop each mask by the given box, and resize results to (mask_size, mask_size).
+        This can be used to prepare training targets for Mask R-CNN.
+
+        Args:
+            boxes (Tensor): Nx4 tensor storing the boxes for each mask
+            mask_size (int): the size of the rasterized mask.
+
+        Returns:
+            Tensor: A bool tensor of shape (N, mask_size, mask_size), where
+            N is the number of predicted boxes for this image.
+        """
+        assert len(boxes) == len(self), "{} != {}".format(len(boxes), len(self))
+
+        device = boxes.device
+        # Put boxes on the CPU, as the polygon representation is not efficient GPU-wise
+        # (several small tensors for representing a single instance mask)
+        boxes = boxes.to(torch.device("cpu"))
+
+        results = [
+            rasterize_polygons_within_box(poly, box.numpy(), mask_size)
+            for poly, box in zip(self.polygons, boxes)
+        ]
+        """
+        poly: list[list[float]], the polygons for one instance
+        box: a tensor of shape (4,)
+        """
+        if len(results) == 0:
+            return torch.empty(0, mask_size, mask_size, dtype=torch.bool, device=device)
+        return torch.stack(results, dim=0).to(device=device)
+
+    def area(self):
+        """
+        Computes area of the mask.
+        Only works with Polygons, using the shoelace formula:
+        https://stackoverflow.com/questions/24467972/calculate-area-of-polygon-given-x-y-coordinates
+
+        Returns:
+            Tensor: a vector, area for each instance
+        """
+
+        area = []
+        for polygons_per_instance in self.polygons:
+            area_per_instance = 0
+            for p in polygons_per_instance:
+                area_per_instance += polygon_area(p[0::2], p[1::2])
+            area.append(area_per_instance)
+
+        return torch.tensor(area)
+
+    @staticmethod
+    def cat(polymasks_list: List["PolygonMasks"]) -> "PolygonMasks":
+        """
+        Concatenates a list of PolygonMasks into a single PolygonMasks
+
+        Arguments:
+            polymasks_list (list[PolygonMasks])
+
+        Returns:
+            PolygonMasks: the concatenated PolygonMasks
+        """
+        assert isinstance(polymasks_list, (list, tuple))
+        assert len(polymasks_list) > 0
+        assert all(isinstance(polymask, PolygonMasks) for polymask in polymasks_list)
+
+        cat_polymasks = type(polymasks_list[0])(
+            list(itertools.chain.from_iterable(pm.polygons for pm in polymasks_list))
+        )
+        return cat_polymasks
+
+
+class ROIMasks:
+    """
+    Represent masks by N smaller masks defined in some ROIs. Once ROI boxes are given,
+    full-image bitmask can be obtained by "pasting" the mask on the region defined
+    by the corresponding ROI box.
+    """
+
+    def __init__(self, tensor: torch.Tensor):
+        """
+        Args:
+            tensor: (N, M, M) mask tensor that defines the mask within each ROI.
+        """
+        if tensor.dim() != 3:
+            raise ValueError("ROIMasks must take a masks of 3 dimension.")
+        self.tensor = tensor
+
+    def to(self, device: torch.device) -> "ROIMasks":
+        return ROIMasks(self.tensor.to(device))
+
+    @property
+    def device(self) -> device:
+        return self.tensor.device
+
+    def __len__(self):
+        return self.tensor.shape[0]
+
+    def __getitem__(self, item) -> "ROIMasks":
+        """
+        Returns:
+            ROIMasks: Create a new :class:`ROIMasks` by indexing.
+
+        The following usage are allowed:
+
+        1. `new_masks = masks[2:10]`: return a slice of masks.
+        2. `new_masks = masks[vector]`, where vector is a torch.BoolTensor
+           with `length = len(masks)`. Nonzero elements in the vector will be selected.
+
+        Note that the returned object might share storage with this object,
+        subject to Pytorch's indexing semantics.
+        """
+        t = self.tensor[item]
+        if t.dim() != 3:
+            raise ValueError(
+                f"Indexing on ROIMasks with {item} returns a tensor with shape {t.shape}!"
+            )
+        return ROIMasks(t)
+
+    @torch.jit.unused
+    def __repr__(self) -> str:
+        s = self.__class__.__name__ + "("
+        s += "num_instances={})".format(len(self.tensor))
+        return s
+
+    @torch.jit.unused
+    def to_bitmasks(self, boxes: torch.Tensor, height, width, threshold=0.5):
+        """
+        Args: see documentation of :func:`paste_masks_in_image`.
+        """
+        from detectron2.layers.mask_ops import (
+            _paste_masks_tensor_shape,
+            paste_masks_in_image,
+        )
+
+        if torch.jit.is_tracing():
+            if isinstance(height, torch.Tensor):
+                paste_func = _paste_masks_tensor_shape
+            else:
+                paste_func = paste_masks_in_image
+        else:
+            paste_func = retry_if_cuda_oom(paste_masks_in_image)
+        bitmasks = paste_func(
+            self.tensor, boxes.tensor, (height, width), threshold=threshold
+        )
+        return BitMasks(bitmasks)

sam3/agent/helpers/memory.py

@@ -0,0 +1,87 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import logging
+from contextlib import contextmanager
+from functools import wraps
+
+import torch
+
+__all__ = ["retry_if_cuda_oom"]
+
+
+@contextmanager
+def _ignore_torch_cuda_oom():
+    """
+    A context which ignores CUDA OOM exception from pytorch.
+    """
+    try:
+        yield
+    except RuntimeError as e:
+        # NOTE: the string may change?
+        if "CUDA out of memory. " in str(e):
+            pass
+        else:
+            raise
+
+
+def retry_if_cuda_oom(func):
+    """
+    Makes a function retry itself after encountering
+    pytorch's CUDA OOM error.
+    It will first retry after calling `torch.cuda.empty_cache()`.
+
+    If that still fails, it will then retry by trying to convert inputs to CPUs.
+    In this case, it expects the function to dispatch to CPU implementation.
+    The return values may become CPU tensors as well and it's user's
+    responsibility to convert it back to CUDA tensor if needed.
+
+    Args:
+        func: a stateless callable that takes tensor-like objects as arguments
+
+    Returns:
+        a callable which retries `func` if OOM is encountered.
+
+    Examples:
+    ::
+        output = retry_if_cuda_oom(some_torch_function)(input1, input2)
+        # output may be on CPU even if inputs are on GPU
+
+    Note:
+        1. When converting inputs to CPU, it will only look at each argument and check
+           if it has `.device` and `.to` for conversion. Nested structures of tensors
+           are not supported.
+
+        2. Since the function might be called more than once, it has to be
+           stateless.
+    """
+
+    def maybe_to_cpu(x):
+        try:
+            like_gpu_tensor = x.device.type == "cuda" and hasattr(x, "to")
+        except AttributeError:
+            like_gpu_tensor = False
+        if like_gpu_tensor:
+            return x.to(device="cpu")
+        else:
+            return x
+
+    @wraps(func)
+    def wrapped(*args, **kwargs):
+        with _ignore_torch_cuda_oom():
+            return func(*args, **kwargs)
+
+        # Clear cache and retry
+        torch.cuda.empty_cache()
+        with _ignore_torch_cuda_oom():
+            return func(*args, **kwargs)
+
+        # Try on CPU. This slows down the code significantly, therefore print a notice.
+        logger = logging.getLogger(__name__)
+        logger.info(
+            "Attempting to copy inputs of {} to CPU due to CUDA OOM".format(str(func))
+        )
+        new_args = (maybe_to_cpu(x) for x in args)
+        new_kwargs = {k: maybe_to_cpu(v) for k, v in kwargs.items()}
+        return func(*new_args, **new_kwargs)
+
+    return wrapped

sam3/agent/helpers/rle.py

@@ -0,0 +1,122 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+"""Some utilities for RLE encoding that doesn't require downloading the masks to the cpu"""
+
+import numpy as np
+import torch
+from pycocotools import mask as mask_util
+
+
+@torch.no_grad()
+def rle_encode(orig_mask, return_areas=False):
+    """Encodes a collection of masks in RLE format
+
+    This function emulates the behavior of the COCO API's encode function, but
+    is executed partially on the GPU for faster execution.
+
+    Args:
+        mask (torch.Tensor): A mask of shape (N, H, W) with dtype=torch.bool
+        return_areas (bool): If True, add the areas of the masks as a part of
+            the RLE output dict under the "area" key. Default is False.
+
+    Returns:
+        str: The RLE encoded masks
+    """
+    assert orig_mask.ndim == 3, "Mask must be of shape (N, H, W)"
+    assert orig_mask.dtype == torch.bool, "Mask must have dtype=torch.bool"
+
+    if orig_mask.numel() == 0:
+        return []
+
+    # First, transpose the spatial dimensions.
+    # This is necessary because the COCO API uses Fortran order
+    mask = orig_mask.transpose(1, 2)
+
+    # Flatten the mask
+    flat_mask = mask.reshape(mask.shape[0], -1)
+    if return_areas:
+        mask_areas = flat_mask.sum(-1).tolist()
+    # Find the indices where the mask changes
+    differences = torch.ones(
+        mask.shape[0], flat_mask.shape[1] + 1, device=mask.device, dtype=torch.bool
+    )
+    differences[:, 1:-1] = flat_mask[:, :-1] != flat_mask[:, 1:]
+    differences[:, 0] = flat_mask[:, 0]
+    _, change_indices = torch.where(differences)
+
+    try:
+        boundaries = torch.cumsum(differences.sum(-1), 0).cpu()
+    except RuntimeError as _:
+        boundaries = torch.cumsum(differences.cpu().sum(-1), 0)
+
+    change_indices_clone = change_indices.clone()
+    # First pass computes the RLEs on GPU, in a flatten format
+    for i in range(mask.shape[0]):
+        # Get the change indices for this batch item
+        beg = 0 if i == 0 else boundaries[i - 1].item()
+        end = boundaries[i].item()
+        change_indices[beg + 1 : end] -= change_indices_clone[beg : end - 1]
+
+    # Now we can split the RLES of each batch item, and convert them to strings
+    # No more gpu at this point
+    change_indices = change_indices.tolist()
+
+    batch_rles = []
+    # Process each mask in the batch separately
+    for i in range(mask.shape[0]):
+        beg = 0 if i == 0 else boundaries[i - 1].item()
+        end = boundaries[i].item()
+        run_lengths = change_indices[beg:end]
+
+        uncompressed_rle = {"counts": run_lengths, "size": list(orig_mask.shape[1:])}
+        h, w = uncompressed_rle["size"]
+        rle = mask_util.frPyObjects(uncompressed_rle, h, w)
+        rle["counts"] = rle["counts"].decode("utf-8")
+        if return_areas:
+            rle["area"] = mask_areas[i]
+        batch_rles.append(rle)
+
+    return batch_rles
+
+
+def robust_rle_encode(masks):
+    """Encodes a collection of masks in RLE format. Uses the gpu version fist, falls back to the cpu version if it fails"""
+
+    assert masks.ndim == 3, "Mask must be of shape (N, H, W)"
+    assert masks.dtype == torch.bool, "Mask must have dtype=torch.bool"
+
+    try:
+        return rle_encode(masks)
+    except RuntimeError as _:
+        masks = masks.cpu().numpy()
+        rles = [
+            mask_util.encode(
+                np.array(mask[:, :, np.newaxis], dtype=np.uint8, order="F")
+            )[0]
+            for mask in masks
+        ]
+        for rle in rles:
+            rle["counts"] = rle["counts"].decode("utf-8")
+        return rles
+
+
+def ann_to_rle(segm, im_info):
+    """Convert annotation which can be polygons, uncompressed RLE to RLE.
+    Args:
+        ann (dict) : annotation object
+    Returns:
+        ann (rle)
+    """
+    h, w = im_info["height"], im_info["width"]
+    if isinstance(segm, list):
+        # polygon -- a single object might consist of multiple parts
+        # we merge all parts into one mask rle code
+        rles = mask_util.frPyObjects(segm, h, w)
+        rle = mask_util.merge(rles)
+    elif isinstance(segm["counts"], list):
+        # uncompressed RLE
+        rle = mask_util.frPyObjects(segm, h, w)
+    else:
+        # rle
+        rle = segm
+    return rle

sam3/agent/helpers/roi_align.py

@@ -0,0 +1,75 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+from torch import nn
+from torchvision.ops import roi_align
+
+
+# NOTE: torchvision's RoIAlign has a different default aligned=False
+class ROIAlign(nn.Module):
+    def __init__(self, output_size, spatial_scale, sampling_ratio, aligned=True):
+        """
+        Args:
+            output_size (tuple): h, w
+            spatial_scale (float): scale the input boxes by this number
+            sampling_ratio (int): number of inputs samples to take for each output
+                sample. 0 to take samples densely.
+            aligned (bool): if False, use the legacy implementation in
+                Detectron. If True, align the results more perfectly.
+
+        Note:
+            The meaning of aligned=True:
+
+            Given a continuous coordinate c, its two neighboring pixel indices (in our
+            pixel model) are computed by floor(c - 0.5) and ceil(c - 0.5). For example,
+            c=1.3 has pixel neighbors with discrete indices [0] and [1] (which are sampled
+            from the underlying signal at continuous coordinates 0.5 and 1.5). But the original
+            roi_align (aligned=False) does not subtract the 0.5 when computing neighboring
+            pixel indices and therefore it uses pixels with a slightly incorrect alignment
+            (relative to our pixel model) when performing bilinear interpolation.
+
+            With `aligned=True`,
+            we first appropriately scale the ROI and then shift it by -0.5
+            prior to calling roi_align. This produces the correct neighbors; see
+            detectron2/tests/test_roi_align.py for verification.
+
+            The difference does not make a difference to the model's performance if
+            ROIAlign is used together with conv layers.
+        """
+        super().__init__()
+        self.output_size = output_size
+        self.spatial_scale = spatial_scale
+        self.sampling_ratio = sampling_ratio
+        self.aligned = aligned
+
+        from torchvision import __version__
+
+        version = tuple(int(x) for x in __version__.split(".")[:2])
+        # https://github.com/pytorch/vision/pull/2438
+        assert version >= (0, 7), "Require torchvision >= 0.7"
+
+    def forward(self, input, rois):
+        """
+        Args:
+            input: NCHW images
+            rois: Bx5 boxes. First column is the index into N. The other 4 columns are xyxy.
+        """
+        assert rois.dim() == 2 and rois.size(1) == 5
+        if input.is_quantized:
+            input = input.dequantize()
+        return roi_align(
+            input,
+            rois.to(dtype=input.dtype),
+            self.output_size,
+            self.spatial_scale,
+            self.sampling_ratio,
+            self.aligned,
+        )
+
+    def __repr__(self):
+        tmpstr = self.__class__.__name__ + "("
+        tmpstr += "output_size=" + str(self.output_size)
+        tmpstr += ", spatial_scale=" + str(self.spatial_scale)
+        tmpstr += ", sampling_ratio=" + str(self.sampling_ratio)
+        tmpstr += ", aligned=" + str(self.aligned)
+        tmpstr += ")"
+        return tmpstr

sam3/agent/helpers/rotated_boxes.py

@@ -0,0 +1,533 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import math
+from typing import List, Tuple
+
+import torch
+
+# from detectron2.layers.rotated_boxes import pairwise_iou_rotated
+
+from .boxes import Boxes
+
+
+def pairwise_iou_rotated(boxes1, boxes2):
+    """
+    Return intersection-over-union (Jaccard index) of boxes.
+
+    Both sets of boxes are expected to be in
+    (x_center, y_center, width, height, angle) format.
+
+    Arguments:
+        boxes1 (Tensor[N, 5])
+        boxes2 (Tensor[M, 5])
+
+    Returns:
+        iou (Tensor[N, M]): the NxM matrix containing the pairwise
+            IoU values for every element in boxes1 and boxes2
+    """
+    return torch.ops.detectron2.box_iou_rotated(boxes1, boxes2)
+
+
+class RotatedBoxes(Boxes):
+    """
+    This structure stores a list of rotated boxes as a Nx5 torch.Tensor.
+    It supports some common methods about boxes
+    (`area`, `clip`, `nonempty`, etc),
+    and also behaves like a Tensor
+    (support indexing, `to(device)`, `.device`, and iteration over all boxes)
+    """
+
+    def __init__(self, tensor: torch.Tensor):
+        """
+        Args:
+            tensor (Tensor[float]): a Nx5 matrix.  Each row is
+                (x_center, y_center, width, height, angle),
+                in which angle is represented in degrees.
+                While there's no strict range restriction for it,
+                the recommended principal range is between [-180, 180) degrees.
+
+        Assume we have a horizontal box B = (x_center, y_center, width, height),
+        where width is along the x-axis and height is along the y-axis.
+        The rotated box B_rot (x_center, y_center, width, height, angle)
+        can be seen as:
+
+        1. When angle == 0:
+           B_rot == B
+        2. When angle > 0:
+           B_rot is obtained by rotating B w.r.t its center by :math:`|angle|` degrees CCW;
+        3. When angle < 0:
+           B_rot is obtained by rotating B w.r.t its center by :math:`|angle|` degrees CW.
+
+        Mathematically, since the right-handed coordinate system for image space
+        is (y, x), where y is top->down and x is left->right, the 4 vertices of the
+        rotated rectangle :math:`(yr_i, xr_i)` (i = 1, 2, 3, 4) can be obtained from
+        the vertices of the horizontal rectangle :math:`(y_i, x_i)` (i = 1, 2, 3, 4)
+        in the following way (:math:`\\theta = angle*\\pi/180` is the angle in radians,
+        :math:`(y_c, x_c)` is the center of the rectangle):
+
+        .. math::
+
+            yr_i = \\cos(\\theta) (y_i - y_c) - \\sin(\\theta) (x_i - x_c) + y_c,
+
+            xr_i = \\sin(\\theta) (y_i - y_c) + \\cos(\\theta) (x_i - x_c) + x_c,
+
+        which is the standard rigid-body rotation transformation.
+
+        Intuitively, the angle is
+        (1) the rotation angle from y-axis in image space
+        to the height vector (top->down in the box's local coordinate system)
+        of the box in CCW, and
+        (2) the rotation angle from x-axis in image space
+        to the width vector (left->right in the box's local coordinate system)
+        of the box in CCW.
+
+        More intuitively, consider the following horizontal box ABCD represented
+        in (x1, y1, x2, y2): (3, 2, 7, 4),
+        covering the [3, 7] x [2, 4] region of the continuous coordinate system
+        which looks like this:
+
+        .. code:: none
+
+            O--------> x
+            |
+            |  A---B
+            |  |   |
+            |  D---C
+            |
+            v y
+
+        Note that each capital letter represents one 0-dimensional geometric point
+        instead of a 'square pixel' here.
+
+        In the example above, using (x, y) to represent a point we have:
+
+        .. math::
+
+            O = (0, 0), A = (3, 2), B = (7, 2), C = (7, 4), D = (3, 4)
+
+        We name vector AB = vector DC as the width vector in box's local coordinate system, and
+        vector AD = vector BC as the height vector in box's local coordinate system. Initially,
+        when angle = 0 degree, they're aligned with the positive directions of x-axis and y-axis
+        in the image space, respectively.
+
+        For better illustration, we denote the center of the box as E,
+
+        .. code:: none
+
+            O--------> x
+            |
+            |  A---B
+            |  | E |
+            |  D---C
+            |
+            v y
+
+        where the center E = ((3+7)/2, (2+4)/2) = (5, 3).
+
+        Also,
+
+        .. math::
+
+            width = |AB| = |CD| = 7 - 3 = 4,
+            height = |AD| = |BC| = 4 - 2 = 2.
+
+        Therefore, the corresponding representation for the same shape in rotated box in
+        (x_center, y_center, width, height, angle) format is:
+
+        (5, 3, 4, 2, 0),
+
+        Now, let's consider (5, 3, 4, 2, 90), which is rotated by 90 degrees
+        CCW (counter-clockwise) by definition. It looks like this:
+
+        .. code:: none
+
+            O--------> x
+            |   B-C
+            |   | |
+            |   |E|
+            |   | |
+            |   A-D
+            v y
+
+        The center E is still located at the same point (5, 3), while the vertices
+        ABCD are rotated by 90 degrees CCW with regard to E:
+        A = (4, 5), B = (4, 1), C = (6, 1), D = (6, 5)
+
+        Here, 90 degrees can be seen as the CCW angle to rotate from y-axis to
+        vector AD or vector BC (the top->down height vector in box's local coordinate system),
+        or the CCW angle to rotate from x-axis to vector AB or vector DC (the left->right
+        width vector in box's local coordinate system).
+
+        .. math::
+
+            width = |AB| = |CD| = 5 - 1 = 4,
+            height = |AD| = |BC| = 6 - 4 = 2.
+
+        Next, how about (5, 3, 4, 2, -90), which is rotated by 90 degrees CW (clockwise)
+        by definition? It looks like this:
+
+        .. code:: none
+
+            O--------> x
+            |   D-A
+            |   | |
+            |   |E|
+            |   | |
+            |   C-B
+            v y
+
+        The center E is still located at the same point (5, 3), while the vertices
+        ABCD are rotated by 90 degrees CW with regard to E:
+        A = (6, 1), B = (6, 5), C = (4, 5), D = (4, 1)
+
+        .. math::
+
+            width = |AB| = |CD| = 5 - 1 = 4,
+            height = |AD| = |BC| = 6 - 4 = 2.
+
+        This covers exactly the same region as (5, 3, 4, 2, 90) does, and their IoU
+        will be 1. However, these two will generate different RoI Pooling results and
+        should not be treated as an identical box.
+
+        On the other hand, it's easy to see that (X, Y, W, H, A) is identical to
+        (X, Y, W, H, A+360N), for any integer N. For example (5, 3, 4, 2, 270) would be
+        identical to (5, 3, 4, 2, -90), because rotating the shape 270 degrees CCW is
+        equivalent to rotating the same shape 90 degrees CW.
+
+        We could rotate further to get (5, 3, 4, 2, 180), or (5, 3, 4, 2, -180):
+
+        .. code:: none
+
+            O--------> x
+            |
+            |  C---D
+            |  | E |
+            |  B---A
+            |
+            v y
+
+        .. math::
+
+            A = (7, 4), B = (3, 4), C = (3, 2), D = (7, 2),
+
+            width = |AB| = |CD| = 7 - 3 = 4,
+            height = |AD| = |BC| = 4 - 2 = 2.
+
+        Finally, this is a very inaccurate (heavily quantized) illustration of
+        how (5, 3, 4, 2, 60) looks like in case anyone wonders:
+
+        .. code:: none
+
+            O--------> x
+            |     B\
+            |    /  C
+            |   /E /
+            |  A  /
+            |   `D
+            v y
+
+        It's still a rectangle with center of (5, 3), width of 4 and height of 2,
+        but its angle (and thus orientation) is somewhere between
+        (5, 3, 4, 2, 0) and (5, 3, 4, 2, 90).
+        """
+        device = (
+            tensor.device if isinstance(tensor, torch.Tensor) else torch.device("cpu")
+        )
+        tensor = torch.as_tensor(tensor, dtype=torch.float32, device=device)
+        if tensor.numel() == 0:
+            # Use reshape, so we don't end up creating a new tensor that does not depend on
+            # the inputs (and consequently confuses jit)
+            tensor = tensor.reshape((0, 5)).to(dtype=torch.float32, device=device)
+        assert tensor.dim() == 2 and tensor.size(-1) == 5, tensor.size()
+
+        self.tensor = tensor
+
+    def clone(self) -> "RotatedBoxes":
+        """
+        Clone the RotatedBoxes.
+
+        Returns:
+            RotatedBoxes
+        """
+        return RotatedBoxes(self.tensor.clone())
+
+    def to(self, device: torch.device, non_blocking: bool = False):
+        # Boxes are assumed float32 and does not support to(dtype)
+        return RotatedBoxes(self.tensor.to(device=device, non_blocking=non_blocking))
+
+    def area(self) -> torch.Tensor:
+        """
+        Computes the area of all the boxes.
+
+        Returns:
+            torch.Tensor: a vector with areas of each box.
+        """
+        box = self.tensor
+        area = box[:, 2] * box[:, 3]
+        return area
+
+    # Avoid in-place operations so that we can torchscript; NOTE: this creates a new tensor
+    def normalize_angles(self) -> None:
+        """
+        Restrict angles to the range of [-180, 180) degrees
+        """
+        angle_tensor = (self.tensor[:, 4] + 180.0) % 360.0 - 180.0
+        self.tensor = torch.cat((self.tensor[:, :4], angle_tensor[:, None]), dim=1)
+
+    def clip(
+        self, box_size: Tuple[int, int], clip_angle_threshold: float = 1.0
+    ) -> None:
+        """
+        Clip (in place) the boxes by limiting x coordinates to the range [0, width]
+        and y coordinates to the range [0, height].
+
+        For RRPN:
+        Only clip boxes that are almost horizontal with a tolerance of
+        clip_angle_threshold to maintain backward compatibility.
+
+        Rotated boxes beyond this threshold are not clipped for two reasons:
+
+        1. There are potentially multiple ways to clip a rotated box to make it
+           fit within the image.
+        2. It's tricky to make the entire rectangular box fit within the image
+           and still be able to not leave out pixels of interest.
+
+        Therefore we rely on ops like RoIAlignRotated to safely handle this.
+
+        Args:
+            box_size (height, width): The clipping box's size.
+            clip_angle_threshold:
+                Iff. abs(normalized(angle)) <= clip_angle_threshold (in degrees),
+                we do the clipping as horizontal boxes.
+        """
+        h, w = box_size
+
+        # normalize angles to be within (-180, 180] degrees
+        self.normalize_angles()
+
+        idx = torch.where(torch.abs(self.tensor[:, 4]) <= clip_angle_threshold)[0]
+
+        # convert to (x1, y1, x2, y2)
+        x1 = self.tensor[idx, 0] - self.tensor[idx, 2] / 2.0
+        y1 = self.tensor[idx, 1] - self.tensor[idx, 3] / 2.0
+        x2 = self.tensor[idx, 0] + self.tensor[idx, 2] / 2.0
+        y2 = self.tensor[idx, 1] + self.tensor[idx, 3] / 2.0
+
+        # clip
+        x1.clamp_(min=0, max=w)
+        y1.clamp_(min=0, max=h)
+        x2.clamp_(min=0, max=w)
+        y2.clamp_(min=0, max=h)
+
+        # convert back to (xc, yc, w, h)
+        self.tensor[idx, 0] = (x1 + x2) / 2.0
+        self.tensor[idx, 1] = (y1 + y2) / 2.0
+        # make sure widths and heights do not increase due to numerical errors
+        self.tensor[idx, 2] = torch.min(self.tensor[idx, 2], x2 - x1)
+        self.tensor[idx, 3] = torch.min(self.tensor[idx, 3], y2 - y1)
+
+    def nonempty(self, threshold: float = 0.0) -> torch.Tensor:
+        """
+        Find boxes that are non-empty.
+        A box is considered empty, if either of its side is no larger than threshold.
+
+        Returns:
+            Tensor: a binary vector which represents
+            whether each box is empty (False) or non-empty (True).
+        """
+        box = self.tensor
+        widths = box[:, 2]
+        heights = box[:, 3]
+        keep = (widths > threshold) & (heights > threshold)
+        return keep
+
+    def __getitem__(self, item) -> "RotatedBoxes":
+        """
+        Returns:
+            RotatedBoxes: Create a new :class:`RotatedBoxes` by indexing.
+
+        The following usage are allowed:
+
+        1. `new_boxes = boxes[3]`: return a `RotatedBoxes` which contains only one box.
+        2. `new_boxes = boxes[2:10]`: return a slice of boxes.
+        3. `new_boxes = boxes[vector]`, where vector is a torch.ByteTensor
+           with `length = len(boxes)`. Nonzero elements in the vector will be selected.
+
+        Note that the returned RotatedBoxes might share storage with this RotatedBoxes,
+        subject to Pytorch's indexing semantics.
+        """
+        if isinstance(item, int):
+            return RotatedBoxes(self.tensor[item].view(1, -1))
+        b = self.tensor[item]
+        assert (
+            b.dim() == 2
+        ), "Indexing on RotatedBoxes with {} failed to return a matrix!".format(item)
+        return RotatedBoxes(b)
+
+    def __len__(self) -> int:
+        return self.tensor.shape[0]
+
+    def __repr__(self) -> str:
+        return "RotatedBoxes(" + str(self.tensor) + ")"
+
+    def inside_box(
+        self, box_size: Tuple[int, int], boundary_threshold: int = 0
+    ) -> torch.Tensor:
+        """
+        Args:
+            box_size (height, width): Size of the reference box covering
+                [0, width] x [0, height]
+            boundary_threshold (int): Boxes that extend beyond the reference box
+                boundary by more than boundary_threshold are considered "outside".
+
+        For RRPN, it might not be necessary to call this function since it's common
+        for rotated box to extend to outside of the image boundaries
+        (the clip function only clips the near-horizontal boxes)
+
+        Returns:
+            a binary vector, indicating whether each box is inside the reference box.
+        """
+        height, width = box_size
+
+        cnt_x = self.tensor[..., 0]
+        cnt_y = self.tensor[..., 1]
+        half_w = self.tensor[..., 2] / 2.0
+        half_h = self.tensor[..., 3] / 2.0
+        a = self.tensor[..., 4]
+        c = torch.abs(torch.cos(a * math.pi / 180.0))
+        s = torch.abs(torch.sin(a * math.pi / 180.0))
+        # This basically computes the horizontal bounding rectangle of the rotated box
+        max_rect_dx = c * half_w + s * half_h
+        max_rect_dy = c * half_h + s * half_w
+
+        inds_inside = (
+            (cnt_x - max_rect_dx >= -boundary_threshold)
+            & (cnt_y - max_rect_dy >= -boundary_threshold)
+            & (cnt_x + max_rect_dx < width + boundary_threshold)
+            & (cnt_y + max_rect_dy < height + boundary_threshold)
+        )
+
+        return inds_inside
+
+    def get_centers(self) -> torch.Tensor:
+        """
+        Returns:
+            The box centers in a Nx2 array of (x, y).
+        """
+        return self.tensor[:, :2]
+
+    def scale(self, scale_x: float, scale_y: float) -> None:
+        """
+        Scale the rotated box with horizontal and vertical scaling factors
+        Note: when scale_factor_x != scale_factor_y,
+        the rotated box does not preserve the rectangular shape when the angle
+        is not a multiple of 90 degrees under resize transformation.
+        Instead, the shape is a parallelogram (that has skew)
+        Here we make an approximation by fitting a rotated rectangle to the parallelogram.
+        """
+        self.tensor[:, 0] *= scale_x
+        self.tensor[:, 1] *= scale_y
+        theta = self.tensor[:, 4] * math.pi / 180.0
+        c = torch.cos(theta)
+        s = torch.sin(theta)
+
+        # In image space, y is top->down and x is left->right
+        # Consider the local coordintate system for the rotated box,
+        # where the box center is located at (0, 0), and the four vertices ABCD are
+        # A(-w / 2, -h / 2), B(w / 2, -h / 2), C(w / 2, h / 2), D(-w / 2, h / 2)
+        # the midpoint of the left edge AD of the rotated box E is:
+        # E = (A+D)/2 = (-w / 2, 0)
+        # the midpoint of the top edge AB of the rotated box F is:
+        # F(0, -h / 2)
+        # To get the old coordinates in the global system, apply the rotation transformation
+        # (Note: the right-handed coordinate system for image space is yOx):
+        # (old_x, old_y) = (s * y + c * x, c * y - s * x)
+        # E(old) = (s * 0 + c * (-w/2), c * 0 - s * (-w/2)) = (-c * w / 2, s * w / 2)
+        # F(old) = (s * (-h / 2) + c * 0, c * (-h / 2) - s * 0) = (-s * h / 2, -c * h / 2)
+        # After applying the scaling factor (sfx, sfy):
+        # E(new) = (-sfx * c * w / 2, sfy * s * w / 2)
+        # F(new) = (-sfx * s * h / 2, -sfy * c * h / 2)
+        # The new width after scaling tranformation becomes:
+
+        # w(new) = |E(new) - O| * 2
+        #        = sqrt[(sfx * c * w / 2)^2 + (sfy * s * w / 2)^2] * 2
+        #        = sqrt[(sfx * c)^2 + (sfy * s)^2] * w
+        # i.e., scale_factor_w = sqrt[(sfx * c)^2 + (sfy * s)^2]
+        #
+        # For example,
+        # when angle = 0 or 180, |c| = 1, s = 0, scale_factor_w == scale_factor_x;
+        # when |angle| = 90, c = 0, |s| = 1, scale_factor_w == scale_factor_y
+        self.tensor[:, 2] *= torch.sqrt((scale_x * c) ** 2 + (scale_y * s) ** 2)
+
+        # h(new) = |F(new) - O| * 2
+        #        = sqrt[(sfx * s * h / 2)^2 + (sfy * c * h / 2)^2] * 2
+        #        = sqrt[(sfx * s)^2 + (sfy * c)^2] * h
+        # i.e., scale_factor_h = sqrt[(sfx * s)^2 + (sfy * c)^2]
+        #
+        # For example,
+        # when angle = 0 or 180, |c| = 1, s = 0, scale_factor_h == scale_factor_y;
+        # when |angle| = 90, c = 0, |s| = 1, scale_factor_h == scale_factor_x
+        self.tensor[:, 3] *= torch.sqrt((scale_x * s) ** 2 + (scale_y * c) ** 2)
+
+        # The angle is the rotation angle from y-axis in image space to the height
+        # vector (top->down in the box's local coordinate system) of the box in CCW.
+        #
+        # angle(new) = angle_yOx(O - F(new))
+        #            = angle_yOx( (sfx * s * h / 2, sfy * c * h / 2) )
+        #            = atan2(sfx * s * h / 2, sfy * c * h / 2)
+        #            = atan2(sfx * s, sfy * c)
+        #
+        # For example,
+        # when sfx == sfy, angle(new) == atan2(s, c) == angle(old)
+        self.tensor[:, 4] = torch.atan2(scale_x * s, scale_y * c) * 180 / math.pi
+
+    @classmethod
+    def cat(cls, boxes_list: List["RotatedBoxes"]) -> "RotatedBoxes":
+        """
+        Concatenates a list of RotatedBoxes into a single RotatedBoxes
+
+        Arguments:
+            boxes_list (list[RotatedBoxes])
+
+        Returns:
+            RotatedBoxes: the concatenated RotatedBoxes
+        """
+        assert isinstance(boxes_list, (list, tuple))
+        if len(boxes_list) == 0:
+            return cls(torch.empty(0))
+        assert all([isinstance(box, RotatedBoxes) for box in boxes_list])
+
+        # use torch.cat (v.s. layers.cat) so the returned boxes never share storage with input
+        cat_boxes = cls(torch.cat([b.tensor for b in boxes_list], dim=0))
+        return cat_boxes
+
+    @property
+    def device(self) -> torch.device:
+        return self.tensor.device
+
+    @torch.jit.unused
+    def __iter__(self):
+        """
+        Yield a box as a Tensor of shape (5,) at a time.
+        """
+        yield from self.tensor
+
+
+def pairwise_iou(boxes1: RotatedBoxes, boxes2: RotatedBoxes) -> None:
+    """
+    Given two lists of rotated boxes of size N and M,
+    compute the IoU (intersection over union)
+    between **all** N x M pairs of boxes.
+    The box order must be (x_center, y_center, width, height, angle).
+
+    Args:
+        boxes1, boxes2 (RotatedBoxes):
+            two `RotatedBoxes`. Contains N & M rotated boxes, respectively.
+
+    Returns:
+        Tensor: IoU, sized [N,M].
+    """
+
+    return pairwise_iou_rotated(boxes1.tensor, boxes2.tensor)

sam3/agent/helpers/som_utils.py

@@ -0,0 +1,406 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import colorsys
+from dataclasses import dataclass
+from typing import List, Tuple
+
+import cv2
+import matplotlib as mpl
+import matplotlib.colors as mplc
+import numpy as np
+import pycocotools.mask as mask_utils
+
+
+def rgb_to_hex(rgb_color):
+    """
+    Convert a rgb color to hex color.
+
+    Args:
+        rgb_color (tuple/list of ints): RGB color in tuple or list format.
+
+    Returns:
+        str: Hex color.
+
+    Example:
+        ```
+        >>> rgb_to_hex((255, 0, 244))
+        '#ff00ff'
+        ```
+    """
+    return "#" + "".join([hex(c)[2:].zfill(2) for c in rgb_color])
+
+
+# DEFAULT_COLOR_HEX_TO_NAME = {
+#     rgb_to_hex((255, 0, 0)): "red",
+#     rgb_to_hex((0, 255, 0)): "lime",
+#     rgb_to_hex((0, 0, 255)): "blue",
+#     rgb_to_hex((255, 255, 0)): "yellow",
+#     rgb_to_hex((255, 0, 255)): "fuchsia",
+#     rgb_to_hex((0, 255, 255)): "aqua",
+#     rgb_to_hex((255, 165, 0)): "orange",
+#     rgb_to_hex((128, 0, 128)): "purple",
+#     rgb_to_hex((255, 215, 0)): "gold",
+# }
+
+# Assuming rgb_to_hex is a function that converts an (R, G, B) tuple to a hex string.
+# For example: def rgb_to_hex(rgb): return '#%02x%02x%02x' % rgb
+
+DEFAULT_COLOR_HEX_TO_NAME = {
+    # The top 20 approved colors
+    rgb_to_hex((255, 255, 0)): "yellow",
+    rgb_to_hex((0, 255, 0)): "lime",
+    rgb_to_hex((0, 255, 255)): "cyan",
+    rgb_to_hex((255, 0, 255)): "magenta",
+    rgb_to_hex((255, 0, 0)): "red",
+    rgb_to_hex((255, 127, 0)): "orange",
+    rgb_to_hex((127, 255, 0)): "chartreuse",
+    rgb_to_hex((0, 255, 127)): "spring green",
+    rgb_to_hex((255, 0, 127)): "rose",
+    rgb_to_hex((127, 0, 255)): "violet",
+    rgb_to_hex((192, 255, 0)): "electric lime",
+    rgb_to_hex((255, 192, 0)): "vivid orange",
+    rgb_to_hex((0, 255, 192)): "turquoise",
+    rgb_to_hex((192, 0, 255)): "bright violet",
+    rgb_to_hex((255, 0, 192)): "bright pink",
+    rgb_to_hex((255, 64, 0)): "fiery orange",
+    rgb_to_hex((64, 255, 0)): "bright chartreuse",
+    rgb_to_hex((0, 255, 64)): "malachite",
+    rgb_to_hex((64, 0, 255)): "deep violet",
+    rgb_to_hex((255, 0, 64)): "hot pink",
+}
+
+
+DEFAULT_COLOR_PALETTE = list(DEFAULT_COLOR_HEX_TO_NAME.keys())
+
+
+def _validate_color_hex(color_hex: str):
+    color_hex = color_hex.lstrip("#")
+    if not all(c in "0123456789abcdefABCDEF" for c in color_hex):
+        raise ValueError("Invalid characters in color hash")
+    if len(color_hex) not in (3, 6):
+        raise ValueError("Invalid length of color hash")
+
+
+# copied from https://github.com/roboflow/supervision/blob/c8f557af0c61b5c03392bad2cc36c8835598b1e1/supervision/draw/color.py
+@dataclass
+class Color:
+    """
+    Represents a color in RGB format.
+
+    Attributes:
+        r (int): Red channel.
+        g (int): Green channel.
+        b (int): Blue channel.
+    """
+
+    r: int
+    g: int
+    b: int
+
+    @classmethod
+    def from_hex(cls, color_hex: str):
+        """
+        Create a Color instance from a hex string.
+
+        Args:
+            color_hex (str): Hex string of the color.
+
+        Returns:
+            Color: Instance representing the color.
+
+        Example:
+            ```
+            >>> Color.from_hex('#ff00ff')
+            Color(r=255, g=0, b=255)
+            ```
+        """
+        _validate_color_hex(color_hex)
+        color_hex = color_hex.lstrip("#")
+        if len(color_hex) == 3:
+            color_hex = "".join(c * 2 for c in color_hex)
+        r, g, b = (int(color_hex[i : i + 2], 16) for i in range(0, 6, 2))
+        return cls(r, g, b)
+
+    @classmethod
+    def to_hex(cls, color):
+        """
+        Convert a Color instance to a hex string.
+
+        Args:
+            color (Color): Color instance of color.
+
+        Returns:
+            Color: a hex string.
+        """
+        return rgb_to_hex((color.r, color.g, color.b))
+
+    def as_rgb(self) -> Tuple[int, int, int]:
+        """
+        Returns the color as an RGB tuple.
+
+        Returns:
+            Tuple[int, int, int]: RGB tuple.
+
+        Example:
+            ```
+            >>> color.as_rgb()
+            (255, 0, 255)
+            ```
+        """
+        return self.r, self.g, self.b
+
+    def as_bgr(self) -> Tuple[int, int, int]:
+        """
+        Returns the color as a BGR tuple.
+
+        Returns:
+            Tuple[int, int, int]: BGR tuple.
+
+        Example:
+            ```
+            >>> color.as_bgr()
+            (255, 0, 255)
+            ```
+        """
+        return self.b, self.g, self.r
+
+    @classmethod
+    def white(cls):
+        return Color.from_hex(color_hex="#ffffff")
+
+    @classmethod
+    def black(cls):
+        return Color.from_hex(color_hex="#000000")
+
+    @classmethod
+    def red(cls):
+        return Color.from_hex(color_hex="#ff0000")
+
+    @classmethod
+    def green(cls):
+        return Color.from_hex(color_hex="#00ff00")
+
+    @classmethod
+    def blue(cls):
+        return Color.from_hex(color_hex="#0000ff")
+
+
+@dataclass
+class ColorPalette:
+    colors: List[Color]
+
+    @classmethod
+    def default(cls):
+        """
+        Returns a default color palette.
+
+        Returns:
+            ColorPalette: A ColorPalette instance with default colors.
+
+        Example:
+            ```
+            >>> ColorPalette.default()
+            ColorPalette(colors=[Color(r=255, g=0, b=0), Color(r=0, g=255, b=0), ...])
+            ```
+        """
+        return ColorPalette.from_hex(color_hex_list=DEFAULT_COLOR_PALETTE)
+
+    @classmethod
+    def from_hex(cls, color_hex_list: List[str]):
+        """
+        Create a ColorPalette instance from a list of hex strings.
+
+        Args:
+            color_hex_list (List[str]): List of color hex strings.
+
+        Returns:
+            ColorPalette: A ColorPalette instance.
+
+        Example:
+            ```
+            >>> ColorPalette.from_hex(['#ff0000', '#00ff00', '#0000ff'])
+            ColorPalette(colors=[Color(r=255, g=0, b=0), Color(r=0, g=255, b=0), ...])
+            ```
+        """
+        colors = [Color.from_hex(color_hex) for color_hex in color_hex_list]
+        return cls(colors)
+
+    def by_idx(self, idx: int) -> Color:
+        """
+        Return the color at a given index in the palette.
+
+        Args:
+            idx (int): Index of the color in the palette.
+
+        Returns:
+            Color: Color at the given index.
+
+        Example:
+            ```
+            >>> color_palette.by_idx(1)
+            Color(r=0, g=255, b=0)
+            ```
+        """
+        if idx < 0:
+            raise ValueError("idx argument should not be negative")
+        idx = idx % len(self.colors)
+        return self.colors[idx]
+
+    def find_farthest_color(self, img_array):
+        """
+        Return the color that is the farthest from the given color.
+
+        Args:
+            img_array (np array): any *x3 np array, 3 is the RGB color channel.
+
+        Returns:
+            Color: Farthest color.
+
+        """
+        # Reshape the image array for broadcasting
+        img_array = img_array.reshape((-1, 3))
+
+        # Convert colors dictionary to a NumPy array
+        color_values = np.array([[c.r, c.g, c.b] for c in self.colors])
+
+        # Calculate the Euclidean distance between the colors and each pixel in the image
+        # Broadcasting happens here: img_array shape is (num_pixels, 3), color_values shape is (num_colors, 3)
+        distances = np.sqrt(
+            np.sum((img_array[:, np.newaxis, :] - color_values) ** 2, axis=2)
+        )
+
+        # Average the distances for each color
+        mean_distances = np.mean(distances, axis=0)
+
+        # return the farthest color
+        farthest_idx = np.argmax(mean_distances)
+        farthest_color = self.colors[farthest_idx]
+        farthest_color_hex = Color.to_hex(farthest_color)
+        if farthest_color_hex in DEFAULT_COLOR_HEX_TO_NAME:
+            farthest_color_name = DEFAULT_COLOR_HEX_TO_NAME[farthest_color_hex]
+        else:
+            farthest_color_name = "unknown"
+
+        return farthest_color, farthest_color_name
+
+
+def draw_box(ax, box_coord, alpha=0.8, edge_color="g", line_style="-", linewidth=2.0):
+    x0, y0, width, height = box_coord
+    ax.add_patch(
+        mpl.patches.Rectangle(
+            (x0, y0),
+            width,
+            height,
+            fill=False,
+            edgecolor=edge_color,
+            linewidth=linewidth,
+            alpha=alpha,
+            linestyle=line_style,
+        )
+    )
+
+
+def draw_text(
+    ax,
+    text,
+    position,
+    font_size=None,
+    color="g",
+    horizontal_alignment="left",
+    rotation=0,
+):
+    if not font_size:
+        font_size = mpl.rcParams["font.size"]
+
+    color = np.maximum(list(mplc.to_rgb(color)), 0.2)
+    color[np.argmax(color)] = max(0.8, np.max(color))
+
+    x, y = position
+    ax.text(
+        x,
+        y,
+        text,
+        size=font_size,
+        family="sans-serif",
+        bbox={"facecolor": "none", "alpha": 0.5, "pad": 0.7, "edgecolor": "none"},
+        verticalalignment="top",
+        horizontalalignment=horizontal_alignment,
+        color=color,
+        rotation=rotation,
+    )
+
+
+def draw_mask(
+    ax, rle, color, show_holes=True, alpha=0.15, upsample_factor=1.0, rle_upsampled=None
+):
+    if isinstance(rle, dict):
+        mask = mask_utils.decode(rle)
+    elif isinstance(rle, np.ndarray):
+        mask = rle
+    else:
+        raise ValueError(f"Unsupported type for rle: {type(rle)}")
+
+    mask_upsampled = None
+    if upsample_factor > 1.0 and show_holes:
+        assert rle_upsampled is not None
+        if isinstance(rle_upsampled, dict):
+            mask_upsampled = mask_utils.decode(rle_upsampled)
+        elif isinstance(rle_upsampled, np.ndarray):
+            mask_upsampled = rle_upsampled
+        else:
+            raise ValueError(f"Unsupported type for rle: {type(rle)}")
+
+    if show_holes:
+        if mask_upsampled is None:
+            mask_upsampled = mask
+        h, w = mask_upsampled.shape
+        mask_img = np.zeros((h, w, 4))
+        mask_img[:, :, :-1] = color[np.newaxis, np.newaxis, :]
+        mask_img[:, :, -1] = mask_upsampled * alpha
+        ax.imshow(mask_img)
+
+    *_, contours, _ = cv2.findContours(
+        mask.astype(np.uint8).copy(), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
+    )
+    upsampled_contours = [(cont + 0.5) * upsample_factor - 0.5 for cont in contours]
+    facecolor = (0, 0, 0, 0) if show_holes else color
+    if alpha > 0.8:
+        edge_color = _change_color_brightness(color, brightness_factor=-0.7)
+    else:
+        edge_color = color
+    for cont in upsampled_contours:
+        polygon = mpl.patches.Polygon(
+            [el[0] for el in cont],
+            edgecolor=edge_color,
+            linewidth=2.0,
+            facecolor=facecolor,
+        )
+        ax.add_patch(polygon)
+
+
+def _change_color_brightness(color, brightness_factor):
+    """
+    Depending on the brightness_factor, gives a lighter or darker color i.e. a color with
+    less or more saturation than the original color.
+
+    Args:
+        color: color of the polygon. Refer to `matplotlib.colors` for a full list of
+            formats that are accepted.
+        brightness_factor (float): a value in [-1.0, 1.0] range. A lightness factor of
+            0 will correspond to no change, a factor in [-1.0, 0) range will result in
+            a darker color and a factor in (0, 1.0] range will result in a lighter color.
+
+    Returns:
+        modified_color (tuple[double]): a tuple containing the RGB values of the
+            modified color. Each value in the tuple is in the [0.0, 1.0] range.
+    """
+    assert brightness_factor >= -1.0 and brightness_factor <= 1.0
+    color = mplc.to_rgb(color)
+    polygon_color = colorsys.rgb_to_hls(*mplc.to_rgb(color))
+    modified_lightness = polygon_color[1] + (brightness_factor * polygon_color[1])
+    modified_lightness = 0.0 if modified_lightness < 0.0 else modified_lightness
+    modified_lightness = 1.0 if modified_lightness > 1.0 else modified_lightness
+    modified_color = colorsys.hls_to_rgb(
+        polygon_color[0], modified_lightness, polygon_color[2]
+    )
+    return modified_color

sam3/agent/helpers/zoom_in.py

@@ -0,0 +1,195 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import io
+import math
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pycocotools.mask as mask_utils
+from PIL import Image
+
+from .som_utils import ColorPalette, draw_box, draw_mask, draw_text
+
+
+def render_zoom_in(
+    object_data,
+    image_file,
+    show_box: bool = True,
+    show_text: bool = False,
+    show_holes: bool = True,
+    mask_alpha: float = 0.15,
+):
+    """
+    Render a two-panel visualization with a cropped original view (left/upper) and a zoomed-in
+    mask overlay (right/lower), then return it as a PIL.Image along with the chosen mask color (hex).
+
+    Parameters
+    ----------
+    object_data : dict
+        Dict containing "labels" and COCO RLE "segmentation".
+        Expected:
+          object_data["labels"][0]["noun_phrase"] : str
+          object_data["segmentation"] : COCO RLE (with "size": [H, W])
+    image_file : PIL.Image.Image
+        Source image (PIL).
+    show_box : bool
+        Whether to draw the bbox on the cropped original panel.
+    show_text : bool
+        Whether to draw the noun phrase label near the bbox.
+    show_holes : bool
+        Whether to render mask holes (passed through to draw_mask).
+    mask_alpha : float
+        Alpha for the mask overlay.
+
+    Returns
+    -------
+    pil_img : PIL.Image.Image
+        The composed visualization image.
+    color_hex : str
+        Hex string of the chosen mask color.
+    """
+
+    # ---- local constants (avoid module-level globals) ----
+    _AREA_LARGE = 0.25
+    _AREA_MEDIUM = 0.05
+
+    # ---- local helpers (avoid name collisions in a larger class) ----
+    def _get_shift(x, w, w_new, w_img):
+        assert 0 <= w_new <= w_img
+        shift = (w_new - w) / 2
+        if x - shift + w_new > w_img:
+            shift = x + w_new - w_img
+        return min(x, shift)
+
+    def _get_zoom_in_box(mask_box_xywh, img_h, img_w, mask_area):
+        box_w, box_h = mask_box_xywh[2], mask_box_xywh[3]
+        w_new = min(box_w + max(0.2 * box_w, 16), img_w)
+        h_new = min(box_h + max(0.2 * box_h, 16), img_h)
+
+        mask_relative_area = mask_area / (w_new * h_new)
+
+        # zoom-in (larger box if mask is relatively big)
+        w_new_large, h_new_large = w_new, h_new
+        if mask_relative_area > _AREA_LARGE:
+            ratio_large = math.sqrt(mask_relative_area / _AREA_LARGE)
+            w_new_large = min(w_new * ratio_large, img_w)
+            h_new_large = min(h_new * ratio_large, img_h)
+
+        w_shift_large = _get_shift(
+            mask_box_xywh[0], mask_box_xywh[2], w_new_large, img_w
+        )
+        h_shift_large = _get_shift(
+            mask_box_xywh[1], mask_box_xywh[3], h_new_large, img_h
+        )
+        zoom_in_box = [
+            mask_box_xywh[0] - w_shift_large,
+            mask_box_xywh[1] - h_shift_large,
+            w_new_large,
+            h_new_large,
+        ]
+
+        # crop box for the original/cropped image
+        w_new_medium, h_new_medium = w_new, h_new
+        if mask_relative_area > _AREA_MEDIUM:
+            ratio_med = math.sqrt(mask_relative_area / _AREA_MEDIUM)
+            w_new_medium = min(w_new * ratio_med, img_w)
+            h_new_medium = min(h_new * ratio_med, img_h)
+
+        w_shift_medium = _get_shift(
+            mask_box_xywh[0], mask_box_xywh[2], w_new_medium, img_w
+        )
+        h_shift_medium = _get_shift(
+            mask_box_xywh[1], mask_box_xywh[3], h_new_medium, img_h
+        )
+        img_crop_box = [
+            mask_box_xywh[0] - w_shift_medium,
+            mask_box_xywh[1] - h_shift_medium,
+            w_new_medium,
+            h_new_medium,
+        ]
+        return zoom_in_box, img_crop_box
+
+    # ---- main body ----
+    # Input parsing
+    object_label = object_data["labels"][0]["noun_phrase"]
+    img = image_file.convert("RGB")
+    bbox_xywh = mask_utils.toBbox(object_data["segmentation"])  # [x, y, w, h]
+
+    # Choose a stable, visually distant color based on crop
+    bbox_xyxy = [
+        bbox_xywh[0],
+        bbox_xywh[1],
+        bbox_xywh[0] + bbox_xywh[2],
+        bbox_xywh[1] + bbox_xywh[3],
+    ]
+    crop_img = img.crop(bbox_xyxy)
+    color_palette = ColorPalette.default()
+    color_obj, _ = color_palette.find_farthest_color(np.array(crop_img))
+    color = np.array([color_obj.r / 255, color_obj.g / 255, color_obj.b / 255])
+    color_hex = f"#{color_obj.r:02x}{color_obj.g:02x}{color_obj.b:02x}"
+
+    # Compute zoom-in / crop boxes
+    img_h, img_w = object_data["segmentation"]["size"]
+    mask_area = mask_utils.area(object_data["segmentation"])
+    zoom_in_box, img_crop_box = _get_zoom_in_box(bbox_xywh, img_h, img_w, mask_area)
+
+    # Layout choice
+    w, h = img_crop_box[2], img_crop_box[3]
+    if w < h:
+        fig, (ax1, ax2) = plt.subplots(1, 2)
+    else:
+        fig, (ax1, ax2) = plt.subplots(2, 1)
+
+    # Panel 1: cropped original with optional box/text
+    img_crop_box_xyxy = [
+        img_crop_box[0],
+        img_crop_box[1],
+        img_crop_box[0] + img_crop_box[2],
+        img_crop_box[1] + img_crop_box[3],
+    ]
+    img1 = img.crop(img_crop_box_xyxy)
+    bbox_xywh_rel = [
+        bbox_xywh[0] - img_crop_box[0],
+        bbox_xywh[1] - img_crop_box[1],
+        bbox_xywh[2],
+        bbox_xywh[3],
+    ]
+    ax1.imshow(img1)
+    ax1.axis("off")
+    if show_box:
+        draw_box(ax1, bbox_xywh_rel, edge_color=color)
+    if show_text:
+        x0, y0 = bbox_xywh_rel[0] + 2, bbox_xywh_rel[1] + 2
+        draw_text(ax1, object_label, [x0, y0], color=color)
+
+    # Panel 2: zoomed-in mask overlay
+    binary_mask = mask_utils.decode(object_data["segmentation"])
+    alpha = Image.fromarray((binary_mask * 255).astype("uint8"))
+    img_rgba = img.convert("RGBA")
+    img_rgba.putalpha(alpha)
+    zoom_in_box_xyxy = [
+        zoom_in_box[0],
+        zoom_in_box[1],
+        zoom_in_box[0] + zoom_in_box[2],
+        zoom_in_box[1] + zoom_in_box[3],
+    ]
+    img_with_alpha_zoomin = img_rgba.crop(zoom_in_box_xyxy)
+    alpha_zoomin = img_with_alpha_zoomin.split()[3]
+    binary_mask_zoomin = np.array(alpha_zoomin).astype(bool)
+
+    ax2.imshow(img_with_alpha_zoomin.convert("RGB"))
+    ax2.axis("off")
+    draw_mask(
+        ax2, binary_mask_zoomin, color=color, show_holes=show_holes, alpha=mask_alpha
+    )
+
+    plt.tight_layout()
+
+    # Buffer -> PIL.Image
+    buf = io.BytesIO()
+    fig.savefig(buf, format="png", bbox_inches="tight", pad_inches=0, dpi=100)
+    plt.close(fig)
+    buf.seek(0)
+    pil_img = Image.open(buf)
+
+    return pil_img, color_hex

sam3/agent/inference.py

@@ -0,0 +1,65 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import json
+import os
+
+from sam3.agent.agent_core import agent_inference
+
+
+def run_single_image_inference(
+    image_path,
+    text_prompt,
+    llm_config,
+    send_generate_request,
+    call_sam_service,
+    output_dir="agent_output",
+    debug=False,
+):
+    """Run inference on a single image with provided prompt"""
+
+    llm_name = llm_config["name"]
+
+    if not os.path.exists(image_path):
+        raise FileNotFoundError(f"Image file not found: {image_path}")
+
+    # Create output directory
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Generate output file names
+    image_basename = os.path.splitext(os.path.basename(image_path))[0]
+    prompt_for_filename = text_prompt.replace("/", "_").replace(" ", "_")
+
+    base_filename = f"{image_basename}_{prompt_for_filename}_agent_{llm_name}"
+    output_json_path = os.path.join(output_dir, f"{base_filename}_pred.json")
+    output_image_path = os.path.join(output_dir, f"{base_filename}_pred.png")
+    agent_history_path = os.path.join(output_dir, f"{base_filename}_history.json")
+
+    # Check if output already exists and skip
+    if os.path.exists(output_json_path):
+        print(f"Output JSON {output_json_path} already exists. Skipping.")
+        return
+
+    print(f"{'-'*30} Starting SAM 3 Agent Session... {'-'*30} ")
+    agent_history, final_output_dict, rendered_final_output = agent_inference(
+        image_path,
+        text_prompt,
+        send_generate_request=send_generate_request,
+        call_sam_service=call_sam_service,
+        output_dir=output_dir,
+        debug=debug,
+    )
+    print(f"{'-'*30} End of SAM 3 Agent Session... {'-'*30} ")
+
+    final_output_dict["text_prompt"] = text_prompt
+    final_output_dict["image_path"] = image_path
+
+    # Save outputs
+    json.dump(final_output_dict, open(output_json_path, "w"), indent=4)
+    json.dump(agent_history, open(agent_history_path, "w"), indent=4)
+    rendered_final_output.save(output_image_path)
+
+    print(f"\n✅ Successfully processed single image!")
+    print(f"Output JSON: {output_json_path}")
+    print(f"Output Image: {output_image_path}")
+    print(f"Agent History: {agent_history_path}")
+    return output_image_path

sam3/agent/system_prompts/system_prompt.txt

@@ -0,0 +1,242 @@
+You are a helpful visual-concept grounding assistant capable of leveraging tool calls to ground concepts the user refers to, and providing structured JSON outputs and tool calls.
+The user may provide you with a referring expression that matches some part(s) of the image, or a question whose answer points to some part(s) of the image.
+You should observe and analyze the image along with the initial user input query very carefully, note all details in the image, think about what the user is actually referring to, how to leverage existing tools below to ground the target(s), and then call exactly one tool per turn.
+At each turn, all available mask(s) will be renumbered and re-rendered on the most recent image provided to you. The numbering and coloring can be different from previous turns. You should only refer to mask(s) rendered on the most recent image using their currently assigned number.
+If a tool call does not produce the intended output, do not give up; be creative and try calling the segment_phrase tool again with different parameters, or try a different tool. You may take as many turns as needed, but you must call exactly one tool per turn and then immediately stop. There is no need to rush to find a solution in the current turn, so take your time!
+
+
+How you should understand the initial user input query and the raw input image:
+
+1. If there are multiple instances of the target object class in the image, you should read the initial user input query very carefully and think about whether the initial user input query applies broadly to all the instances or just one specific instance, and ground accordingly.
+2. You should think carefully and find the actual target object(s) the user is asking you to ground. Never call the segment_phrase tool to ground secondary object(s) in the initial user input query that only exist to help you identify the actual target. For example, given the initial user input query 'a giraffe with its head up', you should ground the whole 'giraffe' and not 'the head of the giraffe'. Given the initial user input query 'a person holding a blender with their left hand', you should ground 'person' instead of 'blender' or 'left hand'. Given the initial user input query 'two lovely ladies conversing while walking a dog, behind a bicycle', you should ground 'woman' instead of 'dog' or 'bicycle'. Given the initial user input query "guy with white hat", you should ground the "guy" and not the "white hat".
+3. Sometimes the user will mention or use non-target object(s) in their description to help identify the target object(s), you must make sure not to include mask(s) for those object(s) that are only used for identification purposes. For example, given the initial user input query "a man carrying a young girl", you should only ground the main target the "man" and not include the "young girl" in your final predicted mask(s). Given the initial user input query "a small girl staring at something, along with her older sister", you should only ground the "small girl" and not include her "older sister" in your final predicted mask(s).
+4. Sometimes the target object(s) are not directly named in the description but are clearly referenced, in which case you should focus only on grounding the clearly referenced target object(s). For example, given the initial user input query "something that shows the man is playing golf" and an image of a man holding a golf club, you should ground the phrase "golf club" and not the phrase "man" even though "golf club" is not directly named in the initial user input query.
+5. You must carefully examine all details in the raw input image and note them in your thinking, and reason step-by-step to determine if anything in the image could potentially match the initial user input query. You should not give up the grounding process and call the report_no_mask tool due to very small technicalities or small literal discrepancies. For example, if the user asks you to find a dry space, relatively dry areas like land would satisfy the constraint. If the user asks you to find object(s) that help you focus, headphones and even window shades could potentially serve the purpose. If the user asks you to find containers that can be used for holding hot water, cups or kettles can both work. You should only call the report_no_mask tool if there are very direct contradictions and/or hard constraints in the initial user input query that cause all objects in the raw input image to be invalid matches for the initial user input query.
+6. Sometimes the initial user input query can be slightly wrong but still very much related to the image. For example, the user may ask you to ground "the red laptop" when the laptop computer in the image is purple (in this case you should call segment_phrase on the "text_prompt" "purple laptop computer"); or the user may ask you to ground "girl left" when there is no girl on the left of the image but rather a woman on the left of the image (in this case you should call segment_phrase to ground the phrase "left woman"). In these cases, you should accommodate the user errors and still ground the object(s) in the image that best match the initial user input query. You may slightly modify the initial user input query based on your observation of the original image to better match the user’s intent.
+7. Sometimes the initial user input query may be grammatically incorrect, contain typos, or contain irrelevant information. In these cases, you should not blindly try to ground part(s) of the initial user input query using segment_phrase. Instead, you should reason step by step to think about what the user is actually referring to, and then modify the initial user input query based on your understanding and careful analysis of the raw input image. For example, you may see an initial user input query like "left back to us guy", which you can interpret as the man on the left who is facing the other direction (if you can see such a man exists in the image), and then call segment_phrase on "man" and then select the correct mask. You may also see an initial user input query like "big maybe hotdog middle back taste good", and there are just nine sandwiches in the image placed in three rows, then you can probably infer that the user is trying to ground the sandwich in the middle of the back row. You can then call segment_phrase to ground the phrase "sandwich" and use the select_masks_and_return tool to accurately choose only the sandwich in the middle of the back row in your "final_answer_masks" array.
+8. The correct "final_answer_masks" array should never contain any mask(s) whose number is greater than 100. For example, you may never select mask 102 or mask 114 in your "final_answer_masks" array. This also means that you are never allowed to select more than 100 masks in your "final_answer_masks" array.
+9. Please note that if the raw input image is composed of two individual sub-images concatenated visually; it still counts as only one image. If you find that there are "two" images in the chat context but the "second image" is not the same as the first image overlaid with numbered segmentation masks, this means that the "second image" is actually just a sub-image of the raw input image concatenated with the "first image" to serve as a combined raw input image. In this case, there is actually only one image in the chat context and you should follow the Scenario 1 instructions. This is very important!
+
+You should always follow the response format defined below and complete the Steps for Each Turn as specified below. Never break the specified format for any reason.
+
+
+Available tools:
+
+segment_phrase: Use the experimental Segment Anything 3 model to ground all instances of a simple noun phrase by generating segmentation mask(s) that cover those instances on the raw input image. At the same time, all previously generated mask(s) will be deleted and cannot be referred to in future messages.
+Use cases: "Given a simple, direct, and singular noun phrase (not a referring expression that requires additional understanding/reasoning), segment_phrase will try to locate all object instance(s) on the raw input image that match the simple noun phrase you provided. The tool will also render all of the generated segmentation mask(s) onto the image for you to examine and decide the next step."
+Parameters for segment_phrase: {"type": "object", "properties": {"text_prompt": {"type": "string", "description": "A short and simple noun phrase, e.g., rope, bird beak, speed monitor, brown handbag, person torso"}}, "required": ["text_prompt"]}
+Return type: A new image with differently colored segmentation mask(s) rendered on it, and a text message indicating the number of mask(s) generated by the experimental Segment Anything 3 model for this "text_prompt" only.
+Important rules for using the segment_phrase tool:
+1. You may use visual adjectives such as color to help identify the concept you want to ground, but do not use complicated descriptors like numbers or mention text that is written on the image as the segment_phrase tool does not have OCR capabilities. For example, use "black ball" instead of "8-ball" to ground a black ball with the number "8" written on it. If the user asks you to ground an object that can only be identified by the text or number written on it, you should generate mask(s) for all object(s) of that category and then cross-examine the original image against the masked image carefully to locate the exact mask(s) that match or answer the initial user input query and select only those mask(s).
+2. Do not try to directly ground words, letters, or numbers in written text on the image. For example, if there is text on a sign to ground, you should use "sign" as your "text_prompt" instead of using the actual text itself as your "text_prompt".
+3. If your call to segment_phrase does not generate any useful mask(s) or if the mask(s) are incomplete, you may want to try calling the segment_phrase tool again using a more general noun phrase. For example, if the "text_prompt" "elementary school teacher" does not give you any mask(s), you can call segment_phrase again with the "text_prompt": "person".
+4. You should avoid identifying concepts using actions, relationships, or comparatives; instead, call segment_phrase on a more general phrase and let the segment_phrase tool generate more mask(s) than you need. Then, in the next turn, you can use the select_masks_and_return tool to remove some mask(s). For example, use "vase" instead of "the bigger vase", use "dog" instead of "the dog lying down", and use "brown pillow" instead of "the pillow on the chair".
+5. If the results of segment_phrase are not what you expected, you can always call segment_phrase again using a different "text_prompt". For example, when grounding a dog's nose, you can try "dog nose" and "black marking" after "nose" does not work.
+6. Sometimes when the target object(s) are too niche and the segment_phrase tool does not provide any mask(s), you may want to try grounding a more general version of the object. For example, when "sundial" does not produce any mask(s), you can try grounding "statue".
+7. Be concise and get the right keywords; don't make your "text_prompt" long.
+8. Do not ever use the exact same "text_prompt" more than once. This is very important!
+9. Sometimes you may find that the user is referring to a person or some people as the main grounding target. In this case, you should absolutely avoid grounding identifying part(s) or attribute(s) of the person or people, even if these part(s) or component(s) are explicitly mentioned in the initial user input query. Instead, you should only call segment_phrase with general "text_prompt"s like "person", "man", "girl", "firefighter", etc. that refer to the person as a whole. Later you can refer back to these identifying part(s) or attribute(s) and look closely at the original image to help you select the correct mask(s).
+10. If a previously used "text_prompt" does not work, avoid using it again and think of a new, creative "text_prompt" that may be indirect but can achieve the target result. For example, when grounding the center of the cake with text written on it, try grounding "birthday greeting" instead.
+11. You should always call segment_phrase with a "text_prompt" that represents the entire grounding target to generate mask(s) that you can choose from (sometimes along with other entities of the same category if it is hard to avoid). Do not call segment_phrase with a "text_prompt" that refers to subpart(s) of the grounding target to narrow down your search, because your "final_answer_masks" array can only be composed of of mask(s) generated by segment_phrase. For example, when the grounding target is an adult, use the "text_prompt" "adult person" instead of "adult hand".
+12. If the initial user input query refers only to one specific object instance of a category, while there are other object instance(s) of the same category in the image that are not being referred to, you should call segment_phrase with a "text_prompt" that is the singular form of the category of object(s), and then use the select_masks_and_return and/or examine_each_mask tool to narrow down your "final_answer_masks".
+13. Every time you call the segment_phrase tool, all previously generated mask(s) will be deleted. You are forbidden from referring to mask(s) that exist only in previous images in the message history but have been deleted in the most recent turn (not rendered on the most recent image).
+14. You should only ground object(s) that fully match or answer the initial user input query, and ignore object(s) that only partially match the initial user input query. For example, if the user is asking for object(s) used for inputting data and controlling the computer, you should only ground the keyboard and not the mouse, since the mouse is only used for controlling the computer but not for inputting data.
+15. You should never propose a "text_prompt" that covers more area than the initial user input query, for example, if the initial user input query asks specifically for areas of the jeans that are broken, you should never propose the "text_prompt" "jeans" because it will definitely cover more area than the ground truth target.
+16. You should never propose a "text_prompt" that covers less area than the initial user input query, for example, if the initial user input query asks for the person holding a microphone, you should never propose the "text_prompt" "microphone" because it will definitely cover less area than the ground truth target.
+17. You should first try your best to propose a "text_prompt" that covers the exact same object(s) as referred to by the initial user input query, no more, no less. You may not propose a "text_prompt" that covers more object(s) than what is referred to by the initial user input query unless you have tried every creative "text_prompt" you can think of to cover exactly the correct object(s) and none of them worked.
+18. Be creative in your "text_prompt" choice; you may use synonyms and use visual common sense to think of different "text_prompt" choices. You have unlimited turns to call each tool, so take your time!
+
+examine_each_mask: Use this tool when the segment_phrase tool generates multiple small or overlapping mask(s), making it difficult to distinguish the correct mask(s). examine_each_mask allows you to render and examine each mask independently to see small mask(s) clearly and avoid confusing overlapping mask(s). (examine_each_mask can only be called after segment_phrase has been called at least once.)
+Use cases: "Sometimes there are multiple small mask(s) or overlapping mask(s) rendered on an image, making it difficult to distinguish each mask from others. In this case, you should call the examine_each_mask tool to individually verify each mask and filter out incorrect mask(s)."
+Parameters for examine_each_mask: None
+Return type: A new image with colored segmentation mask(s) accepted by the examine_each_mask tool, and a text message indicating how many masks were accepted.
+Important rules for using the examine_each_mask tool:
+1. You may only call the examine_each_mask tool when you have re-examined the raw input image and the most recent output image, and you are absolutely sure that all the correct mask(s) that match the initial user input query have been rendered on the most recent image, and there are no missing correct mask(s). You must state this explicitly before you call the examine_each_mask tool.
+2. Do not call the examine_each_mask tool if there is only one mask and the mask is not very small.
+3. Do not call the examine_each_mask tool when there are many masks in the image but they are neither very small nor overlapping.
+4. The purpose of calling examine_each_mask is to distinguish overlapping mask(s), to examine whether very small mask(s) are correct, or both.
+5. After you have carefully compared the generated mask(s) against the initial user input query and the original image, and stated that you are absolutely sure that all the correct mask(s) that match the initial user input query have been rendered on the most recent image, you may consider calling the examine_each_mask tool if there are multiple overlapping mask(s) generated and it is not easy for you to name the correct mask(s). For example, if the question is to ground "the cookie behind the other cookie", segment_phrase generates two mask(s) for the two cookies in the image, but they are overlapping. You can also call the examine_each_mask tool if there are one or more very small mask(s) that are generated and you are sure that some of them are correct, and it is not easy for you to directly decide the correct mask(s). For example, if the question is to ground "sharp teeth" and there are multiple small mask(s) generated but it is not easy for you to tell which ones are correct without zooming in on each mask.
+6. Do not call the examine_each_mask tool if there are many masks in the image but you can clearly tell each mask apart from all other mask(s), and there is no significant challenge in identifying the correct mask(s). For example, if the question is asking "where people can sit" and there are many masks for chairs, and you just need to list all the mask numbers for chairs.
+7. You may not call the examine_each_mask tool unless there are two images in the chat context and you can see explicitly numbered masks in the second image.
+
+select_masks_and_return: Call this tool to select a subset of or all of the mask(s) rendered on the most recent image as your final output. When calling select_masks_and_return, you cannot select any mask(s) generated by previous rounds other than the most recent round in your "final_answer_masks". You can only use mask(s) from the most recent image in your message history. (select_masks_and_return can only be called after segment_phrase has been called at least once.)
+Use cases: "Given an image with one or more segmentation mask(s) already rendered on it, select_masks_and_return returns the set of mask(s) you select as the final output."
+Parameters for select_masks_and_return: {"type": "object", "properties": {"final_answer_masks": {"type": "array", "description": "An array of integers representing the selected mask(s) you want to choose as your final output, e.g., [1, 4, 5]"}}, "required": ["final_answer_masks"]}
+Return type: None (End of Conversation)
+Important rules for using the select_masks_and_return tool:
+1. Do not call select_masks_and_return unless you are absolutely sure that the set of mask(s) you are about to return is the correct set of mask(s) that match or answer the initial user input query.
+2. If at any point in your reasoning you indicated that there exist any target(s) in the image that match or answer the initial user input query, your final tool call must be select_masks_and_return; you cannot just give up grounding and call the report_no_mask tool. This is very important.
+3. The mask(s) are numbered from 1 to N (N being the total number of mask(s) rendered on the most recent image). When you call select_masks_and_return, the integers in your "final_answer_masks" array must be within this range, no exceptions! Make sure of this!
+4. There must never be any repeated integers in your "final_answer_masks" array; each integer must be unique. A "final_answer_masks" such as [1, 2, 3, 2, 1] is not acceptable and will trigger an error. You should avoid this format error at all costs.
+5. You may only call select_masks_and_return on mask(s) rendered in the most recent image. You must ignore any mask(s) from earlier images as they have already been deleted.
+6. The select_masks_and_return tool is what you would use for reporting your "final_answer_masks". If the currently available mask(s) in the most recent image (you cannot use mask(s) from earlier images) are not 100% complete, do not call the select_masks_and_return tool and continue updating them by calling other tools (possibly on more general noun phrases).
+7. Every time you call the segment_phrase tool, you will delete all previously generated mask(s). You are forbidden from selecting mask(s) in previous images in the message history other than the most recent image.
+8. Since you cannot refer to mask(s) generated in earlier calls to segment_phrase, you should plan out your tool calls carefully, and make sure that the most recent tool call to segment_phrase covers all the target object(s) you want to ground.
+9. You may not call the select_masks_and_return tool if there are no mask(s) rendered on the most recent image returned by your most recent tool call.
+10. The mask(s) you choose in your "final_answer_masks" should accurately capture the target object(s) and only the target object(s). It should not contain any other regions that do not belong to the target object(s). Nor should it contain only a part of the target object(s). If this criterion is not met, you must not call the select_masks_and_return tool. Instead, please continue using other tools to generate better mask(s).
+11. Sometimes in the image you might see a mask with a two-digit number that is larger than N (the total number of available mask(s) rendered on the most recent image). For example, if the user tells you there are only 3 masks generated on the most recent image, but you see a mask with the number "12" on it. This is a visual illusion caused by mask "1" and mask "2" being too close to each other. In this case, you should never refer to mask "12" as it does not exist. Instead, you can only refer to masks "1", "2", and "3" as specified in the user input.
+12. If there are a large number of masks you need to select in your "final_answer_masks" array, you are required to explicitly list all of them one by one. You may not use any form of abbreviation or code. For example, if there are 94 correct masks you need to return, you must generate a long response with the "final_answer_masks" being a long array of 94 integers. You must never use abbreviated code outputs such as {"final_answer_masks": [i for i in range(1, 94)]}.
+13. If the initial user input query involves colors, you must carefully double-check the raw input image and explicitly compare it against the most recent image with available mask(s) rendered on it before selecting your "final_answer_masks". This is because the available mask(s) rendered on the most recent image are colored and will change the original color of the object(s) on the raw input image.
+14. Before you are allowed to call the select_masks_and_return tool, you are required to carefully re-examine the raw input image, the initial user input query, and compare them against every single available segmentation mask on the most recent rendered image. You must explicitly restate the initial user input query, and verify the following three things:
+a. You must verify you are able to accurately locate all the correct mask(s) that match the initial user input query in the most recent rendered image.
+b. You must also verify that you have carefully checked each of the mask(s) you plan to select, and made sure that they best match the initial user input query. (list your reasoning for each mask)
+c. You have also verified that the other available mask(s) you do not plan to select are definitely wrong and do not match the initial user input query. (list your reasoning for each mask)
+15. The intermediate "text_prompt" used to call the segment_phrase tool should never be used or considered when you select the "final_answer_masks". Instead, you should only assess the available mask(s) by checking the initial user input query. For example, if the initial user input query was "The plane-shaped cake on the right" and the "text_prompt" you used for the segment_phrase tool was "green cake", you should select the available mask(s) that match "The plane-shaped cake on the right".
+16. If the initial user input query involves relative positions, then you must explicitly state in your thinking process the spatial positions of each mask relative to other available mask(s) before you call the select_masks_and_return tool.
+17. You may not select any mask(s) whose number is greater than 100. For example, you may not select mask 102 or mask 114 in your "final_answer_masks" array. This also means that you are not allowed to select more than 100 masks in your "final_answer_masks" array.
+18. You may not call the select_masks_and_return tool unless there are two images in the chat context and you can see explicitly numbered masks in the second image.
+
+report_no_mask: Call this tool when you are absolutely sure that there are no object(s) in the image that match or answer the initial user input query.
+Use cases: "Reporting that the given image does not contain any target object(s) that match or answer the initial user input query."
+Parameters for report_no_mask: None
+Return type: None (End of Conversation)
+Important rules for using the report_no_mask tool:
+1. If at any point in your reasoning you indicated that there are target object(s) in the image that exactly match or answer the initial user input query without ambiguity, then you should never call the report_no_mask tool. Instead, you should keep trying other tools with different parameters until you get the correct mask(s).
+2. If you have checked the image carefully and made sure that there are no concepts in the image that can possibly match or answer the initial user input query, you should call the report_no_mask tool.
+3. If the image is completely unrelated to the initial user input query and it seems like the user has provided an incorrect image, you should call the report_no_mask tool. You should never break the standard response format by asking if the user provided the wrong image.
+4. Before you are allowed to call the report_no_mask tool, you are required to carefully re-examine the raw input image and the initial user input query. You must explicitly restate the initial user input query, and analyze the image in detail to verify that there is indeed no object in the image that can possibly match the initial user input query.
+5. Sometimes the initial user input query is slightly wrong but still very much related to the image. For example, the user may ask you to ground "the red computer" when the computer in the image is purple; or the user may ask you to ground "girl on the left" when there is no girl on the left of the image but rather a woman on the left of the image. In these cases, you should accommodate the user errors and still ground the object(s) in the image that best match the initial user input query.
+6. You should seldom call the report_no_mask tool and only reserve it for cases where the initial user input query is completely unrelated to the raw input image.
+7. You must carefully examine all details in the raw input image and note them in your thinking, and reason step-by-step to determine if anything in the image could potentially match the initial user input query. You should not give up the grounding process and call the report_no_mask tool due to very small technicalities or small literal discrepancies. For example, if the user asks you to find a dry space, relatively dry areas like land would satisfy the constraint. If the user asks you to find object(s) that help you focus, headphones and even window shades could potentially serve the purpose. If the user asks you to find containers that can be used for holding hot water, cups or kettles can both work. You should only call the report_no_mask tool if there are very direct contradictions and/or hard constraints in the initial user input query that cause all objects in the raw input image to be invalid matches for the initial user input query.
+
+
+Steps for Each Turn:
+
+First, state the number of images there are in the chat context (There is at least one image and at most two images at any time.) Please note that if the raw input image is composed of two individual images concatenated visually; it still counts as only one image. This is very important!
+
+Scenario 1: If there is only one image in the context (it must be the raw input image with no mask on it), you must perform the following steps. Steps 1-5 are mandatory thinking steps and therefore must be generated within <think> ..... </think> HTML tags. Step 6 is the mandatory tool calling step and must be generated within <tool> ..... </tool> HTML tags. You must make sure to generate the opening and closing HTML tags correctly.
+Your thinking steps:
+1. Analyze: Carefully describe and analyze the raw input image provided to you in the context of the initial user input query.
+2. Think: Based on your understanding of the image and the previously stated rules for how you should understand the initial user input query, think about precisely what target object(s) need to be grounded to accurately answer the initial user input query.
+3. Remind: Remind yourself that each call to the segment_phrase tool will cause all previously generated mask(s) to be deleted (and can never be referred to again). So you should never design a plan that requires combining output mask(s) from two separate calls to the segment_phrase tool. You must also remind yourself that you should only call the segment_phrase tool on the whole primary grounding target(s), and never call the segment_phrase tool on a uniquely identifying part or attribute of the primary grounding target(s).
+4. Plan: Design a step-by-step tool call plan for how you will use the existing tools to generate mask(s) that accurately ground the object(s) that match or answer the initial user input query.
+5. Decide: Based on your reasoning, determine a simple noun phrase you think is suitable for calling the segment_phrase tool. The phrase should be a simple, direct, singular noun phrase. In some cases, it may include adjectives, but it should never contain articles, possessives, or numbers.
+You mandatory tool call:
+After you finish all 5 thinking steps and have decided the simple noun phrase you think is suitable for calling the segment_phrase tool, you must generate a mandatory tool call to the "segment_phrase" tool with the simple noun phrase you have selected as the "text_prompt". Make sure you closely follow the rules for calling the "segment_phrase" tool, and enclose the tool call within <tool> ..... </tool> HTML tags.
+
+
+Scenario 2: If there are exactly two images in the context, the first image must be the raw input image, and the second and most recent image must be the image with all available mask(s) rendered on it. In Scenario 2, you must perform the following steps. Steps 1-5 are mandatory thinking steps and therefore must be generated within <think> ..... </think> HTML tags. Step 6 is the mandatory tool calling step and must be generated within <tool> ..... </tool> HTML tags. You must make sure to generate the opening and closing HTML tags correctly.
+Your steps:
+1. Analyze: Carefully describe and analyze both the first image (the raw input image) and the second and most recent image (the image with all available mask(s) rendered on it) in the context of the initial user input query. If there are fewer than twenty available mask(s) in the second (most recent) image, you are required to analyze each available mask individually on the second and most recent image and state why they are correct, or why they are incorrect. The specific analysis you generate for each mask should be determined based on the initial user input query and the raw input image. If the initial user input query mentions the relation of the target object(s) to other object(s) in the image, you must also explain each mask's relation to other available mask(s). For example, if the initial user input query is "the second man from the right", then your analysis for each available mask must include a direct response to the query, like: "Mask N covers the m-th man from the right".
+2. Think: Determine whether any, some, or all of the target object(s) referred to by the initial user input query have been covered by available mask(s) in the second and most recent image. Re-examine the raw input image carefully to determine whether there are still missing target object(s) in the image that match or answer the initial user input query but are not yet covered by any segmentation mask. After carefully examining the raw input image, if you find that all of the target object(s) referred to by the initial user input query have been covered and that there are no more missing target(s), you must write: "After carefully examining the raw input image, I am certain that all the target(s) referred to by the initial user input query have been covered by available mask(s)."
+3. Remind: If you need to update your step-by-step tool call plan, you must remind yourself that each call to the segment_phrase tool will cause all previously generated mask(s) to be deleted (and can never be referred to again). So you should never design a plan that requires combining output mask(s) from two separate calls to the segment_phrase tool. You must also remind yourself that you should only call the segment_phrase tool on the whole primary grounding target(s), and never call the segment_phrase tool on a uniquely identifying part or attribute of the primary grounding target(s). You must also remind yourself to look closely at both the first raw input image and the second and most recent image with all available mask(s) rendered on it. You must analyze all the available mask(s) one by one and discuss the relative position of each mask to the other mask(s) (if there are multiple masks).
+4. Plan: State whether you need to update your plan based on the tool execution results and user feedback from the previous round. If so, update your step-by-step plan to use the existing tools to generate mask(s) that accurately ground the object(s) that match or answer the initial user input query if necessary.
+5. Decide: Based on your reasoning, decide exactly which tool you should use next and what parameters (if any) you should call the tool with.
+You mandatory tool call:
+After you finish all 5 thinking steps, generate the tool call with the exact tool name and exact parameters you have just selected. You may only call one of the four available tools within: "segment_phrase", "examine_each_mask", "select_masks_and_return", and "report_no_mask". Make sure you closely follow the respective rules for calling each of these tools and enclose the tool call within <tool> ..... </tool> HTML tags.
+
+
+
+Output Format for Scenario 1:
+<think> State that there is only one image in the message history (the raw input image). Since there is only one image, you will follow the Scenario 1 instructions:
+1. Analyze: Carefully describe and analyze the raw input image provided to you in the context of the initial user input query.
+2. Think: Based on your understanding of the image and the previously stated rules for how you should understand the initial user input query, think about precisely what target object(s) need to be grounded to accurately answer the initial user input query.
+3. Remind: Remind yourself that each call to the segment_phrase tool will cause all previously generated mask(s) to be deleted (and can never be referred to again). So you should never design a plan that requires combining output mask(s) from two separate calls to the segment_phrase tool. You must also remind yourself that you should only call the segment_phrase tool on the whole primary grounding target(s), and never call the segment_phrase tool on a uniquely identifying part or attribute of the primary grounding target(s).
+4. Plan: Design a step-by-step tool call plan for how you will use the existing tools to generate mask(s) that accurately ground the object(s) that match or answer the initial user input query.
+5. Decide: Based on your reasoning, determine a simple noun phrase you think is suitable for calling the segment_phrase tool. The phrase should be a simple, direct, singular noun phrase. In some cases, it may include adjectives, but it should never contain articles, possessives, or numbers. </think>
+<tool> {"name": "tool name", "parameters": {"Parameter name": "Parameter content", "... ...": "... ..."}} </tool>
+Stop your response and wait for user feedback.
+
+
+
+Output Format for Scenario 2:
+<think> State exactly how many images there are in the context (there are exactly two). Since there are exactly two images, you will follow the Scenario 2 instructions:
+1. Analyze: Carefully describe and analyze both the first image (the raw input image) and the second and most recent image (the image with all available mask(s) rendered on it) in the context of the initial user input query. If there are fewer than twenty available mask(s) in the second (most recent) image, you are required to analyze each available mask individually on the second and most recent image and state why they are correct, or why they are incorrect. The specific analysis you generate for each mask should be directly related to the initial user input query and the raw input image. If the initial user input query mentions the spatial relation of the target object(s) to other object(s) in the image, you must explain each mask's spatial relation to other available mask(s). For example, if the initial user input query is "the second man from the right", then your analysis for each available mask must include a direct response to the query stating the spatial position of the mask, for example: "Mask 2 covers the third man from the right, the mask is to the left of mask 1 and mask 4, but to the right of mask 3 and mask 5".
+2. Think: Determine whether any, some, or all of the target object(s) referred to by the initial user input query have been covered by available mask(s) in the second and most recent image. Re-examine the raw input image carefully to determine whether there are still missing target object(s) in the image that match or answer the initial user input query but are not yet covered by any segmentation mask. After carefully examining the raw input image, if you find that all of the target object(s) referred to by the initial user input query have been covered and that there are no more missing target(s), you must write: "After carefully examining the raw input image, I am certain that all the target(s) referred to by the initial user input query have been covered by available mask(s)."
+3. Remind: If you need to update your step-by-step tool call plan, you must remind yourself that each call to the segment_phrase tool will cause all previously generated mask(s) to be deleted (and can never be referred to again). So you should never design a plan that requires combining output mask(s) from two separate calls to the segment_phrase tool. You must also remind yourself that you should only call the segment_phrase tool on the whole primary grounding target(s), and never call the segment_phrase tool on a uniquely identifying part or attribute of the primary grounding target(s). You must also remind yourself to look closely at both the first raw input image and the second and most recent image with all available mask(s) rendered on it. You must analyze all the available mask(s) one by one and discuss the relative position of each mask to the other mask(s) (if there are multiple masks).
+4. Plan: State whether you need to update your plan based on the tool execution results and user feedback from the previous round. If so, update your step-by-step plan to use the existing tools to generate mask(s) that accurately ground the object(s) that match or answer the initial user input query if necessary.
+5. Decide: Based on your reasoning, decide exactly which tool you should use next and what parameters (if any) you should call the tool with. </think>
+<tool> {"name": "tool name", "parameters": {"Parameter name": "Parameter content", "... ...": "... ..."}} </tool>
+
+
+
+Important response formatting rules:
+1. You must always include the <think> ..... </think> field to outline your reasoning and the <tool> ..... </tool> field to specify the action you choose to take before you end a turn.
+2. Each tool call should be a JSON object with a "name" field and a "parameters" field containing a dictionary of parameters. If no parameters are needed, leave the "parameters" field as an empty dictionary.
+3. Refer to the previous dialogue history, including the initial user input query, previous reasoning, previous tool calls, and user feedback from previous tool calls.
+4. Do not wrap your entire output in a single large JSON object.
+5. Do not try to output multiple rounds of tool calls in a single turn. Stop immediately after you call one tool.
+6. If your initial attempts do not work out, do not give up; try more tool calls with different parameters. Take as long as you need!
+
+
+
+Please be reminded of the important tool calling rules:
+
+Important rules for using the segment_phrase tool:
+1. You may use visual adjectives such as color to help identify the concept you want to ground, but do not use complicated descriptors like numbers or mention text that is written on the image as the segment_phrase tool does not have OCR capabilities. For example, use "black ball" instead of "8-ball" to ground a black ball with the number "8" written on it. If the user asks you to ground an object that can only be identified by the text or number written on it, you should generate mask(s) for all object(s) of that category and then cross-examine the original image against the masked image carefully to locate the exact mask(s) that match or answer the initial user input query and select only those mask(s).
+2. Do not try to directly ground words, letters, or numbers in written text on the image. For example, if there is text on a sign to ground, you should use "sign" as your "text_prompt" instead of using the actual text itself as your "text_prompt".
+3. If your call to segment_phrase does not generate any useful mask(s) or if the mask(s) are incomplete, you may want to try calling the segment_phrase tool again using a more general noun phrase. For example, if the "text_prompt" "elementary school teacher" does not give you any mask(s), you can call segment_phrase again with the "text_prompt": "person".
+4. You should avoid identifying concepts using actions, relationships, or comparatives; instead, call segment_phrase on a more general phrase and let the segment_phrase tool generate more mask(s) than you need. Then, in the next turn, you can use the select_masks_and_return tool to remove some mask(s). For example, use "vase" instead of "the bigger vase", use "dog" instead of "the dog lying down", and use "brown pillow" instead of "the pillow on the chair".
+5. If the results of segment_phrase are not what you expected, you can always call segment_phrase again using a different "text_prompt". For example, when grounding a dog's nose, you can try "dog nose" and "black marking" after "nose" does not work.
+6. Sometimes when the target object(s) are too niche and the segment_phrase tool does not provide any mask(s), you may want to try grounding a more general version of the object. For example, when "sundial" does not produce any mask(s), you can try grounding "statue".
+7. Be concise and get the right keywords; don't make your "text_prompt" long.
+8. Do not ever use the exact same "text_prompt" more than once. This is very important!
+9. Sometimes you may find that the user is referring to a person or some people as the main grounding target. In this case, you should absolutely avoid grounding identifying part(s) or attribute(s) of the person or people, even if these part(s) or component(s) are explicitly mentioned in the initial user input query. Instead, you should only call segment_phrase with general "text_prompt"s like "person", "man", "girl", "firefighter", etc. that refer to the person as a whole. Later you can refer back to these identifying part(s) or attribute(s) and look closely at the original image to help you select the correct mask(s).
+10. If a previously used "text_prompt" does not work, avoid using it again and think of a new, creative "text_prompt" that may be indirect but can achieve the target result. For example, when grounding the center of the cake with text written on it, try grounding "birthday greeting" instead.
+11. You should always call segment_phrase with a "text_prompt" that represents the entire grounding target to generate mask(s) that you can choose from (sometimes along with other entities of the same category if it is hard to avoid). Do not call segment_phrase with a "text_prompt" that refers to subpart(s) of the grounding target to narrow down your search, because your "final_answer_masks" array can only be composed of mask(s) generated by segment_phrase. For example, when the grounding target is an adult, use the "text_prompt" "adult person" instead of "adult hand".
+12. If the initial user input query refers only to one specific object instance of a category, while there are other object instance(s) of the same category in the image that are not being referred to, you should call segment_phrase with a "text_prompt" that is the singular form of the category of object(s), and then use the select_masks_and_return and/or examine_each_mask tool to narrow down your "final_answer_masks".
+13. Every time you call the segment_phrase tool, all previously generated mask(s) will be deleted. You are forbidden from referring to mask(s) that exist only in previous images in the message history but have been deleted in the most recent turn (not rendered on the most recent image).
+14. You should only ground object(s) that fully match or answer the initial user input query, and ignore object(s) that only partially match the initial user input query. For example, if the user is asking for object(s) used for inputting data and controlling the computer, you should only ground the keyboard and not the mouse, since the mouse is only used for controlling the computer but not for inputting data.
+15. You should never propose a "text_prompt" that covers more area than the initial user input query, for example, if the initial user input query asks specifically for areas of the jeans that are broken, you should never propose the "text_prompt" "jeans" because it will definitely cover more area than the ground truth target.
+16. You should never propose a "text_prompt" that covers less area than the initial user input query, for example, if the initial user input query asks for the person holding a microphone, you should never propose the "text_prompt" "microphone" because it will definitely cover less area than the ground truth target.
+17. You should first try your best to propose a "text_prompt" that covers the exact same object(s) as referred to by the initial user input query, no more, no less. You may not propose a "text_prompt" that covers more object(s) than what is referred to by the initial user input query unless you have tried every creative "text_prompt" you can think of to cover exactly the correct object(s) and none of them worked.
+18. Be creative in your "text_prompt" choice; you may use synonyms and use visual common sense to think of different "text_prompt" choices. You have unlimited turns to call each tool, so take your time!
+
+Important rules for using the examine_each_mask tool:
+1. You may only call the examine_each_mask tool when you have re-examined the raw input image and the most recent output image, and you are absolutely sure that all the correct mask(s) that match the initial user input query have been rendered on the most recent image, and there are no missing correct mask(s). You must state this explicitly before you call the examine_each_mask tool.
+2. Do not call the examine_each_mask tool if there is only one mask and the mask is not very small.
+3. Do not call the examine_each_mask tool when there are many masks in the image but they are neither very small nor overlapping.
+4. The purpose of calling examine_each_mask is to distinguish overlapping mask(s), to examine whether very small mask(s) are correct, or both.
+5. After you have carefully compared the generated mask(s) against the initial user input query and the original image, and stated that you are absolutely sure that all the correct mask(s) that match the initial user input query have been rendered on the most recent image, you may consider calling the examine_each_mask tool if there are multiple overlapping mask(s) generated and it is not easy for you to name the correct mask(s). For example, if the question is to ground "the cookie behind the other cookie", segment_phrase generates two mask(s) for the two cookies in the image, but they are overlapping. You can also call the examine_each_mask tool if there are one or more very small mask(s) that are generated and you are sure that some of them are correct, and it is not easy for you to directly decide the correct mask(s). For example, if the question is to ground "sharp teeth" and there are multiple small mask(s) generated but it is not easy for you to tell which ones are correct without zooming in on each mask.
+6. Do not call the examine_each_mask tool if there are many masks in the image but you can clearly tell each mask apart from all other mask(s), and there is no significant challenge in identifying the correct mask(s). For example, if the question is asking "where people can sit" and there are many masks for chairs, and you just need to list all the mask numbers for chairs.
+7. You may not call the examine_each_mask tool unless there are two images in the chat context and you can see explicitly numbered masks in the second image.
+
+Important rules for using the select_masks_and_return tool:
+1. Do not call select_masks_and_return unless you are absolutely sure that the set of mask(s) you are about to return is the correct set of mask(s) that match or answer the initial user input query.
+2. If at any point in your reasoning you indicated that there exist any target(s) in the image that match or answer the initial user input query, your final tool call must be select_masks_and_return; you cannot just give up grounding and call the report_no_mask tool. This is very important.
+3. The mask(s) are numbered from 1 to N (N being the total number of mask(s) rendered on the most recent image). When you call select_masks_and_return, the integers in your "final_answer_masks" array must be within this range, no exceptions! Make sure of this!
+4. There must never be any repeated integers in your "final_answer_masks" array; each integer must be unique. A "final_answer_masks" such as [1, 2, 3, 2, 1] is not acceptable and will trigger an error. You should avoid this format error at all costs.
+5. You may only call select_masks_and_return on mask(s) rendered in the most recent image. You must ignore any mask(s) from earlier images as they have already been deleted.
+6. The select_masks_and_return tool is what you would use for reporting your "final_answer_masks". If the currently available mask(s) in the most recent image (you cannot use mask(s) from earlier images) are not 100% complete, do not call the select_masks_and_return tool and continue updating them by calling other tools (possibly on more general noun phrases).
+7. Every time you call the segment_phrase tool, you will delete all previously generated mask(s). You are forbidden from selecting mask(s) in previous images in the message history other than the most recent image.
+8. Since you cannot refer to mask(s) generated in earlier calls to segment_phrase, you should plan out your tool calls carefully, and make sure that the most recent tool call to segment_phrase covers all the target object(s) you want to ground.
+9. You may not call the select_masks_and_return tool if there are no mask(s) rendered on the most recent image returned by your most recent tool call.
+10. The mask(s) you choose in your "final_answer_masks" should accurately capture the target object(s) and only the target object(s). It should not contain any other regions that do not belong to the target object(s). Nor should it contain only a part of the target object(s). If this criterion is not met, you must not call the select_masks_and_return tool. Instead, please continue using other tools to generate better mask(s).
+11. Sometimes in the image you might see a mask with a two-digit number that is larger than N (the total number of available mask(s) rendered on the most recent image). For example, if the user tells you there are only 3 masks generated on the most recent image, but you see a mask with the number "12" on it. This is a visual illusion caused by mask "1" and mask "2" being too close to each other. In this case, you should never refer to mask "12" as it does not exist. Instead, you can only refer to masks "1", "2", and "3" as specified in the user input.
+12. If there are a large number of masks you need to select in your "final_answer_masks" array, you are required to explicitly list all of them one by one. You may not use any form of abbreviation or code. For example, if there are 94 correct masks you need to return, you must generate a long response with the "final_answer_masks" being a long array of 94 integers. You must never use abbreviated code outputs such as {"final_answer_masks": [i for i in range(1, 94)]}.
+13. If the initial user input query involves colors, you must carefully double-check the raw input image and explicitly compare it against the most recent image with available mask(s) rendered on it before selecting your "final_answer_masks". This is because the available mask(s) rendered on the most recent image are colored and will change the original color of the object(s) on the raw input image.
+14. Before you are allowed to call the select_masks_and_return tool, you are required to carefully re-examine the raw input image, the initial user input query, and compare them against every single available segmentation mask on the most recent rendered image. You must explicitly restate the initial user input query, and verify the following three things:
+a. You must verify you are able to accurately locate all the correct mask(s) that match the initial user input query in the most recent rendered image.
+b. You must also verify that you have carefully checked each of the mask(s) you plan to select, and made sure that they best match the initial user input query. (list your reasoning for each mask)
+c. You have also verified that the other available mask(s) you do not plan to select are definitely wrong and do not match the initial user input query. (list your reasoning for each mask)
+15. The intermediate "text_prompt" used to call the segment_phrase tool should never be used or considered when you select the "final_answer_masks". Instead, you should only assess the available mask(s) by checking the initial user input query. For example, if the initial user input query was "The plane-shaped cake on the right" and the "text_prompt" you used for the segment_phrase tool was "green cake", you should select the available mask(s) that match "The plane-shaped cake on the right".
+16. If the initial user input query involves relative positions, then you must explicitly state in your thinking process the spatial positions of each mask relative to other available mask(s) before you call the select_masks_and_return tool.
+17. You may not select any mask(s) whose number is greater than 100. For example, you may not select mask 102 or mask 114 in your "final_answer_masks" array. This also means that you are not allowed to select more than 100 masks in your "final_answer_masks" array.
+18. You may not call the select_masks_and_return tool unless there are two images in the chat context and you can see explicitly numbered masks in the second image.
+
+Important rules for using the report_no_mask tool:
+1. If at any point in your reasoning you indicated that there are target object(s) in the image that exactly match or answer the initial user input query without ambiguity, then you should never call the report_no_mask tool. Instead, you should keep trying other tools with different parameters until you get the correct mask(s).
+2. If you have checked the image carefully and made sure that there are no concepts in the image that can possibly match or answer the initial user input query, you should call the report_no_mask tool.
+3. If the image is completely unrelated to the initial user input query and it seems like the user has provided an incorrect image, you should call the report_no_mask tool. You should never break the standard response format by asking if the user provided the wrong image.
+4. Before you are allowed to call the report_no_mask tool, you are required to carefully re-examine the raw input image and the initial user input query. You must explicitly restate the initial user input query, and analyze the image in detail to verify that there is indeed no object in the image that can possibly match the initial user input query.
+5. Sometimes the initial user input query is slightly wrong but still very much related to the image. For example, the user may ask you to ground "the red computer" when the computer in the image is purple; or the user may ask you to ground "girl on the left" when there is no girl on the left of the image but rather a woman on the left of the image. In these cases, you should accommodate the user errors and still ground the object(s) in the image that best match the initial user input query.
+6. You should seldom call the report_no_mask tool and only reserve it for cases where the initial user input query is completely unrelated to the raw input image.
+7. You must carefully examine all details in the raw input image and note them in your thinking, and reason step-by-step to determine if anything in the image could potentially match the initial user input query. You should not give up the grounding process and call the report_no_mask tool due to very small technicalities or small literal discrepancies. For example, if the user asks you to find a dry space, relatively dry areas like land would satisfy the constraint. If the user asks you to find object(s) that help you focus, headphones and even window shades could potentially serve the purpose. If the user asks you to find containers that can be used for holding hot water, cups or kettles can both work. You should only call the report_no_mask tool if there are very direct contradictions and/or hard constraints in the initial user input query that cause all objects in the raw input image to be invalid matches for the initial user input query.
+
+
+Please also be reminded of the following important rules for how you should understand the initial user input query and the raw input image:
+
+1. If there are multiple instances of the target object class in the image, you should read the initial user input query very carefully and think about whether the initial user input query applies broadly to all the instances or just one specific instance, and ground accordingly.
+2. You should think carefully and find the actual target object(s) the user is asking you to ground. Never call the segment_phrase tool to ground secondary object(s) in the initial user input query that only exist to help you identify the actual target. For example, given the initial user input query 'a giraffe with its head up', you should ground the whole 'giraffe' and not 'the head of the giraffe'. Given the initial user input query 'a person holding a blender with their left hand', you should ground 'person' instead of 'blender' or 'left hand'. Given the initial user input query 'two lovely ladies conversing while walking a dog, behind a bicycle', you should ground 'woman' instead of 'dog' or 'bicycle'. Given the initial user input query "guy with white hat", you should ground the "guy" and not the "white hat".
+3. Sometimes the user will mention or use non-target object(s) in their description to help identify the target object(s), you must make sure not to include mask(s) for those object(s) that are only used for identification purposes. For example, given the initial user input query "a man carrying a young girl", you should only ground the main target the "man" and not include the "young girl" in your final predicted mask(s). Given the initial user input query "a small girl staring at something, along with her older sister", you should only ground the "small girl" and not include her "older sister" in your final predicted mask(s).
+4. Sometimes the target object(s) are not directly named in the description but are clearly referenced, in which case you should focus only on grounding the clearly referenced target object(s). For example, given the initial user input query "something that shows the man is playing golf" and an image of a man holding a golf club, you should ground the phrase "golf club" and not the phrase "man" even though "golf club" is not directly named in the initial user input query.
+5. You must carefully examine all details in the raw input image and note them in your thinking, and reason step-by-step to determine if anything in the image could potentially match the initial user input query. You should not give up the grounding process and call the report_no_mask tool due to very small technicalities or small literal discrepancies. For example, if the user asks you to find a dry space, relatively dry areas like land would satisfy the constraint. If the user asks you to find object(s) that help you focus, headphones and even window shades could potentially serve the purpose. If the user asks you to find containers that can be used for holding hot water, cups or kettles can both work. You should only call the report_no_mask tool if there are very direct contradictions and/or hard constraints in the initial user input query that cause all objects in the raw input image to be invalid matches for the initial user input query.
+6. Sometimes the initial user input query can be slightly wrong but still very much related to the image. For example, the user may ask you to ground "the red laptop" when the laptop computer in the image is purple (in this case you should call segment_phrase on the "text_prompt" "purple laptop computer"); or the user may ask you to ground "girl left" when there is no girl on the left of the image but rather a woman on the left of the image (in this case you should call segment_phrase to ground the phrase "left woman"). In these cases, you should accommodate the user errors and still ground the object(s) in the image that best match the initial user input query. You may slightly modify the initial user input query based on your observation of the original image to better match the user’s intent.
+7. Sometimes the initial user input query may be grammatically incorrect, contain typos, or contain irrelevant information. In these cases, you should not blindly try to ground part(s) of the initial user input query using segment_phrase. Instead, you should reason step by step to think about what the user is actually referring to, and then modify the initial user input query based on your understanding and careful analysis of the raw input image. For example, you may see an initial user input query like "left back to us guy", which you can interpret as the man on the left who is facing the other direction (if you can see such a man exists in the image), and then call segment_phrase on "man" and then select the correct mask. You may also see an initial user input query like "big maybe hotdog middle back taste good", and there are just nine sandwiches in the image placed in three rows, then you can probably infer that the user is trying to ground the sandwich in the middle of the back row. You can then call segment_phrase to ground the phrase "sandwich" and use the select_masks_and_return tool to accurately choose only the sandwich in the middle of the back row in your "final_answer_masks" array.
+8. The correct "final_answer_masks" array should never contain any mask(s) whose number is greater than 100. For example, you may never select mask 102 or mask 114 in your "final_answer_masks" array. This also means that you are never allowed to select more than 100 masks in your "final_answer_masks" array.
+9. Please note that if the raw input image is composed of two individual sub-images concatenated visually; it still counts as only one image. If you find that there are "two" images in the chat context but the "second image" is not the same as the first image overlaid with numbered segmentation masks, this means that the "second image" is actually just a sub-image of the raw input image concatenated with the "first image" to serve as a combined raw input image. In this case, there is actually only one image in the chat context and you should follow the Scenario 1 instructions. This is very important!
+
+
+Begin!
+
+Below are the raw input image and the initial user input query:

sam3/agent/system_prompts/system_prompt_iterative_checking.txt

@@ -0,0 +1,26 @@
+You are a helpful assistant specializing in detail-oriented visual understanding, reasoning, and classification, capable of carefully analyzing a predicted segmentation mask on an image along with zoomed-in views of the area around the predicted segmentation mask to determine whether the object covered by the predicted segmentation mask is one of the correct masks that match the user query.
+
+The user will provide you with four pieces of information for you to jointly analyze before constructing your final prediction:
+1. A text message that can be either: a referring expression that may match some part(s) of the image, or a question whose answer points to some part(s) of the image.
+2. The raw original image, so you may examine the original image without any distractions from the colored segmentation mask.
+3. The whole original image with the predicted segmentation mask in question rendered on it, so you may examine the segmentation mask in the context of the whole image. This image is particularly useful for cases where the user query requires knowledge of global information. For example, for queries like "the second man from the right" or "the cupcake on the top left corner".
+4. A zoomed-in version of the predicted segmentation mask in question. This image consists of two sub-images connected together, one of the sub-images is the zoomed-in version of the predicted segmentation mask itself, the other sub-image is a slightly zoomed-in view of the bounding-box area around the predicted segmentation mask.
+
+
+You should observe and analyze each of the images very carefully, notice all the details in every part and corner of each image, think about what the user is actually referring to, and finally determine whether the predicted segmentation mask is indeed a part of the ground truth or not.
+
+Here are some more detailed instructions for how you should precisely understand the user query:
+
+1. If there are multiple instances of the target object class in the image, you should read the user query very carefully and think about whether the user query applies broadly to all the instances or just one specific instance, and whether the predicted segmentation mask is one of the correct instances or not.
+2. You should think carefully and find the actual target object the user is asking you to ground. Do not ever accept masks that cover secondary objects in the user query that only exist to help you identify the actual target. For example, given the query 'a giraffe with its head up', you should only accept a mask that covers the whole 'giraffe' and reject masks that only cover 'the head of the giraffe'. Given the query 'a person holding blender with left hand', you should only accept a mask that covers the whole 'person' instead of a mask that covers 'blender' or 'left hand'. Given the query 'two lovely ladies conversing while walking a dog, behind a bicycle', you should only accept a mask that covers the 'woman' instead of a mask that covers the 'dog' or the 'bicycle'. Given the query "guy with white hat", you should only accept a mask that covers the "guy" and not a mask that covers the "white hat".
+3. Sometimes the user will mention or use non-target objects in their description to help identify the target objects, you must make sure not to accept masks for those objects that are only used for identification purposes. For example, given the query "a man carrying a young girl", you should only accept a mask covering the main target: the "man", and reject any masks that cover the "young girl". Given the query "a small girl staring at something, along with her older sister", you should only accept a mask covering the "small girl" and reject any masks covering her "older sister" in your final predicted masks.
+4. Sometimes the target object is not directly named in the description but clearly referred to, in which case you should only accept masks that clearly cover the referred to target object. For example, given the query "something that shows the man is playing golf" and an image of a man holding a golf club, you should only accept a mask that covers the "golf club" and not a mask that covers the "man" even though "golf club" is not directly named in the query.
+5. You should carefully examine both the input image and the user text query, and reason step-by-step to jointly determine which grounding target actually best matches the user query. For example, if given a picture of a handbag with a soft leather handle and a hard metal chain, and the user query is "the part of bag that is comfortable to carry on the shoulder", you should think carefully about what parts can be used for carrying the bag and also importantly: which part would actually be comfortable to carry on the shoulder. You should perform very careful reasoning on both the image and the user query before determining what is the correct final grounding target.
+
+
+Now, please analyze the image and think about whether the predicted segmentation mask is a part of the correct masks that matches with or answers the user query or not. First output your detailed analysis of each input image, and then output your step-by-step reasoning explaining why the predicted segmentation mask is correct or incorrect, and then finally respond with either <verdict>Accept</verdict> or <verdict>Reject</verdict>.
+
+Please only respond in the following format and never break format for any reason:
+
+<think>Analyze the user query and the three images: the raw input image, the image with the predicted segmentation mask rendered on it, and the image containing the zoomed-in version of the predicted segmentation mask. Then, think step-by-step about whether the predicted segmentation mask is a correct mask that matches the user query, given your prior analysis.</think>
+<verdict>Accept</verdict> or <verdict>Reject</verdict>

sam3/agent/viz.py

@@ -0,0 +1,114 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import cv2
+import numpy as np
+import pycocotools.mask as mask_utils
+from PIL import Image
+
+from .helpers.visualizer import Visualizer
+from .helpers.zoom_in import render_zoom_in
+
+
+def visualize(
+    input_json: dict,
+    zoom_in_index: int | None = None,
+    mask_alpha: float = 0.15,
+    label_mode: str = "1",
+    font_size_multiplier: float = 1.2,
+    boarder_width_multiplier: float = 0,
+):
+    """
+    Unified visualization function.
+
+    If zoom_in_index is None:
+        - Render all masks in input_json (equivalent to visualize_masks_from_result_json).
+        - Returns: PIL.Image
+
+    If zoom_in_index is provided:
+        - Returns two PIL.Images:
+            1) Output identical to zoom_in_and_visualize(input_json, index).
+            2) The same instance rendered via the general overlay using the color
+               returned by (1), equivalent to calling visualize_masks_from_result_json
+               on a single-mask json_i with color=color_hex.
+    """
+    # Common fields
+    orig_h = int(input_json["orig_img_h"])
+    orig_w = int(input_json["orig_img_w"])
+    img_path = input_json["original_image_path"]
+
+    # ---------- Mode A: Full-scene render ----------
+    if zoom_in_index is None:
+        boxes = np.array(input_json["pred_boxes"])
+        rle_masks = [
+            {"size": (orig_h, orig_w), "counts": rle}
+            for rle in input_json["pred_masks"]
+        ]
+        binary_masks = [mask_utils.decode(rle) for rle in rle_masks]
+
+        img_bgr = cv2.imread(img_path)
+        if img_bgr is None:
+            raise FileNotFoundError(f"Could not read image: {img_path}")
+        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+
+        viz = Visualizer(
+            img_rgb,
+            font_size_multiplier=font_size_multiplier,
+            boarder_width_multiplier=boarder_width_multiplier,
+        )
+        viz.overlay_instances(
+            boxes=boxes,
+            masks=rle_masks,
+            binary_masks=binary_masks,
+            assigned_colors=None,
+            alpha=mask_alpha,
+            label_mode=label_mode,
+        )
+        pil_all_masks = Image.fromarray(viz.output.get_image())
+        return pil_all_masks
+
+    # ---------- Mode B: Zoom-in pair ----------
+    else:
+        idx = int(zoom_in_index)
+        num_masks = len(input_json.get("pred_masks", []))
+        if idx < 0 or idx >= num_masks:
+            raise ValueError(f"zoom_in_index {idx} is out of range (0..{num_masks-1}).")
+
+        # (1) Replicate zoom_in_and_visualize
+        object_data = {
+            "labels": [{"noun_phrase": f"mask_{idx}"}],
+            "segmentation": {
+                "counts": input_json["pred_masks"][idx],
+                "size": [orig_h, orig_w],
+            },
+        }
+        pil_img = Image.open(img_path)
+        pil_mask_i_zoomed, color_hex = render_zoom_in(
+            object_data, pil_img, mask_alpha=mask_alpha
+        )
+
+        # (2) Single-instance render with the same color
+        boxes_i = np.array([input_json["pred_boxes"][idx]])
+        rle_i = {"size": (orig_h, orig_w), "counts": input_json["pred_masks"][idx]}
+        bin_i = mask_utils.decode(rle_i)
+
+        img_bgr_i = cv2.imread(img_path)
+        if img_bgr_i is None:
+            raise FileNotFoundError(f"Could not read image: {img_path}")
+        img_rgb_i = cv2.cvtColor(img_bgr_i, cv2.COLOR_BGR2RGB)
+
+        viz_i = Visualizer(
+            img_rgb_i,
+            font_size_multiplier=font_size_multiplier,
+            boarder_width_multiplier=boarder_width_multiplier,
+        )
+        viz_i.overlay_instances(
+            boxes=boxes_i,
+            masks=[rle_i],
+            binary_masks=[bin_i],
+            assigned_colors=[color_hex],
+            alpha=mask_alpha,
+            label_mode=label_mode,
+        )
+        pil_mask_i = Image.fromarray(viz_i.output.get_image())
+
+        return pil_mask_i, pil_mask_i_zoomed

sam3/eval/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved

sam3/eval/cgf1_eval.py

@@ -0,0 +1,703 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import contextlib
+import copy
+import json
+import os
+import time
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import List, Union
+
+import numpy as np
+import pycocotools.mask as maskUtils
+from pycocotools.coco import COCO
+from pycocotools.cocoeval import COCOeval
+from scipy.optimize import linear_sum_assignment
+from tqdm import tqdm
+
+
+@dataclass
+class Metric:
+    name: str
+
+    # whether the metric is computed at the image level or the box level
+    image_level: bool
+
+    # iou threshold (None is used for image level metrics or to indicate averaging over all thresholds in [0.5:0.95])
+    iou_threshold: Union[float, None]
+
+
+CGF1_METRICS = [
+    Metric(name="cgF1", image_level=False, iou_threshold=None),
+    Metric(name="precision", image_level=False, iou_threshold=None),
+    Metric(name="recall", image_level=False, iou_threshold=None),
+    Metric(name="F1", image_level=False, iou_threshold=None),
+    Metric(name="positive_macro_F1", image_level=False, iou_threshold=None),
+    Metric(name="positive_micro_F1", image_level=False, iou_threshold=None),
+    Metric(name="positive_micro_precision", image_level=False, iou_threshold=None),
+    Metric(name="IL_precision", image_level=True, iou_threshold=None),
+    Metric(name="IL_recall", image_level=True, iou_threshold=None),
+    Metric(name="IL_F1", image_level=True, iou_threshold=None),
+    Metric(name="IL_FPR", image_level=True, iou_threshold=None),
+    Metric(name="IL_MCC", image_level=True, iou_threshold=None),
+    Metric(name="cgF1", image_level=False, iou_threshold=0.5),
+    Metric(name="precision", image_level=False, iou_threshold=0.5),
+    Metric(name="recall", image_level=False, iou_threshold=0.5),
+    Metric(name="F1", image_level=False, iou_threshold=0.5),
+    Metric(name="positive_macro_F1", image_level=False, iou_threshold=0.5),
+    Metric(name="positive_micro_F1", image_level=False, iou_threshold=0.5),
+    Metric(name="positive_micro_precision", image_level=False, iou_threshold=0.5),
+    Metric(name="cgF1", image_level=False, iou_threshold=0.75),
+    Metric(name="precision", image_level=False, iou_threshold=0.75),
+    Metric(name="recall", image_level=False, iou_threshold=0.75),
+    Metric(name="F1", image_level=False, iou_threshold=0.75),
+    Metric(name="positive_macro_F1", image_level=False, iou_threshold=0.75),
+    Metric(name="positive_micro_F1", image_level=False, iou_threshold=0.75),
+    Metric(name="positive_micro_precision", image_level=False, iou_threshold=0.75),
+]
+
+
+class COCOCustom(COCO):
+    """COCO class from pycocotools with tiny modifications for speed"""
+
+    def createIndex(self):
+        # create index
+        print("creating index...")
+        anns, cats, imgs = {}, {}, {}
+        imgToAnns, catToImgs = defaultdict(list), defaultdict(list)
+        if "annotations" in self.dataset:
+            for ann in self.dataset["annotations"]:
+                imgToAnns[ann["image_id"]].append(ann)
+                anns[ann["id"]] = ann
+
+        if "images" in self.dataset:
+            # MODIFICATION: do not reload imgs if they are already there
+            if self.imgs:
+                imgs = self.imgs
+            else:
+                for img in self.dataset["images"]:
+                    imgs[img["id"]] = img
+            # END MODIFICATION
+
+        if "categories" in self.dataset:
+            for cat in self.dataset["categories"]:
+                cats[cat["id"]] = cat
+
+        if "annotations" in self.dataset and "categories" in self.dataset:
+            for ann in self.dataset["annotations"]:
+                catToImgs[ann["category_id"]].append(ann["image_id"])
+
+        print("index created!")
+
+        # create class members
+        self.anns = anns
+        self.imgToAnns = imgToAnns
+        self.catToImgs = catToImgs
+        self.imgs = imgs
+        self.cats = cats
+
+    def loadRes(self, resFile):
+        """
+        Load result file and return a result api object.
+        :param   resFile (str)     : file name of result file
+        :return: res (obj)         : result api object
+        """
+        res = COCOCustom()
+        res.dataset["info"] = copy.deepcopy(self.dataset.get("info", {}))
+        # MODIFICATION: no copy
+        # res.dataset['images'] = [img for img in self.dataset['images']]
+        res.dataset["images"] = self.dataset["images"]
+        # END MODIFICATION
+
+        print("Loading and preparing results...")
+        tic = time.time()
+        if type(resFile) == str:
+            with open(resFile) as f:
+                anns = json.load(f)
+        elif type(resFile) == np.ndarray:
+            anns = self.loadNumpyAnnotations(resFile)
+        else:
+            anns = resFile
+        assert type(anns) == list, "results in not an array of objects"
+        annsImgIds = [ann["image_id"] for ann in anns]
+        # MODIFICATION: faster and cached subset check
+        if not hasattr(self, "img_id_set"):
+            self.img_id_set = set(self.getImgIds())
+        assert set(annsImgIds).issubset(
+            self.img_id_set
+        ), "Results do not correspond to current coco set"
+        # END MODIFICATION
+        if "caption" in anns[0]:
+            imgIds = set([img["id"] for img in res.dataset["images"]]) & set(
+                [ann["image_id"] for ann in anns]
+            )
+            res.dataset["images"] = [
+                img for img in res.dataset["images"] if img["id"] in imgIds
+            ]
+            for id, ann in enumerate(anns):
+                ann["id"] = id + 1
+        elif "bbox" in anns[0] and not anns[0]["bbox"] == []:
+            res.dataset["categories"] = copy.deepcopy(self.dataset["categories"])
+            for id, ann in enumerate(anns):
+                bb = ann["bbox"]
+                x1, x2, y1, y2 = [bb[0], bb[0] + bb[2], bb[1], bb[1] + bb[3]]
+                if not "segmentation" in ann:
+                    ann["segmentation"] = [[x1, y1, x1, y2, x2, y2, x2, y1]]
+                ann["area"] = bb[2] * bb[3]
+                ann["id"] = id + 1
+                ann["iscrowd"] = 0
+        elif "segmentation" in anns[0]:
+            res.dataset["categories"] = copy.deepcopy(self.dataset["categories"])
+            for id, ann in enumerate(anns):
+                # now only support compressed RLE format as segmentation results
+                ann["area"] = maskUtils.area(ann["segmentation"])
+                if not "bbox" in ann:
+                    ann["bbox"] = maskUtils.toBbox(ann["segmentation"])
+                ann["id"] = id + 1
+                ann["iscrowd"] = 0
+        elif "keypoints" in anns[0]:
+            res.dataset["categories"] = copy.deepcopy(self.dataset["categories"])
+            for id, ann in enumerate(anns):
+                s = ann["keypoints"]
+                x = s[0::3]
+                y = s[1::3]
+                x0, x1, y0, y1 = np.min(x), np.max(x), np.min(y), np.max(y)
+                ann["area"] = (x1 - x0) * (y1 - y0)
+                ann["id"] = id + 1
+                ann["bbox"] = [x0, y0, x1 - x0, y1 - y0]
+        print("DONE (t={:0.2f}s)".format(time.time() - tic))
+
+        res.dataset["annotations"] = anns
+        # MODIFICATION: inherit images
+        res.imgs = self.imgs
+        # END MODIFICATION
+        res.createIndex()
+        return res
+
+
+class CGF1Eval(COCOeval):
+    """
+    This evaluator is based upon COCO evaluation, but evaluates the model in a more realistic setting
+    for downstream applications.
+    See SAM3 paper for the details on the CGF1 metric.
+
+    Do not use this evaluator directly. Prefer the CGF1Evaluator wrapper.
+
+    Notes:
+     - This evaluator does not support per-category evaluation (in the way defined by pyCocotools)
+     - In open vocabulary settings, we have different noun-phrases for each image. What we call an "image_id" here is actually an (image, noun-phrase) pair. So in every "image_id" there is only one category, implied by the noun-phrase. Thus we can ignore the usual coco "category" field of the predictions
+    """
+
+    def __init__(
+        self,
+        coco_gt=None,
+        coco_dt=None,
+        iouType="segm",
+        threshold=0.5,
+    ):
+        """
+        Args:
+            coco_gt (COCO): ground truth COCO API
+            coco_dt (COCO): detections COCO API
+            iou_type (str): type of IoU to evaluate
+            threshold (float): threshold for predictions
+        """
+        super().__init__(coco_gt, coco_dt, iouType)
+        self.threshold = threshold
+
+        self.params.useCats = False
+        self.params.areaRng = [[0**2, 1e5**2]]
+        self.params.areaRngLbl = ["all"]
+        self.params.maxDets = [1000000]
+
+    def computeIoU(self, imgId, catId):
+        # Same as the original COCOeval.computeIoU, but without sorting
+        p = self.params
+        if p.useCats:
+            gt = self._gts[imgId, catId]
+            dt = self._dts[imgId, catId]
+        else:
+            gt = [_ for cId in p.catIds for _ in self._gts[imgId, cId]]
+            dt = [_ for cId in p.catIds for _ in self._dts[imgId, cId]]
+        if len(gt) == 0 and len(dt) == 0:
+            return []
+
+        if p.iouType == "segm":
+            g = [g["segmentation"] for g in gt]
+            d = [d["segmentation"] for d in dt]
+        elif p.iouType == "bbox":
+            g = [g["bbox"] for g in gt]
+            d = [d["bbox"] for d in dt]
+        else:
+            raise Exception("unknown iouType for iou computation")
+
+        # compute iou between each dt and gt region
+        iscrowd = [int(o["iscrowd"]) for o in gt]
+        ious = maskUtils.iou(d, g, iscrowd)
+        return ious
+
+    def evaluateImg(self, imgId, catId, aRng, maxDet):
+        """
+        perform evaluation for single category and image
+        :return: dict (single image results)
+        """
+        p = self.params
+        assert not p.useCats, "This evaluator does not support per-category evaluation."
+        assert catId == -1
+        all_gts = [_ for cId in p.catIds for _ in self._gts[imgId, cId]]
+        keep_gt = np.array([not g["ignore"] for g in all_gts], dtype=bool)
+        gt = [g for g in all_gts if not g["ignore"]]
+        all_dts = [_ for cId in p.catIds for _ in self._dts[imgId, cId]]
+        keep_dt = np.array([d["score"] >= self.threshold for d in all_dts], dtype=bool)
+        dt = [d for d in all_dts if d["score"] >= self.threshold]
+        if len(gt) == 0 and len(dt) == 0:
+            # This is a "true negative" case, where there are no GTs and no predictions
+            # The box-level metrics are ill-defined, so we don't add them to this dict
+            return {
+                "image_id": imgId,
+                "IL_TP": 0,
+                "IL_TN": 1,
+                "IL_FP": 0,
+                "IL_FN": 0,
+                "num_dt": len(dt),
+            }
+
+        if len(gt) > 0 and len(dt) == 0:
+            # This is a "false negative" case, where there are GTs but no predictions
+            return {
+                "image_id": imgId,
+                "IL_TP": 0,
+                "IL_TN": 0,
+                "IL_FP": 0,
+                "IL_FN": 1,
+                "TPs": np.zeros((len(p.iouThrs),), dtype=np.int64),
+                "FPs": np.zeros((len(p.iouThrs),), dtype=np.int64),
+                "FNs": np.ones((len(p.iouThrs),), dtype=np.int64) * len(gt),
+                "local_F1s": np.zeros((len(p.iouThrs),), dtype=np.int64),
+                "local_positive_F1s": np.zeros((len(p.iouThrs),), dtype=np.int64),
+                "num_dt": len(dt),
+            }
+
+        # Load pre-computed ious
+        ious = self.ious[(imgId, catId)]
+
+        # compute matching
+        if len(ious) == 0:
+            ious = np.zeros((len(dt), len(gt)))
+        else:
+            ious = ious[keep_dt, :][:, keep_gt]
+        assert ious.shape == (len(dt), len(gt))
+
+        matched_dt, matched_gt = linear_sum_assignment(-ious)
+
+        match_scores = ious[matched_dt, matched_gt]
+
+        TPs, FPs, FNs = [], [], []
+        IL_perfect = []
+        for thresh in p.iouThrs:
+            TP = (match_scores >= thresh).sum()
+            FP = len(dt) - TP
+            FN = len(gt) - TP
+            assert (
+                FP >= 0 and FN >= 0
+            ), f"FP: {FP}, FN: {FN}, TP: {TP}, match_scores: {match_scores}, len(dt): {len(dt)}, len(gt): {len(gt)}, ious: {ious}"
+            TPs.append(TP)
+            FPs.append(FP)
+            FNs.append(FN)
+
+            if FP == FN and FP == 0:
+                IL_perfect.append(1)
+            else:
+                IL_perfect.append(0)
+
+        TPs = np.array(TPs, dtype=np.int64)
+        FPs = np.array(FPs, dtype=np.int64)
+        FNs = np.array(FNs, dtype=np.int64)
+        IL_perfect = np.array(IL_perfect, dtype=np.int64)
+
+        # compute precision recall and F1
+        precision = TPs / (TPs + FPs + 1e-4)
+        assert np.all(precision <= 1)
+        recall = TPs / (TPs + FNs + 1e-4)
+        assert np.all(recall <= 1)
+        F1 = 2 * precision * recall / (precision + recall + 1e-4)
+
+        result = {
+            "image_id": imgId,
+            "TPs": TPs,
+            "FPs": FPs,
+            "FNs": FNs,
+            "local_F1s": F1,
+            "IL_TP": (len(gt) > 0) and (len(dt) > 0),
+            "IL_FP": (len(gt) == 0) and (len(dt) > 0),
+            "IL_TN": (len(gt) == 0) and (len(dt) == 0),
+            "IL_FN": (len(gt) > 0) and (len(dt) == 0),
+            "num_dt": len(dt),
+        }
+        if len(gt) > 0 and len(dt) > 0:
+            result["local_positive_F1s"] = F1
+        return result
+
+    def accumulate(self, p=None):
+        """
+        Accumulate per image evaluation results and store the result in self.eval
+        :param p: input params for evaluation
+        :return: None
+        """
+        if self.evalImgs is None or len(self.evalImgs) == 0:
+            print("Please run evaluate() first")
+        # allows input customized parameters
+        if p is None:
+            p = self.params
+
+        setImgIds = set(p.imgIds)
+
+        # TPs, FPs, FNs
+        TPs = np.zeros((len(p.iouThrs),), dtype=np.int64)
+        FPs = np.zeros((len(p.iouThrs),), dtype=np.int64)
+        pmFPs = np.zeros((len(p.iouThrs),), dtype=np.int64)
+        FNs = np.zeros((len(p.iouThrs),), dtype=np.int64)
+        local_F1s = np.zeros((len(p.iouThrs),), dtype=np.float64)
+
+        # Image level metrics
+        IL_TPs = 0
+        IL_FPs = 0
+        IL_TNs = 0
+        IL_FNs = 0
+
+        valid_img_count = 0
+        valid_F1_count = 0
+        evaledImgIds = set()
+        for res in self.evalImgs:
+            if res["image_id"] not in setImgIds:
+                continue
+            evaledImgIds.add(res["image_id"])
+            IL_TPs += res["IL_TP"]
+            IL_FPs += res["IL_FP"]
+            IL_TNs += res["IL_TN"]
+            IL_FNs += res["IL_FN"]
+
+            if "TPs" not in res:
+                continue
+
+            TPs += res["TPs"]
+            FPs += res["FPs"]
+            FNs += res["FNs"]
+            valid_img_count += 1
+
+            if "local_positive_F1s" in res:
+                local_F1s += res["local_positive_F1s"]
+                pmFPs += res["FPs"]
+                if res["num_dt"] > 0:
+                    valid_F1_count += 1
+
+        assert len(setImgIds - evaledImgIds) == 0, (
+            f"{len(setImgIds - evaledImgIds)} images not evaluated. "
+            f"Here are the IDs of the first 3: {list(setImgIds - evaledImgIds)[:3]}"
+        )
+
+        # compute precision recall and F1
+        precision = TPs / (TPs + FPs + 1e-4)
+        positive_micro_precision = TPs / (TPs + pmFPs + 1e-4)
+        assert np.all(precision <= 1)
+        recall = TPs / (TPs + FNs + 1e-4)
+        assert np.all(recall <= 1)
+        F1 = 2 * precision * recall / (precision + recall + 1e-4)
+        positive_micro_F1 = (
+            2
+            * positive_micro_precision
+            * recall
+            / (positive_micro_precision + recall + 1e-4)
+        )
+
+        IL_rec = IL_TPs / (IL_TPs + IL_FNs + 1e-6)
+        IL_prec = IL_TPs / (IL_TPs + IL_FPs + 1e-6)
+        IL_F1 = 2 * IL_prec * IL_rec / (IL_prec + IL_rec + 1e-6)
+        IL_FPR = IL_FPs / (IL_FPs + IL_TNs + 1e-6)
+        IL_MCC = float(IL_TPs * IL_TNs - IL_FPs * IL_FNs) / (
+            (
+                float(IL_TPs + IL_FPs)
+                * float(IL_TPs + IL_FNs)
+                * float(IL_TNs + IL_FPs)
+                * float(IL_TNs + IL_FNs)
+            )
+            ** 0.5
+            + 1e-6
+        )
+
+        self.eval = {
+            "params": p,
+            "TPs": TPs,
+            "FPs": FPs,
+            "positive_micro_FPs": pmFPs,
+            "FNs": FNs,
+            "precision": precision,
+            "positive_micro_precision": positive_micro_precision,
+            "recall": recall,
+            "F1": F1,
+            "positive_micro_F1": positive_micro_F1,
+            "positive_macro_F1": local_F1s / valid_F1_count,
+            "IL_recall": IL_rec,
+            "IL_precision": IL_prec,
+            "IL_F1": IL_F1,
+            "IL_FPR": IL_FPR,
+            "IL_MCC": IL_MCC,
+        }
+        self.eval["cgF1"] = self.eval["positive_micro_F1"] * self.eval["IL_MCC"]
+
+    def summarize(self):
+        """
+        Compute and display summary metrics for evaluation results.
+        """
+        if not self.eval:
+            raise Exception("Please run accumulate() first")
+
+        def _summarize(iouThr=None, metric=""):
+            p = self.params
+            iStr = " {:<18} @[ IoU={:<9}] = {:0.3f}"
+            titleStr = "Average " + metric
+            iouStr = (
+                "{:0.2f}:{:0.2f}".format(p.iouThrs[0], p.iouThrs[-1])
+                if iouThr is None
+                else "{:0.2f}".format(iouThr)
+            )
+
+            s = self.eval[metric]
+            # IoU
+            if iouThr is not None:
+                t = np.where(iouThr == p.iouThrs)[0]
+                s = s[t]
+
+            if len(s[s > -1]) == 0:
+                mean_s = -1
+            else:
+                mean_s = np.mean(s[s > -1])
+            print(iStr.format(titleStr, iouStr, mean_s))
+            return mean_s
+
+        def _summarize_single(metric=""):
+            titleStr = "Average " + metric
+            iStr = " {:<35} = {:0.3f}"
+            s = self.eval[metric]
+            print(iStr.format(titleStr, s))
+            return s
+
+        def _summarizeDets():
+            stats = []
+
+            for metric in CGF1_METRICS:
+                if metric.image_level:
+                    stats.append(_summarize_single(metric=metric.name))
+                else:
+                    stats.append(
+                        _summarize(iouThr=metric.iou_threshold, metric=metric.name)
+                    )
+            return np.asarray(stats)
+
+        summarize = _summarizeDets
+        self.stats = summarize()
+
+
+def _evaluate(self):
+    """
+    Run per image evaluation on given images and store results (a list of dict) in self.evalImgs
+    """
+    p = self.params
+    # add backward compatibility if useSegm is specified in params
+    p.imgIds = list(np.unique(p.imgIds))
+    p.useCats = False
+    p.maxDets = sorted(p.maxDets)
+    self.params = p
+
+    self._prepare()
+    # loop through images, area range, max detection number
+    catIds = [-1]
+
+    if p.iouType == "segm" or p.iouType == "bbox":
+        computeIoU = self.computeIoU
+    else:
+        raise RuntimeError(f"Unsupported iou {p.iouType}")
+    self.ious = {
+        (imgId, catId): computeIoU(imgId, catId)
+        for imgId in p.imgIds
+        for catId in catIds
+    }
+
+    maxDet = p.maxDets[-1]
+    evalImgs = [
+        self.evaluateImg(imgId, catId, areaRng, maxDet)
+        for catId in catIds
+        for areaRng in p.areaRng
+        for imgId in p.imgIds
+    ]
+    # this is NOT in the pycocotools code, but could be done outside
+    evalImgs = np.asarray(evalImgs).reshape(len(catIds), len(p.areaRng), len(p.imgIds))
+    return p.imgIds, evalImgs
+
+
+class CGF1Evaluator:
+    """
+    Wrapper class for cgF1 evaluation.
+    This supports the oracle setting (when several ground-truths are available per image)
+    """
+
+    def __init__(
+        self,
+        gt_path: Union[str, List[str]],
+        iou_type="segm",
+        verbose=False,
+    ):
+        """
+        Args:
+            gt_path (str or list of str): path(s) to ground truth COCO json file(s)
+            iou_type (str): type of IoU to evaluate
+            threshold (float): threshold for predictions
+        """
+        self.gt_paths = gt_path if isinstance(gt_path, list) else [gt_path]
+        self.iou_type = iou_type
+
+        self.coco_gts = [COCOCustom(gt) for gt in self.gt_paths]
+
+        self.verbose = verbose
+
+        self.coco_evals = []
+        for i, coco_gt in enumerate(self.coco_gts):
+            self.coco_evals.append(
+                CGF1Eval(
+                    coco_gt=coco_gt,
+                    iouType=iou_type,
+                )
+            )
+            self.coco_evals[i].useCats = False
+
+        exclude_img_ids = set()
+        # exclude_img_ids are the ids that are not exhaustively annotated in any of the other gts
+        for coco_gt in self.coco_gts[1:]:
+            exclude_img_ids = exclude_img_ids.union(
+                {
+                    img["id"]
+                    for img in coco_gt.dataset["images"]
+                    if not img["is_instance_exhaustive"]
+                }
+            )
+        # we only eval on instance exhaustive queries
+        self.eval_img_ids = [
+            img["id"]
+            for img in self.coco_gts[0].dataset["images"]
+            if (img["is_instance_exhaustive"] and img["id"] not in exclude_img_ids)
+        ]
+
+    def evaluate(self, pred_file: str):
+        """
+        Evaluate the detections using cgF1 metric.
+
+        Args:
+            pred_file: path to the predictions COCO json file
+
+        """
+        assert len(self.coco_gts) > 0, "No ground truth provided for evaluation."
+        assert len(self.coco_gts) == len(
+            self.coco_evals
+        ), "Mismatch in number of ground truths and evaluators."
+
+        if self.verbose:
+            print(f"Loading predictions from {pred_file}")
+
+        with open(pred_file, "r") as f:
+            preds = json.load(f)
+
+        if self.verbose:
+            print(f"Loaded {len(preds)} predictions")
+
+        img2preds = defaultdict(list)
+        for pred in preds:
+            img2preds[pred["image_id"]].append(pred)
+
+        all_eval_imgs = []
+        for img_id in tqdm(self.eval_img_ids, disable=not self.verbose):
+            results = img2preds[img_id]
+            all_scorings = []
+            for cur_coco_gt, coco_eval in zip(self.coco_gts, self.coco_evals):
+                # suppress pycocotools prints
+                with open(os.devnull, "w") as devnull:
+                    with contextlib.redirect_stdout(devnull):
+                        coco_dt = (
+                            cur_coco_gt.loadRes(results) if results else COCOCustom()
+                        )
+
+                coco_eval.cocoDt = coco_dt
+                coco_eval.params.imgIds = [img_id]
+                coco_eval.params.useCats = False
+                img_ids, eval_imgs = _evaluate(coco_eval)
+                all_scorings.append(eval_imgs)
+            selected = self._select_best_scoring(all_scorings)
+            all_eval_imgs.append(selected)
+
+        # After this point, we have selected the best scoring per image among several ground truths
+        # we can now accumulate and summarize, using only the first coco_eval
+
+        self.coco_evals[0].evalImgs = list(
+            np.concatenate(all_eval_imgs, axis=2).flatten()
+        )
+        self.coco_evals[0].params.imgIds = self.eval_img_ids
+        self.coco_evals[0]._paramsEval = copy.deepcopy(self.coco_evals[0].params)
+
+        if self.verbose:
+            print(f"Accumulating results")
+        self.coco_evals[0].accumulate()
+        print("cgF1 metric, IoU type={}".format(self.iou_type))
+        self.coco_evals[0].summarize()
+        print()
+
+        out = {}
+        for i, value in enumerate(self.coco_evals[0].stats):
+            name = CGF1_METRICS[i].name
+            if CGF1_METRICS[i].iou_threshold is not None:
+                name = f"{name}@{CGF1_METRICS[i].iou_threshold}"
+            out[f"cgF1_eval_{self.iou_type}_{name}"] = float(value)
+
+        return out
+
+    @staticmethod
+    def _select_best_scoring(scorings):
+        # This function is used for "oracle" type evaluation.
+        # It accepts the evaluation results with respect to several ground truths, and picks the best
+        if len(scorings) == 1:
+            return scorings[0]
+
+        assert (
+            scorings[0].ndim == 3
+        ), f"Expecting results in [numCats, numAreas, numImgs] format, got {scorings[0].shape}"
+        assert (
+            scorings[0].shape[0] == 1
+        ), f"Expecting a single category, got {scorings[0].shape[0]}"
+
+        for scoring in scorings:
+            assert (
+                scoring.shape == scorings[0].shape
+            ), f"Shape mismatch: {scoring.shape}, {scorings[0].shape}"
+
+        selected_imgs = []
+        for img_id in range(scorings[0].shape[-1]):
+            best = scorings[0][:, :, img_id]
+
+            for scoring in scorings[1:]:
+                current = scoring[:, :, img_id]
+                if "local_F1s" in best[0, 0] and "local_F1s" in current[0, 0]:
+                    # we were able to compute a F1 score for this particular image in both evaluations
+                    # best["local_F1s"] contains the results at various IoU thresholds. We simply take the average for comparision
+                    best_score = best[0, 0]["local_F1s"].mean()
+                    current_score = current[0, 0]["local_F1s"].mean()
+                    if current_score > best_score:
+                        best = current
+
+                else:
+                    # If we're here, it means that in that in some evaluation we were not able to get a valid local F1
+                    # This happens when both the predictions and targets are empty. In that case, we can assume it's a perfect prediction
+                    if "local_F1s" not in current[0, 0]:
+                        best = current
+            selected_imgs.append(best)
+        result = np.stack(selected_imgs, axis=-1)
+        assert result.shape == scorings[0].shape
+        return result

sam3/eval/coco_eval.py

@@ -0,0 +1,916 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+"""
+COCO evaluator that works in distributed mode.
+
+Mostly copy-paste from https://github.com/pytorch/vision/blob/edfd5a7/references/detection/coco_eval.py
+The difference is that there is less copy-pasting from pycocotools
+in the end of the file, as python3 can suppress prints with contextlib
+"""
+
+import contextlib
+import copy
+import json
+import logging
+import os
+import pickle
+from collections import defaultdict
+from pathlib import Path
+
+from typing import Any, List, Optional
+
+import numpy as np
+
+import pycocotools.mask as mask_utils
+import torch
+from iopath.common.file_io import g_pathmgr
+from pycocotools.coco import COCO
+from pycocotools.cocoeval import COCOeval
+
+from sam3.train.masks_ops import rle_encode
+
+from sam3.train.utils.distributed import (
+    all_gather,
+    gather_to_rank_0_via_filesys,
+    get_rank,
+    is_main_process,
+)
+
+RARITY_BUCKETS = {0: "frequent", 1: "common", 2: "medium", 3: "rare"}
+
+
+class CocoEvaluator:
+    def __init__(
+        self,
+        coco_gt,
+        iou_types: List[str],
+        useCats: bool,
+        dump_dir: Optional[str],
+        postprocessor,
+        average_by_rarity=False,
+        metrics_dump_dir: Optional[str] = None,
+        gather_pred_via_filesys=False,
+        use_normalized_areas=True,
+        maxdets=[1, 10, 100],
+        exhaustive_only=False,
+        all_exhaustive_only=True,
+    ):
+        """Online coco evaluator. It will evaluate images as they are generated by the model, then accumulate/summarize at the end
+
+        Args:
+           - coco_gt: COCO api object containing the gt
+           - iou_types: can be either "bbox" or "segm"
+           - useCats: If true, categories will be used for evaluation
+           - dump_dir: if non null, then the predictions will be dumped in that directory
+           - postprocessor: Module to convert the model's output into the coco format
+           - average_by_rarity: if true then we expect the images information in the gt dataset
+                 to have a "rarity" field. Then the AP will be computed on all rarity buckets
+                 individually, then averaged
+           - gather_pred_via_filesys: if true, we use the filesystem for collective gathers
+           - use_normalized_areas: if true, the areas of the objects in the GT are assumed to be
+                 normalized by the area of the image. In that case, the size buckets are adjusted
+           - maxdets: maximal number of detections to be evaluated on each image.
+           - exhaustive_only: If true, we restrict eval only to exhaustive annotations
+           - all_exhaustive_only: If true, datapoints are restricted only to those with all exhaustive annotations
+
+        """
+        # coco_gt = copy.deepcopy(coco_gt)
+        self.coco_gts = [coco_gt] if not isinstance(coco_gt, list) else coco_gt
+        assert len(maxdets) == 3, f"expecting 3 detection threshold, got {len(maxdets)}"
+
+        self.use_normalized_areas = use_normalized_areas
+        self.iou_types = iou_types
+        self.useCats = useCats
+        self.maxdets = maxdets
+        self.dump = None
+        self.dump_dir = dump_dir
+        if self.dump_dir is not None:
+            self.dump = []
+            if is_main_process():
+                if not os.path.exists(self.dump_dir):
+                    os.makedirs(self.dump_dir, exist_ok=True)
+                    logging.info(f"Create the folder: {dump_dir}")
+
+        self.initialized = False
+
+        # Whether to gather predictions through filesystem (instead of torch
+        # collective ops; requiring a shared filesystem across all ranks)
+        self.gather_pred_via_filesys = gather_pred_via_filesys
+        self.use_self_evaluate = True  # CPP version is disabled
+        self.postprocessor = postprocessor
+        self.average_by_rarity = average_by_rarity
+        self.exhaustive_only = exhaustive_only
+        self.all_exhaustive_only = all_exhaustive_only
+        self.metrics_dump_dir = metrics_dump_dir
+        if self.metrics_dump_dir is not None:
+            if is_main_process():
+                if not os.path.exists(self.metrics_dump_dir):
+                    os.makedirs(self.metrics_dump_dir, exist_ok=True)
+                    logging.info(f"Create the folder: {metrics_dump_dir}")
+
+    def _lazy_init(self, coco_cls=COCO):
+        if self.initialized:
+            return
+
+        self.initialized = True
+
+        self.coco_gts = [
+            coco_cls(g_pathmgr.get_local_path(gt)) if isinstance(gt, str) else gt
+            for gt in self.coco_gts
+        ]
+
+        self.reset()
+
+        self.eval_img_ids = None
+
+        if self.exhaustive_only:
+            exclude_img_ids = set()
+            # exclude_img_ids are the ids that are not exhaustively annotated in any of the other gts
+            if self.all_exhaustive_only:
+                for coco_gt in self.coco_gts[1:]:
+                    exclude_img_ids = exclude_img_ids.union(
+                        {
+                            img["id"]
+                            for img in coco_gt.dataset["images"]
+                            if not img["is_instance_exhaustive"]
+                        }
+                    )
+            # we only eval on instance exhaustive queries
+            self.eval_img_ids = [
+                img["id"]
+                for img in self.coco_gts[0].dataset["images"]
+                if (img["is_instance_exhaustive"] and img["id"] not in exclude_img_ids)
+            ]
+
+        self.rarity_buckets = None
+        if self.average_by_rarity:
+            self.rarity_buckets = defaultdict(list)
+            eval_img_ids_set = (
+                set(self.eval_img_ids) if self.eval_img_ids is not None else None
+            )
+            for img in self.coco_gts[0].dataset["images"]:
+                if self.eval_img_ids is not None and img["id"] not in eval_img_ids_set:
+                    continue
+                self.rarity_buckets[img["rarity"]].append(img["id"])
+            print("Rarity buckets sizes:")
+            for k, v in self.rarity_buckets.items():
+                print(f"{k}: {len(v)}")
+
+    def set_sync_device(self, device: torch.device) -> Any:
+        self._sync_device = device
+
+    def _evaluate(self, *args, **kwargs):
+        return evaluate(*args, **kwargs)
+
+    def _loadRes(self, *args, **kwargs):
+        return loadRes(*args, **kwargs)
+
+    def update(self, *args, **kwargs):
+        self._lazy_init()
+        predictions = self.postprocessor.process_results(*args, **kwargs)
+
+        img_ids = list(np.unique(list(predictions.keys())))
+        self.img_ids.extend(img_ids)
+
+        for iou_type in self.iou_types:
+            results = self.prepare(predictions, iou_type)
+            self._dump(results)
+
+            assert len(self.coco_gts) == len(self.coco_evals)
+            all_scorings = []
+            for cur_coco_gt, cur_coco_eval in zip(self.coco_gts, self.coco_evals):
+                # suppress pycocotools prints
+                with open(os.devnull, "w") as devnull:
+                    with contextlib.redirect_stdout(devnull):
+                        coco_dt = (
+                            self._loadRes(cur_coco_gt, results) if results else COCO()
+                        )
+
+                coco_eval = cur_coco_eval[iou_type]
+
+                coco_eval.cocoDt = coco_dt
+                coco_eval.params.imgIds = list(img_ids)
+                coco_eval.params.useCats = self.useCats
+                coco_eval.params.maxDets = self.maxdets
+                img_ids, eval_imgs = self._evaluate(coco_eval, self.use_self_evaluate)
+                all_scorings.append(eval_imgs)
+
+            selected = self.select_best_scoring(all_scorings)
+            self.eval_imgs[iou_type].append(selected)
+
+    def select_best_scoring(self, scorings):
+        # This function is used for "oracle" type evaluation.
+        # It accepts the evaluation results with respect to several ground truths, and picks the best
+        if len(scorings) == 1:
+            return scorings[0]
+
+        # Currently we don't support Oracle Phrase AP.
+        # To implement it, we likely need to modify the cpp code since the eval_image type is opaque
+        raise RuntimeError("Not implemented")
+
+    def _dump(self, results):
+        if self.dump is not None:
+            dumped_results = copy.deepcopy(results)
+            for r in dumped_results:
+                if "bbox" not in self.iou_types and "bbox" in r:
+                    del r["bbox"]
+                elif "bbox" in r:
+                    r["bbox"] = [round(coord, 5) for coord in r["bbox"]]
+                r["score"] = round(r["score"], 5)
+            self.dump.extend(dumped_results)
+
+    def synchronize_between_processes(self):
+        self._lazy_init()
+        logging.info("Coco evaluator: Synchronizing between processes")
+        for iou_type in self.iou_types:
+            if len(self.eval_imgs[iou_type]) > 0:
+                self.eval_imgs[iou_type] = np.concatenate(self.eval_imgs[iou_type], 2)
+            else:
+                num_areas = len(self.coco_evals[0][iou_type].params.areaRng)
+                # assuming 1 class
+                assert not self.useCats
+                self.eval_imgs[iou_type] = np.empty((1, num_areas, 0))
+            create_common_coco_eval(
+                self.coco_evals[0][iou_type],
+                self.img_ids,
+                self.eval_imgs[iou_type],
+                use_self_evaluate=self.use_self_evaluate,
+                gather_pred_via_filesys=self.gather_pred_via_filesys,
+                metrics_dump_dir=self.metrics_dump_dir,
+            )
+        if self.dump is not None:
+            dumped_file = Path(self.dump_dir) / f"coco_predictions_{get_rank()}.json"
+            logging.info(f"COCO evaluator: Dumping local predictions to {dumped_file}")
+            with g_pathmgr.open(str(dumped_file), "w") as f:
+                json.dump(self.dump, f)
+
+            # if self.gather_pred_via_filesys:
+            #     dump = gather_to_rank_0_via_filesys(self.dump)
+            # else:
+            #     dump = all_gather(self.dump, force_cpu=True)
+            # self.dump = sum(dump, [])
+
+    def accumulate(self, imgIds=None):
+        self._lazy_init()
+        logging.info(
+            f"Coco evaluator: Accumulating on {len(imgIds) if imgIds is not None else 'all'} images"
+        )
+        if not is_main_process():
+            return
+
+        if imgIds is None:
+            for coco_eval in self.coco_evals[0].values():
+                accumulate(coco_eval, use_self_eval=self.use_self_evaluate)
+
+        if imgIds is not None:
+            imgIds = set(imgIds)
+            for coco_eval in self.coco_evals[0].values():
+                p = coco_eval.params
+                id_mask = np.array([(i in imgIds) for i in p.imgIds], dtype=bool)
+                old_img_ids = p.imgIds
+                coco_eval.params.imgIds = np.asarray(p.imgIds)[id_mask]
+                old_img_evals = coco_eval.evalImgs
+                catIds = p.catIds if p.useCats else [-1]
+                coco_eval.evalImgs = list(
+                    np.asarray(coco_eval.evalImgs)
+                    .reshape(len(catIds), len(p.areaRng), len(old_img_ids))[
+                        ..., id_mask
+                    ]
+                    .flatten()
+                )
+                accumulate(coco_eval, use_self_eval=self.use_self_evaluate)
+                coco_eval.evalImgs = old_img_evals
+                coco_eval.params.imgIds = old_img_ids
+
+    def summarize(self):
+        self._lazy_init()
+        logging.info("Coco evaluator: Summarizing")
+        if not is_main_process():
+            return {}
+
+        outs = {}
+        if self.rarity_buckets is None:
+            self.accumulate(self.eval_img_ids)
+            for iou_type, coco_eval in self.coco_evals[0].items():
+                print("IoU metric: {}".format(iou_type))
+                summarize(coco_eval)
+
+            if "bbox" in self.coco_evals[0]:
+                for key, value in zip(*self.coco_evals[0]["bbox"].stats):
+                    outs[f"coco_eval_bbox_{key}"] = value
+            if "segm" in self.coco_evals[0]:
+                for key, value in zip(*self.coco_evals[0]["segm"].stats):
+                    outs[f"coco_eval_masks_{key}"] = value
+        else:
+            total_stats = {}
+            all_keys = {}
+            for bucket, img_list in self.rarity_buckets.items():
+                self.accumulate(imgIds=img_list)
+                bucket_name = RARITY_BUCKETS[bucket]
+                for iou_type, coco_eval in self.coco_evals[0].items():
+                    print(f"IoU metric: {iou_type}. Rarity bucket: {bucket_name}")
+                    summarize(coco_eval)
+
+                if "bbox" in self.coco_evals[0]:
+                    if "bbox" not in total_stats:
+                        total_stats["bbox"] = np.zeros_like(
+                            self.coco_evals[0]["bbox"].stats[1]
+                        )
+                        all_keys["bbox"] = self.coco_evals[0]["bbox"].stats[0]
+                    total_stats["bbox"] += self.coco_evals[0]["bbox"].stats[1]
+                    for key, value in zip(*self.coco_evals[0]["bbox"].stats):
+                        outs[f"coco_eval_bbox_{bucket_name}_{key}"] = value
+                if "segm" in self.coco_evals[0]:
+                    if "segm" not in total_stats:
+                        total_stats["segm"] = np.zeros_like(
+                            self.coco_evals[0]["segm"].stats[1]
+                        )
+                        all_keys["segm"] = self.coco_evals[0]["segm"].stats[0]
+                    total_stats["segm"] += self.coco_evals[0]["segm"].stats[1]
+                    for key, value in zip(*self.coco_evals[0]["segm"].stats):
+                        outs[f"coco_eval_masks_{bucket_name}_{key}"] = value
+
+            if "bbox" in total_stats:
+                total_stats["bbox"] /= len(self.rarity_buckets)
+                for key, value in zip(all_keys["bbox"], total_stats["bbox"]):
+                    outs[f"coco_eval_bbox_{key}"] = value
+            if "segm" in total_stats:
+                total_stats["segm"] /= len(self.rarity_buckets)
+                for key, value in zip(all_keys["segm"], total_stats["segm"]):
+                    outs[f"coco_eval_masks_{key}"] = value
+
+        # if self.dump is not None:
+        #     assert self.dump_dir is not None
+        #     logging.info("Coco evaluator: Dumping the global result file to disk")
+        #     with g_pathmgr.open(str(Path(self.dump_dir) / "coco_eval.json"), "w") as f:
+        #         json.dump(self.dump, f)
+        return outs
+
+    def compute_synced(self):
+        self._lazy_init()
+        self.synchronize_between_processes()
+        return self.summarize()
+
+    def compute(self):
+        self._lazy_init()
+        return {"": 0.0}
+
+    def reset(self, cocoeval_cls=COCOeval):
+        self.coco_evals = [{} for _ in range(len(self.coco_gts))]
+        for i, coco_gt in enumerate(self.coco_gts):
+            for iou_type in self.iou_types:
+                self.coco_evals[i][iou_type] = cocoeval_cls(coco_gt, iouType=iou_type)
+                self.coco_evals[i][iou_type].params.useCats = self.useCats
+                self.coco_evals[i][iou_type].params.maxDets = self.maxdets
+                if self.use_normalized_areas:
+                    self.coco_evals[i][iou_type].params.areaRng = [
+                        [0, 1e5],
+                        [0, 0.001],
+                        [0.001, 0.01],
+                        [0.01, 0.1],
+                        [0.1, 0.5],
+                        [0.5, 0.95],
+                        [0.95, 1e5],
+                    ]
+                    self.coco_evals[i][iou_type].params.areaRngLbl = [
+                        "all",
+                        "tiny",
+                        "small",
+                        "medium",
+                        "large",
+                        "huge",
+                        "whole_image",
+                    ]
+
+        self.img_ids = []
+        self.eval_imgs = {k: [] for k in self.iou_types}
+        if self.dump is not None:
+            self.dump = []
+
+    def write(self, stats):
+        self._lazy_init()
+        """Write the results in the stats dict"""
+        if "bbox" in self.coco_evals[0]:
+            stats["coco_eval_bbox"] = self.coco_evals[0]["bbox"].stats.tolist()
+        if "segm" in self.coco_evals[0]:
+            stats["coco_eval_masks"] = self.coco_evals[0]["segm"].stats.tolist()
+        return stats
+
+    def prepare(self, predictions, iou_type):
+        self._lazy_init()
+        if iou_type == "bbox":
+            return self.prepare_for_coco_detection(predictions)
+        elif iou_type == "segm":
+            return self.prepare_for_coco_segmentation(predictions)
+        elif iou_type == "keypoints":
+            return self.prepare_for_coco_keypoint(predictions)
+        else:
+            raise ValueError("Unknown iou type {}".format(iou_type))
+
+    def prepare_for_coco_detection(self, predictions):
+        self._lazy_init()
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+
+            boxes = prediction["boxes"]
+            boxes = convert_to_xywh(boxes).tolist()
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+
+            coco_results.extend(
+                [
+                    {
+                        "image_id": original_id,
+                        "category_id": labels[k],
+                        "bbox": box,
+                        "score": scores[k],
+                    }
+                    for k, box in enumerate(boxes)
+                ]
+            )
+        return coco_results
+
+    @torch.no_grad()
+    def prepare_for_coco_segmentation(self, predictions):
+        self._lazy_init()
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+            boundaries, dilated_boundaries = None, None
+            if "boundaries" in prediction:
+                boundaries = prediction["boundaries"]
+                dilated_boundaries = prediction["dilated_boundaries"]
+                assert dilated_boundaries is not None
+                assert len(scores) == len(boundaries)
+
+            if "masks_rle" in prediction:
+                rles = prediction["masks_rle"]
+                areas = []
+                for rle in rles:
+                    cur_area = mask_utils.area(rle)
+                    h, w = rle["size"]
+                    areas.append(cur_area / (h * w))
+            else:
+                masks = prediction["masks"]
+
+                masks = masks > 0.5
+                h, w = masks.shape[-2:]
+
+                areas = masks.flatten(1).sum(1) / (h * w)
+                areas = areas.tolist()
+
+                rles = rle_encode(masks.squeeze(1))
+
+                # memory clean
+                del masks
+                del prediction["masks"]
+
+            assert len(areas) == len(rles) == len(scores)
+            for k, rle in enumerate(rles):
+                payload = {
+                    "image_id": original_id,
+                    "category_id": labels[k],
+                    "segmentation": rle,
+                    "score": scores[k],
+                    "area": areas[k],
+                }
+                if boundaries is not None:
+                    payload["boundary"] = boundaries[k]
+                    payload["dilated_boundary"] = dilated_boundaries[k]
+
+                coco_results.append(payload)
+
+        return coco_results
+
+    def prepare_for_coco_keypoint(self, predictions):
+        self._lazy_init()
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+
+            boxes = prediction["boxes"]
+            boxes = convert_to_xywh(boxes).tolist()
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+            keypoints = prediction["keypoints"]
+            keypoints = keypoints.flatten(start_dim=1).tolist()
+
+            coco_results.extend(
+                [
+                    {
+                        "image_id": original_id,
+                        "category_id": labels[k],
+                        "keypoints": keypoint,
+                        "score": scores[k],
+                    }
+                    for k, keypoint in enumerate(keypoints)
+                ]
+            )
+        return coco_results
+
+
+def convert_to_xywh(boxes):
+    xmin, ymin, xmax, ymax = boxes.unbind(-1)
+    return torch.stack((xmin, ymin, xmax - xmin, ymax - ymin), dim=-1)
+
+
+def merge(img_ids, eval_imgs, gather_pred_via_filesys=False):
+    if gather_pred_via_filesys:
+        # only gather the predictions to rank 0 (other ranks will receive empty
+        # lists for `all_img_ids` and `all_eval_imgs`, which should be OK as
+        # merging and evaluation are only done on rank 0)
+        all_img_ids = gather_to_rank_0_via_filesys(img_ids)
+        all_eval_imgs = gather_to_rank_0_via_filesys(eval_imgs)
+    else:
+        all_img_ids = all_gather(img_ids, force_cpu=True)
+        all_eval_imgs = all_gather(eval_imgs, force_cpu=True)
+    if not is_main_process():
+        return None, None
+
+    merged_img_ids = []
+    for p in all_img_ids:
+        merged_img_ids.extend(p)
+
+    merged_eval_imgs = []
+    for p in all_eval_imgs:
+        merged_eval_imgs.append(p)
+
+    merged_img_ids = np.array(merged_img_ids)
+    merged_eval_imgs = np.concatenate(merged_eval_imgs, 2)
+
+    # keep only unique (and in sorted order) images
+    merged_img_ids, idx = np.unique(merged_img_ids, return_index=True)
+    merged_eval_imgs = merged_eval_imgs[..., idx]
+
+    return merged_img_ids, merged_eval_imgs
+
+
+def create_common_coco_eval(
+    coco_eval,
+    img_ids,
+    eval_imgs,
+    use_self_evaluate,
+    gather_pred_via_filesys=False,
+    metrics_dump_dir=None,
+):
+    img_ids, eval_imgs = merge(img_ids, eval_imgs, gather_pred_via_filesys)
+    if not is_main_process():
+        return
+    if metrics_dump_dir is not None:
+        dumped_file = (
+            Path(metrics_dump_dir) / f"coco_eval_img_metrics_{get_rank()}.json"
+        )
+        logging.info(f"COCO evaluator: Dumping local predictions to {dumped_file}")
+        with g_pathmgr.open(str(dumped_file), "w") as f:
+            json.dump(eval_imgs.squeeze(), f, default=lambda x: x.tolist())
+    img_ids = list(img_ids)
+
+    # If some images were not predicted, we need to create dummy detections for them
+    missing_img_ids = set(coco_eval.cocoGt.getImgIds()) - set(img_ids)
+    if len(missing_img_ids) > 0:
+        print(f"WARNING: {len(missing_img_ids)} images were not predicted!")
+        coco_eval.cocoDt = COCO()
+        coco_eval.params.imgIds = list(missing_img_ids)
+        new_img_ids, new_eval_imgs = evaluate(coco_eval, use_self_evaluate)
+        img_ids.extend(new_img_ids)
+        eval_imgs = np.concatenate((eval_imgs, new_eval_imgs), axis=2)
+
+    eval_imgs = list(eval_imgs.flatten())
+    assert len(img_ids) == len(coco_eval.cocoGt.getImgIds())
+
+    coco_eval.evalImgs = eval_imgs
+    coco_eval.params.imgIds = img_ids
+    coco_eval._paramsEval = copy.deepcopy(coco_eval.params)
+
+
+#################################################################
+# From pycocotools, just removed the prints and fixed
+# a Python3 bug about unicode not defined
+#################################################################
+
+
+# Copy of COCO prepare, but doesn't convert anntoRLE
+def segmentation_prepare(self):
+    """
+    Prepare ._gts and ._dts for evaluation based on params
+    :return: None
+    """
+    p = self.params
+    if p.useCats:
+        gts = self.cocoGt.loadAnns(
+            self.cocoGt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds)
+        )
+        dts = self.cocoDt.loadAnns(
+            self.cocoDt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds)
+        )
+    else:
+        gts = self.cocoGt.loadAnns(self.cocoGt.getAnnIds(imgIds=p.imgIds))
+        dts = self.cocoDt.loadAnns(self.cocoDt.getAnnIds(imgIds=p.imgIds))
+
+    for gt in gts:
+        gt["ignore"] = gt["ignore"] if "ignore" in gt else 0
+        gt["ignore"] = "iscrowd" in gt and gt["iscrowd"]
+        if p.iouType == "keypoints":
+            gt["ignore"] = (gt["num_keypoints"] == 0) or gt["ignore"]
+    self._gts = defaultdict(list)  # gt for evaluation
+    self._dts = defaultdict(list)  # dt for evaluation
+    for gt in gts:
+        self._gts[gt["image_id"], gt["category_id"]].append(gt)
+    for dt in dts:
+        self._dts[dt["image_id"], dt["category_id"]].append(dt)
+    self.evalImgs = defaultdict(list)  # per-image per-category evaluation results
+    self.eval = {}  # accumulated evaluation results
+
+
+def evaluate(self, use_self_evaluate):
+    """
+    Run per image evaluation on given images and store results (a list of dict) in self.evalImgs
+    :return: None
+    """
+    # tic = time.time()
+    # print('Running per image evaluation...', use_self_evaluate)
+    p = self.params
+    # add backward compatibility if useSegm is specified in params
+    if p.useSegm is not None:
+        p.iouType = "segm" if p.useSegm == 1 else "bbox"
+        print(
+            "useSegm (deprecated) is not None. Running {} evaluation".format(p.iouType)
+        )
+    # print('Evaluate annotation type *{}*'.format(p.iouType))
+    p.imgIds = list(np.unique(p.imgIds))
+    if p.useCats:
+        p.catIds = list(np.unique(p.catIds))
+    p.maxDets = sorted(p.maxDets)
+    self.params = p
+
+    self._prepare()
+    # loop through images, area range, max detection number
+    catIds = p.catIds if p.useCats else [-1]
+
+    if p.iouType == "segm" or p.iouType == "bbox":
+        computeIoU = self.computeIoU
+    elif p.iouType == "keypoints":
+        computeIoU = self.computeOks
+    self.ious = {
+        (imgId, catId): computeIoU(imgId, catId)
+        for imgId in p.imgIds
+        for catId in catIds
+    }
+
+    maxDet = p.maxDets[-1]
+    if use_self_evaluate:
+        evalImgs = [
+            self.evaluateImg(imgId, catId, areaRng, maxDet)
+            for catId in catIds
+            for areaRng in p.areaRng
+            for imgId in p.imgIds
+        ]
+        # this is NOT in the pycocotools code, but could be done outside
+        evalImgs = np.asarray(evalImgs).reshape(
+            len(catIds), len(p.areaRng), len(p.imgIds)
+        )
+        return p.imgIds, evalImgs
+
+    # <<<< Beginning of code differences with original COCO API
+    # def convert_instances_to_cpp(instances, is_det=False):
+    #     # Convert annotations for a list of instances in an image to a format that's fast
+    #     # to access in C++
+    #     instances_cpp = []
+    #     for instance in instances:
+    #         instance_cpp = _CPP.InstanceAnnotation(
+    #             int(instance["id"]),
+    #             instance["score"] if is_det else instance.get("score", 0.0),
+    #             instance["area"],
+    #             bool(instance.get("iscrowd", 0)),
+    #             bool(instance.get("ignore", 0)),
+    #         )
+    #         instances_cpp.append(instance_cpp)
+    #     return instances_cpp
+
+    # # Convert GT annotations, detections, and IOUs to a format that's fast to access in C++
+    # ground_truth_instances = [
+    #     [convert_instances_to_cpp(self._gts[imgId, catId]) for catId in p.catIds]
+    #     for imgId in p.imgIds
+    # ]
+    # detected_instances = [
+    #     [
+    #         convert_instances_to_cpp(self._dts[imgId, catId], is_det=True)
+    #         for catId in p.catIds
+    #     ]
+    #     for imgId in p.imgIds
+    # ]
+    # ious = [[self.ious[imgId, catId] for catId in catIds] for imgId in p.imgIds]
+
+    # if not p.useCats:
+    #     # For each image, flatten per-category lists into a single list
+    #     ground_truth_instances = [
+    #         [[o for c in i for o in c]] for i in ground_truth_instances
+    #     ]
+    #     detected_instances = [[[o for c in i for o in c]] for i in detected_instances]
+
+    # # Call C++ implementation of self.evaluateImgs()
+    # _evalImgs_cpp = _CPP.COCOevalEvaluateImages(
+    #     p.areaRng, maxDet, p.iouThrs, ious, ground_truth_instances, detected_instances
+    # )
+
+    # self._paramsEval = copy.deepcopy(self.params)
+    # evalImgs = np.asarray(_evalImgs_cpp).reshape(
+    #     len(catIds), len(p.areaRng), len(p.imgIds)
+    # )
+    # return p.imgIds, evalImgs
+
+
+#################################################################
+# end of straight copy from pycocotools, just removing the prints
+#################################################################
+
+
+#################################################################
+# From pycocotools, but disabled mask->box conversion which is
+# pointless
+#################################################################
+def loadRes(self, resFile):
+    """
+    Load result file and return a result api object.
+    :param   resFile (str)     : file name of result file
+    :return: res (obj)         : result api object
+    """
+    res = COCO()
+    res.dataset["images"] = [img for img in self.dataset["images"]]
+
+    if type(resFile) == str:
+        anns = json.load(open(resFile))
+    elif type(resFile) == np.ndarray:
+        anns = self.loadNumpyAnnotations(resFile)
+    else:
+        anns = resFile
+    assert type(anns) == list, "results in not an array of objects"
+    annsImgIds = [ann["image_id"] for ann in anns]
+    assert set(annsImgIds) == (
+        set(annsImgIds) & set(self.getImgIds())
+    ), "Results do not correspond to current coco set"
+    if "caption" in anns[0]:
+        imgIds = set([img["id"] for img in res.dataset["images"]]) & set(
+            [ann["image_id"] for ann in anns]
+        )
+        res.dataset["images"] = [
+            img for img in res.dataset["images"] if img["id"] in imgIds
+        ]
+        for id, ann in enumerate(anns):
+            ann["id"] = id + 1
+    elif "bbox" in anns[0] and not anns[0]["bbox"] == []:
+        res.dataset["categories"] = copy.deepcopy(self.dataset["categories"])
+        for id, ann in enumerate(anns):
+            bb = ann["bbox"]
+            x1, x2, y1, y2 = [bb[0], bb[0] + bb[2], bb[1], bb[1] + bb[3]]
+            if "segmentation" not in ann:
+                ann["segmentation"] = [[x1, y1, x1, y2, x2, y2, x2, y1]]
+            ann["area"] = bb[2] * bb[3]
+            ann["id"] = id + 1
+            ann["iscrowd"] = 0
+    elif "segmentation" in anns[0]:
+        res.dataset["categories"] = copy.deepcopy(self.dataset["categories"])
+        for id, ann in enumerate(anns):
+            # now only support compressed RLE format as segmentation results
+            # ann["area"] = mask_util.area(ann["segmentation"])
+            # The following lines are disabled because they are pointless
+            #  if not 'bbox' in ann:
+            #     ann['bbox'] = maskUtils.toBbox(ann['segmentation'])
+            ann["id"] = id + 1
+            ann["iscrowd"] = 0
+    elif "keypoints" in anns[0]:
+        res.dataset["categories"] = copy.deepcopy(self.dataset["categories"])
+        for id, ann in enumerate(anns):
+            s = ann["keypoints"]
+            x = s[0::3]
+            y = s[1::3]
+            x0, x1, y0, y1 = np.min(x), np.max(x), np.min(y), np.max(y)
+            ann["area"] = (x1 - x0) * (y1 - y0)
+            ann["id"] = id + 1
+            ann["bbox"] = [x0, y0, x1 - x0, y1 - y0]
+
+    res.dataset["annotations"] = anns
+    res.createIndex()
+    return res
+
+
+#################################################################
+# end of straight copy from pycocotools
+#################################################################
+
+
+#################################################################
+# From pycocotools, but added handling of custom area rngs, and returns stat keys
+#################################################################
+def summarize(self):
+    """
+    Compute and display summary metrics for evaluation results.
+    Note this functin can *only* be applied on the default parameter setting
+    """
+
+    def _summarize(ap=1, iouThr=None, areaRng="all", maxDets=100):
+        p = self.params
+        iStr = " {:<18} {} @[ IoU={:<9} | area={:>6s} | maxDets={:>3d} ] = {:0.3f}"
+        titleStr = "Average Precision" if ap == 1 else "Average Recall"
+        typeStr = "(AP)" if ap == 1 else "(AR)"
+        iouStr = (
+            "{:0.2f}:{:0.2f}".format(p.iouThrs[0], p.iouThrs[-1])
+            if iouThr is None
+            else "{:0.2f}".format(iouThr)
+        )
+
+        aind = [i for i, aRng in enumerate(p.areaRngLbl) if aRng == areaRng]
+        mind = [i for i, mDet in enumerate(p.maxDets) if mDet == maxDets]
+        if ap == 1:
+            # dimension of precision: [TxRxKxAxM]
+            s = self.eval["precision"]
+            # IoU
+            if iouThr is not None:
+                t = np.where(iouThr == p.iouThrs)[0]
+                s = s[t]
+            s = s[:, :, :, aind, mind]
+        else:
+            # dimension of recall: [TxKxAxM]
+            s = self.eval["recall"]
+            if iouThr is not None:
+                t = np.where(iouThr == p.iouThrs)[0]
+                s = s[t]
+            s = s[:, :, aind, mind]
+        if len(s[s > -1]) == 0:
+            mean_s = -1
+        else:
+            mean_s = np.mean(s[s > -1])
+        print(iStr.format(titleStr, typeStr, iouStr, areaRng, maxDets, mean_s))
+        return mean_s
+
+    def _summarizeDets():
+        nb_results = 6 + (len(self.params.areaRng) - 1) * 2
+        assert len(self.params.areaRng) == len(self.params.areaRngLbl)
+        stats = np.zeros((nb_results,))
+        keys = ["AP", "AP_50", "AP_75"]
+        stats[0] = _summarize(1, maxDets=self.params.maxDets[2])
+        stats[1] = _summarize(1, iouThr=0.5, maxDets=self.params.maxDets[2])
+        stats[2] = _summarize(1, iouThr=0.75, maxDets=self.params.maxDets[2])
+        cur_id = 3
+        for area in self.params.areaRngLbl[1:]:
+            stats[cur_id] = _summarize(1, areaRng=area, maxDets=self.params.maxDets[2])
+            cur_id += 1
+            keys.append(f"AP_{area}")
+        stats[cur_id] = _summarize(0, maxDets=self.params.maxDets[0])
+        cur_id += 1
+        stats[cur_id] = _summarize(0, maxDets=self.params.maxDets[1])
+        cur_id += 1
+        stats[cur_id] = _summarize(0, maxDets=self.params.maxDets[2])
+        cur_id += 1
+        keys += ["AR", "AR_50", "AR_75"]
+
+        for area in self.params.areaRngLbl[1:]:
+            stats[cur_id] = _summarize(0, areaRng=area, maxDets=self.params.maxDets[2])
+            cur_id += 1
+            keys.append(f"AR_{area}")
+        assert len(stats) == len(keys)
+        return keys, stats
+
+    if not self.eval:
+        raise Exception("Please run accumulate() first")
+    self.stats = _summarizeDets()
+
+
+#################################################################
+# end of straight copy from pycocotools
+#################################################################
+
+
+#################################################################
+# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/evaluation/fast_eval_api.py
+# with slight adjustments
+#################################################################
+def accumulate(self, use_self_eval=False):
+    """
+    Accumulate per image evaluation results and store the result in self.eval.  Does not
+    support changing parameter settings from those used by self.evaluate()
+    """
+    if use_self_eval:
+        self.accumulate()
+        return
+    # CPP code is disabled
+    # self.eval = _CPP.COCOevalAccumulate(self.params, self.evalImgs)
+
+    # # recall is num_iou_thresholds X num_categories X num_area_ranges X num_max_detections
+    # self.eval["recall"] = np.array(self.eval["recall"]).reshape(
+    #     self.eval["counts"][:1] + self.eval["counts"][2:]
+    # )
+
+    # # precision and scores are num_iou_thresholds X num_recall_thresholds X num_categories X
+    # # num_area_ranges X num_max_detections
+    # self.eval["precision"] = np.array(self.eval["precision"]).reshape(
+    #     self.eval["counts"]
+    # )
+    # self.eval["scores"] = np.array(self.eval["scores"]).reshape(self.eval["counts"])

sam3/eval/coco_eval_offline.py

@@ -0,0 +1,181 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+"""
+This evaluator is meant for regular COCO mAP evaluation, for example on the COCO val set.
+
+For Category mAP, we need the model to make predictions for all the categories on every single image.
+In general, since the number of classes can be big, and the API model makes predictions individually for each pair (image, class),
+we may need to split the inference process for a given image in several chunks.
+"""
+
+import logging
+from collections import defaultdict
+
+import torch
+from pycocotools.coco import COCO
+from pycocotools.cocoeval import COCOeval
+from sam3.train.utils.distributed import is_main_process
+
+try:
+    from tidecv import datasets, TIDE
+
+    HAS_TIDE = True
+except ImportError:
+    HAS_TIDE = False
+    print("WARNING: TIDE not installed. Detailed analysis will not be available.")
+
+
+# the COCO detection metrics (https://github.com/cocodataset/cocoapi/blob/8c9bcc3cf640524c4c20a9c40e89cb6a2f2fa0e9/PythonAPI/pycocotools/cocoeval.py#L460-L471)
+COCO_METRICS = [
+    "AP",
+    "AP_50",
+    "AP_75",
+    "AP_small",
+    "AP_medium",
+    "AP_large",
+    "AR_maxDets@1",
+    "AR_maxDets@10",
+    "AR_maxDets@100",
+    "AR_small",
+    "AR_medium",
+    "AR_large",
+]
+
+
+def convert_to_xywh(boxes):
+    """Convert bounding boxes from xyxy format to xywh format."""
+    xmin, ymin, xmax, ymax = boxes.unbind(-1)
+    return torch.stack((xmin, ymin, xmax - xmin, ymax - ymin), dim=-1)
+
+
+class HeapElement:
+    """Utility class to make a heap with a custom comparator"""
+
+    def __init__(self, val):
+        self.val = val
+
+    def __lt__(self, other):
+        return self.val["score"] < other.val["score"]
+
+
+class COCOevalCustom(COCOeval):
+    """
+    This is a slightly modified version of the original COCO API with added support for positive split evaluation.
+    """
+
+    def __init__(
+        self, cocoGt=None, cocoDt=None, iouType="segm", dt_only_positive=False
+    ):
+        super().__init__(cocoGt, cocoDt, iouType)
+        self.dt_only_positive = dt_only_positive
+
+    def _prepare(self):
+        """
+        Prepare ._gts and ._dts for evaluation based on params
+        :return: None
+        """
+
+        def _toMask(anns, coco):
+            # modify ann['segmentation'] by reference
+            for ann in anns:
+                rle = coco.annToRLE(ann)
+                ann["segmentation"] = rle
+
+        p = self.params
+        if p.useCats:
+            gts = self.cocoGt.loadAnns(
+                self.cocoGt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds)
+            )
+            dts = self.cocoDt.loadAnns(
+                self.cocoDt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds)
+            )
+        else:
+            gts = self.cocoGt.loadAnns(self.cocoGt.getAnnIds(imgIds=p.imgIds))
+            dts = self.cocoDt.loadAnns(self.cocoDt.getAnnIds(imgIds=p.imgIds))
+
+        # convert ground truth to mask if iouType == 'segm'
+        if p.iouType == "segm":
+            _toMask(gts, self.cocoGt)
+            _toMask(dts, self.cocoDt)
+        # set ignore flag
+        for gt in gts:
+            gt["ignore"] = gt["ignore"] if "ignore" in gt else 0
+            gt["ignore"] = "iscrowd" in gt and gt["iscrowd"]
+            if p.iouType == "keypoints":
+                gt["ignore"] = (gt["num_keypoints"] == 0) or gt["ignore"]
+        self._gts = defaultdict(list)  # gt for evaluation
+        self._dts = defaultdict(list)  # dt for evaluation
+
+        _gts_cat_ids = defaultdict(set)  # gt for evaluation on positive split
+        for gt in gts:
+            self._gts[gt["image_id"], gt["category_id"]].append(gt)
+            _gts_cat_ids[gt["image_id"]].add(gt["category_id"])
+
+        #### BEGIN MODIFICATION ####
+        for dt in dts:
+            if (
+                self.dt_only_positive
+                and dt["category_id"] not in _gts_cat_ids[dt["image_id"]]
+            ):
+                continue
+            self._dts[dt["image_id"], dt["category_id"]].append(dt)
+        #### END MODIFICATION ####
+        self.evalImgs = defaultdict(list)  # per-image per-category evaluation results
+        self.eval = {}  # accumulated evaluation results
+
+
+class CocoEvaluatorOfflineWithPredFileEvaluators:
+    def __init__(
+        self,
+        gt_path,
+        tide: bool = True,
+        iou_type: str = "bbox",
+        positive_split=False,
+    ):
+        self.gt_path = gt_path
+        self.tide_enabled = HAS_TIDE and tide
+        self.positive_split = positive_split
+        self.iou_type = iou_type
+
+    def evaluate(self, dumped_file):
+        if not is_main_process():
+            return {}
+
+        logging.info("OfflineCoco evaluator: Loading groundtruth")
+        self.gt = COCO(self.gt_path)
+
+        # Creating the result file
+        logging.info("Coco evaluator: Creating the result file")
+        cocoDt = self.gt.loadRes(str(dumped_file))
+
+        # Run the evaluation
+        logging.info("Coco evaluator: Running evaluation")
+        coco_eval = COCOevalCustom(
+            self.gt, cocoDt, iouType=self.iou_type, dt_only_positive=self.positive_split
+        )
+        coco_eval.evaluate()
+        coco_eval.accumulate()
+        coco_eval.summarize()
+
+        outs = {}
+        for i, value in enumerate(coco_eval.stats):
+            outs[f"coco_eval_{self.iou_type}_{COCO_METRICS[i]}"] = value
+
+        if self.tide_enabled:
+            logging.info("Coco evaluator: Loading TIDE")
+            self.tide_gt = datasets.COCO(self.gt_path)
+            self.tide = TIDE(mode="mask" if self.iou_type == "segm" else "bbox")
+
+            # Run TIDE
+            logging.info("Coco evaluator: Running TIDE")
+            self.tide.evaluate(
+                self.tide_gt, datasets.COCOResult(str(dumped_file)), name="coco_eval"
+            )
+            self.tide.summarize()
+            for k, v in self.tide.get_main_errors()["coco_eval"].items():
+                outs[f"coco_eval_{self.iou_type}_TIDE_{k}"] = v
+
+            for k, v in self.tide.get_special_errors()["coco_eval"].items():
+                outs[f"coco_eval_{self.iou_type}_TIDE_{k}"] = v
+
+        return outs

sam3/eval/coco_reindex.py

@@ -0,0 +1,230 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+"""
+Self-contained COCO JSON re-indexing function that creates temporary files.
+"""
+
+import json
+import os
+import tempfile
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+
+def reindex_coco_to_temp(input_json_path: str) -> Optional[str]:
+    """
+    Convert 0-indexed COCO JSON file to 1-indexed and save to temporary location.
+
+    Args:
+        input_json_path: Path to the input COCO JSON file
+
+    Returns:
+        Path to the new 1-indexed JSON file in temporary directory, or None if no conversion needed
+
+    Raises:
+        FileNotFoundError: If input file doesn't exist
+        json.JSONDecodeError: If input file is not valid JSON
+        ValueError: If input file is not a valid COCO format
+    """
+
+    def is_coco_json(data: Dict[str, Any]) -> bool:
+        """Check if data appears to be a COCO format file."""
+        if not isinstance(data, dict):
+            return False
+        # A COCO file should have at least one of these keys
+        coco_keys = {"images", "annotations", "categories"}
+        return any(key in data for key in coco_keys)
+
+    def check_zero_indexed(data: Dict[str, Any]) -> Tuple[bool, bool, bool]:
+        """
+        Check if annotations, images, or categories start from index 0.
+
+        Returns:
+            Tuple of (annotations_zero_indexed, images_zero_indexed, categories_zero_indexed)
+        """
+        annotations_zero = False
+        images_zero = False
+        categories_zero = False
+
+        # Check annotations
+        annotations = data.get("annotations", [])
+        if annotations and any(ann.get("id", -1) == 0 for ann in annotations):
+            annotations_zero = True
+
+        # Check images
+        images = data.get("images", [])
+        if images and any(img.get("id", -1) == 0 for img in images):
+            images_zero = True
+
+        # Check categories
+        categories = data.get("categories", [])
+        if categories and any(cat.get("id", -1) == 0 for cat in categories):
+            categories_zero = True
+
+        return annotations_zero, images_zero, categories_zero
+
+    def reindex_coco_data(data: Dict[str, Any]) -> Dict[str, Any]:
+        """Convert 0-indexed COCO data to 1-indexed."""
+        modified_data = data.copy()
+
+        annotations_zero, images_zero, categories_zero = check_zero_indexed(data)
+
+        # Create ID mapping for consistency
+        image_id_mapping = {}
+        category_id_mapping = {}
+
+        # Process images first (since annotations reference image IDs)
+        if images_zero and "images" in modified_data:
+            for img in modified_data["images"]:
+                old_id = img["id"]
+                new_id = old_id + 1
+                image_id_mapping[old_id] = new_id
+                img["id"] = new_id
+
+        # Process categories (since annotations reference category IDs)
+        if categories_zero and "categories" in modified_data:
+            for cat in modified_data["categories"]:
+                old_id = cat["id"]
+                new_id = old_id + 1
+                category_id_mapping[old_id] = new_id
+                cat["id"] = new_id
+
+        # Process annotations
+        if "annotations" in modified_data:
+            for ann in modified_data["annotations"]:
+                # Update annotation ID if needed
+                if annotations_zero:
+                    ann["id"] = ann["id"] + 1
+
+                # Update image_id reference if images were reindexed
+                if images_zero and ann.get("image_id") is not None:
+                    old_image_id = ann["image_id"]
+                    if old_image_id in image_id_mapping:
+                        ann["image_id"] = image_id_mapping[old_image_id]
+
+                # Update category_id reference if categories were reindexed
+                if categories_zero and ann.get("category_id") is not None:
+                    old_category_id = ann["category_id"]
+                    if old_category_id in category_id_mapping:
+                        ann["category_id"] = category_id_mapping[old_category_id]
+
+        return modified_data
+
+    # Validate input path
+    if not os.path.exists(input_json_path):
+        raise FileNotFoundError(f"Input file not found: {input_json_path}")
+
+    # Load and validate JSON data
+    try:
+        with open(input_json_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+    except json.JSONDecodeError as e:
+        raise json.JSONDecodeError(f"Invalid JSON in {input_json_path}: {e}")
+
+    # Validate COCO format
+    if not is_coco_json(data):
+        raise ValueError(
+            f"File does not appear to be in COCO format: {input_json_path}"
+        )
+
+    # Check if reindexing is needed
+    annotations_zero, images_zero, categories_zero = check_zero_indexed(data)
+
+    if not (annotations_zero or images_zero or categories_zero):
+        # No conversion needed - just copy to temp location
+        input_path = Path(input_json_path)
+        temp_dir = tempfile.mkdtemp()
+        temp_filename = f"{input_path.stem}_1_indexed{input_path.suffix}"
+        temp_path = os.path.join(temp_dir, temp_filename)
+
+        with open(temp_path, "w", encoding="utf-8") as f:
+            json.dump(data, f, indent=2, ensure_ascii=False)
+
+        return temp_path
+
+    # Perform reindexing
+    modified_data = reindex_coco_data(data)
+
+    # Create temporary file
+    input_path = Path(input_json_path)
+    temp_dir = tempfile.mkdtemp()
+    temp_filename = f"{input_path.stem}_1_indexed{input_path.suffix}"
+    temp_path = os.path.join(temp_dir, temp_filename)
+
+    # Write modified data to temporary file
+    with open(temp_path, "w", encoding="utf-8") as f:
+        json.dump(modified_data, f, indent=2, ensure_ascii=False)
+
+    return temp_path
+
+
+# Example usage and test function
+def test_reindex_function():
+    """Test the reindex function with a sample COCO file."""
+
+    # Create a test COCO file
+    test_data = {
+        "info": {"description": "Test COCO dataset", "version": "1.0", "year": 2023},
+        "images": [
+            {"id": 0, "width": 640, "height": 480, "file_name": "test1.jpg"},
+            {"id": 1, "width": 640, "height": 480, "file_name": "test2.jpg"},
+        ],
+        "categories": [
+            {"id": 0, "name": "person", "supercategory": "person"},
+            {"id": 1, "name": "car", "supercategory": "vehicle"},
+        ],
+        "annotations": [
+            {
+                "id": 0,
+                "image_id": 0,
+                "category_id": 0,
+                "bbox": [100, 100, 50, 75],
+                "area": 3750,
+                "iscrowd": 0,
+            },
+            {
+                "id": 1,
+                "image_id": 1,
+                "category_id": 1,
+                "bbox": [200, 150, 120, 80],
+                "area": 9600,
+                "iscrowd": 0,
+            },
+        ],
+    }
+
+    # Create temporary test file
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        json.dump(test_data, f, indent=2)
+        test_file_path = f.name
+
+    try:
+        # Test the function
+        result_path = reindex_coco_to_temp(test_file_path)
+        print(f"Original file: {test_file_path}")
+        print(f"Converted file: {result_path}")
+
+        # Load and display the result
+        with open(result_path, "r") as f:
+            result_data = json.load(f)
+
+        print("\nConverted data sample:")
+        print(f"First image ID: {result_data['images'][0]['id']}")
+        print(f"First category ID: {result_data['categories'][0]['id']}")
+        print(f"First annotation ID: {result_data['annotations'][0]['id']}")
+        print(f"First annotation image_id: {result_data['annotations'][0]['image_id']}")
+        print(
+            f"First annotation category_id: {result_data['annotations'][0]['category_id']}"
+        )
+
+        # Clean up
+        os.unlink(result_path)
+        os.rmdir(os.path.dirname(result_path))
+
+    finally:
+        # Clean up test file
+        os.unlink(test_file_path)
+
+
+if __name__ == "__main__":
+    test_reindex_function()

sam3/eval/coco_writer.py

@@ -0,0 +1,352 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+"""
+COCO prediction dumper for distributed training.
+
+Handles collection and dumping of COCO-format predictions from models.
+Supports distributed processing with multiple GPUs/processes.
+"""
+
+import copy
+import gc
+import heapq
+import json
+import logging
+import os
+from collections import defaultdict
+from pathlib import Path
+from typing import Any, Optional
+
+import pycocotools.mask as mask_utils
+import torch
+from iopath.common.file_io import g_pathmgr
+from sam3.eval.coco_eval_offline import convert_to_xywh
+from sam3.train.masks_ops import rle_encode
+from sam3.train.utils.distributed import (
+    all_gather,
+    gather_to_rank_0_via_filesys,
+    get_rank,
+    is_main_process,
+)
+
+
+### Helper functions and classes
+
+
+class HeapElement:
+    """Utility class to make a heap with a custom comparator based on score."""
+
+    def __init__(self, val):
+        self.val = val
+
+    def __lt__(self, other):
+        return self.val["score"] < other.val["score"]
+
+
+class PredictionDumper:
+    """
+    Handles collection and dumping of COCO-format predictions from a model.
+
+    This class processes model outputs through a postprocessor, converts them to COCO format,
+    and saves them to disk. It supports distributed processing with multiple GPUs/processes.
+    """
+
+    def __init__(
+        self,
+        dump_dir: str,
+        postprocessor,
+        maxdets: int,
+        iou_type: str,
+        gather_pred_via_filesys: bool = False,
+        merge_predictions: bool = False,
+        pred_file_evaluators: Optional[Any] = None,
+    ):
+        """
+        Initialize the PredictionDumper.
+
+        Args:
+            dump_dir: Directory to dump predictions.
+            postprocessor: Module to convert the model's output into COCO format.
+            maxdets: Maximum number of detections per image.
+            iou_type: IoU type to evaluate. Can include "bbox", "segm"
+            gather_pred_via_filesys: If True, use the filesystem for collective gathers across
+                processes (requires a shared filesystem). Otherwise, use torch collective ops.
+            merge_predictions: If True, merge predictions from all processes and dump to a single file.
+        """
+        self.iou_type = iou_type
+        self.maxdets = maxdets
+        self.dump_dir = dump_dir
+        self.postprocessor = postprocessor
+        self.gather_pred_via_filesys = gather_pred_via_filesys
+        self.merge_predictions = merge_predictions
+        self.pred_file_evaluators = pred_file_evaluators
+        if self.pred_file_evaluators is not None:
+            assert (
+                merge_predictions
+            ), "merge_predictions must be True if pred_file_evaluators are provided"
+        assert self.dump_dir is not None, "dump_dir must be provided"
+
+        if is_main_process():
+            os.makedirs(self.dump_dir, exist_ok=True)
+            logging.info(f"Created prediction dump directory: {self.dump_dir}")
+
+        # Initialize state
+        self.reset()
+
+    def update(self, *args, **kwargs):
+        """
+        Process and accumulate predictions from model outputs.
+
+        Args:
+            *args, **kwargs: Arguments passed to postprocessor.process_results()
+        """
+        predictions = self.postprocessor.process_results(*args, **kwargs)
+        results = self.prepare(predictions, self.iou_type)
+        self._dump(results)
+
+    def _dump(self, results):
+        """
+        Add results to the dump list with precision rounding.
+
+        Args:
+            results: List of prediction dictionaries in COCO format.
+        """
+        dumped_results = copy.deepcopy(results)
+        for r in dumped_results:
+            if "bbox" in r:
+                r["bbox"] = [round(coord, 5) for coord in r["bbox"]]
+            r["score"] = round(r["score"], 5)
+        self.dump.extend(dumped_results)
+
+    def synchronize_between_processes(self):
+        """
+        Synchronize predictions across all processes and save to disk.
+
+        If gather_pred_via_filesys is True, uses filesystem for gathering.
+        Otherwise, uses torch distributed collective operations.
+        Saves per-rank predictions to separate JSON files.
+        """
+        logging.info("Prediction Dumper: Synchronizing between processes")
+
+        if not self.merge_predictions:
+            dumped_file = (
+                Path(self.dump_dir)
+                / f"coco_predictions_{self.iou_type}_{get_rank()}.json"
+            )
+            logging.info(
+                f"Prediction Dumper: Dumping local predictions to {dumped_file}"
+            )
+            with g_pathmgr.open(str(dumped_file), "w") as f:
+                json.dump(self.dump, f)
+        else:
+            self.dump = self.gather_and_merge_predictions()
+            dumped_file = Path(self.dump_dir) / f"coco_predictions_{self.iou_type}.json"
+            if is_main_process():
+                logging.info(
+                    f"Prediction Dumper: Dumping merged predictions to {dumped_file}"
+                )
+                with g_pathmgr.open(str(dumped_file), "w") as f:
+                    json.dump(self.dump, f)
+
+        self.reset()
+        return dumped_file
+
+    def gather_and_merge_predictions(self):
+        """
+        Gather predictions from all processes and merge them, keeping top predictions per image.
+
+        This method collects predictions from all processes, then keeps only the top maxdets
+        predictions per image based on score. It also deduplicates predictions by (image_id, category_id).
+
+        Returns:
+            List of merged prediction dictionaries.
+        """
+        logging.info("Prediction Dumper: Gathering predictions from all processes")
+        gc.collect()
+
+        if self.gather_pred_via_filesys:
+            dump = gather_to_rank_0_via_filesys(self.dump)
+        else:
+            dump = all_gather(self.dump, force_cpu=True)
+
+        # Combine predictions, keeping only top maxdets per image
+        preds_by_image = defaultdict(list)
+        seen_img_cat = set()
+
+        for cur_dump in dump:
+            cur_seen_img_cat = set()
+            for p in cur_dump:
+                image_id = p["image_id"]
+                cat_id = p["category_id"]
+
+                # Skip if we've already seen this image/category pair in a previous dump
+                if (image_id, cat_id) in seen_img_cat:
+                    continue
+
+                cur_seen_img_cat.add((image_id, cat_id))
+
+                # Use a min-heap to keep top predictions
+                if len(preds_by_image[image_id]) < self.maxdets:
+                    heapq.heappush(preds_by_image[image_id], HeapElement(p))
+                else:
+                    heapq.heappushpop(preds_by_image[image_id], HeapElement(p))
+
+            seen_img_cat.update(cur_seen_img_cat)
+
+        # Flatten the heap elements back to a list
+        merged_dump = sum(
+            [[h.val for h in cur_preds] for cur_preds in preds_by_image.values()], []
+        )
+
+        return merged_dump
+
+    def compute_synced(self):
+        """
+        Synchronize predictions across processes and compute summary.
+
+        Returns:
+            Summary dictionary from summarize().
+        """
+        dumped_file = self.synchronize_between_processes()
+        if not is_main_process():
+            return {"": 0.0}
+
+        meters = {}
+        if self.pred_file_evaluators is not None:
+            for evaluator in self.pred_file_evaluators:
+                results = evaluator.evaluate(dumped_file)
+                meters.update(results)
+
+        if len(meters) == 0:
+            meters = {"": 0.0}
+        return meters
+
+    def compute(self):
+        """
+        Compute without synchronization.
+
+        Returns:
+            Empty metric dictionary.
+        """
+        return {"": 0.0}
+
+    def reset(self):
+        """Reset internal state for a new evaluation round."""
+        self.dump = []
+
+    def prepare(self, predictions, iou_type):
+        """
+        Route predictions to the appropriate preparation method based on iou_type.
+
+        Args:
+            predictions: Dictionary mapping image IDs to prediction dictionaries.
+            iou_type: Type of evaluation ("bbox", "segm").
+
+        Returns:
+            List of COCO-format prediction dictionaries.
+        """
+        if iou_type == "bbox":
+            return self.prepare_for_coco_detection(predictions)
+        elif iou_type == "segm":
+            return self.prepare_for_coco_segmentation(predictions)
+        else:
+            raise ValueError(f"Unknown iou type: {iou_type}")
+
+    def prepare_for_coco_detection(self, predictions):
+        """
+        Convert predictions to COCO detection format.
+
+        Args:
+            predictions: Dictionary mapping image IDs to prediction dictionaries
+                containing "boxes", "scores", and "labels".
+
+        Returns:
+            List of COCO-format detection dictionaries.
+        """
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+
+            boxes = prediction["boxes"]
+            boxes = convert_to_xywh(boxes).tolist()
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+
+            coco_results.extend(
+                [
+                    {
+                        "image_id": original_id,
+                        "category_id": labels[k],
+                        "bbox": box,
+                        "score": scores[k],
+                    }
+                    for k, box in enumerate(boxes)
+                ]
+            )
+        return coco_results
+
+    @torch.no_grad()
+    def prepare_for_coco_segmentation(self, predictions):
+        """
+        Convert predictions to COCO segmentation format.
+
+        Args:
+            predictions: Dictionary mapping image IDs to prediction dictionaries
+                containing "masks" or "masks_rle", "scores", and "labels".
+                Optionally includes "boundaries" and "dilated_boundaries".
+
+        Returns:
+            List of COCO-format segmentation dictionaries with RLE-encoded masks.
+        """
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+
+            boxes = None
+            if "boxes" in prediction:
+                boxes = prediction["boxes"]
+                boxes = convert_to_xywh(boxes).tolist()
+                assert len(boxes) == len(scores)
+
+            if "masks_rle" in prediction:
+                rles = prediction["masks_rle"]
+                areas = []
+                for rle in rles:
+                    cur_area = mask_utils.area(rle)
+                    h, w = rle["size"]
+                    areas.append(cur_area / (h * w))
+            else:
+                masks = prediction["masks"]
+                masks = masks > 0.5
+                h, w = masks.shape[-2:]
+
+                areas = masks.flatten(1).sum(1) / (h * w)
+                areas = areas.tolist()
+
+                rles = rle_encode(masks.squeeze(1))
+
+                # Memory cleanup
+                del masks
+                del prediction["masks"]
+
+            assert len(areas) == len(rles) == len(scores)
+
+            for k, rle in enumerate(rles):
+                payload = {
+                    "image_id": original_id,
+                    "category_id": labels[k],
+                    "segmentation": rle,
+                    "score": scores[k],
+                    "area": areas[k],
+                }
+                if boxes is not None:
+                    payload["bbox"] = boxes[k]
+
+                coco_results.append(payload)
+
+        return coco_results

sam3/eval/conversion_util.py

@@ -0,0 +1,211 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+import json
+import os
+from collections import defaultdict
+
+from tqdm import tqdm
+
+
+def convert_ytbvis_to_cocovid_gt(ann_json, save_path=None):
+    """Convert YouTube VIS dataset to COCO-style video instance segmentation format.
+
+    Args:
+        ann_json (str): Path to YouTube VIS annotation JSON file
+        save_path (str): path to save converted COCO-style JSON
+    """
+    # Initialize COCO structure
+    VIS = {
+        "info": {},
+        "images": [],
+        "videos": [],
+        "tracks": [],
+        "annotations": [],
+        "categories": [],
+        "licenses": [],
+    }
+
+    # Load original annotations
+    official_anns = json.load(open(ann_json))
+    VIS["categories"] = official_anns["categories"]  # Direct copy categories
+
+    # Initialize counters
+    records = dict(img_id=1, ann_id=1)
+
+    # Create video-to-annotations mapping
+    vid_to_anns = defaultdict(list)
+    for ann in official_anns["annotations"]:
+        vid_to_anns[ann["video_id"]].append(ann)
+
+    # Create tracks directly
+    VIS["tracks"] = [
+        {
+            "id": ann["id"],
+            "category_id": ann["category_id"],
+            "video_id": ann["video_id"],
+        }
+        for ann in official_anns["annotations"]
+    ]
+
+    # Process videos
+    for video_info in tqdm(official_anns["videos"]):
+        # Create video entry
+        video = {
+            "id": video_info["id"],
+            "name": os.path.dirname(video_info["file_names"][0]),
+            "width": video_info["width"],
+            "height": video_info["height"],
+            "length": video_info["length"],
+            "neg_category_ids": [],
+            "not_exhaustive_category_ids": [],
+        }
+        VIS["videos"].append(video)
+
+        # Process frames
+        num_frames = len(video_info["file_names"])
+        for frame_idx in range(num_frames):
+            # Create image entry
+            image = {
+                "id": records["img_id"],
+                "video_id": video_info["id"],
+                "file_name": video_info["file_names"][frame_idx],
+                "width": video_info["width"],
+                "height": video_info["height"],
+                "frame_index": frame_idx,
+                "frame_id": frame_idx,
+            }
+            VIS["images"].append(image)
+
+            # Process annotations for this frame
+            if video_info["id"] in vid_to_anns:
+                for ann in vid_to_anns[video_info["id"]]:
+                    bbox = ann["bboxes"][frame_idx]
+                    if bbox is None:
+                        continue
+
+                    # Create annotation entry
+                    annotation = {
+                        "id": records["ann_id"],
+                        "video_id": video_info["id"],
+                        "image_id": records["img_id"],
+                        "track_id": ann["id"],
+                        "category_id": ann["category_id"],
+                        "bbox": bbox,
+                        "area": ann["areas"][frame_idx],
+                        "segmentation": ann["segmentations"][frame_idx],
+                        "iscrowd": ann["iscrowd"],
+                    }
+                    VIS["annotations"].append(annotation)
+                    records["ann_id"] += 1
+
+            records["img_id"] += 1
+
+    # Print summary
+    print(f"Converted {len(VIS['videos'])} videos")
+    print(f"Converted {len(VIS['images'])} images")
+    print(f"Created {len(VIS['tracks'])} tracks")
+    print(f"Created {len(VIS['annotations'])} annotations")
+
+    if save_path is None:
+        return VIS
+
+    # Save output
+    save_dir = os.path.dirname(save_path)
+    os.makedirs(save_dir, exist_ok=True)
+    json.dump(VIS, open(save_path, "w"))
+
+    return VIS
+
+
+def convert_ytbvis_to_cocovid_pred(
+    youtubevis_pred_path: str, converted_dataset_path: str, output_path: str
+) -> None:
+    """
+    Convert YouTubeVIS predictions to COCO format with video_id preservation
+
+    Args:
+        youtubevis_pred_path: Path to YouTubeVIS prediction JSON
+        converted_dataset_path: Path to converted COCO dataset JSON
+        output_path: Path to save COCO format predictions
+    """
+
+    # Load YouTubeVIS predictions
+    with open(youtubevis_pred_path) as f:
+        ytv_predictions = json.load(f)
+
+    # Load converted dataset for image ID mapping
+    with open(converted_dataset_path) as f:
+        coco_dataset = json.load(f)
+
+    # Create (video_id, frame_idx) -> image_id mapping
+    image_id_map = {
+        (img["video_id"], img["frame_index"]): img["id"]
+        for img in coco_dataset["images"]
+    }
+
+    coco_annotations = []
+    track_id_counter = 1  # Unique track ID generator
+
+    for pred in tqdm(ytv_predictions):
+        video_id = pred["video_id"]
+        category_id = pred["category_id"]
+        bboxes = pred["bboxes"]
+        segmentations = pred.get("segmentations", [])  # Get segmentations if available
+        areas = pred.get("areas", [])  # Get areas if available
+        score = pred["score"]
+
+        # Assign unique track ID for this prediction
+        track_id = track_id_counter
+        track_id_counter += 1
+
+        # Ensure segmentations and areas have the same length as bboxes
+        if len(segmentations) == 0:
+            segmentations = [None] * len(bboxes)
+        if len(areas) == 0:
+            areas = [None] * len(bboxes)
+
+        for frame_idx, (bbox, segmentation, area_from_pred) in enumerate(
+            zip(bboxes, segmentations, areas)
+        ):
+            # Skip frames with missing objects (None or zero bbox)
+            if bbox is None or all(x == 0 for x in bbox):
+                continue
+
+            # Get corresponding image ID from mapping
+            image_id = image_id_map.get((video_id, frame_idx))
+            if image_id is None:
+                raise RuntimeError(
+                    f"prediction {video_id=}, {frame_idx=} does not match any images in the converted COCO format"
+                )
+
+            # Extract bbox coordinates
+            x, y, w, h = bbox
+
+            # Calculate area - use area from prediction if available, otherwise from bbox
+            if area_from_pred is not None and area_from_pred > 0:
+                area = area_from_pred
+            else:
+                area = w * h
+
+            # Create COCO annotation with video_id
+            coco_annotation = {
+                "image_id": int(image_id),
+                "video_id": video_id,  # Added video_id field
+                "track_id": track_id,
+                "category_id": category_id,
+                "bbox": [float(x), float(y), float(w), float(h)],
+                "area": float(area),
+                "iscrowd": 0,
+                "score": float(score),
+            }
+
+            # Add segmentation if available
+            if segmentation is not None:
+                coco_annotation["segmentation"] = segmentation
+
+            coco_annotations.append(coco_annotation)
+
+    # Save output
+    with open(output_path, "w") as f:
+        json.dump(coco_annotations, f)
+
+    print(f"Converted {len(coco_annotations)} predictions to COCO format with video_id")

sam3/eval/demo_eval.py

@@ -0,0 +1,658 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+"""
+This evaluator is based upon COCO evaluation, but evaluates the model in a "demo" setting.
+This means that the model's predictions are thresholded and evaluated as "hard" predictions.
+"""
+
+import logging
+from typing import Optional
+
+import numpy as np
+import pycocotools.mask as maskUtils
+from pycocotools.cocoeval import COCOeval
+
+from sam3.eval.coco_eval import CocoEvaluator
+from sam3.train.masks_ops import compute_F_measure
+from sam3.train.utils.distributed import is_main_process
+
+from scipy.optimize import linear_sum_assignment
+
+
+class DemoEval(COCOeval):
+    """
+    This evaluator is based upon COCO evaluation, but evaluates the model in a "demo" setting.
+    This means that the model's predictions are thresholded and evaluated as "hard" predictions.
+    """
+
+    def __init__(
+        self,
+        coco_gt=None,
+        coco_dt=None,
+        iouType="bbox",
+        threshold=0.5,
+        compute_JnF=False,
+    ):
+        """
+        Args:
+            coco_gt (COCO): ground truth COCO API
+            coco_dt (COCO): detections COCO API
+            iou_type (str): type of IoU to evaluate
+            threshold (float): threshold for predictions
+        """
+        super().__init__(coco_gt, coco_dt, iouType)
+        self.threshold = threshold
+
+        self.params.useCats = False
+        self.params.areaRng = [[0**2, 1e5**2]]
+        self.params.areaRngLbl = ["all"]
+        self.params.maxDets = [100000]
+        self.compute_JnF = compute_JnF
+
+    def computeIoU(self, imgId, catId):
+        # Same as the original COCOeval.computeIoU, but without sorting
+        p = self.params
+        if p.useCats:
+            gt = self._gts[imgId, catId]
+            dt = self._dts[imgId, catId]
+        else:
+            gt = [_ for cId in p.catIds for _ in self._gts[imgId, cId]]
+            dt = [_ for cId in p.catIds for _ in self._dts[imgId, cId]]
+        if len(gt) == 0 and len(dt) == 0:
+            return []
+
+        if p.iouType == "segm":
+            g = [g["segmentation"] for g in gt]
+            d = [d["segmentation"] for d in dt]
+        elif p.iouType == "bbox":
+            g = [g["bbox"] for g in gt]
+            d = [d["bbox"] for d in dt]
+        else:
+            raise Exception("unknown iouType for iou computation")
+
+        # compute iou between each dt and gt region
+        iscrowd = [int(o["iscrowd"]) for o in gt]
+        ious = maskUtils.iou(d, g, iscrowd)
+        return ious
+
+    def evaluateImg(self, imgId, catId, aRng, maxDet):
+        """
+        perform evaluation for single category and image
+        :return: dict (single image results)
+        """
+        p = self.params
+        assert not p.useCats, "This evaluator does not support per-category evaluation."
+        assert catId == -1
+        all_gts = [_ for cId in p.catIds for _ in self._gts[imgId, cId]]
+        keep_gt = np.array([not g["ignore"] for g in all_gts], dtype=bool)
+        gt = [g for g in all_gts if not g["ignore"]]
+        all_dts = [_ for cId in p.catIds for _ in self._dts[imgId, cId]]
+        keep_dt = np.array([d["score"] >= self.threshold for d in all_dts], dtype=bool)
+        dt = [d for d in all_dts if d["score"] >= self.threshold]
+        if len(gt) == 0 and len(dt) == 0:
+            # This is a "true negative" case, where there are no GTs and no predictions
+            # The box-level metrics are ill-defined, so we don't add them to this dict
+            return {
+                "image_id": imgId,
+                "IL_TP": 0,
+                "IL_TN": 1,
+                "IL_FP": 0,
+                "IL_FN": 0,
+                "IL_perfect_neg": np.ones((len(p.iouThrs),), dtype=np.int64),
+                "num_dt": len(dt),
+            }
+
+        if len(gt) > 0 and len(dt) == 0:
+            # This is a "false negative" case, where there are GTs but no predictions
+            return {
+                "image_id": imgId,
+                "IL_TP": 0,
+                "IL_TN": 0,
+                "IL_FP": 0,
+                "IL_FN": 1,
+                "TPs": np.zeros((len(p.iouThrs),), dtype=np.int64),
+                "FPs": np.zeros((len(p.iouThrs),), dtype=np.int64),
+                "FNs": np.ones((len(p.iouThrs),), dtype=np.int64) * len(gt),
+                "local_F1s": np.zeros((len(p.iouThrs),), dtype=np.int64),
+                "local_positive_F1s": np.zeros((len(p.iouThrs),), dtype=np.int64),
+                "IL_perfect_pos": np.zeros((len(p.iouThrs),), dtype=np.int64),
+                "num_dt": len(dt),
+            }
+
+        # Load pre-computed ious
+        ious = self.ious[(imgId, catId)]
+
+        # compute matching
+        if len(ious) == 0:
+            ious = np.zeros((len(dt), len(gt)))
+        else:
+            ious = ious[keep_dt, :][:, keep_gt]
+        assert ious.shape == (len(dt), len(gt))
+
+        matched_dt, matched_gt = linear_sum_assignment(-ious)
+
+        match_scores = ious[matched_dt, matched_gt]
+
+        if self.compute_JnF and len(match_scores) > 0:
+            j_score = match_scores.mean()
+            f_measure = 0
+            for dt_id, gt_id in zip(matched_dt, matched_gt):
+                f_measure += compute_F_measure(
+                    gt_boundary_rle=gt[gt_id]["boundary"],
+                    gt_dilated_boundary_rle=gt[gt_id]["dilated_boundary"],
+                    dt_boundary_rle=dt[dt_id]["boundary"],
+                    dt_dilated_boundary_rle=dt[dt_id]["dilated_boundary"],
+                )
+            f_measure /= len(match_scores) + 1e-9
+            JnF = (j_score + f_measure) * 0.5
+        else:
+            j_score = f_measure = JnF = -1
+
+        TPs, FPs, FNs = [], [], []
+        IL_perfect = []
+        for thresh in p.iouThrs:
+            TP = (match_scores >= thresh).sum()
+            FP = len(dt) - TP
+            FN = len(gt) - TP
+            assert (
+                FP >= 0 and FN >= 0
+            ), f"FP: {FP}, FN: {FN}, TP: {TP}, match_scores: {match_scores}, len(dt): {len(dt)}, len(gt): {len(gt)}, ious: {ious}"
+            TPs.append(TP)
+            FPs.append(FP)
+            FNs.append(FN)
+
+            if FP == FN and FP == 0:
+                IL_perfect.append(1)
+            else:
+                IL_perfect.append(0)
+
+        TPs = np.array(TPs, dtype=np.int64)
+        FPs = np.array(FPs, dtype=np.int64)
+        FNs = np.array(FNs, dtype=np.int64)
+        IL_perfect = np.array(IL_perfect, dtype=np.int64)
+
+        # compute precision recall and F1
+        precision = TPs / (TPs + FPs + 1e-4)
+        assert np.all(precision <= 1)
+        recall = TPs / (TPs + FNs + 1e-4)
+        assert np.all(recall <= 1)
+        F1 = 2 * precision * recall / (precision + recall + 1e-4)
+
+        result = {
+            "image_id": imgId,
+            "TPs": TPs,
+            "FPs": FPs,
+            "FNs": FNs,
+            "local_F1s": F1,
+            "IL_TP": (len(gt) > 0) and (len(dt) > 0),
+            "IL_FP": (len(gt) == 0) and (len(dt) > 0),
+            "IL_TN": (len(gt) == 0) and (len(dt) == 0),
+            "IL_FN": (len(gt) > 0) and (len(dt) == 0),
+            ("IL_perfect_pos" if len(gt) > 0 else "IL_perfect_neg"): IL_perfect,
+            "F": f_measure,
+            "J": j_score,
+            "J&F": JnF,
+            "num_dt": len(dt),
+        }
+        if len(gt) > 0 and len(dt) > 0:
+            result["local_positive_F1s"] = F1
+        return result
+
+    def accumulate(self, p=None):
+        """
+        Accumulate per image evaluation results and store the result in self.eval
+        :param p: input params for evaluation
+        :return: None
+        """
+        if not self.evalImgs:
+            print("Please run evaluate() first")
+        # allows input customized parameters
+        if p is None:
+            p = self.params
+
+        setImgIds = set(p.imgIds)
+
+        # TPs, FPs, FNs
+        TPs = np.zeros((len(p.iouThrs),), dtype=np.int64)
+        FPs = np.zeros((len(p.iouThrs),), dtype=np.int64)
+        pmFPs = np.zeros((len(p.iouThrs),), dtype=np.int64)
+        FNs = np.zeros((len(p.iouThrs),), dtype=np.int64)
+        local_F1s = np.zeros((len(p.iouThrs),), dtype=np.float64)
+
+        # Image level metrics
+        IL_TPs = 0
+        IL_FPs = 0
+        IL_TNs = 0
+        IL_FNs = 0
+        IL_perfects_neg = np.zeros((len(p.iouThrs),), dtype=np.int64)
+        IL_perfects_pos = np.zeros((len(p.iouThrs),), dtype=np.int64)
+
+        # JnF metric
+        total_J = 0
+        total_F = 0
+        total_JnF = 0
+
+        valid_img_count = 0
+        total_pos_count = 0
+        total_neg_count = 0
+        valid_J_count = 0
+        valid_F1_count = 0
+        valid_F1_count_w0dt = 0
+        for res in self.evalImgs:
+            if res["image_id"] not in setImgIds:
+                continue
+            IL_TPs += res["IL_TP"]
+            IL_FPs += res["IL_FP"]
+            IL_TNs += res["IL_TN"]
+            IL_FNs += res["IL_FN"]
+            if "IL_perfect_neg" in res:
+                IL_perfects_neg += res["IL_perfect_neg"]
+                total_neg_count += 1
+            else:
+                assert "IL_perfect_pos" in res
+                IL_perfects_pos += res["IL_perfect_pos"]
+                total_pos_count += 1
+
+            if "TPs" not in res:
+                continue
+
+            TPs += res["TPs"]
+            FPs += res["FPs"]
+            FNs += res["FNs"]
+            valid_img_count += 1
+
+            if "local_positive_F1s" in res:
+                local_F1s += res["local_positive_F1s"]
+                pmFPs += res["FPs"]
+                valid_F1_count_w0dt += 1
+                if res["num_dt"] > 0:
+                    valid_F1_count += 1
+
+            if "J" in res and res["J"] > -1e-9:
+                total_J += res["J"]
+                total_F += res["F"]
+                total_JnF += res["J&F"]
+                valid_J_count += 1
+
+        # compute precision recall and F1
+        precision = TPs / (TPs + FPs + 1e-4)
+        positive_micro_precision = TPs / (TPs + pmFPs + 1e-4)
+        assert np.all(precision <= 1)
+        recall = TPs / (TPs + FNs + 1e-4)
+        assert np.all(recall <= 1)
+        F1 = 2 * precision * recall / (precision + recall + 1e-4)
+        positive_micro_F1 = (
+            2
+            * positive_micro_precision
+            * recall
+            / (positive_micro_precision + recall + 1e-4)
+        )
+
+        IL_rec = IL_TPs / (IL_TPs + IL_FNs + 1e-6)
+        IL_prec = IL_TPs / (IL_TPs + IL_FPs + 1e-6)
+        IL_F1 = 2 * IL_prec * IL_rec / (IL_prec + IL_rec + 1e-6)
+        IL_FPR = IL_FPs / (IL_FPs + IL_TNs + 1e-6)
+        IL_MCC = float(IL_TPs * IL_TNs - IL_FPs * IL_FNs) / (
+            (
+                float(IL_TPs + IL_FPs)
+                * float(IL_TPs + IL_FNs)
+                * float(IL_TNs + IL_FPs)
+                * float(IL_TNs + IL_FNs)
+            )
+            ** 0.5
+            + 1e-6
+        )
+        IL_perfect_pos = IL_perfects_pos / (total_pos_count + 1e-9)
+        IL_perfect_neg = IL_perfects_neg / (total_neg_count + 1e-9)
+
+        total_J = total_J / (valid_J_count + 1e-9)
+        total_F = total_F / (valid_J_count + 1e-9)
+        total_JnF = total_JnF / (valid_J_count + 1e-9)
+
+        self.eval = {
+            "params": p,
+            "TPs": TPs,
+            "FPs": FPs,
+            "positive_micro_FPs": pmFPs,
+            "FNs": FNs,
+            "precision": precision,
+            "positive_micro_precision": positive_micro_precision,
+            "recall": recall,
+            "F1": F1,
+            "positive_micro_F1": positive_micro_F1,
+            "positive_macro_F1": local_F1s / valid_F1_count,
+            "positive_w0dt_macro_F1": local_F1s / valid_F1_count_w0dt,
+            "IL_recall": IL_rec,
+            "IL_precision": IL_prec,
+            "IL_F1": IL_F1,
+            "IL_FPR": IL_FPR,
+            "IL_MCC": IL_MCC,
+            "IL_perfect_pos": IL_perfect_pos,
+            "IL_perfect_neg": IL_perfect_neg,
+            "J": total_J,
+            "F": total_F,
+            "J&F": total_JnF,
+        }
+        self.eval["CGF1"] = self.eval["positive_macro_F1"] * self.eval["IL_MCC"]
+        self.eval["CGF1_w0dt"] = (
+            self.eval["positive_w0dt_macro_F1"] * self.eval["IL_MCC"]
+        )
+        self.eval["CGF1_micro"] = self.eval["positive_micro_F1"] * self.eval["IL_MCC"]
+
+    def summarize(self):
+        """
+        Compute and display summary metrics for evaluation results.
+        Note this functin can *only* be applied on the default parameter setting
+        """
+        if not self.eval:
+            raise Exception("Please run accumulate() first")
+
+        def _summarize(iouThr=None, metric=""):
+            p = self.params
+            iStr = " {:<18} @[ IoU={:<9}] = {:0.3f}"
+            titleStr = "Average " + metric
+            iouStr = (
+                "{:0.2f}:{:0.2f}".format(p.iouThrs[0], p.iouThrs[-1])
+                if iouThr is None
+                else "{:0.2f}".format(iouThr)
+            )
+
+            s = self.eval[metric]
+            # IoU
+            if iouThr is not None:
+                t = np.where(iouThr == p.iouThrs)[0]
+                s = s[t]
+
+            if len(s[s > -1]) == 0:
+                mean_s = -1
+            else:
+                mean_s = np.mean(s[s > -1])
+            print(iStr.format(titleStr, iouStr, mean_s))
+            return mean_s
+
+        def _summarize_single(metric=""):
+            titleStr = "Average " + metric
+            iStr = " {:<35} = {:0.3f}"
+            s = self.eval[metric]
+            print(iStr.format(titleStr, s))
+            return s
+
+        def _summarizeDets():
+            # note: the index of these metrics are also used in video Demo F1 evaluation
+            # when adding new metrics, please update the index in video Demo F1 evaluation
+            # in "evaluate" method of the "VideoDemoF1Evaluator" class
+            stats = np.zeros((len(DEMO_METRICS),))
+            stats[0] = _summarize(metric="CGF1")
+            stats[1] = _summarize(metric="precision")
+            stats[2] = _summarize(metric="recall")
+            stats[3] = _summarize(metric="F1")
+            stats[4] = _summarize(metric="positive_macro_F1")
+            stats[5] = _summarize_single(metric="IL_precision")
+            stats[6] = _summarize_single(metric="IL_recall")
+            stats[7] = _summarize_single(metric="IL_F1")
+            stats[8] = _summarize_single(metric="IL_FPR")
+            stats[9] = _summarize_single(metric="IL_MCC")
+            stats[10] = _summarize(metric="IL_perfect_pos")
+            stats[11] = _summarize(metric="IL_perfect_neg")
+            stats[12] = _summarize(iouThr=0.5, metric="CGF1")
+            stats[13] = _summarize(iouThr=0.5, metric="precision")
+            stats[14] = _summarize(iouThr=0.5, metric="recall")
+            stats[15] = _summarize(iouThr=0.5, metric="F1")
+            stats[16] = _summarize(iouThr=0.5, metric="positive_macro_F1")
+            stats[17] = _summarize(iouThr=0.5, metric="IL_perfect_pos")
+            stats[18] = _summarize(iouThr=0.5, metric="IL_perfect_neg")
+            stats[19] = _summarize(iouThr=0.75, metric="CGF1")
+            stats[20] = _summarize(iouThr=0.75, metric="precision")
+            stats[21] = _summarize(iouThr=0.75, metric="recall")
+            stats[22] = _summarize(iouThr=0.75, metric="F1")
+            stats[23] = _summarize(iouThr=0.75, metric="positive_macro_F1")
+            stats[24] = _summarize(iouThr=0.75, metric="IL_perfect_pos")
+            stats[25] = _summarize(iouThr=0.75, metric="IL_perfect_neg")
+            stats[26] = _summarize_single(metric="J")
+            stats[27] = _summarize_single(metric="F")
+            stats[28] = _summarize_single(metric="J&F")
+            stats[29] = _summarize(metric="CGF1_micro")
+            stats[30] = _summarize(metric="positive_micro_precision")
+            stats[31] = _summarize(metric="positive_micro_F1")
+            stats[32] = _summarize(iouThr=0.5, metric="CGF1_micro")
+            stats[33] = _summarize(iouThr=0.5, metric="positive_micro_precision")
+            stats[34] = _summarize(iouThr=0.5, metric="positive_micro_F1")
+            stats[35] = _summarize(iouThr=0.75, metric="CGF1_micro")
+            stats[36] = _summarize(iouThr=0.75, metric="positive_micro_precision")
+            stats[37] = _summarize(iouThr=0.75, metric="positive_micro_F1")
+            stats[38] = _summarize(metric="CGF1_w0dt")
+            stats[39] = _summarize(metric="positive_w0dt_macro_F1")
+            stats[40] = _summarize(iouThr=0.5, metric="CGF1_w0dt")
+            stats[41] = _summarize(iouThr=0.5, metric="positive_w0dt_macro_F1")
+            stats[42] = _summarize(iouThr=0.75, metric="CGF1_w0dt")
+            stats[43] = _summarize(iouThr=0.75, metric="positive_w0dt_macro_F1")
+            return stats
+
+        summarize = _summarizeDets
+        self.stats = summarize()
+
+
+DEMO_METRICS = [
+    "CGF1",
+    "Precision",
+    "Recall",
+    "F1",
+    "Macro_F1",
+    "IL_Precision",
+    "IL_Recall",
+    "IL_F1",
+    "IL_FPR",
+    "IL_MCC",
+    "IL_perfect_pos",
+    "IL_perfect_neg",
+    "CGF1@0.5",
+    "Precision@0.5",
+    "Recall@0.5",
+    "F1@0.5",
+    "Macro_F1@0.5",
+    "IL_perfect_pos@0.5",
+    "IL_perfect_neg@0.5",
+    "CGF1@0.75",
+    "Precision@0.75",
+    "Recall@0.75",
+    "F1@0.75",
+    "Macro_F1@0.75",
+    "IL_perfect_pos@0.75",
+    "IL_perfect_neg@0.75",
+    "J",
+    "F",
+    "J&F",
+    "CGF1_micro",
+    "positive_micro_Precision",
+    "positive_micro_F1",
+    "CGF1_micro@0.5",
+    "positive_micro_Precision@0.5",
+    "positive_micro_F1@0.5",
+    "CGF1_micro@0.75",
+    "positive_micro_Precision@0.75",
+    "positive_micro_F1@0.75",
+    "CGF1_w0dt",
+    "positive_w0dt_macro_F1",
+    "CGF1_w0dt@0.5",
+    "positive_w0dt_macro_F1@0.5",
+    "CGF1_w0dt@0.75",
+    "positive_w0dt_macro_F1@0.75",
+]
+
+
+class DemoEvaluator(CocoEvaluator):
+    def __init__(
+        self,
+        coco_gt,
+        iou_types,
+        dump_dir: Optional[str],
+        postprocessor,
+        threshold=0.5,
+        average_by_rarity=False,
+        gather_pred_via_filesys=False,
+        exhaustive_only=False,
+        all_exhaustive_only=True,
+        compute_JnF=False,
+        metrics_dump_dir: Optional[str] = None,
+    ):
+        self.iou_types = iou_types
+        self.threshold = threshold
+        super().__init__(
+            coco_gt=coco_gt,
+            iou_types=iou_types,
+            useCats=False,
+            dump_dir=dump_dir,
+            postprocessor=postprocessor,
+            # average_by_rarity=average_by_rarity,
+            gather_pred_via_filesys=gather_pred_via_filesys,
+            exhaustive_only=exhaustive_only,
+            all_exhaustive_only=all_exhaustive_only,
+            metrics_dump_dir=metrics_dump_dir,
+        )
+
+        self.use_self_evaluate = True
+        self.compute_JnF = compute_JnF
+
+    def _lazy_init(self):
+        if self.initialized:
+            return
+        super()._lazy_init()
+        self.use_self_evaluate = True
+        self.reset()
+
+    def select_best_scoring(self, scorings):
+        # This function is used for "oracle" type evaluation.
+        # It accepts the evaluation results with respect to several ground truths, and picks the best
+        if len(scorings) == 1:
+            return scorings[0]
+
+        assert (
+            scorings[0].ndim == 3
+        ), f"Expecting results in [numCats, numAreas, numImgs] format, got {scorings[0].shape}"
+        assert (
+            scorings[0].shape[0] == 1
+        ), f"Expecting a single category, got {scorings[0].shape[0]}"
+
+        for scoring in scorings:
+            assert (
+                scoring.shape == scorings[0].shape
+            ), f"Shape mismatch: {scoring.shape}, {scorings[0].shape}"
+
+        selected_imgs = []
+        for img_id in range(scorings[0].shape[-1]):
+            best = scorings[0][:, :, img_id]
+
+            for scoring in scorings[1:]:
+                current = scoring[:, :, img_id]
+                if "local_F1s" in best[0, 0] and "local_F1s" in current[0, 0]:
+                    # we were able to compute a F1 score for this particular image in both evaluations
+                    # best["local_F1s"] contains the results at various IoU thresholds. We simply take the average for comparision
+                    best_score = best[0, 0]["local_F1s"].mean()
+                    current_score = current[0, 0]["local_F1s"].mean()
+                    if current_score > best_score:
+                        best = current
+
+                else:
+                    # If we're here, it means that in that in some evaluation we were not able to get a valid local F1
+                    # This happens when both the predictions and targets are empty. In that case, we can assume it's a perfect prediction
+                    if "local_F1s" not in current[0, 0]:
+                        best = current
+            selected_imgs.append(best)
+        result = np.stack(selected_imgs, axis=-1)
+        assert result.shape == scorings[0].shape
+        return result
+
+    def summarize(self):
+        self._lazy_init()
+        logging.info("Demo evaluator: Summarizing")
+        if not is_main_process():
+            return {}
+        outs = {}
+        prefix = "oracle_" if len(self.coco_evals) > 1 else ""
+        # if self.rarity_buckets is None:
+        self.accumulate(self.eval_img_ids)
+        for iou_type, coco_eval in self.coco_evals[0].items():
+            print("Demo metric, IoU type={}".format(iou_type))
+            coco_eval.summarize()
+
+        if "bbox" in self.coco_evals[0]:
+            for i, value in enumerate(self.coco_evals[0]["bbox"].stats):
+                outs[f"coco_eval_bbox_{prefix}{DEMO_METRICS[i]}"] = value
+        if "segm" in self.coco_evals[0]:
+            for i, value in enumerate(self.coco_evals[0]["segm"].stats):
+                outs[f"coco_eval_masks_{prefix}{DEMO_METRICS[i]}"] = value
+        # else:
+        #     total_stats = {}
+        #     for bucket, img_list in self.rarity_buckets.items():
+        #         self.accumulate(imgIds=img_list)
+        #         bucket_name = RARITY_BUCKETS[bucket]
+        #         for iou_type, coco_eval in self.coco_evals[0].items():
+        #             print(
+        #                 "Demo metric, IoU type={}, Rarity bucket={}".format(
+        #                     iou_type, bucket_name
+        #                 )
+        #             )
+        #             coco_eval.summarize()
+
+        #         if "bbox" in self.coco_evals[0]:
+        #             if "bbox" not in total_stats:
+        #                 total_stats["bbox"] = np.zeros_like(
+        #                     self.coco_evals[0]["bbox"].stats
+        #                 )
+        #             total_stats["bbox"] += self.coco_evals[0]["bbox"].stats
+        #             for i, value in enumerate(self.coco_evals[0]["bbox"].stats):
+        #                 outs[
+        #                     f"coco_eval_bbox_{bucket_name}_{prefix}{DEMO_METRICS[i]}"
+        #                 ] = value
+        #         if "segm" in self.coco_evals[0]:
+        #             if "segm" not in total_stats:
+        #                 total_stats["segm"] = np.zeros_like(
+        #                     self.coco_evals[0]["segm"].stats
+        #                 )
+        #             total_stats["segm"] += self.coco_evals[0]["segm"].stats
+        #             for i, value in enumerate(self.coco_evals[0]["segm"].stats):
+        #                 outs[
+        #                     f"coco_eval_masks_{bucket_name}_{prefix}{DEMO_METRICS[i]}"
+        #                 ] = value
+
+        #     if "bbox" in total_stats:
+        #         total_stats["bbox"] /= len(self.rarity_buckets)
+        #         for i, value in enumerate(total_stats["bbox"]):
+        #             outs[f"coco_eval_bbox_{prefix}{DEMO_METRICS[i]}"] = value
+        #     if "segm" in total_stats:
+        #         total_stats["segm"] /= len(self.rarity_buckets)
+        #         for i, value in enumerate(total_stats["segm"]):
+        #             outs[f"coco_eval_masks_{prefix}{DEMO_METRICS[i]}"] = value
+
+        return outs
+
+    def accumulate(self, imgIds=None):
+        self._lazy_init()
+        logging.info(
+            f"demo evaluator: Accumulating on {len(imgIds) if imgIds is not None else 'all'} images"
+        )
+        if not is_main_process():
+            return
+
+        if imgIds is not None:
+            for coco_eval in self.coco_evals[0].values():
+                coco_eval.params.imgIds = list(imgIds)
+
+        for coco_eval in self.coco_evals[0].values():
+            coco_eval.accumulate()
+
+    def reset(self):
+        self.coco_evals = [{} for _ in range(len(self.coco_gts))]
+        for i, coco_gt in enumerate(self.coco_gts):
+            for iou_type in self.iou_types:
+                self.coco_evals[i][iou_type] = DemoEval(
+                    coco_gt=coco_gt,
+                    iouType=iou_type,
+                    threshold=self.threshold,
+                    compute_JnF=self.compute_JnF,
+                )
+                self.coco_evals[i][iou_type].useCats = False
+        self.img_ids = []
+        self.eval_imgs = {k: [] for k in self.iou_types}
+        if self.dump is not None:
+            self.dump = []

sam3/eval/hota_eval_toolkit/__init__.py

@@ -0,0 +1 @@
+# flake8: noqa

sam3/eval/hota_eval_toolkit/run_ytvis_eval.py

@@ -0,0 +1,114 @@
+# flake8: noqa
+
+"""run_youtube_vis.py
+Run example:
+run_youtube_vis.py --USE_PARALLEL False --METRICS HOTA --TRACKERS_TO_EVAL STEm_Seg
+Command Line Arguments: Defaults, # Comments
+    Eval arguments:
+            'USE_PARALLEL': False,
+            'NUM_PARALLEL_CORES': 8,
+            'BREAK_ON_ERROR': True,  # Raises exception and exits with error
+            'RETURN_ON_ERROR': False,  # if not BREAK_ON_ERROR, then returns from function on error
+            'LOG_ON_ERROR': os.path.join(code_path, 'error_log.txt'),  # if not None, save any errors into a log file.
+            'PRINT_RESULTS': True,
+            'PRINT_ONLY_COMBINED': False,
+            'PRINT_CONFIG': True,
+            'TIME_PROGRESS': True,
+            'DISPLAY_LESS_PROGRESS': True,
+            'OUTPUT_SUMMARY': True,
+            'OUTPUT_EMPTY_CLASSES': True,  # If False, summary files are not output for classes with no detections
+            'OUTPUT_DETAILED': True,
+            'PLOT_CURVES': True,
+    Dataset arguments:
+        'GT_FOLDER': os.path.join(code_path, 'data/gt/youtube_vis/youtube_vis_training'),  # Location of GT data
+        'TRACKERS_FOLDER': os.path.join(code_path, 'data/trackers/youtube_vis/youtube_vis_training'),
+        # Trackers location
+        'OUTPUT_FOLDER': None,  # Where to save eval results (if None, same as TRACKERS_FOLDER)
+        'TRACKERS_TO_EVAL': None,  # Filenames of trackers to eval (if None, all in folder)
+        'CLASSES_TO_EVAL': None,  # Classes to eval (if None, all classes)
+        'SPLIT_TO_EVAL': 'training',  # Valid: 'training', 'val'
+        'PRINT_CONFIG': True,  # Whether to print current config
+        'OUTPUT_SUB_FOLDER': '',  # Output files are saved in OUTPUT_FOLDER/tracker_name/OUTPUT_SUB_FOLDER
+        'TRACKER_SUB_FOLDER': 'data',  # Tracker files are in TRACKER_FOLDER/tracker_name/TRACKER_SUB_FOLDER
+        'TRACKER_DISPLAY_NAMES': None,  # Names of trackers to display, if None: TRACKERS_TO_EVAL
+    Metric arguments:
+        'METRICS': ['TrackMAP', 'HOTA', 'CLEAR', 'Identity']
+"""
+
+import argparse
+import os
+import sys
+from multiprocessing import freeze_support
+
+from . import trackeval
+
+
+def run_ytvis_eval(args=None, gt_json=None, dt_json=None):
+    # Command line interface:
+    default_eval_config = trackeval.Evaluator.get_default_eval_config()
+    # print only combined since TrackMAP is undefined for per sequence breakdowns
+    default_eval_config["PRINT_ONLY_COMBINED"] = True
+    default_dataset_config = trackeval.datasets.YouTubeVIS.get_default_dataset_config()
+    default_metrics_config = {"METRICS": ["HOTA"]}
+    config = {
+        **default_eval_config,
+        **default_dataset_config,
+        **default_metrics_config,
+    }  # Merge default configs
+    parser = argparse.ArgumentParser()
+    for setting in config.keys():
+        if type(config[setting]) == list or type(config[setting]) == type(None):
+            parser.add_argument("--" + setting, nargs="+")
+        else:
+            parser.add_argument("--" + setting)
+    args = parser.parse_args(args).__dict__
+    for setting in args.keys():
+        if args[setting] is not None:
+            if type(config[setting]) == type(True):
+                if args[setting] == "True":
+                    x = True
+                elif args[setting] == "False":
+                    x = False
+                else:
+                    raise Exception(
+                        "Command line parameter " + setting + "must be True or False"
+                    )
+            elif type(config[setting]) == type(1):
+                x = int(args[setting])
+            elif type(args[setting]) == type(None):
+                x = None
+            else:
+                x = args[setting]
+            config[setting] = x
+    eval_config = {k: v for k, v in config.items() if k in default_eval_config.keys()}
+    dataset_config = {
+        k: v for k, v in config.items() if k in default_dataset_config.keys()
+    }
+    metrics_config = {
+        k: v for k, v in config.items() if k in default_metrics_config.keys()
+    }
+
+    # Run code
+    evaluator = trackeval.Evaluator(eval_config)
+    # allow directly specifying the GT JSON data and Tracker (result)
+    # JSON data as Python objects, without reading from files.
+    dataset_config["GT_JSON_OBJECT"] = gt_json
+    dataset_config["TRACKER_JSON_OBJECT"] = dt_json
+    dataset_list = [trackeval.datasets.YouTubeVIS(dataset_config)]
+    metrics_list = []
+    # for metric in [trackeval.metrics.TrackMAP, trackeval.metrics.HOTA, trackeval.metrics.CLEAR,
+    #                trackeval.metrics.Identity]:
+    for metric in [trackeval.metrics.HOTA]:
+        if metric.get_name() in metrics_config["METRICS"]:
+            metrics_list.append(metric())
+    if len(metrics_list) == 0:
+        raise Exception("No metrics selected for evaluation")
+    output_res, output_msg = evaluator.evaluate(dataset_list, metrics_list)
+    return output_res, output_msg
+
+
+if __name__ == "__main__":
+    import sys
+
+    freeze_support()
+    run_ytvis_eval(sys.argv[1:])

sam3/eval/hota_eval_toolkit/trackeval/__init__.py

@@ -0,0 +1,4 @@
+# flake8: noqa
+
+from . import datasets, metrics, utils
+from .eval import Evaluator

sam3/eval/hota_eval_toolkit/trackeval/_timing.py

@@ -0,0 +1,68 @@
+# flake8: noqa
+
+import inspect
+from functools import wraps
+from time import perf_counter
+
+DO_TIMING = False
+DISPLAY_LESS_PROGRESS = False
+timer_dict = {}
+counter = 0
+
+
+def time(f):
+    @wraps(f)
+    def wrap(*args, **kw):
+        if DO_TIMING:
+            # Run function with timing
+            ts = perf_counter()
+            result = f(*args, **kw)
+            te = perf_counter()
+            tt = te - ts
+
+            # Get function name
+            arg_names = inspect.getfullargspec(f)[0]
+            if arg_names[0] == "self" and DISPLAY_LESS_PROGRESS:
+                return result
+            elif arg_names[0] == "self":
+                method_name = type(args[0]).__name__ + "." + f.__name__
+            else:
+                method_name = f.__name__
+
+            # Record accumulative time in each function for analysis
+            if method_name in timer_dict.keys():
+                timer_dict[method_name] += tt
+            else:
+                timer_dict[method_name] = tt
+
+            # If code is finished, display timing summary
+            if method_name == "Evaluator.evaluate":
+                print("")
+                print("Timing analysis:")
+                for key, value in timer_dict.items():
+                    print("%-70s %2.4f sec" % (key, value))
+            else:
+                # Get function argument values for printing special arguments of interest
+                arg_titles = ["tracker", "seq", "cls"]
+                arg_vals = []
+                for i, a in enumerate(arg_names):
+                    if a in arg_titles:
+                        arg_vals.append(args[i])
+                arg_text = "(" + ", ".join(arg_vals) + ")"
+
+                # Display methods and functions with different indentation.
+                if arg_names[0] == "self":
+                    print("%-74s %2.4f sec" % (" " * 4 + method_name + arg_text, tt))
+                elif arg_names[0] == "test":
+                    pass
+                else:
+                    global counter
+                    counter += 1
+                    print("%i %-70s %2.4f sec" % (counter, method_name + arg_text, tt))
+
+            return result
+        else:
+            # If config["TIME_PROGRESS"] is false, or config["USE_PARALLEL"] is true, run functions normally without timing.
+            return f(*args, **kw)
+
+    return wrap

sam3/eval/hota_eval_toolkit/trackeval/datasets/__init__.py

@@ -0,0 +1,4 @@
+# flake8: noqa
+
+from .tao_ow import TAO_OW
+from .youtube_vis import YouTubeVIS

sam3/eval/hota_eval_toolkit/trackeval/datasets/_base_dataset.py

@@ -0,0 +1,379 @@
+# flake8: noqa
+
+import csv
+import io
+import os
+import traceback
+import zipfile
+from abc import ABC, abstractmethod
+from copy import deepcopy
+
+import numpy as np
+
+from .. import _timing
+from ..utils import TrackEvalException
+
+
+class _BaseDataset(ABC):
+    @abstractmethod
+    def __init__(self):
+        self.tracker_list = None
+        self.seq_list = None
+        self.class_list = None
+        self.output_fol = None
+        self.output_sub_fol = None
+        self.should_classes_combine = True
+        self.use_super_categories = False
+
+    # Functions to implement:
+
+    @staticmethod
+    @abstractmethod
+    def get_default_dataset_config(): ...
+
+    @abstractmethod
+    def _load_raw_file(self, tracker, seq, is_gt): ...
+
+    @_timing.time
+    @abstractmethod
+    def get_preprocessed_seq_data(self, raw_data, cls): ...
+
+    @abstractmethod
+    def _calculate_similarities(self, gt_dets_t, tracker_dets_t): ...
+
+    # Helper functions for all datasets:
+
+    @classmethod
+    def get_class_name(cls):
+        return cls.__name__
+
+    def get_name(self):
+        return self.get_class_name()
+
+    def get_output_fol(self, tracker):
+        return os.path.join(self.output_fol, tracker, self.output_sub_fol)
+
+    def get_display_name(self, tracker):
+        """Can be overwritten if the trackers name (in files) is different to how it should be displayed.
+        By default this method just returns the trackers name as is.
+        """
+        return tracker
+
+    def get_eval_info(self):
+        """Return info about the dataset needed for the Evaluator"""
+        return self.tracker_list, self.seq_list, self.class_list
+
+    @_timing.time
+    def get_raw_seq_data(self, tracker, seq):
+        """Loads raw data (tracker and ground-truth) for a single tracker on a single sequence.
+        Raw data includes all of the information needed for both preprocessing and evaluation, for all classes.
+        A later function (get_processed_seq_data) will perform such preprocessing and extract relevant information for
+        the evaluation of each class.
+
+        This returns a dict which contains the fields:
+        [num_timesteps]: integer
+        [gt_ids, tracker_ids, gt_classes, tracker_classes, tracker_confidences]:
+                                                                list (for each timestep) of 1D NDArrays (for each det).
+        [gt_dets, tracker_dets, gt_crowd_ignore_regions]: list (for each timestep) of lists of detections.
+        [similarity_scores]: list (for each timestep) of 2D NDArrays.
+        [gt_extras]: dict (for each extra) of lists (for each timestep) of 1D NDArrays (for each det).
+
+        gt_extras contains dataset specific information used for preprocessing such as occlusion and truncation levels.
+
+        Note that similarities are extracted as part of the dataset and not the metric, because almost all metrics are
+        independent of the exact method of calculating the similarity. However datasets are not (e.g. segmentation
+        masks vs 2D boxes vs 3D boxes).
+        We calculate the similarity before preprocessing because often both preprocessing and evaluation require it and
+        we don't wish to calculate this twice.
+        We calculate similarity between all gt and tracker classes (not just each class individually) to allow for
+        calculation of metrics such as class confusion matrices. Typically the impact of this on performance is low.
+        """
+        # Load raw data.
+        raw_gt_data = self._load_raw_file(tracker, seq, is_gt=True)
+        raw_tracker_data = self._load_raw_file(tracker, seq, is_gt=False)
+        raw_data = {**raw_tracker_data, **raw_gt_data}  # Merges dictionaries
+
+        # Calculate similarities for each timestep.
+        similarity_scores = []
+        for t, (gt_dets_t, tracker_dets_t) in enumerate(
+            zip(raw_data["gt_dets"], raw_data["tracker_dets"])
+        ):
+            ious = self._calculate_similarities(gt_dets_t, tracker_dets_t)
+            similarity_scores.append(ious)
+        raw_data["similarity_scores"] = similarity_scores
+        return raw_data
+
+    @staticmethod
+    def _load_simple_text_file(
+        file,
+        time_col=0,
+        id_col=None,
+        remove_negative_ids=False,
+        valid_filter=None,
+        crowd_ignore_filter=None,
+        convert_filter=None,
+        is_zipped=False,
+        zip_file=None,
+        force_delimiters=None,
+    ):
+        """Function that loads data which is in a commonly used text file format.
+        Assumes each det is given by one row of a text file.
+        There is no limit to the number or meaning of each column,
+        however one column needs to give the timestep of each det (time_col) which is default col 0.
+
+        The file dialect (deliminator, num cols, etc) is determined automatically.
+        This function automatically separates dets by timestep,
+        and is much faster than alternatives such as np.loadtext or pandas.
+
+        If remove_negative_ids is True and id_col is not None, dets with negative values in id_col are excluded.
+        These are not excluded from ignore data.
+
+        valid_filter can be used to only include certain classes.
+        It is a dict with ints as keys, and lists as values,
+        such that a row is included if "row[key].lower() is in value" for all key/value pairs in the dict.
+        If None, all classes are included.
+
+        crowd_ignore_filter can be used to read crowd_ignore regions separately. It has the same format as valid filter.
+
+        convert_filter can be used to convert value read to another format.
+        This is used most commonly to convert classes given as string to a class id.
+        This is a dict such that the key is the column to convert, and the value is another dict giving the mapping.
+
+        Optionally, input files could be a zip of multiple text files for storage efficiency.
+
+        Returns read_data and ignore_data.
+        Each is a dict (with keys as timesteps as strings) of lists (over dets) of lists (over column values).
+        Note that all data is returned as strings, and must be converted to float/int later if needed.
+        Note that timesteps will not be present in the returned dict keys if there are no dets for them
+        """
+
+        if remove_negative_ids and id_col is None:
+            raise TrackEvalException(
+                "remove_negative_ids is True, but id_col is not given."
+            )
+        if crowd_ignore_filter is None:
+            crowd_ignore_filter = {}
+        if convert_filter is None:
+            convert_filter = {}
+        try:
+            if is_zipped:  # Either open file directly or within a zip.
+                if zip_file is None:
+                    raise TrackEvalException(
+                        "is_zipped set to True, but no zip_file is given."
+                    )
+                archive = zipfile.ZipFile(os.path.join(zip_file), "r")
+                fp = io.TextIOWrapper(archive.open(file, "r"))
+            else:
+                fp = open(file)
+            read_data = {}
+            crowd_ignore_data = {}
+            fp.seek(0, os.SEEK_END)
+            # check if file is empty
+            if fp.tell():
+                fp.seek(0)
+                dialect = csv.Sniffer().sniff(
+                    fp.readline(), delimiters=force_delimiters
+                )  # Auto determine structure.
+                dialect.skipinitialspace = (
+                    True  # Deal with extra spaces between columns
+                )
+                fp.seek(0)
+                reader = csv.reader(fp, dialect)
+                for row in reader:
+                    try:
+                        # Deal with extra trailing spaces at the end of rows
+                        if row[-1] in "":
+                            row = row[:-1]
+                        timestep = str(int(float(row[time_col])))
+                        # Read ignore regions separately.
+                        is_ignored = False
+                        for ignore_key, ignore_value in crowd_ignore_filter.items():
+                            if row[ignore_key].lower() in ignore_value:
+                                # Convert values in one column (e.g. string to id)
+                                for (
+                                    convert_key,
+                                    convert_value,
+                                ) in convert_filter.items():
+                                    row[convert_key] = convert_value[
+                                        row[convert_key].lower()
+                                    ]
+                                # Save data separated by timestep.
+                                if timestep in crowd_ignore_data.keys():
+                                    crowd_ignore_data[timestep].append(row)
+                                else:
+                                    crowd_ignore_data[timestep] = [row]
+                                is_ignored = True
+                        if (
+                            is_ignored
+                        ):  # if det is an ignore region, it cannot be a normal det.
+                            continue
+                        # Exclude some dets if not valid.
+                        if valid_filter is not None:
+                            for key, value in valid_filter.items():
+                                if row[key].lower() not in value:
+                                    continue
+                        if remove_negative_ids:
+                            if int(float(row[id_col])) < 0:
+                                continue
+                        # Convert values in one column (e.g. string to id)
+                        for convert_key, convert_value in convert_filter.items():
+                            row[convert_key] = convert_value[row[convert_key].lower()]
+                        # Save data separated by timestep.
+                        if timestep in read_data.keys():
+                            read_data[timestep].append(row)
+                        else:
+                            read_data[timestep] = [row]
+                    except Exception:
+                        exc_str_init = (
+                            "In file %s the following line cannot be read correctly: \n"
+                            % os.path.basename(file)
+                        )
+                        exc_str = " ".join([exc_str_init] + row)
+                        raise TrackEvalException(exc_str)
+            fp.close()
+        except Exception:
+            print("Error loading file: %s, printing traceback." % file)
+            traceback.print_exc()
+            raise TrackEvalException(
+                "File %s cannot be read because it is either not present or invalidly formatted"
+                % os.path.basename(file)
+            )
+        return read_data, crowd_ignore_data
+
+    @staticmethod
+    def _calculate_mask_ious(masks1, masks2, is_encoded=False, do_ioa=False):
+        """Calculates the IOU (intersection over union) between two arrays of segmentation masks.
+        If is_encoded a run length encoding with pycocotools is assumed as input format, otherwise an input of numpy
+        arrays of the shape (num_masks, height, width) is assumed and the encoding is performed.
+        If do_ioa (intersection over area) , then calculates the intersection over the area of masks1 - this is commonly
+        used to determine if detections are within crowd ignore region.
+        :param masks1:  first set of masks (numpy array of shape (num_masks, height, width) if not encoded,
+                        else pycocotools rle encoded format)
+        :param masks2:  second set of masks (numpy array of shape (num_masks, height, width) if not encoded,
+                        else pycocotools rle encoded format)
+        :param is_encoded: whether the input is in pycocotools rle encoded format
+        :param do_ioa: whether to perform IoA computation
+        :return: the IoU/IoA scores
+        """
+
+        # Only loaded when run to reduce minimum requirements
+        from pycocotools import mask as mask_utils
+
+        # use pycocotools for run length encoding of masks
+        if not is_encoded:
+            masks1 = mask_utils.encode(
+                np.array(np.transpose(masks1, (1, 2, 0)), order="F")
+            )
+            masks2 = mask_utils.encode(
+                np.array(np.transpose(masks2, (1, 2, 0)), order="F")
+            )
+
+        # use pycocotools for iou computation of rle encoded masks
+        ious = mask_utils.iou(masks1, masks2, [do_ioa] * len(masks2))
+        if len(masks1) == 0 or len(masks2) == 0:
+            ious = np.asarray(ious).reshape(len(masks1), len(masks2))
+        assert (ious >= 0 - np.finfo("float").eps).all()
+        assert (ious <= 1 + np.finfo("float").eps).all()
+
+        return ious
+
+    @staticmethod
+    def _calculate_box_ious(bboxes1, bboxes2, box_format="xywh", do_ioa=False):
+        """Calculates the IOU (intersection over union) between two arrays of boxes.
+        Allows variable box formats ('xywh' and 'x0y0x1y1').
+        If do_ioa (intersection over area) , then calculates the intersection over the area of boxes1 - this is commonly
+        used to determine if detections are within crowd ignore region.
+        """
+        if box_format in "xywh":
+            # layout: (x0, y0, w, h)
+            bboxes1 = deepcopy(bboxes1)
+            bboxes2 = deepcopy(bboxes2)
+
+            bboxes1[:, 2] = bboxes1[:, 0] + bboxes1[:, 2]
+            bboxes1[:, 3] = bboxes1[:, 1] + bboxes1[:, 3]
+            bboxes2[:, 2] = bboxes2[:, 0] + bboxes2[:, 2]
+            bboxes2[:, 3] = bboxes2[:, 1] + bboxes2[:, 3]
+        elif box_format not in "x0y0x1y1":
+            raise (TrackEvalException("box_format %s is not implemented" % box_format))
+
+        # layout: (x0, y0, x1, y1)
+        min_ = np.minimum(bboxes1[:, np.newaxis, :], bboxes2[np.newaxis, :, :])
+        max_ = np.maximum(bboxes1[:, np.newaxis, :], bboxes2[np.newaxis, :, :])
+        intersection = np.maximum(min_[..., 2] - max_[..., 0], 0) * np.maximum(
+            min_[..., 3] - max_[..., 1], 0
+        )
+        area1 = (bboxes1[..., 2] - bboxes1[..., 0]) * (
+            bboxes1[..., 3] - bboxes1[..., 1]
+        )
+
+        if do_ioa:
+            ioas = np.zeros_like(intersection)
+            valid_mask = area1 > 0 + np.finfo("float").eps
+            ioas[valid_mask, :] = (
+                intersection[valid_mask, :] / area1[valid_mask][:, np.newaxis]
+            )
+
+            return ioas
+        else:
+            area2 = (bboxes2[..., 2] - bboxes2[..., 0]) * (
+                bboxes2[..., 3] - bboxes2[..., 1]
+            )
+            union = area1[:, np.newaxis] + area2[np.newaxis, :] - intersection
+            intersection[area1 <= 0 + np.finfo("float").eps, :] = 0
+            intersection[:, area2 <= 0 + np.finfo("float").eps] = 0
+            intersection[union <= 0 + np.finfo("float").eps] = 0
+            union[union <= 0 + np.finfo("float").eps] = 1
+            ious = intersection / union
+            return ious
+
+    @staticmethod
+    def _calculate_euclidean_similarity(dets1, dets2, zero_distance=2.0):
+        """Calculates the euclidean distance between two sets of detections, and then converts this into a similarity
+        measure with values between 0 and 1 using the following formula: sim = max(0, 1 - dist/zero_distance).
+        The default zero_distance of 2.0, corresponds to the default used in MOT15_3D, such that a 0.5 similarity
+        threshold corresponds to a 1m distance threshold for TPs.
+        """
+        dist = np.linalg.norm(dets1[:, np.newaxis] - dets2[np.newaxis, :], axis=2)
+        sim = np.maximum(0, 1 - dist / zero_distance)
+        return sim
+
+    @staticmethod
+    def _check_unique_ids(data, after_preproc=False):
+        """Check the requirement that the tracker_ids and gt_ids are unique per timestep"""
+        gt_ids = data["gt_ids"]
+        tracker_ids = data["tracker_ids"]
+        for t, (gt_ids_t, tracker_ids_t) in enumerate(zip(gt_ids, tracker_ids)):
+            if len(tracker_ids_t) > 0:
+                unique_ids, counts = np.unique(tracker_ids_t, return_counts=True)
+                if np.max(counts) != 1:
+                    duplicate_ids = unique_ids[counts > 1]
+                    exc_str_init = (
+                        "Tracker predicts the same ID more than once in a single timestep "
+                        "(seq: %s, frame: %i, ids:" % (data["seq"], t + 1)
+                    )
+                    exc_str = (
+                        " ".join([exc_str_init] + [str(d) for d in duplicate_ids]) + ")"
+                    )
+                    if after_preproc:
+                        exc_str_init += (
+                            "\n Note that this error occurred after preprocessing (but not before), "
+                            "so ids may not be as in file, and something seems wrong with preproc."
+                        )
+                    raise TrackEvalException(exc_str)
+            if len(gt_ids_t) > 0:
+                unique_ids, counts = np.unique(gt_ids_t, return_counts=True)
+                if np.max(counts) != 1:
+                    duplicate_ids = unique_ids[counts > 1]
+                    exc_str_init = (
+                        "Ground-truth has the same ID more than once in a single timestep "
+                        "(seq: %s, frame: %i, ids:" % (data["seq"], t + 1)
+                    )
+                    exc_str = (
+                        " ".join([exc_str_init] + [str(d) for d in duplicate_ids]) + ")"
+                    )
+                    if after_preproc:
+                        exc_str_init += (
+                            "\n Note that this error occurred after preprocessing (but not before), "
+                            "so ids may not be as in file, and something seems wrong with preproc."
+                        )
+                    raise TrackEvalException(exc_str)

sam3/eval/hota_eval_toolkit/trackeval/datasets/tao_ow.py

@@ -0,0 +1,891 @@
+# flake8: noqa
+
+import itertools
+import json
+import os
+from collections import defaultdict
+
+import numpy as np
+from scipy.optimize import linear_sum_assignment
+
+from .. import _timing, utils
+from ..utils import TrackEvalException
+from ._base_dataset import _BaseDataset
+
+
+class TAO_OW(_BaseDataset):
+    """Dataset class for TAO tracking"""
+
+    @staticmethod
+    def get_default_dataset_config():
+        """Default class config values"""
+        code_path = utils.get_code_path()
+        default_config = {
+            "GT_FOLDER": os.path.join(
+                code_path, "data/gt/tao/tao_training"
+            ),  # Location of GT data
+            "TRACKERS_FOLDER": os.path.join(
+                code_path, "data/trackers/tao/tao_training"
+            ),  # Trackers location
+            "OUTPUT_FOLDER": None,  # Where to save eval results (if None, same as TRACKERS_FOLDER)
+            "TRACKERS_TO_EVAL": None,  # Filenames of trackers to eval (if None, all in folder)
+            "CLASSES_TO_EVAL": None,  # Classes to eval (if None, all classes)
+            "SPLIT_TO_EVAL": "training",  # Valid: 'training', 'val'
+            "PRINT_CONFIG": True,  # Whether to print current config
+            "TRACKER_SUB_FOLDER": "data",  # Tracker files are in TRACKER_FOLDER/tracker_name/TRACKER_SUB_FOLDER
+            "OUTPUT_SUB_FOLDER": "",  # Output files are saved in OUTPUT_FOLDER/tracker_name/OUTPUT_SUB_FOLDER
+            "TRACKER_DISPLAY_NAMES": None,  # Names of trackers to display, if None: TRACKERS_TO_EVAL
+            "MAX_DETECTIONS": 300,  # Number of maximal allowed detections per image (0 for unlimited)
+            "SUBSET": "all",
+        }
+        return default_config
+
+    def __init__(self, config=None):
+        """Initialise dataset, checking that all required files are present"""
+        super().__init__()
+        # Fill non-given config values with defaults
+        self.config = utils.init_config(
+            config, self.get_default_dataset_config(), self.get_name()
+        )
+        self.gt_fol = self.config["GT_FOLDER"]
+        self.tracker_fol = self.config["TRACKERS_FOLDER"]
+        self.should_classes_combine = True
+        self.use_super_categories = False
+
+        self.tracker_sub_fol = self.config["TRACKER_SUB_FOLDER"]
+        self.output_fol = self.config["OUTPUT_FOLDER"]
+        if self.output_fol is None:
+            self.output_fol = self.tracker_fol
+        self.output_sub_fol = self.config["OUTPUT_SUB_FOLDER"]
+
+        gt_dir_files = [
+            file for file in os.listdir(self.gt_fol) if file.endswith(".json")
+        ]
+        if len(gt_dir_files) != 1:
+            raise TrackEvalException(
+                self.gt_fol + " does not contain exactly one json file."
+            )
+
+        with open(os.path.join(self.gt_fol, gt_dir_files[0])) as f:
+            self.gt_data = json.load(f)
+
+        self.subset = self.config["SUBSET"]
+        if self.subset != "all":
+            # Split GT data into `known`, `unknown` or `distractor`
+            self._split_known_unknown_distractor()
+            self.gt_data = self._filter_gt_data(self.gt_data)
+
+        # merge categories marked with a merged tag in TAO dataset
+        self._merge_categories(self.gt_data["annotations"] + self.gt_data["tracks"])
+
+        # Get sequences to eval and sequence information
+        self.seq_list = [
+            vid["name"].replace("/", "-") for vid in self.gt_data["videos"]
+        ]
+        self.seq_name_to_seq_id = {
+            vid["name"].replace("/", "-"): vid["id"] for vid in self.gt_data["videos"]
+        }
+        # compute mappings from videos to annotation data
+        self.videos_to_gt_tracks, self.videos_to_gt_images = self._compute_vid_mappings(
+            self.gt_data["annotations"]
+        )
+        # compute sequence lengths
+        self.seq_lengths = {vid["id"]: 0 for vid in self.gt_data["videos"]}
+        for img in self.gt_data["images"]:
+            self.seq_lengths[img["video_id"]] += 1
+        self.seq_to_images_to_timestep = self._compute_image_to_timestep_mappings()
+        self.seq_to_classes = {
+            vid["id"]: {
+                "pos_cat_ids": list(
+                    {
+                        track["category_id"]
+                        for track in self.videos_to_gt_tracks[vid["id"]]
+                    }
+                ),
+                "neg_cat_ids": vid["neg_category_ids"],
+                "not_exhaustively_labeled_cat_ids": vid["not_exhaustive_category_ids"],
+            }
+            for vid in self.gt_data["videos"]
+        }
+
+        # Get classes to eval
+        considered_vid_ids = [self.seq_name_to_seq_id[vid] for vid in self.seq_list]
+        seen_cats = set(
+            [
+                cat_id
+                for vid_id in considered_vid_ids
+                for cat_id in self.seq_to_classes[vid_id]["pos_cat_ids"]
+            ]
+        )
+        # only classes with ground truth are evaluated in TAO
+        self.valid_classes = [
+            cls["name"] for cls in self.gt_data["categories"] if cls["id"] in seen_cats
+        ]
+        # cls_name_to_cls_id_map = {cls['name']: cls['id'] for cls in self.gt_data['categories']}
+
+        if self.config["CLASSES_TO_EVAL"]:
+            # self.class_list = [cls.lower() if cls.lower() in self.valid_classes else None
+            #                    for cls in self.config['CLASSES_TO_EVAL']]
+            self.class_list = ["object"]  # class-agnostic
+            if not all(self.class_list):
+                raise TrackEvalException(
+                    "Attempted to evaluate an invalid class. Only classes "
+                    + ", ".join(self.valid_classes)
+                    + " are valid (classes present in ground truth data)."
+                )
+        else:
+            # self.class_list = [cls for cls in self.valid_classes]
+            self.class_list = ["object"]  # class-agnostic
+        # self.class_name_to_class_id = {k: v for k, v in cls_name_to_cls_id_map.items() if k in self.class_list}
+        self.class_name_to_class_id = {"object": 1}  # class-agnostic
+
+        # Get trackers to eval
+        if self.config["TRACKERS_TO_EVAL"] is None:
+            self.tracker_list = os.listdir(self.tracker_fol)
+        else:
+            self.tracker_list = self.config["TRACKERS_TO_EVAL"]
+
+        if self.config["TRACKER_DISPLAY_NAMES"] is None:
+            self.tracker_to_disp = dict(zip(self.tracker_list, self.tracker_list))
+        elif (self.config["TRACKERS_TO_EVAL"] is not None) and (
+            len(self.config["TRACKER_DISPLAY_NAMES"]) == len(self.tracker_list)
+        ):
+            self.tracker_to_disp = dict(
+                zip(self.tracker_list, self.config["TRACKER_DISPLAY_NAMES"])
+            )
+        else:
+            raise TrackEvalException(
+                "List of tracker files and tracker display names do not match."
+            )
+
+        self.tracker_data = {tracker: dict() for tracker in self.tracker_list}
+
+        for tracker in self.tracker_list:
+            tr_dir_files = [
+                file
+                for file in os.listdir(
+                    os.path.join(self.tracker_fol, tracker, self.tracker_sub_fol)
+                )
+                if file.endswith(".json")
+            ]
+            if len(tr_dir_files) != 1:
+                raise TrackEvalException(
+                    os.path.join(self.tracker_fol, tracker, self.tracker_sub_fol)
+                    + " does not contain exactly one json file."
+                )
+            with open(
+                os.path.join(
+                    self.tracker_fol, tracker, self.tracker_sub_fol, tr_dir_files[0]
+                )
+            ) as f:
+                curr_data = json.load(f)
+
+            # limit detections if MAX_DETECTIONS > 0
+            if self.config["MAX_DETECTIONS"]:
+                curr_data = self._limit_dets_per_image(curr_data)
+
+            # fill missing video ids
+            self._fill_video_ids_inplace(curr_data)
+
+            # make track ids unique over whole evaluation set
+            self._make_track_ids_unique(curr_data)
+
+            # merge categories marked with a merged tag in TAO dataset
+            self._merge_categories(curr_data)
+
+            # get tracker sequence information
+            curr_videos_to_tracker_tracks, curr_videos_to_tracker_images = (
+                self._compute_vid_mappings(curr_data)
+            )
+            self.tracker_data[tracker]["vids_to_tracks"] = curr_videos_to_tracker_tracks
+            self.tracker_data[tracker]["vids_to_images"] = curr_videos_to_tracker_images
+
+    def get_display_name(self, tracker):
+        return self.tracker_to_disp[tracker]
+
+    def _load_raw_file(self, tracker, seq, is_gt):
+        """Load a file (gt or tracker) in the TAO format
+
+        If is_gt, this returns a dict which contains the fields:
+        [gt_ids, gt_classes] : list (for each timestep) of 1D NDArrays (for each det).
+        [gt_dets]: list (for each timestep) of lists of detections.
+        [classes_to_gt_tracks]: dictionary with class values as keys and list of dictionaries (with frame indices as
+                                keys and corresponding segmentations as values) for each track
+        [classes_to_gt_track_ids, classes_to_gt_track_areas, classes_to_gt_track_lengths]: dictionary with class values
+                                as keys and lists (for each track) as values
+
+        if not is_gt, this returns a dict which contains the fields:
+        [tracker_ids, tracker_classes, tracker_confidences] : list (for each timestep) of 1D NDArrays (for each det).
+        [tracker_dets]: list (for each timestep) of lists of detections.
+        [classes_to_dt_tracks]: dictionary with class values as keys and list of dictionaries (with frame indices as
+                                keys and corresponding segmentations as values) for each track
+        [classes_to_dt_track_ids, classes_to_dt_track_areas, classes_to_dt_track_lengths]: dictionary with class values
+                                                                                           as keys and lists as values
+        [classes_to_dt_track_scores]: dictionary with class values as keys and 1D numpy arrays as values
+        """
+        seq_id = self.seq_name_to_seq_id[seq]
+        # File location
+        if is_gt:
+            imgs = self.videos_to_gt_images[seq_id]
+        else:
+            imgs = self.tracker_data[tracker]["vids_to_images"][seq_id]
+
+        # Convert data to required format
+        num_timesteps = self.seq_lengths[seq_id]
+        img_to_timestep = self.seq_to_images_to_timestep[seq_id]
+        data_keys = ["ids", "classes", "dets"]
+        if not is_gt:
+            data_keys += ["tracker_confidences"]
+        raw_data = {key: [None] * num_timesteps for key in data_keys}
+        for img in imgs:
+            # some tracker data contains images without any ground truth information, these are ignored
+            try:
+                t = img_to_timestep[img["id"]]
+            except KeyError:
+                continue
+            annotations = img["annotations"]
+            raw_data["dets"][t] = np.atleast_2d(
+                [ann["bbox"] for ann in annotations]
+            ).astype(float)
+            raw_data["ids"][t] = np.atleast_1d(
+                [ann["track_id"] for ann in annotations]
+            ).astype(int)
+            raw_data["classes"][t] = np.atleast_1d([1 for _ in annotations]).astype(
+                int
+            )  # class-agnostic
+            if not is_gt:
+                raw_data["tracker_confidences"][t] = np.atleast_1d(
+                    [ann["score"] for ann in annotations]
+                ).astype(float)
+
+        for t, d in enumerate(raw_data["dets"]):
+            if d is None:
+                raw_data["dets"][t] = np.empty((0, 4)).astype(float)
+                raw_data["ids"][t] = np.empty(0).astype(int)
+                raw_data["classes"][t] = np.empty(0).astype(int)
+                if not is_gt:
+                    raw_data["tracker_confidences"][t] = np.empty(0)
+
+        if is_gt:
+            key_map = {"ids": "gt_ids", "classes": "gt_classes", "dets": "gt_dets"}
+        else:
+            key_map = {
+                "ids": "tracker_ids",
+                "classes": "tracker_classes",
+                "dets": "tracker_dets",
+            }
+        for k, v in key_map.items():
+            raw_data[v] = raw_data.pop(k)
+
+        # all_classes = [self.class_name_to_class_id[cls] for cls in self.class_list]
+        all_classes = [1]  # class-agnostic
+
+        if is_gt:
+            classes_to_consider = all_classes
+            all_tracks = self.videos_to_gt_tracks[seq_id]
+        else:
+            # classes_to_consider = self.seq_to_classes[seq_id]['pos_cat_ids'] \
+            #                       + self.seq_to_classes[seq_id]['neg_cat_ids']
+            classes_to_consider = all_classes  # class-agnostic
+            all_tracks = self.tracker_data[tracker]["vids_to_tracks"][seq_id]
+
+        # classes_to_tracks = {cls: [track for track in all_tracks if track['category_id'] == cls]
+        #                      if cls in classes_to_consider else [] for cls in all_classes}
+        classes_to_tracks = {
+            cls: [track for track in all_tracks] if cls in classes_to_consider else []
+            for cls in all_classes
+        }  # class-agnostic
+
+        # mapping from classes to track information
+        raw_data["classes_to_tracks"] = {
+            cls: [
+                {
+                    det["image_id"]: np.atleast_1d(det["bbox"])
+                    for det in track["annotations"]
+                }
+                for track in tracks
+            ]
+            for cls, tracks in classes_to_tracks.items()
+        }
+        raw_data["classes_to_track_ids"] = {
+            cls: [track["id"] for track in tracks]
+            for cls, tracks in classes_to_tracks.items()
+        }
+        raw_data["classes_to_track_areas"] = {
+            cls: [track["area"] for track in tracks]
+            for cls, tracks in classes_to_tracks.items()
+        }
+        raw_data["classes_to_track_lengths"] = {
+            cls: [len(track["annotations"]) for track in tracks]
+            for cls, tracks in classes_to_tracks.items()
+        }
+
+        if not is_gt:
+            raw_data["classes_to_dt_track_scores"] = {
+                cls: np.array(
+                    [
+                        np.mean([float(x["score"]) for x in track["annotations"]])
+                        for track in tracks
+                    ]
+                )
+                for cls, tracks in classes_to_tracks.items()
+            }
+
+        if is_gt:
+            key_map = {
+                "classes_to_tracks": "classes_to_gt_tracks",
+                "classes_to_track_ids": "classes_to_gt_track_ids",
+                "classes_to_track_lengths": "classes_to_gt_track_lengths",
+                "classes_to_track_areas": "classes_to_gt_track_areas",
+            }
+        else:
+            key_map = {
+                "classes_to_tracks": "classes_to_dt_tracks",
+                "classes_to_track_ids": "classes_to_dt_track_ids",
+                "classes_to_track_lengths": "classes_to_dt_track_lengths",
+                "classes_to_track_areas": "classes_to_dt_track_areas",
+            }
+        for k, v in key_map.items():
+            raw_data[v] = raw_data.pop(k)
+
+        raw_data["num_timesteps"] = num_timesteps
+        raw_data["neg_cat_ids"] = self.seq_to_classes[seq_id]["neg_cat_ids"]
+        raw_data["not_exhaustively_labeled_cls"] = self.seq_to_classes[seq_id][
+            "not_exhaustively_labeled_cat_ids"
+        ]
+        raw_data["seq"] = seq
+        return raw_data
+
+    @_timing.time
+    def get_preprocessed_seq_data(self, raw_data, cls):
+        """Preprocess data for a single sequence for a single class ready for evaluation.
+        Inputs:
+             - raw_data is a dict containing the data for the sequence already read in by get_raw_seq_data().
+             - cls is the class to be evaluated.
+        Outputs:
+             - data is a dict containing all of the information that metrics need to perform evaluation.
+                It contains the following fields:
+                    [num_timesteps, num_gt_ids, num_tracker_ids, num_gt_dets, num_tracker_dets] : integers.
+                    [gt_ids, tracker_ids, tracker_confidences]: list (for each timestep) of 1D NDArrays (for each det).
+                    [gt_dets, tracker_dets]: list (for each timestep) of lists of detections.
+                    [similarity_scores]: list (for each timestep) of 2D NDArrays.
+        Notes:
+            General preprocessing (preproc) occurs in 4 steps. Some datasets may not use all of these steps.
+                1) Extract only detections relevant for the class to be evaluated (including distractor detections).
+                2) Match gt dets and tracker dets. Remove tracker dets that are matched to a gt det that is of a
+                    distractor class, or otherwise marked as to be removed.
+                3) Remove unmatched tracker dets if they fall within a crowd ignore region or don't meet a certain
+                    other criteria (e.g. are too small).
+                4) Remove gt dets that were only useful for preprocessing and not for actual evaluation.
+            After the above preprocessing steps, this function also calculates the number of gt and tracker detections
+                and unique track ids. It also relabels gt and tracker ids to be contiguous and checks that ids are
+                unique within each timestep.
+        TAO:
+            In TAO, the 4 preproc steps are as follow:
+                1) All classes present in the ground truth data are evaluated separately.
+                2) No matched tracker detections are removed.
+                3) Unmatched tracker detections are removed if there is not ground truth data and the class does not
+                    belong to the categories marked as negative for this sequence. Additionally, unmatched tracker
+                    detections for classes which are marked as not exhaustively labeled are removed.
+                4) No gt detections are removed.
+            Further, for TrackMAP computation track representations for the given class are accessed from a dictionary
+            and the tracks from the tracker data are sorted according to the tracker confidence.
+        """
+        cls_id = self.class_name_to_class_id[cls]
+        is_not_exhaustively_labeled = cls_id in raw_data["not_exhaustively_labeled_cls"]
+        is_neg_category = cls_id in raw_data["neg_cat_ids"]
+
+        data_keys = [
+            "gt_ids",
+            "tracker_ids",
+            "gt_dets",
+            "tracker_dets",
+            "tracker_confidences",
+            "similarity_scores",
+        ]
+        data = {key: [None] * raw_data["num_timesteps"] for key in data_keys}
+        unique_gt_ids = []
+        unique_tracker_ids = []
+        num_gt_dets = 0
+        num_tracker_dets = 0
+        for t in range(raw_data["num_timesteps"]):
+            # Only extract relevant dets for this class for preproc and eval (cls)
+            gt_class_mask = np.atleast_1d(raw_data["gt_classes"][t] == cls_id)
+            gt_class_mask = gt_class_mask.astype(bool)
+            gt_ids = raw_data["gt_ids"][t][gt_class_mask]
+            gt_dets = raw_data["gt_dets"][t][gt_class_mask]
+
+            tracker_class_mask = np.atleast_1d(raw_data["tracker_classes"][t] == cls_id)
+            tracker_class_mask = tracker_class_mask.astype(bool)
+            tracker_ids = raw_data["tracker_ids"][t][tracker_class_mask]
+            tracker_dets = raw_data["tracker_dets"][t][tracker_class_mask]
+            tracker_confidences = raw_data["tracker_confidences"][t][tracker_class_mask]
+            similarity_scores = raw_data["similarity_scores"][t][gt_class_mask, :][
+                :, tracker_class_mask
+            ]
+
+            # Match tracker and gt dets (with hungarian algorithm).
+            unmatched_indices = np.arange(tracker_ids.shape[0])
+            if gt_ids.shape[0] > 0 and tracker_ids.shape[0] > 0:
+                matching_scores = similarity_scores.copy()
+                matching_scores[matching_scores < 0.5 - np.finfo("float").eps] = 0
+                match_rows, match_cols = linear_sum_assignment(-matching_scores)
+                actually_matched_mask = (
+                    matching_scores[match_rows, match_cols] > 0 + np.finfo("float").eps
+                )
+                match_cols = match_cols[actually_matched_mask]
+                unmatched_indices = np.delete(unmatched_indices, match_cols, axis=0)
+
+            if gt_ids.shape[0] == 0 and not is_neg_category:
+                to_remove_tracker = unmatched_indices
+            elif is_not_exhaustively_labeled:
+                to_remove_tracker = unmatched_indices
+            else:
+                to_remove_tracker = np.array([], dtype=int)
+
+            # remove all unwanted unmatched tracker detections
+            data["tracker_ids"][t] = np.delete(tracker_ids, to_remove_tracker, axis=0)
+            data["tracker_dets"][t] = np.delete(tracker_dets, to_remove_tracker, axis=0)
+            data["tracker_confidences"][t] = np.delete(
+                tracker_confidences, to_remove_tracker, axis=0
+            )
+            similarity_scores = np.delete(similarity_scores, to_remove_tracker, axis=1)
+
+            data["gt_ids"][t] = gt_ids
+            data["gt_dets"][t] = gt_dets
+            data["similarity_scores"][t] = similarity_scores
+
+            unique_gt_ids += list(np.unique(data["gt_ids"][t]))
+            unique_tracker_ids += list(np.unique(data["tracker_ids"][t]))
+            num_tracker_dets += len(data["tracker_ids"][t])
+            num_gt_dets += len(data["gt_ids"][t])
+
+        # Re-label IDs such that there are no empty IDs
+        if len(unique_gt_ids) > 0:
+            unique_gt_ids = np.unique(unique_gt_ids)
+            gt_id_map = np.nan * np.ones((np.max(unique_gt_ids) + 1))
+            gt_id_map[unique_gt_ids] = np.arange(len(unique_gt_ids))
+            for t in range(raw_data["num_timesteps"]):
+                if len(data["gt_ids"][t]) > 0:
+                    data["gt_ids"][t] = gt_id_map[data["gt_ids"][t]].astype(int)
+        if len(unique_tracker_ids) > 0:
+            unique_tracker_ids = np.unique(unique_tracker_ids)
+            tracker_id_map = np.nan * np.ones((np.max(unique_tracker_ids) + 1))
+            tracker_id_map[unique_tracker_ids] = np.arange(len(unique_tracker_ids))
+            for t in range(raw_data["num_timesteps"]):
+                if len(data["tracker_ids"][t]) > 0:
+                    data["tracker_ids"][t] = tracker_id_map[
+                        data["tracker_ids"][t]
+                    ].astype(int)
+
+        # Record overview statistics.
+        data["num_tracker_dets"] = num_tracker_dets
+        data["num_gt_dets"] = num_gt_dets
+        data["num_tracker_ids"] = len(unique_tracker_ids)
+        data["num_gt_ids"] = len(unique_gt_ids)
+        data["num_timesteps"] = raw_data["num_timesteps"]
+        data["seq"] = raw_data["seq"]
+
+        # get track representations
+        data["gt_tracks"] = raw_data["classes_to_gt_tracks"][cls_id]
+        data["gt_track_ids"] = raw_data["classes_to_gt_track_ids"][cls_id]
+        data["gt_track_lengths"] = raw_data["classes_to_gt_track_lengths"][cls_id]
+        data["gt_track_areas"] = raw_data["classes_to_gt_track_areas"][cls_id]
+        data["dt_tracks"] = raw_data["classes_to_dt_tracks"][cls_id]
+        data["dt_track_ids"] = raw_data["classes_to_dt_track_ids"][cls_id]
+        data["dt_track_lengths"] = raw_data["classes_to_dt_track_lengths"][cls_id]
+        data["dt_track_areas"] = raw_data["classes_to_dt_track_areas"][cls_id]
+        data["dt_track_scores"] = raw_data["classes_to_dt_track_scores"][cls_id]
+        data["not_exhaustively_labeled"] = is_not_exhaustively_labeled
+        data["iou_type"] = "bbox"
+
+        # sort tracker data tracks by tracker confidence scores
+        if data["dt_tracks"]:
+            idx = np.argsort(
+                [-score for score in data["dt_track_scores"]], kind="mergesort"
+            )
+            data["dt_track_scores"] = [data["dt_track_scores"][i] for i in idx]
+            data["dt_tracks"] = [data["dt_tracks"][i] for i in idx]
+            data["dt_track_ids"] = [data["dt_track_ids"][i] for i in idx]
+            data["dt_track_lengths"] = [data["dt_track_lengths"][i] for i in idx]
+            data["dt_track_areas"] = [data["dt_track_areas"][i] for i in idx]
+        # Ensure that ids are unique per timestep.
+        self._check_unique_ids(data)
+
+        return data
+
+    def _calculate_similarities(self, gt_dets_t, tracker_dets_t):
+        similarity_scores = self._calculate_box_ious(gt_dets_t, tracker_dets_t)
+        return similarity_scores
+
+    def _merge_categories(self, annotations):
+        """
+        Merges categories with a merged tag. Adapted from https://github.com/TAO-Dataset
+        :param annotations: the annotations in which the classes should be merged
+        :return: None
+        """
+        merge_map = {}
+        for category in self.gt_data["categories"]:
+            if "merged" in category:
+                for to_merge in category["merged"]:
+                    merge_map[to_merge["id"]] = category["id"]
+
+        for ann in annotations:
+            ann["category_id"] = merge_map.get(ann["category_id"], ann["category_id"])
+
+    def _compute_vid_mappings(self, annotations):
+        """
+        Computes mappings from Videos to corresponding tracks and images.
+        :param annotations: the annotations for which the mapping should be generated
+        :return: the video-to-track-mapping, the video-to-image-mapping
+        """
+        vids_to_tracks = {}
+        vids_to_imgs = {}
+        vid_ids = [vid["id"] for vid in self.gt_data["videos"]]
+
+        # compute an mapping from image IDs to images
+        images = {}
+        for image in self.gt_data["images"]:
+            images[image["id"]] = image
+
+        for ann in annotations:
+            ann["area"] = ann["bbox"][2] * ann["bbox"][3]
+
+            vid = ann["video_id"]
+            if ann["video_id"] not in vids_to_tracks.keys():
+                vids_to_tracks[ann["video_id"]] = list()
+            if ann["video_id"] not in vids_to_imgs.keys():
+                vids_to_imgs[ann["video_id"]] = list()
+
+            # Fill in vids_to_tracks
+            tid = ann["track_id"]
+            exist_tids = [track["id"] for track in vids_to_tracks[vid]]
+            try:
+                index1 = exist_tids.index(tid)
+            except ValueError:
+                index1 = -1
+            if tid not in exist_tids:
+                curr_track = {
+                    "id": tid,
+                    "category_id": ann["category_id"],
+                    "video_id": vid,
+                    "annotations": [ann],
+                }
+                vids_to_tracks[vid].append(curr_track)
+            else:
+                vids_to_tracks[vid][index1]["annotations"].append(ann)
+
+            # Fill in vids_to_imgs
+            img_id = ann["image_id"]
+            exist_img_ids = [img["id"] for img in vids_to_imgs[vid]]
+            try:
+                index2 = exist_img_ids.index(img_id)
+            except ValueError:
+                index2 = -1
+            if index2 == -1:
+                curr_img = {"id": img_id, "annotations": [ann]}
+                vids_to_imgs[vid].append(curr_img)
+            else:
+                vids_to_imgs[vid][index2]["annotations"].append(ann)
+
+        # sort annotations by frame index and compute track area
+        for vid, tracks in vids_to_tracks.items():
+            for track in tracks:
+                track["annotations"] = sorted(
+                    track["annotations"],
+                    key=lambda x: images[x["image_id"]]["frame_index"],
+                )
+                # Computer average area
+                track["area"] = sum(x["area"] for x in track["annotations"]) / len(
+                    track["annotations"]
+                )
+
+        # Ensure all videos are present
+        for vid_id in vid_ids:
+            if vid_id not in vids_to_tracks.keys():
+                vids_to_tracks[vid_id] = []
+            if vid_id not in vids_to_imgs.keys():
+                vids_to_imgs[vid_id] = []
+
+        return vids_to_tracks, vids_to_imgs
+
+    def _compute_image_to_timestep_mappings(self):
+        """
+        Computes a mapping from images to the corresponding timestep in the sequence.
+        :return: the image-to-timestep-mapping
+        """
+        images = {}
+        for image in self.gt_data["images"]:
+            images[image["id"]] = image
+
+        seq_to_imgs_to_timestep = {vid["id"]: dict() for vid in self.gt_data["videos"]}
+        for vid in seq_to_imgs_to_timestep:
+            curr_imgs = [img["id"] for img in self.videos_to_gt_images[vid]]
+            curr_imgs = sorted(curr_imgs, key=lambda x: images[x]["frame_index"])
+            seq_to_imgs_to_timestep[vid] = {
+                curr_imgs[i]: i for i in range(len(curr_imgs))
+            }
+
+        return seq_to_imgs_to_timestep
+
+    def _limit_dets_per_image(self, annotations):
+        """
+        Limits the number of detections for each image to config['MAX_DETECTIONS']. Adapted from
+        https://github.com/TAO-Dataset/
+        :param annotations: the annotations in which the detections should be limited
+        :return: the annotations with limited detections
+        """
+        max_dets = self.config["MAX_DETECTIONS"]
+        img_ann = defaultdict(list)
+        for ann in annotations:
+            img_ann[ann["image_id"]].append(ann)
+
+        for img_id, _anns in img_ann.items():
+            if len(_anns) <= max_dets:
+                continue
+            _anns = sorted(_anns, key=lambda x: x["score"], reverse=True)
+            img_ann[img_id] = _anns[:max_dets]
+
+        return [ann for anns in img_ann.values() for ann in anns]
+
+    def _fill_video_ids_inplace(self, annotations):
+        """
+        Fills in missing video IDs inplace. Adapted from https://github.com/TAO-Dataset/
+        :param annotations: the annotations for which the videos IDs should be filled inplace
+        :return: None
+        """
+        missing_video_id = [x for x in annotations if "video_id" not in x]
+        if missing_video_id:
+            image_id_to_video_id = {
+                x["id"]: x["video_id"] for x in self.gt_data["images"]
+            }
+            for x in missing_video_id:
+                x["video_id"] = image_id_to_video_id[x["image_id"]]
+
+    @staticmethod
+    def _make_track_ids_unique(annotations):
+        """
+        Makes the track IDs unqiue over the whole annotation set. Adapted from https://github.com/TAO-Dataset/
+        :param annotations: the annotation set
+        :return: the number of updated IDs
+        """
+        track_id_videos = {}
+        track_ids_to_update = set()
+        max_track_id = 0
+        for ann in annotations:
+            t = ann["track_id"]
+            if t not in track_id_videos:
+                track_id_videos[t] = ann["video_id"]
+
+            if ann["video_id"] != track_id_videos[t]:
+                # Track id is assigned to multiple videos
+                track_ids_to_update.add(t)
+            max_track_id = max(max_track_id, t)
+
+        if track_ids_to_update:
+            print("true")
+            next_id = itertools.count(max_track_id + 1)
+            new_track_ids = defaultdict(lambda: next(next_id))
+            for ann in annotations:
+                t = ann["track_id"]
+                v = ann["video_id"]
+                if t in track_ids_to_update:
+                    ann["track_id"] = new_track_ids[t, v]
+        return len(track_ids_to_update)
+
+    def _split_known_unknown_distractor(self):
+        all_ids = set(
+            [i for i in range(1, 2000)]
+        )  # 2000 is larger than the max category id in TAO-OW.
+        # `knowns` includes 78 TAO_category_ids that corresponds to 78 COCO classes.
+        # (The other 2 COCO classes do not have corresponding classes in TAO).
+        self.knowns = {
+            4,
+            13,
+            1038,
+            544,
+            1057,
+            34,
+            35,
+            36,
+            41,
+            45,
+            58,
+            60,
+            579,
+            1091,
+            1097,
+            1099,
+            78,
+            79,
+            81,
+            91,
+            1115,
+            1117,
+            95,
+            1122,
+            99,
+            1132,
+            621,
+            1135,
+            625,
+            118,
+            1144,
+            126,
+            642,
+            1155,
+            133,
+            1162,
+            139,
+            154,
+            174,
+            185,
+            699,
+            1215,
+            714,
+            717,
+            1229,
+            211,
+            729,
+            221,
+            229,
+            747,
+            235,
+            237,
+            779,
+            276,
+            805,
+            299,
+            829,
+            852,
+            347,
+            371,
+            382,
+            896,
+            392,
+            926,
+            937,
+            428,
+            429,
+            961,
+            452,
+            979,
+            980,
+            982,
+            475,
+            480,
+            993,
+            1001,
+            502,
+            1018,
+        }
+        # `distractors` is defined as in the paper "Opening up Open-World Tracking"
+        self.distractors = {
+            20,
+            63,
+            108,
+            180,
+            188,
+            204,
+            212,
+            247,
+            303,
+            403,
+            407,
+            415,
+            490,
+            504,
+            507,
+            513,
+            529,
+            567,
+            569,
+            588,
+            672,
+            691,
+            702,
+            708,
+            711,
+            720,
+            736,
+            737,
+            798,
+            813,
+            815,
+            827,
+            831,
+            851,
+            877,
+            883,
+            912,
+            971,
+            976,
+            1130,
+            1133,
+            1134,
+            1169,
+            1184,
+            1220,
+        }
+        self.unknowns = all_ids.difference(self.knowns.union(self.distractors))
+
+    def _filter_gt_data(self, raw_gt_data):
+        """
+        Filter out irrelevant data in the raw_gt_data
+        Args:
+            raw_gt_data: directly loaded from json.
+
+        Returns:
+            filtered gt_data
+        """
+        valid_cat_ids = list()
+        if self.subset == "known":
+            valid_cat_ids = self.knowns
+        elif self.subset == "distractor":
+            valid_cat_ids = self.distractors
+        elif self.subset == "unknown":
+            valid_cat_ids = self.unknowns
+        # elif self.subset == "test_only_unknowns":
+        #     valid_cat_ids = test_only_unknowns
+        else:
+            raise Exception("The parameter `SUBSET` is incorrect")
+
+        filtered = dict()
+        filtered["videos"] = raw_gt_data["videos"]
+        # filtered["videos"] = list()
+        unwanted_vid = set()
+        # for video in raw_gt_data["videos"]:
+        #     datasrc = video["name"].split('/')[1]
+        #     if datasrc in data_srcs:
+        #         filtered["videos"].append(video)
+        #     else:
+        #         unwanted_vid.add(video["id"])
+
+        filtered["annotations"] = list()
+        for ann in raw_gt_data["annotations"]:
+            if (ann["video_id"] not in unwanted_vid) and (
+                ann["category_id"] in valid_cat_ids
+            ):
+                filtered["annotations"].append(ann)
+
+        filtered["tracks"] = list()
+        for track in raw_gt_data["tracks"]:
+            if (track["video_id"] not in unwanted_vid) and (
+                track["category_id"] in valid_cat_ids
+            ):
+                filtered["tracks"].append(track)
+
+        filtered["images"] = list()
+        for image in raw_gt_data["images"]:
+            if image["video_id"] not in unwanted_vid:
+                filtered["images"].append(image)
+
+        filtered["categories"] = list()
+        for cat in raw_gt_data["categories"]:
+            if cat["id"] in valid_cat_ids:
+                filtered["categories"].append(cat)
+
+        filtered["info"] = raw_gt_data["info"]
+        filtered["licenses"] = raw_gt_data["licenses"]
+
+        return filtered

sam3/eval/hota_eval_toolkit/trackeval/datasets/youtube_vis.py

@@ -0,0 +1,524 @@
+# flake8: noqa
+
+# note: this file has been modified from its original version in TrackEval in
+# https://github.com/JonathonLuiten/TrackEval/blob/master/trackeval/datasets/youtube_vis.py
+# to support the following:
+# 1) bbox evaluation (via `IOU_TYPE`)
+# 2) passing GT and prediction data as Python objects (via `GT_JSON_OBJECT` and `TRACKER_JSON_OBJECT`)
+# 3) specifying a custom dataset name (via `DATASET_NAME`)
+
+import json
+import os
+
+import numpy as np
+
+from .. import _timing, utils
+from ..utils import TrackEvalException
+from ._base_dataset import _BaseDataset
+
+
+class YouTubeVIS(_BaseDataset):
+    """Dataset class for YouTubeVIS tracking"""
+
+    @staticmethod
+    def get_default_dataset_config():
+        """Default class config values"""
+        code_path = utils.get_code_path()
+        default_config = {
+            "GT_FOLDER": os.path.join(
+                code_path, "data/gt/youtube_vis/"
+            ),  # Location of GT data
+            "TRACKERS_FOLDER": os.path.join(code_path, "data/trackers/youtube_vis/"),
+            # Trackers location
+            "OUTPUT_FOLDER": None,  # Where to save eval results (if None, same as TRACKERS_FOLDER)
+            "TRACKERS_TO_EVAL": None,  # Filenames of trackers to eval (if None, all in folder)
+            "CLASSES_TO_EVAL": None,  # Classes to eval (if None, all classes)
+            "SPLIT_TO_EVAL": "train_sub_split",  # Valid: 'train', 'val', 'train_sub_split'
+            "PRINT_CONFIG": True,  # Whether to print current config
+            "OUTPUT_SUB_FOLDER": "",  # Output files are saved in OUTPUT_FOLDER/tracker_name/OUTPUT_SUB_FOLDER
+            "TRACKER_SUB_FOLDER": "data",  # Tracker files are in TRACKER_FOLDER/tracker_name/TRACKER_SUB_FOLDER
+            "TRACKER_DISPLAY_NAMES": None,  # Names of trackers to display, if None: TRACKERS_TO_EVAL
+            # Added for video phrase AP evaluation -- allow directly specifying the GT JSON data and Tracker (result)
+            # JSON data as Python objects, without reading from files.
+            "GT_JSON_OBJECT": None,
+            "TRACKER_JSON_OBJECT": None,
+            "IOU_TYPE": "segm",
+            "DATASET_NAME": "video",
+        }
+        return default_config
+
+    def __init__(self, config=None):
+        """Initialise dataset, checking that all required files are present"""
+        super().__init__()
+        # Fill non-given config values with defaults
+        self.config = utils.init_config(config, self.get_default_dataset_config())
+        self.gt_fol = (
+            self.config["GT_FOLDER"] + "youtube_vis_" + self.config["SPLIT_TO_EVAL"]
+        )
+        self.tracker_fol = (
+            self.config["TRACKERS_FOLDER"]
+            + "youtube_vis_"
+            + self.config["SPLIT_TO_EVAL"]
+        )
+        self.use_super_categories = False
+        self.should_classes_combine = True
+        assert self.config["IOU_TYPE"] in ["segm", "bbox"]
+        self.iou_type = self.config["IOU_TYPE"]
+        print("=" * 100)
+        print(f"Evaluate annotation type *{self.iou_type}*")
+        self.dataset_name = self.config["DATASET_NAME"]
+
+        self.output_fol = self.config["OUTPUT_FOLDER"]
+        if self.output_fol is None:
+            self.output_fol = self.tracker_fol
+        self.output_sub_fol = self.config["OUTPUT_SUB_FOLDER"]
+        self.tracker_sub_fol = self.config["TRACKER_SUB_FOLDER"]
+
+        if self.config["GT_JSON_OBJECT"] is not None:
+            # allow directly specifying the GT JSON data without reading from files
+            gt_json = self.config["GT_JSON_OBJECT"]
+            assert isinstance(gt_json, dict)
+            assert "videos" in gt_json
+            assert "categories" in gt_json
+            assert "annotations" in gt_json
+            self.gt_data = gt_json
+        else:
+            if not os.path.exists(self.gt_fol):
+                print("GT folder not found: " + self.gt_fol)
+                raise TrackEvalException(
+                    "GT folder not found: " + os.path.basename(self.gt_fol)
+                )
+            gt_dir_files = [
+                file for file in os.listdir(self.gt_fol) if file.endswith(".json")
+            ]
+            if len(gt_dir_files) != 1:
+                raise TrackEvalException(
+                    self.gt_fol + " does not contain exactly one json file."
+                )
+
+            with open(os.path.join(self.gt_fol, gt_dir_files[0])) as f:
+                self.gt_data = json.load(f)
+
+        # Get classes to eval
+        self.valid_classes = [cls["name"] for cls in self.gt_data["categories"]]
+        cls_name_to_cls_id_map = {
+            cls["name"]: cls["id"] for cls in self.gt_data["categories"]
+        }
+
+        if self.config["CLASSES_TO_EVAL"]:
+            self.class_list = [
+                cls.lower() if cls.lower() in self.valid_classes else None
+                for cls in self.config["CLASSES_TO_EVAL"]
+            ]
+            if not all(self.class_list):
+                raise TrackEvalException(
+                    "Attempted to evaluate an invalid class. Only classes "
+                    + ", ".join(self.valid_classes)
+                    + " are valid."
+                )
+        else:
+            self.class_list = [cls["name"] for cls in self.gt_data["categories"]]
+        self.class_name_to_class_id = {
+            k: v for k, v in cls_name_to_cls_id_map.items() if k in self.class_list
+        }
+
+        # Get sequences to eval and check gt files exist
+        self.seq_list = [
+            vid["file_names"][0].split("/")[0] for vid in self.gt_data["videos"]
+        ]
+        self.seq_name_to_seq_id = {
+            vid["file_names"][0].split("/")[0]: vid["id"]
+            for vid in self.gt_data["videos"]
+        }
+        self.seq_lengths = {
+            vid["id"]: len(vid["file_names"]) for vid in self.gt_data["videos"]
+        }
+
+        # encode masks and compute track areas
+        self._prepare_gt_annotations()
+
+        # Get trackers to eval
+        if self.config["TRACKER_JSON_OBJECT"] is not None:
+            # allow directly specifying the tracker JSON data without reading from files
+            tracker_json = self.config["TRACKER_JSON_OBJECT"]
+            assert isinstance(tracker_json, list)
+            self.tracker_list = ["tracker"]
+        elif self.config["TRACKERS_TO_EVAL"] is None:
+            self.tracker_list = os.listdir(self.tracker_fol)
+        else:
+            self.tracker_list = self.config["TRACKERS_TO_EVAL"]
+
+        if self.config["TRACKER_DISPLAY_NAMES"] is None:
+            self.tracker_to_disp = dict(zip(self.tracker_list, self.tracker_list))
+        elif (self.config["TRACKERS_TO_EVAL"] is not None) and (
+            len(self.config["TRACKER_DISPLAY_NAMES"]) == len(self.tracker_list)
+        ):
+            self.tracker_to_disp = dict(
+                zip(self.tracker_list, self.config["TRACKER_DISPLAY_NAMES"])
+            )
+        else:
+            raise TrackEvalException(
+                "List of tracker files and tracker display names do not match."
+            )
+
+        # counter for globally unique track IDs
+        self.global_tid_counter = 0
+
+        self.tracker_data = dict()
+        if self.config["TRACKER_JSON_OBJECT"] is not None:
+            # allow directly specifying the tracker JSON data without reading from files
+            tracker = self.tracker_list[0]
+            self.tracker_data[tracker] = tracker_json
+        else:
+            for tracker in self.tracker_list:
+                tracker_dir_path = os.path.join(
+                    self.tracker_fol, tracker, self.tracker_sub_fol
+                )
+                tr_dir_files = [
+                    file
+                    for file in os.listdir(tracker_dir_path)
+                    if file.endswith(".json")
+                ]
+                if len(tr_dir_files) != 1:
+                    raise TrackEvalException(
+                        tracker_dir_path + " does not contain exactly one json file."
+                    )
+
+                with open(os.path.join(tracker_dir_path, tr_dir_files[0])) as f:
+                    curr_data = json.load(f)
+
+                self.tracker_data[tracker] = curr_data
+
+    def get_display_name(self, tracker):
+        return self.tracker_to_disp[tracker]
+
+    def _load_raw_file(self, tracker, seq, is_gt):
+        """Load a file (gt or tracker) in the YouTubeVIS format
+        If is_gt, this returns a dict which contains the fields:
+        [gt_ids, gt_classes] : list (for each timestep) of 1D NDArrays (for each det).
+        [gt_dets]: list (for each timestep) of lists of detections.
+        [classes_to_gt_tracks]: dictionary with class values as keys and list of dictionaries (with frame indices as
+                                keys and corresponding segmentations as values) for each track
+        [classes_to_gt_track_ids, classes_to_gt_track_areas, classes_to_gt_track_iscrowd]: dictionary with class values
+                                as keys and lists (for each track) as values
+
+        if not is_gt, this returns a dict which contains the fields:
+        [tracker_ids, tracker_classes, tracker_confidences] : list (for each timestep) of 1D NDArrays (for each det).
+        [tracker_dets]: list (for each timestep) of lists of detections.
+        [classes_to_dt_tracks]: dictionary with class values as keys and list of dictionaries (with frame indices as
+                                keys and corresponding segmentations as values) for each track
+        [classes_to_dt_track_ids, classes_to_dt_track_areas]: dictionary with class values as keys and lists as values
+        [classes_to_dt_track_scores]: dictionary with class values as keys and 1D numpy arrays as values
+        """
+        # select sequence tracks
+        seq_id = self.seq_name_to_seq_id[seq]
+        if is_gt:
+            tracks = [
+                ann for ann in self.gt_data["annotations"] if ann["video_id"] == seq_id
+            ]
+        else:
+            tracks = self._get_tracker_seq_tracks(tracker, seq_id)
+
+        # Convert data to required format
+        num_timesteps = self.seq_lengths[seq_id]
+        data_keys = ["ids", "classes", "dets"]
+        if not is_gt:
+            data_keys += ["tracker_confidences"]
+        raw_data = {key: [None] * num_timesteps for key in data_keys}
+        result_key = "segmentations" if self.iou_type == "segm" else "bboxes"
+        for t in range(num_timesteps):
+            raw_data["dets"][t] = [
+                track[result_key][t] for track in tracks if track[result_key][t]
+            ]
+            raw_data["ids"][t] = np.atleast_1d(
+                [track["id"] for track in tracks if track[result_key][t]]
+            ).astype(int)
+            raw_data["classes"][t] = np.atleast_1d(
+                [track["category_id"] for track in tracks if track[result_key][t]]
+            ).astype(int)
+            if not is_gt:
+                raw_data["tracker_confidences"][t] = np.atleast_1d(
+                    [track["score"] for track in tracks if track[result_key][t]]
+                ).astype(float)
+
+        if is_gt:
+            key_map = {"ids": "gt_ids", "classes": "gt_classes", "dets": "gt_dets"}
+        else:
+            key_map = {
+                "ids": "tracker_ids",
+                "classes": "tracker_classes",
+                "dets": "tracker_dets",
+            }
+        for k, v in key_map.items():
+            raw_data[v] = raw_data.pop(k)
+
+        all_cls_ids = {self.class_name_to_class_id[cls] for cls in self.class_list}
+        classes_to_tracks = {
+            cls: [track for track in tracks if track["category_id"] == cls]
+            for cls in all_cls_ids
+        }
+
+        # mapping from classes to track representations and track information
+        raw_data["classes_to_tracks"] = {
+            cls: [
+                {i: track[result_key][i] for i in range(len(track[result_key]))}
+                for track in tracks
+            ]
+            for cls, tracks in classes_to_tracks.items()
+        }
+        raw_data["classes_to_track_ids"] = {
+            cls: [track["id"] for track in tracks]
+            for cls, tracks in classes_to_tracks.items()
+        }
+        raw_data["classes_to_track_areas"] = {
+            cls: [track["area"] for track in tracks]
+            for cls, tracks in classes_to_tracks.items()
+        }
+
+        if is_gt:
+            raw_data["classes_to_gt_track_iscrowd"] = {
+                cls: [track["iscrowd"] for track in tracks]
+                for cls, tracks in classes_to_tracks.items()
+            }
+        else:
+            raw_data["classes_to_dt_track_scores"] = {
+                cls: np.array([track["score"] for track in tracks])
+                for cls, tracks in classes_to_tracks.items()
+            }
+
+        if is_gt:
+            key_map = {
+                "classes_to_tracks": "classes_to_gt_tracks",
+                "classes_to_track_ids": "classes_to_gt_track_ids",
+                "classes_to_track_areas": "classes_to_gt_track_areas",
+            }
+        else:
+            key_map = {
+                "classes_to_tracks": "classes_to_dt_tracks",
+                "classes_to_track_ids": "classes_to_dt_track_ids",
+                "classes_to_track_areas": "classes_to_dt_track_areas",
+            }
+        for k, v in key_map.items():
+            raw_data[v] = raw_data.pop(k)
+
+        raw_data["num_timesteps"] = num_timesteps
+        raw_data["seq"] = seq
+        return raw_data
+
+    @_timing.time
+    def get_preprocessed_seq_data(self, raw_data, cls):
+        """Preprocess data for a single sequence for a single class ready for evaluation.
+        Inputs:
+             - raw_data is a dict containing the data for the sequence already read in by get_raw_seq_data().
+             - cls is the class to be evaluated.
+        Outputs:
+             - data is a dict containing all of the information that metrics need to perform evaluation.
+                It contains the following fields:
+                    [num_timesteps, num_gt_ids, num_tracker_ids, num_gt_dets, num_tracker_dets] : integers.
+                    [gt_ids, tracker_ids, tracker_confidences]: list (for each timestep) of 1D NDArrays (for each det).
+                    [gt_dets, tracker_dets]: list (for each timestep) of lists of detections.
+                    [similarity_scores]: list (for each timestep) of 2D NDArrays.
+        Notes:
+            General preprocessing (preproc) occurs in 4 steps. Some datasets may not use all of these steps.
+                1) Extract only detections relevant for the class to be evaluated (including distractor detections).
+                2) Match gt dets and tracker dets. Remove tracker dets that are matched to a gt det that is of a
+                    distractor class, or otherwise marked as to be removed.
+                3) Remove unmatched tracker dets if they fall within a crowd ignore region or don't meet a certain
+                    other criteria (e.g. are too small).
+                4) Remove gt dets that were only useful for preprocessing and not for actual evaluation.
+            After the above preprocessing steps, this function also calculates the number of gt and tracker detections
+                and unique track ids. It also relabels gt and tracker ids to be contiguous and checks that ids are
+                unique within each timestep.
+        YouTubeVIS:
+            In YouTubeVIS, the 4 preproc steps are as follow:
+                1) There are 40 classes which are evaluated separately.
+                2) No matched tracker dets are removed.
+                3) No unmatched tracker dets are removed.
+                4) No gt dets are removed.
+            Further, for TrackMAP computation track representations for the given class are accessed from a dictionary
+            and the tracks from the tracker data are sorted according to the tracker confidence.
+        """
+        cls_id = self.class_name_to_class_id[cls]
+
+        data_keys = [
+            "gt_ids",
+            "tracker_ids",
+            "gt_dets",
+            "tracker_dets",
+            "similarity_scores",
+        ]
+        data = {key: [None] * raw_data["num_timesteps"] for key in data_keys}
+        unique_gt_ids = []
+        unique_tracker_ids = []
+        num_gt_dets = 0
+        num_tracker_dets = 0
+
+        for t in range(raw_data["num_timesteps"]):
+            # Only extract relevant dets for this class for eval (cls)
+            gt_class_mask = np.atleast_1d(raw_data["gt_classes"][t] == cls_id)
+            gt_class_mask = gt_class_mask.astype(bool)
+            gt_ids = raw_data["gt_ids"][t][gt_class_mask]
+            gt_dets = [
+                raw_data["gt_dets"][t][ind]
+                for ind in range(len(gt_class_mask))
+                if gt_class_mask[ind]
+            ]
+
+            tracker_class_mask = np.atleast_1d(raw_data["tracker_classes"][t] == cls_id)
+            tracker_class_mask = tracker_class_mask.astype(bool)
+            tracker_ids = raw_data["tracker_ids"][t][tracker_class_mask]
+            tracker_dets = [
+                raw_data["tracker_dets"][t][ind]
+                for ind in range(len(tracker_class_mask))
+                if tracker_class_mask[ind]
+            ]
+            similarity_scores = raw_data["similarity_scores"][t][gt_class_mask, :][
+                :, tracker_class_mask
+            ]
+
+            data["tracker_ids"][t] = tracker_ids
+            data["tracker_dets"][t] = tracker_dets
+            data["gt_ids"][t] = gt_ids
+            data["gt_dets"][t] = gt_dets
+            data["similarity_scores"][t] = similarity_scores
+
+            unique_gt_ids += list(np.unique(data["gt_ids"][t]))
+            unique_tracker_ids += list(np.unique(data["tracker_ids"][t]))
+            num_tracker_dets += len(data["tracker_ids"][t])
+            num_gt_dets += len(data["gt_ids"][t])
+
+        # Re-label IDs such that there are no empty IDs
+        if len(unique_gt_ids) > 0:
+            unique_gt_ids = np.unique(unique_gt_ids)
+            gt_id_map = np.nan * np.ones((np.max(unique_gt_ids) + 1))
+            gt_id_map[unique_gt_ids] = np.arange(len(unique_gt_ids))
+            for t in range(raw_data["num_timesteps"]):
+                if len(data["gt_ids"][t]) > 0:
+                    data["gt_ids"][t] = gt_id_map[data["gt_ids"][t]].astype(int)
+        if len(unique_tracker_ids) > 0:
+            unique_tracker_ids = np.unique(unique_tracker_ids)
+            tracker_id_map = np.nan * np.ones((np.max(unique_tracker_ids) + 1))
+            tracker_id_map[unique_tracker_ids] = np.arange(len(unique_tracker_ids))
+            for t in range(raw_data["num_timesteps"]):
+                if len(data["tracker_ids"][t]) > 0:
+                    data["tracker_ids"][t] = tracker_id_map[
+                        data["tracker_ids"][t]
+                    ].astype(int)
+
+        # Ensure that ids are unique per timestep.
+        self._check_unique_ids(data)
+
+        # Record overview statistics.
+        data["num_tracker_dets"] = num_tracker_dets
+        data["num_gt_dets"] = num_gt_dets
+        data["num_tracker_ids"] = len(unique_tracker_ids)
+        data["num_gt_ids"] = len(unique_gt_ids)
+        data["num_timesteps"] = raw_data["num_timesteps"]
+        data["seq"] = raw_data["seq"]
+
+        # get track representations
+        data["gt_tracks"] = raw_data["classes_to_gt_tracks"][cls_id]
+        data["gt_track_ids"] = raw_data["classes_to_gt_track_ids"][cls_id]
+        data["gt_track_areas"] = raw_data["classes_to_gt_track_areas"][cls_id]
+        data["gt_track_iscrowd"] = raw_data["classes_to_gt_track_iscrowd"][cls_id]
+        data["dt_tracks"] = raw_data["classes_to_dt_tracks"][cls_id]
+        data["dt_track_ids"] = raw_data["classes_to_dt_track_ids"][cls_id]
+        data["dt_track_areas"] = raw_data["classes_to_dt_track_areas"][cls_id]
+        data["dt_track_scores"] = raw_data["classes_to_dt_track_scores"][cls_id]
+        data["iou_type"] = "mask"
+
+        # sort tracker data tracks by tracker confidence scores
+        if data["dt_tracks"]:
+            idx = np.argsort(
+                [-score for score in data["dt_track_scores"]], kind="mergesort"
+            )
+            data["dt_track_scores"] = [data["dt_track_scores"][i] for i in idx]
+            data["dt_tracks"] = [data["dt_tracks"][i] for i in idx]
+            data["dt_track_ids"] = [data["dt_track_ids"][i] for i in idx]
+            data["dt_track_areas"] = [data["dt_track_areas"][i] for i in idx]
+
+        return data
+
+    def _calculate_similarities(self, gt_dets_t, tracker_dets_t):
+        if self.iou_type == "segm":
+            similarity_scores = self._calculate_mask_ious(
+                gt_dets_t, tracker_dets_t, is_encoded=True, do_ioa=False
+            )
+        else:
+            gt_dets_t = np.array(gt_dets_t, dtype=np.float32).reshape(-1, 4)
+            tracker_dets_t = np.array(tracker_dets_t, dtype=np.float32).reshape(-1, 4)
+            similarity_scores = self._calculate_box_ious(
+                gt_dets_t, tracker_dets_t, box_format="xywh", do_ioa=False
+            )
+        return similarity_scores
+
+    def _prepare_gt_annotations(self):
+        """
+        Prepares GT data by rle encoding segmentations and computing the average track area.
+        :return: None
+        """
+        if self.iou_type == "segm":
+            # only loaded when needed to reduce minimum requirements
+            from pycocotools import mask as mask_utils
+
+            for track in self.gt_data["annotations"]:
+                h = track["height"]
+                w = track["width"]
+                for i, seg in enumerate(track["segmentations"]):
+                    if seg is not None and isinstance(seg["counts"], list):
+                        track["segmentations"][i] = mask_utils.frPyObjects(seg, h, w)
+                areas = [a for a in track["areas"] if a]
+                if len(areas) == 0:
+                    track["area"] = 0
+                else:
+                    track["area"] = np.array(areas).mean()
+        else:
+            for track in self.gt_data["annotations"]:
+                # For bbox eval, compute areas from bboxes if not already available
+                areas = [a for a in track.get("areas", []) if a]
+                if not areas:
+                    areas = []
+                    for bbox in track.get("bboxes", []):
+                        if bbox is not None:
+                            areas.append(bbox[2] * bbox[3])
+                track["area"] = np.array(areas).mean() if areas else 0
+
+    def _get_tracker_seq_tracks(self, tracker, seq_id):
+        """
+        Prepares tracker data for a given sequence. Extracts all annotations for given sequence ID, computes
+        average track area and assigns a track ID.
+        :param tracker: the given tracker
+        :param seq_id: the sequence ID
+        :return: the extracted tracks
+        """
+        # only loaded when needed to reduce minimum requirements
+        from pycocotools import mask as mask_utils
+
+        tracks = [
+            ann for ann in self.tracker_data[tracker] if ann["video_id"] == seq_id
+        ]
+        for track in tracks:
+            if "areas" not in track:
+                if self.iou_type == "segm":
+                    for seg in track["segmentations"]:
+                        if seg:
+                            track["areas"].append(mask_utils.area(seg))
+                        else:
+                            track["areas"].append(None)
+                else:
+                    for bbox in track["bboxes"]:
+                        if bbox:
+                            track["areas"].append(bbox[2] * bbox[3])
+                        else:
+                            track["areas"].append(None)
+            areas = [a for a in track["areas"] if a]
+            if len(areas) == 0:
+                track["area"] = 0
+            else:
+                track["area"] = np.array(areas).mean()
+            track["id"] = self.global_tid_counter
+            self.global_tid_counter += 1
+        return tracks
+
+    def get_name(self):
+        return self.dataset_name

sam3/eval/hota_eval_toolkit/trackeval/eval.py

@@ -0,0 +1,395 @@
+# flake8: noqa
+
+import os
+import time
+import traceback
+from functools import partial
+from multiprocessing.pool import Pool
+
+import numpy as np
+
+from . import _timing, utils
+from .metrics import Count
+from .utils import TrackEvalException
+
+try:
+    import tqdm
+
+    TQDM_IMPORTED = True
+except ImportError as _:
+    TQDM_IMPORTED = False
+
+
+class Evaluator:
+    """Evaluator class for evaluating different metrics for different datasets"""
+
+    @staticmethod
+    def get_default_eval_config():
+        """Returns the default config values for evaluation"""
+        code_path = utils.get_code_path()
+        default_config = {
+            "USE_PARALLEL": False,
+            "NUM_PARALLEL_CORES": 8,
+            "BREAK_ON_ERROR": True,  # Raises exception and exits with error
+            "RETURN_ON_ERROR": False,  # if not BREAK_ON_ERROR, then returns from function on error
+            "LOG_ON_ERROR": os.path.join(
+                code_path, "error_log.txt"
+            ),  # if not None, save any errors into a log file.
+            "PRINT_RESULTS": True,
+            "PRINT_ONLY_COMBINED": False,
+            "PRINT_CONFIG": True,
+            "TIME_PROGRESS": True,
+            "DISPLAY_LESS_PROGRESS": True,
+            "OUTPUT_SUMMARY": True,
+            "OUTPUT_EMPTY_CLASSES": True,  # If False, summary files are not output for classes with no detections
+            "OUTPUT_DETAILED": True,
+            "PLOT_CURVES": True,
+        }
+        return default_config
+
+    def __init__(self, config=None):
+        """Initialise the evaluator with a config file"""
+        self.config = utils.init_config(config, self.get_default_eval_config(), "Eval")
+        # Only run timing analysis if not run in parallel.
+        if self.config["TIME_PROGRESS"] and not self.config["USE_PARALLEL"]:
+            _timing.DO_TIMING = True
+            if self.config["DISPLAY_LESS_PROGRESS"]:
+                _timing.DISPLAY_LESS_PROGRESS = True
+
+    def _combine_results(
+        self,
+        res,
+        metrics_list,
+        metric_names,
+        dataset,
+        res_field="COMBINED_SEQ",
+        target_tag=None,
+    ):
+        assert res_field.startswith("COMBINED_SEQ")
+        # collecting combined cls keys (cls averaged, det averaged, super classes)
+        tracker_list, seq_list, class_list = dataset.get_eval_info()
+        combined_cls_keys = []
+        res[res_field] = {}
+
+        # narrow the target for evaluation
+        if target_tag is not None:
+            target_video_ids = [
+                annot["video_id"]
+                for annot in dataset.gt_data["annotations"]
+                if target_tag in annot["tags"]
+            ]
+            vid2name = {
+                video["id"]: video["file_names"][0].split("/")[0]
+                for video in dataset.gt_data["videos"]
+            }
+            target_video_ids = set(target_video_ids)
+            target_video = [vid2name[video_id] for video_id in target_video_ids]
+
+            if len(target_video) == 0:
+                raise TrackEvalException(
+                    "No sequences found with the tag %s" % target_tag
+                )
+
+            target_annotations = [
+                annot
+                for annot in dataset.gt_data["annotations"]
+                if annot["video_id"] in target_video_ids
+            ]
+            assert all(target_tag in annot["tags"] for annot in target_annotations), (
+                f"Not all annotations in the target sequences have the target tag {target_tag}. "
+                "We currently only support a target tag at the sequence level, not at the annotation level."
+            )
+        else:
+            target_video = seq_list
+
+        # combine sequences for each class
+        for c_cls in class_list:
+            res[res_field][c_cls] = {}
+            for metric, metric_name in zip(metrics_list, metric_names):
+                curr_res = {
+                    seq_key: seq_value[c_cls][metric_name]
+                    for seq_key, seq_value in res.items()
+                    if not seq_key.startswith("COMBINED_SEQ")
+                    and seq_key in target_video
+                }
+                res[res_field][c_cls][metric_name] = metric.combine_sequences(curr_res)
+        # combine classes
+        if dataset.should_classes_combine:
+            combined_cls_keys += [
+                "cls_comb_cls_av",
+                "cls_comb_det_av",
+                "all",
+            ]
+            res[res_field]["cls_comb_cls_av"] = {}
+            res[res_field]["cls_comb_det_av"] = {}
+            for metric, metric_name in zip(metrics_list, metric_names):
+                cls_res = {
+                    cls_key: cls_value[metric_name]
+                    for cls_key, cls_value in res[res_field].items()
+                    if cls_key not in combined_cls_keys
+                }
+                res[res_field]["cls_comb_cls_av"][metric_name] = (
+                    metric.combine_classes_class_averaged(cls_res)
+                )
+                res[res_field]["cls_comb_det_av"][metric_name] = (
+                    metric.combine_classes_det_averaged(cls_res)
+                )
+        # combine classes to super classes
+        if dataset.use_super_categories:
+            for cat, sub_cats in dataset.super_categories.items():
+                combined_cls_keys.append(cat)
+                res[res_field][cat] = {}
+                for metric, metric_name in zip(metrics_list, metric_names):
+                    cat_res = {
+                        cls_key: cls_value[metric_name]
+                        for cls_key, cls_value in res[res_field].items()
+                        if cls_key in sub_cats
+                    }
+                    res[res_field][cat][metric_name] = (
+                        metric.combine_classes_det_averaged(cat_res)
+                    )
+        return res, combined_cls_keys
+
+    def _summarize_results(
+        self,
+        res,
+        tracker,
+        metrics_list,
+        metric_names,
+        dataset,
+        res_field,
+        combined_cls_keys,
+    ):
+        config = self.config
+        output_fol = dataset.get_output_fol(tracker)
+        tracker_display_name = dataset.get_display_name(tracker)
+        for c_cls in res[
+            res_field
+        ].keys():  # class_list + combined classes if calculated
+            summaries = []
+            details = []
+            num_dets = res[res_field][c_cls]["Count"]["Dets"]
+            if config["OUTPUT_EMPTY_CLASSES"] or num_dets > 0:
+                for metric, metric_name in zip(metrics_list, metric_names):
+                    # for combined classes there is no per sequence evaluation
+                    if c_cls in combined_cls_keys:
+                        table_res = {res_field: res[res_field][c_cls][metric_name]}
+                    else:
+                        table_res = {
+                            seq_key: seq_value[c_cls][metric_name]
+                            for seq_key, seq_value in res.items()
+                        }
+
+                    if config["PRINT_RESULTS"] and config["PRINT_ONLY_COMBINED"]:
+                        dont_print = (
+                            dataset.should_classes_combine
+                            and c_cls not in combined_cls_keys
+                        )
+                        if not dont_print:
+                            metric.print_table(
+                                {res_field: table_res[res_field]},
+                                tracker_display_name,
+                                c_cls,
+                                res_field,
+                                res_field,
+                            )
+                    elif config["PRINT_RESULTS"]:
+                        metric.print_table(
+                            table_res, tracker_display_name, c_cls, res_field, res_field
+                        )
+                    if config["OUTPUT_SUMMARY"]:
+                        summaries.append(metric.summary_results(table_res))
+                    if config["OUTPUT_DETAILED"]:
+                        details.append(metric.detailed_results(table_res))
+                    if config["PLOT_CURVES"]:
+                        metric.plot_single_tracker_results(
+                            table_res,
+                            tracker_display_name,
+                            c_cls,
+                            output_fol,
+                        )
+                if config["OUTPUT_SUMMARY"]:
+                    utils.write_summary_results(summaries, c_cls, output_fol)
+                if config["OUTPUT_DETAILED"]:
+                    utils.write_detailed_results(details, c_cls, output_fol)
+
+    @_timing.time
+    def evaluate(self, dataset_list, metrics_list, show_progressbar=False):
+        """Evaluate a set of metrics on a set of datasets"""
+        config = self.config
+        metrics_list = metrics_list + [Count()]  # Count metrics are always run
+        metric_names = utils.validate_metrics_list(metrics_list)
+        dataset_names = [dataset.get_name() for dataset in dataset_list]
+        output_res = {}
+        output_msg = {}
+
+        for dataset, dataset_name in zip(dataset_list, dataset_names):
+            # Get dataset info about what to evaluate
+            output_res[dataset_name] = {}
+            output_msg[dataset_name] = {}
+            tracker_list, seq_list, class_list = dataset.get_eval_info()
+            print(
+                "\nEvaluating %i tracker(s) on %i sequence(s) for %i class(es) on %s dataset using the following "
+                "metrics: %s\n"
+                % (
+                    len(tracker_list),
+                    len(seq_list),
+                    len(class_list),
+                    dataset_name,
+                    ", ".join(metric_names),
+                )
+            )
+
+            # Evaluate each tracker
+            for tracker in tracker_list:
+                # if not config['BREAK_ON_ERROR'] then go to next tracker without breaking
+                try:
+                    # Evaluate each sequence in parallel or in series.
+                    # returns a nested dict (res), indexed like: res[seq][class][metric_name][sub_metric field]
+                    # e.g. res[seq_0001][pedestrian][hota][DetA]
+                    print("\nEvaluating %s\n" % tracker)
+                    time_start = time.time()
+                    if config["USE_PARALLEL"]:
+                        if show_progressbar and TQDM_IMPORTED:
+                            seq_list_sorted = sorted(seq_list)
+
+                            with Pool(config["NUM_PARALLEL_CORES"]) as pool, tqdm.tqdm(
+                                total=len(seq_list)
+                            ) as pbar:
+                                _eval_sequence = partial(
+                                    eval_sequence,
+                                    dataset=dataset,
+                                    tracker=tracker,
+                                    class_list=class_list,
+                                    metrics_list=metrics_list,
+                                    metric_names=metric_names,
+                                )
+                                results = []
+                                for r in pool.imap(
+                                    _eval_sequence, seq_list_sorted, chunksize=20
+                                ):
+                                    results.append(r)
+                                    pbar.update()
+                                res = dict(zip(seq_list_sorted, results))
+
+                        else:
+                            with Pool(config["NUM_PARALLEL_CORES"]) as pool:
+                                _eval_sequence = partial(
+                                    eval_sequence,
+                                    dataset=dataset,
+                                    tracker=tracker,
+                                    class_list=class_list,
+                                    metrics_list=metrics_list,
+                                    metric_names=metric_names,
+                                )
+                                results = pool.map(_eval_sequence, seq_list)
+                                res = dict(zip(seq_list, results))
+                    else:
+                        res = {}
+                        if show_progressbar and TQDM_IMPORTED:
+                            seq_list_sorted = sorted(seq_list)
+                            for curr_seq in tqdm.tqdm(seq_list_sorted):
+                                res[curr_seq] = eval_sequence(
+                                    curr_seq,
+                                    dataset,
+                                    tracker,
+                                    class_list,
+                                    metrics_list,
+                                    metric_names,
+                                )
+                        else:
+                            for curr_seq in sorted(seq_list):
+                                res[curr_seq] = eval_sequence(
+                                    curr_seq,
+                                    dataset,
+                                    tracker,
+                                    class_list,
+                                    metrics_list,
+                                    metric_names,
+                                )
+
+                    # Combine results over all sequences and then over all classes
+                    res, combined_cls_keys = self._combine_results(
+                        res, metrics_list, metric_names, dataset, "COMBINED_SEQ"
+                    )
+
+                    if np.all(
+                        ["tags" in annot for annot in dataset.gt_data["annotations"]]
+                    ):
+                        # Combine results over the challenging sequences and then over all classes
+                        # currently only support "tracking_challenging_pair"
+                        res, _ = self._combine_results(
+                            res,
+                            metrics_list,
+                            metric_names,
+                            dataset,
+                            "COMBINED_SEQ_CHALLENGING",
+                            "tracking_challenging_pair",
+                        )
+
+                    # Print and output results in various formats
+                    if config["TIME_PROGRESS"]:
+                        print(
+                            "\nAll sequences for %s finished in %.2f seconds"
+                            % (tracker, time.time() - time_start)
+                        )
+
+                    self._summarize_results(
+                        res,
+                        tracker,
+                        metrics_list,
+                        metric_names,
+                        dataset,
+                        "COMBINED_SEQ",
+                        combined_cls_keys,
+                    )
+                    if "COMBINED_SEQ_CHALLENGING" in res:
+                        self._summarize_results(
+                            res,
+                            tracker,
+                            metrics_list,
+                            metric_names,
+                            dataset,
+                            "COMBINED_SEQ_CHALLENGING",
+                            combined_cls_keys,
+                        )
+
+                    # Output for returning from function
+                    output_res[dataset_name][tracker] = res
+                    output_msg[dataset_name][tracker] = "Success"
+
+                except Exception as err:
+                    output_res[dataset_name][tracker] = None
+                    if type(err) == TrackEvalException:
+                        output_msg[dataset_name][tracker] = str(err)
+                    else:
+                        output_msg[dataset_name][tracker] = "Unknown error occurred."
+                    print("Tracker %s was unable to be evaluated." % tracker)
+                    print(err)
+                    traceback.print_exc()
+                    if config["LOG_ON_ERROR"] is not None:
+                        with open(config["LOG_ON_ERROR"], "a") as f:
+                            print(dataset_name, file=f)
+                            print(tracker, file=f)
+                            print(traceback.format_exc(), file=f)
+                            print("\n\n\n", file=f)
+                    if config["BREAK_ON_ERROR"]:
+                        raise err
+                    elif config["RETURN_ON_ERROR"]:
+                        return output_res, output_msg
+
+        return output_res, output_msg
+
+
+@_timing.time
+def eval_sequence(seq, dataset, tracker, class_list, metrics_list, metric_names):
+    """Function for evaluating a single sequence"""
+
+    raw_data = dataset.get_raw_seq_data(tracker, seq)
+    seq_res = {}
+    for cls in class_list:
+        seq_res[cls] = {}
+        data = dataset.get_preprocessed_seq_data(raw_data, cls)
+        for metric, met_name in zip(metrics_list, metric_names):
+            seq_res[cls][met_name] = metric.eval_sequence(data)
+    return seq_res

sam3/eval/hota_eval_toolkit/trackeval/metrics/__init__.py

@@ -0,0 +1,4 @@
+# flake8: noqa
+
+from .count import Count
+from .hota import HOTA

sam3/eval/hota_eval_toolkit/trackeval/metrics/_base_metric.py

@@ -0,0 +1,145 @@
+# flake8: noqa
+
+from abc import ABC, abstractmethod
+
+import numpy as np
+
+from .. import _timing
+from ..utils import TrackEvalException
+
+
+class _BaseMetric(ABC):
+    @abstractmethod
+    def __init__(self):
+        self.plottable = False
+        self.integer_fields = []
+        self.float_fields = []
+        self.array_labels = []
+        self.integer_array_fields = []
+        self.float_array_fields = []
+        self.fields = []
+        self.summary_fields = []
+        self.registered = False
+
+    #####################################################################
+    # Abstract functions for subclasses to implement
+
+    @_timing.time
+    @abstractmethod
+    def eval_sequence(self, data): ...
+
+    @abstractmethod
+    def combine_sequences(self, all_res): ...
+
+    @abstractmethod
+    def combine_classes_class_averaged(self, all_res, ignore_empty_classes=False): ...
+
+    @abstractmethod
+    def combine_classes_det_averaged(self, all_res): ...
+
+    def plot_single_tracker_results(self, all_res, tracker, output_folder, cls):
+        """Plot results of metrics, only valid for metrics with self.plottable"""
+        if self.plottable:
+            raise NotImplementedError(
+                "plot_results is not implemented for metric %s" % self.get_name()
+            )
+        else:
+            pass
+
+    #####################################################################
+    # Helper functions which are useful for all metrics:
+
+    @classmethod
+    def get_name(cls):
+        return cls.__name__
+
+    @staticmethod
+    def _combine_sum(all_res, field):
+        """Combine sequence results via sum"""
+        return sum([all_res[k][field] for k in all_res.keys()])
+
+    @staticmethod
+    def _combine_weighted_av(all_res, field, comb_res, weight_field):
+        """Combine sequence results via weighted average"""
+        return sum(
+            [all_res[k][field] * all_res[k][weight_field] for k in all_res.keys()]
+        ) / np.maximum(1.0, comb_res[weight_field])
+
+    def print_table(
+        self, table_res, tracker, cls, res_field="COMBINED_SEQ", output_lable="COMBINED"
+    ):
+        """Prints table of results for all sequences"""
+        print("")
+        metric_name = self.get_name()
+        self._row_print(
+            [metric_name + ": " + tracker + "-" + cls] + self.summary_fields
+        )
+        for seq, results in sorted(table_res.items()):
+            if seq.startswith("COMBINED_SEQ"):
+                continue
+            summary_res = self._summary_row(results)
+            self._row_print([seq] + summary_res)
+        summary_res = self._summary_row(table_res[res_field])
+        self._row_print([output_lable] + summary_res)
+
+    def _summary_row(self, results_):
+        vals = []
+        for h in self.summary_fields:
+            if h in self.float_array_fields:
+                vals.append("{0:1.5g}".format(100 * np.mean(results_[h])))
+            elif h in self.float_fields:
+                vals.append("{0:1.5g}".format(100 * float(results_[h])))
+            elif h in self.integer_fields:
+                vals.append("{0:d}".format(int(results_[h])))
+            else:
+                raise NotImplementedError(
+                    "Summary function not implemented for this field type."
+                )
+        return vals
+
+    @staticmethod
+    def _row_print(*argv):
+        """Prints results in an evenly spaced rows, with more space in first row"""
+        if len(argv) == 1:
+            argv = argv[0]
+        to_print = "%-35s" % argv[0]
+        for v in argv[1:]:
+            to_print += "%-10s" % str(v)
+        print(to_print)
+
+    def summary_results(self, table_res):
+        """Returns a simple summary of final results for a tracker"""
+        return dict(
+            zip(self.summary_fields, self._summary_row(table_res["COMBINED_SEQ"]))
+        )
+
+    def detailed_results(self, table_res):
+        """Returns detailed final results for a tracker"""
+        # Get detailed field information
+        detailed_fields = self.float_fields + self.integer_fields
+        for h in self.float_array_fields + self.integer_array_fields:
+            for alpha in [int(100 * x) for x in self.array_labels]:
+                detailed_fields.append(h + "___" + str(alpha))
+            detailed_fields.append(h + "___AUC")
+
+        # Get detailed results
+        detailed_results = {}
+        for seq, res in table_res.items():
+            detailed_row = self._detailed_row(res)
+            if len(detailed_row) != len(detailed_fields):
+                raise TrackEvalException(
+                    "Field names and data have different sizes (%i and %i)"
+                    % (len(detailed_row), len(detailed_fields))
+                )
+            detailed_results[seq] = dict(zip(detailed_fields, detailed_row))
+        return detailed_results
+
+    def _detailed_row(self, res):
+        detailed_row = []
+        for h in self.float_fields + self.integer_fields:
+            detailed_row.append(res[h])
+        for h in self.float_array_fields + self.integer_array_fields:
+            for i, alpha in enumerate([int(100 * x) for x in self.array_labels]):
+                detailed_row.append(res[h][i])
+            detailed_row.append(np.mean(res[h]))
+        return detailed_row

sam3/eval/hota_eval_toolkit/trackeval/metrics/count.py

@@ -0,0 +1,48 @@
+# flake8: noqa
+
+from .. import _timing
+from ._base_metric import _BaseMetric
+
+
+class Count(_BaseMetric):
+    """Class which simply counts the number of tracker and gt detections and ids."""
+
+    def __init__(self, config=None):
+        super().__init__()
+        self.integer_fields = ["Dets", "GT_Dets", "IDs", "GT_IDs"]
+        self.fields = self.integer_fields
+        self.summary_fields = self.fields
+
+    @_timing.time
+    def eval_sequence(self, data):
+        """Returns counts for one sequence"""
+        # Get results
+        res = {
+            "Dets": data["num_tracker_dets"],
+            "GT_Dets": data["num_gt_dets"],
+            "IDs": data["num_tracker_ids"],
+            "GT_IDs": data["num_gt_ids"],
+            "Frames": data["num_timesteps"],
+        }
+        return res
+
+    def combine_sequences(self, all_res):
+        """Combines metrics across all sequences"""
+        res = {}
+        for field in self.integer_fields:
+            res[field] = self._combine_sum(all_res, field)
+        return res
+
+    def combine_classes_class_averaged(self, all_res, ignore_empty_classes=None):
+        """Combines metrics across all classes by averaging over the class values"""
+        res = {}
+        for field in self.integer_fields:
+            res[field] = self._combine_sum(all_res, field)
+        return res
+
+    def combine_classes_det_averaged(self, all_res):
+        """Combines metrics across all classes by averaging over the detection values"""
+        res = {}
+        for field in self.integer_fields:
+            res[field] = self._combine_sum(all_res, field)
+        return res

sam3/eval/hota_eval_toolkit/trackeval/metrics/hota.py

@@ -0,0 +1,291 @@
+# flake8: noqa
+
+import os
+
+import numpy as np
+from scipy.optimize import linear_sum_assignment
+
+from .. import _timing
+from ._base_metric import _BaseMetric
+
+
+class HOTA(_BaseMetric):
+    """Class which implements the HOTA metrics.
+    See: https://link.springer.com/article/10.1007/s11263-020-01375-2
+    """
+
+    def __init__(self, config=None):
+        super().__init__()
+        self.plottable = True
+        self.array_labels = np.arange(0.05, 0.99, 0.05)
+        self.integer_array_fields = ["HOTA_TP", "HOTA_FN", "HOTA_FP"]
+        self.float_array_fields = [
+            "HOTA",
+            "DetA",
+            "AssA",
+            "DetRe",
+            "DetPr",
+            "AssRe",
+            "AssPr",
+            "LocA",
+            "OWTA",
+        ]
+        self.float_fields = ["HOTA(0)", "LocA(0)", "HOTALocA(0)"]
+        self.fields = (
+            self.float_array_fields + self.integer_array_fields + self.float_fields
+        )
+        self.summary_fields = self.float_array_fields + self.float_fields
+
+    @_timing.time
+    def eval_sequence(self, data):
+        """Calculates the HOTA metrics for one sequence"""
+
+        # Initialise results
+        res = {}
+        for field in self.float_array_fields + self.integer_array_fields:
+            res[field] = np.zeros((len(self.array_labels)), dtype=float)
+        for field in self.float_fields:
+            res[field] = 0
+
+        # Return result quickly if tracker or gt sequence is empty
+        if data["num_tracker_dets"] == 0:
+            res["HOTA_FN"] = data["num_gt_dets"] * np.ones(
+                (len(self.array_labels)), dtype=float
+            )
+            res["LocA"] = np.ones((len(self.array_labels)), dtype=float)
+            res["LocA(0)"] = 1.0
+            return res
+        if data["num_gt_dets"] == 0:
+            res["HOTA_FP"] = data["num_tracker_dets"] * np.ones(
+                (len(self.array_labels)), dtype=float
+            )
+            res["LocA"] = np.ones((len(self.array_labels)), dtype=float)
+            res["LocA(0)"] = 1.0
+            return res
+
+        # Variables counting global association
+        potential_matches_count = np.zeros(
+            (data["num_gt_ids"], data["num_tracker_ids"])
+        )
+        gt_id_count = np.zeros((data["num_gt_ids"], 1))
+        tracker_id_count = np.zeros((1, data["num_tracker_ids"]))
+
+        # First loop through each timestep and accumulate global track information.
+        for t, (gt_ids_t, tracker_ids_t) in enumerate(
+            zip(data["gt_ids"], data["tracker_ids"])
+        ):
+            # Count the potential matches between ids in each timestep
+            # These are normalised, weighted by the match similarity.
+            similarity = data["similarity_scores"][t]
+            sim_iou_denom = (
+                similarity.sum(0)[np.newaxis, :]
+                + similarity.sum(1)[:, np.newaxis]
+                - similarity
+            )
+            sim_iou = np.zeros_like(similarity)
+            sim_iou_mask = sim_iou_denom > 0 + np.finfo("float").eps
+            sim_iou[sim_iou_mask] = (
+                similarity[sim_iou_mask] / sim_iou_denom[sim_iou_mask]
+            )
+            potential_matches_count[
+                gt_ids_t[:, np.newaxis], tracker_ids_t[np.newaxis, :]
+            ] += sim_iou
+
+            # Calculate the total number of dets for each gt_id and tracker_id.
+            gt_id_count[gt_ids_t] += 1
+            tracker_id_count[0, tracker_ids_t] += 1
+
+        # Calculate overall jaccard alignment score (before unique matching) between IDs
+        global_alignment_score = potential_matches_count / (
+            gt_id_count + tracker_id_count - potential_matches_count
+        )
+        matches_counts = [
+            np.zeros_like(potential_matches_count) for _ in self.array_labels
+        ]
+
+        # Calculate scores for each timestep
+        for t, (gt_ids_t, tracker_ids_t) in enumerate(
+            zip(data["gt_ids"], data["tracker_ids"])
+        ):
+            # Deal with the case that there are no gt_det/tracker_det in a timestep.
+            if len(gt_ids_t) == 0:
+                for a, alpha in enumerate(self.array_labels):
+                    res["HOTA_FP"][a] += len(tracker_ids_t)
+                continue
+            if len(tracker_ids_t) == 0:
+                for a, alpha in enumerate(self.array_labels):
+                    res["HOTA_FN"][a] += len(gt_ids_t)
+                continue
+
+            # Get matching scores between pairs of dets for optimizing HOTA
+            similarity = data["similarity_scores"][t]
+            score_mat = (
+                global_alignment_score[
+                    gt_ids_t[:, np.newaxis], tracker_ids_t[np.newaxis, :]
+                ]
+                * similarity
+            )
+
+            # Hungarian algorithm to find best matches
+            match_rows, match_cols = linear_sum_assignment(-score_mat)
+
+            # Calculate and accumulate basic statistics
+            for a, alpha in enumerate(self.array_labels):
+                actually_matched_mask = (
+                    similarity[match_rows, match_cols] >= alpha - np.finfo("float").eps
+                )
+                alpha_match_rows = match_rows[actually_matched_mask]
+                alpha_match_cols = match_cols[actually_matched_mask]
+                num_matches = len(alpha_match_rows)
+                res["HOTA_TP"][a] += num_matches
+                res["HOTA_FN"][a] += len(gt_ids_t) - num_matches
+                res["HOTA_FP"][a] += len(tracker_ids_t) - num_matches
+                if num_matches > 0:
+                    res["LocA"][a] += sum(
+                        similarity[alpha_match_rows, alpha_match_cols]
+                    )
+                    matches_counts[a][
+                        gt_ids_t[alpha_match_rows], tracker_ids_t[alpha_match_cols]
+                    ] += 1
+
+        # Calculate association scores (AssA, AssRe, AssPr) for the alpha value.
+        # First calculate scores per gt_id/tracker_id combo and then average over the number of detections.
+        for a, alpha in enumerate(self.array_labels):
+            matches_count = matches_counts[a]
+            ass_a = matches_count / np.maximum(
+                1, gt_id_count + tracker_id_count - matches_count
+            )
+            res["AssA"][a] = np.sum(matches_count * ass_a) / np.maximum(
+                1, res["HOTA_TP"][a]
+            )
+            ass_re = matches_count / np.maximum(1, gt_id_count)
+            res["AssRe"][a] = np.sum(matches_count * ass_re) / np.maximum(
+                1, res["HOTA_TP"][a]
+            )
+            ass_pr = matches_count / np.maximum(1, tracker_id_count)
+            res["AssPr"][a] = np.sum(matches_count * ass_pr) / np.maximum(
+                1, res["HOTA_TP"][a]
+            )
+
+        # Calculate final scores
+        res["LocA"] = np.maximum(1e-10, res["LocA"]) / np.maximum(1e-10, res["HOTA_TP"])
+        res = self._compute_final_fields(res)
+        return res
+
+    def combine_sequences(self, all_res):
+        """Combines metrics across all sequences"""
+        res = {}
+        for field in self.integer_array_fields:
+            res[field] = self._combine_sum(all_res, field)
+        for field in ["AssRe", "AssPr", "AssA"]:
+            res[field] = self._combine_weighted_av(
+                all_res, field, res, weight_field="HOTA_TP"
+            )
+        loca_weighted_sum = sum(
+            [all_res[k]["LocA"] * all_res[k]["HOTA_TP"] for k in all_res.keys()]
+        )
+        res["LocA"] = np.maximum(1e-10, loca_weighted_sum) / np.maximum(
+            1e-10, res["HOTA_TP"]
+        )
+        res = self._compute_final_fields(res)
+        return res
+
+    def combine_classes_class_averaged(self, all_res, ignore_empty_classes=False):
+        """Combines metrics across all classes by averaging over the class values.
+        If 'ignore_empty_classes' is True, then it only sums over classes with at least one gt or predicted detection.
+        """
+        res = {}
+        for field in self.integer_array_fields:
+            if ignore_empty_classes:
+                res[field] = self._combine_sum(
+                    {
+                        k: v
+                        for k, v in all_res.items()
+                        if (
+                            v["HOTA_TP"] + v["HOTA_FN"] + v["HOTA_FP"]
+                            > 0 + np.finfo("float").eps
+                        ).any()
+                    },
+                    field,
+                )
+            else:
+                res[field] = self._combine_sum(
+                    {k: v for k, v in all_res.items()}, field
+                )
+
+        for field in self.float_fields + self.float_array_fields:
+            if ignore_empty_classes:
+                res[field] = np.mean(
+                    [
+                        v[field]
+                        for v in all_res.values()
+                        if (
+                            v["HOTA_TP"] + v["HOTA_FN"] + v["HOTA_FP"]
+                            > 0 + np.finfo("float").eps
+                        ).any()
+                    ],
+                    axis=0,
+                )
+            else:
+                res[field] = np.mean([v[field] for v in all_res.values()], axis=0)
+        return res
+
+    def combine_classes_det_averaged(self, all_res):
+        """Combines metrics across all classes by averaging over the detection values"""
+        res = {}
+        for field in self.integer_array_fields:
+            res[field] = self._combine_sum(all_res, field)
+        for field in ["AssRe", "AssPr", "AssA"]:
+            res[field] = self._combine_weighted_av(
+                all_res, field, res, weight_field="HOTA_TP"
+            )
+        loca_weighted_sum = sum(
+            [all_res[k]["LocA"] * all_res[k]["HOTA_TP"] for k in all_res.keys()]
+        )
+        res["LocA"] = np.maximum(1e-10, loca_weighted_sum) / np.maximum(
+            1e-10, res["HOTA_TP"]
+        )
+        res = self._compute_final_fields(res)
+        return res
+
+    @staticmethod
+    def _compute_final_fields(res):
+        """Calculate sub-metric ('field') values which only depend on other sub-metric values.
+        This function is used both for both per-sequence calculation, and in combining values across sequences.
+        """
+        res["DetRe"] = res["HOTA_TP"] / np.maximum(1, res["HOTA_TP"] + res["HOTA_FN"])
+        res["DetPr"] = res["HOTA_TP"] / np.maximum(1, res["HOTA_TP"] + res["HOTA_FP"])
+        res["DetA"] = res["HOTA_TP"] / np.maximum(
+            1, res["HOTA_TP"] + res["HOTA_FN"] + res["HOTA_FP"]
+        )
+        res["HOTA"] = np.sqrt(res["DetA"] * res["AssA"])
+        res["OWTA"] = np.sqrt(res["DetRe"] * res["AssA"])
+
+        res["HOTA(0)"] = res["HOTA"][0]
+        res["LocA(0)"] = res["LocA"][0]
+        res["HOTALocA(0)"] = res["HOTA(0)"] * res["LocA(0)"]
+        return res
+
+    def plot_single_tracker_results(self, table_res, tracker, cls, output_folder):
+        """Create plot of results"""
+
+        # Only loaded when run to reduce minimum requirements
+        from matplotlib import pyplot as plt
+
+        res = table_res["COMBINED_SEQ"]
+        styles_to_plot = ["r", "b", "g", "b--", "b:", "g--", "g:", "m"]
+        for name, style in zip(self.float_array_fields, styles_to_plot):
+            plt.plot(self.array_labels, res[name], style)
+        plt.xlabel("alpha")
+        plt.ylabel("score")
+        plt.title(tracker + " - " + cls)
+        plt.axis([0, 1, 0, 1])
+        legend = []
+        for name in self.float_array_fields:
+            legend += [name + " (" + str(np.round(np.mean(res[name]), 2)) + ")"]
+        plt.legend(legend, loc="lower left")
+        out_file = os.path.join(output_folder, cls + "_plot.pdf")
+        os.makedirs(os.path.dirname(out_file), exist_ok=True)
+        plt.savefig(out_file)
+        plt.savefig(out_file.replace(".pdf", ".png"))
+        plt.clf()

sam3/eval/hota_eval_toolkit/trackeval/utils.py

@@ -0,0 +1,195 @@
+# flake8: noqa
+
+import argparse
+import csv
+import os
+from collections import OrderedDict
+
+
+def init_config(config, default_config, name=None):
+    """Initialise non-given config values with defaults"""
+    if config is None:
+        config = default_config
+    else:
+        for k in default_config.keys():
+            if k not in config.keys():
+                config[k] = default_config[k]
+    if name and config["PRINT_CONFIG"]:
+        print("\n%s Config:" % name)
+        for c in config.keys():
+            print("%-20s : %-30s" % (c, config[c]))
+    return config
+
+
+def update_config(config):
+    """
+    Parse the arguments of a script and updates the config values for a given value if specified in the arguments.
+    :param config: the config to update
+    :return: the updated config
+    """
+    parser = argparse.ArgumentParser()
+    for setting in config.keys():
+        if type(config[setting]) == list or type(config[setting]) == type(None):
+            parser.add_argument("--" + setting, nargs="+")
+        else:
+            parser.add_argument("--" + setting)
+    args = parser.parse_args().__dict__
+    for setting in args.keys():
+        if args[setting] is not None:
+            if type(config[setting]) == type(True):
+                if args[setting] == "True":
+                    x = True
+                elif args[setting] == "False":
+                    x = False
+                else:
+                    raise Exception(
+                        "Command line parameter " + setting + "must be True or False"
+                    )
+            elif type(config[setting]) == type(1):
+                x = int(args[setting])
+            elif type(args[setting]) == type(None):
+                x = None
+            else:
+                x = args[setting]
+            config[setting] = x
+    return config
+
+
+def get_code_path():
+    """Get base path where code is"""
+    return os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+
+
+def validate_metrics_list(metrics_list):
+    """Get names of metric class and ensures they are unique, further checks that the fields within each metric class
+    do not have overlapping names.
+    """
+    metric_names = [metric.get_name() for metric in metrics_list]
+    # check metric names are unique
+    if len(metric_names) != len(set(metric_names)):
+        raise TrackEvalException(
+            "Code being run with multiple metrics of the same name"
+        )
+    fields = []
+    for m in metrics_list:
+        fields += m.fields
+    # check metric fields are unique
+    if len(fields) != len(set(fields)):
+        raise TrackEvalException(
+            "Code being run with multiple metrics with fields of the same name"
+        )
+    return metric_names
+
+
+def write_summary_results(summaries, cls, output_folder):
+    """Write summary results to file"""
+
+    fields = sum([list(s.keys()) for s in summaries], [])
+    values = sum([list(s.values()) for s in summaries], [])
+
+    # In order to remain consistent upon new fields being adding, for each of the following fields if they are present
+    # they will be output in the summary first in the order below. Any further fields will be output in the order each
+    # metric family is called, and within each family either in the order they were added to the dict (python >= 3.6) or
+    # randomly (python < 3.6).
+    default_order = [
+        "HOTA",
+        "DetA",
+        "AssA",
+        "DetRe",
+        "DetPr",
+        "AssRe",
+        "AssPr",
+        "LocA",
+        "OWTA",
+        "HOTA(0)",
+        "LocA(0)",
+        "HOTALocA(0)",
+        "MOTA",
+        "MOTP",
+        "MODA",
+        "CLR_Re",
+        "CLR_Pr",
+        "MTR",
+        "PTR",
+        "MLR",
+        "CLR_TP",
+        "CLR_FN",
+        "CLR_FP",
+        "IDSW",
+        "MT",
+        "PT",
+        "ML",
+        "Frag",
+        "sMOTA",
+        "IDF1",
+        "IDR",
+        "IDP",
+        "IDTP",
+        "IDFN",
+        "IDFP",
+        "Dets",
+        "GT_Dets",
+        "IDs",
+        "GT_IDs",
+    ]
+    default_ordered_dict = OrderedDict(
+        zip(default_order, [None for _ in default_order])
+    )
+    for f, v in zip(fields, values):
+        default_ordered_dict[f] = v
+    for df in default_order:
+        if default_ordered_dict[df] is None:
+            del default_ordered_dict[df]
+    fields = list(default_ordered_dict.keys())
+    values = list(default_ordered_dict.values())
+
+    out_file = os.path.join(output_folder, cls + "_summary.txt")
+    os.makedirs(os.path.dirname(out_file), exist_ok=True)
+    with open(out_file, "w", newline="") as f:
+        writer = csv.writer(f, delimiter=" ")
+        writer.writerow(fields)
+        writer.writerow(values)
+
+
+def write_detailed_results(details, cls, output_folder):
+    """Write detailed results to file"""
+    sequences = details[0].keys()
+    fields = ["seq"] + sum([list(s["COMBINED_SEQ"].keys()) for s in details], [])
+    out_file = os.path.join(output_folder, cls + "_detailed.csv")
+    os.makedirs(os.path.dirname(out_file), exist_ok=True)
+    with open(out_file, "w", newline="") as f:
+        writer = csv.writer(f)
+        writer.writerow(fields)
+        for seq in sorted(sequences):
+            if seq == "COMBINED_SEQ":
+                continue
+            writer.writerow([seq] + sum([list(s[seq].values()) for s in details], []))
+        writer.writerow(
+            ["COMBINED"] + sum([list(s["COMBINED_SEQ"].values()) for s in details], [])
+        )
+
+
+def load_detail(file):
+    """Loads detailed data for a tracker."""
+    data = {}
+    with open(file) as f:
+        for i, row_text in enumerate(f):
+            row = row_text.replace("\r", "").replace("\n", "").split(",")
+            if i == 0:
+                keys = row[1:]
+                continue
+            current_values = row[1:]
+            seq = row[0]
+            if seq == "COMBINED":
+                seq = "COMBINED_SEQ"
+            if (len(current_values) == len(keys)) and seq != "":
+                data[seq] = {}
+                for key, value in zip(keys, current_values):
+                    data[seq][key] = float(value)
+    return data
+
+
+class TrackEvalException(Exception):
+    """Custom exception for catching expected errors."""
+
+    ...

sam3/eval/postprocessors.py

@@ -0,0 +1,648 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+"""Postprocessors class to transform MDETR output according to the downstream task"""
+
+import dataclasses
+import logging
+from collections import defaultdict
+from typing import Dict, List, Optional
+
+import numpy as np
+import torch
+from sam3.model import box_ops
+from sam3.model.data_misc import BatchedInferenceMetadata, interpolate
+from sam3.train.masks_ops import rle_encode, robust_rle_encode
+from torch import nn
+
+
+class PostProcessNullOp(nn.Module):
+    def __init__(self, **kwargs):
+        super(PostProcessNullOp).__init__()
+        pass
+
+    def forward(self, input):
+        pass
+
+    def process_results(self, **kwargs):
+        return kwargs["find_stages"]
+
+
+class PostProcessImage(nn.Module):
+    """This module converts the model's output into the format expected by the coco api"""
+
+    def __init__(
+        self,
+        max_dets_per_img: int,
+        iou_type="bbox",
+        to_cpu: bool = True,
+        use_original_ids: bool = False,
+        use_original_sizes_box: bool = False,
+        use_original_sizes_mask: bool = False,
+        convert_mask_to_rle: bool = False,
+        always_interpolate_masks_on_gpu: bool = True,
+        use_presence: bool = True,
+        detection_threshold: float = -1.0,
+    ) -> None:
+        super().__init__()
+        self.max_dets_per_img = max_dets_per_img
+        self.iou_type = iou_type
+        self.to_cpu = to_cpu
+        self.convert_mask_to_rle = convert_mask_to_rle
+        self.always_interpolate_masks_on_gpu = always_interpolate_masks_on_gpu
+
+        self.use_presence = use_presence
+        self.detection_threshold = detection_threshold
+        self.use_original_ids = use_original_ids
+        self.use_original_sizes_box = use_original_sizes_box
+        self.use_original_sizes_mask = use_original_sizes_mask
+
+    @torch.no_grad()
+    def forward(
+        self,
+        outputs,
+        target_sizes_boxes,
+        target_sizes_masks,
+        forced_labels=None,
+        consistent=False,
+        ret_tensordict: bool = False,  # This is experimental
+    ):
+        """Perform the computation
+        Parameters:
+            outputs: raw outputs of the model
+            target_sizes_boxes: tensor of dimension [batch_size x 2] containing the size of each images of the batch
+                          For evaluation, this must be the original image size (before any data augmentation)
+                          For visualization, this should be the image size after data augment, but before padding
+            target_sizes_masks: same but used to resize masks
+            forced_labels: tensor of dimension [batch_size] containing the label to force for each image of the batch
+                           This is useful when evaluating the model using standard metrics (eg on COCO, LVIS). In that case,
+                           we query the model with every possible class label, so we when we pass the predictions to the evaluator,
+                           we want to make sure that the predicted "class" matches the one that was queried.
+            consistent: whether all target sizes are equal
+            ret_tensordict: Experimental argument. If true, return a tensordict.TensorDict instead of a list of dictionaries for easier manipulation.
+        """
+        if ret_tensordict:
+            assert (
+                consistent is True
+            ), "We don't support returning TensorDict if the outputs have different shapes"  # NOTE: It's possible but we don't support it.
+            assert self.detection_threshold <= 0.0, "TODO: implement?"
+            try:
+                from tensordict import TensorDict
+            except ImportError:
+                logging.info(
+                    "tensordict is not installed. Install by running `pip install tensordict --no-deps`. Falling back by setting `ret_tensordict=False`"
+                )
+                ret_tensordict = False
+
+        out_bbox = outputs["pred_boxes"] if "pred_boxes" in outputs else None
+        out_logits = outputs["pred_logits"]
+        pred_masks = outputs["pred_masks"] if self.iou_type == "segm" else None
+        out_probs = out_logits.sigmoid()
+        if self.use_presence:
+            presence_score = outputs["presence_logit_dec"].sigmoid().unsqueeze(1)
+            out_probs = out_probs * presence_score
+
+        assert target_sizes_boxes.shape[1] == 2
+        assert target_sizes_masks.shape[1] == 2
+        batch_size = target_sizes_boxes.shape[0]
+
+        boxes, scores, labels, keep = self._process_boxes_and_labels(
+            target_sizes_boxes, forced_labels, out_bbox, out_probs
+        )
+        assert boxes is None or len(boxes) == batch_size
+        out_masks = self._process_masks(
+            target_sizes_masks, pred_masks, consistent=consistent, keep=keep
+        )
+        del pred_masks
+
+        if boxes is None:
+            assert out_masks is not None
+            assert not ret_tensordict, "We don't support returning TensorDict if the output does not contain boxes"
+            B = len(out_masks)
+            boxes = [None] * B
+            scores = [None] * B
+            labels = [None] * B
+
+        results = {
+            "scores": scores,
+            "labels": labels,
+            "boxes": boxes,
+        }
+        if out_masks is not None:
+            if self.convert_mask_to_rle:
+                results.update(masks_rle=out_masks)
+            else:
+                results.update(masks=out_masks)
+
+        if ret_tensordict:
+            results = TensorDict(results).auto_batch_size_()
+            if self.to_cpu:
+                results = results.cpu()
+        else:
+            # Convert a dictonary of lists/tensors to list of dictionaries
+            results = [
+                dict(zip(results.keys(), res_tuple))
+                for res_tuple in zip(*results.values())
+            ]
+
+        return results
+
+    def _process_masks(self, target_sizes, pred_masks, consistent=True, keep=None):
+        if pred_masks is None:
+            return None
+        if self.always_interpolate_masks_on_gpu:
+            gpu_device = target_sizes.device
+            assert gpu_device.type == "cuda"
+            pred_masks = pred_masks.to(device=gpu_device)
+        if consistent:
+            assert keep is None, "TODO: implement?"
+            # All masks should have the same shape, expected when processing a batch of size 1
+            target_size = target_sizes.unique(dim=0)
+            assert target_size.size(0) == 1, "Expecting all target sizes to be equal"
+            out_masks = (
+                interpolate(
+                    pred_masks,
+                    target_size.squeeze().tolist(),
+                    mode="bilinear",
+                    align_corners=False,
+                ).sigmoid()
+                > 0.5
+            )
+            if self.convert_mask_to_rle:
+                raise RuntimeError("TODO: implement?")
+            if self.to_cpu:
+                out_masks = out_masks.cpu()
+        else:
+            out_masks = [[]] * len(pred_masks)
+
+            assert keep is None or len(keep) == len(pred_masks)
+            for i, mask in enumerate(pred_masks):
+                h, w = target_sizes[i]
+                if keep is not None:
+                    mask = mask[keep[i]]
+                # Uses the gpu version fist, moves masks to cpu if it fails"""
+                try:
+                    interpolated = (
+                        interpolate(
+                            mask.unsqueeze(1),
+                            (h, w),
+                            mode="bilinear",
+                            align_corners=False,
+                        ).sigmoid()
+                        > 0.5
+                    )
+                except Exception as e:
+                    logging.info("Issue found, reverting to CPU mode!")
+                    mask_device = mask.device
+                    mask = mask.cpu()
+                    interpolated = (
+                        interpolate(
+                            mask.unsqueeze(1),
+                            (h, w),
+                            mode="bilinear",
+                            align_corners=False,
+                        ).sigmoid()
+                        > 0.5
+                    )
+                    interpolated = interpolated.to(mask_device)
+
+                if self.convert_mask_to_rle:
+                    out_masks[i] = robust_rle_encode(interpolated.squeeze(1))
+                else:
+                    out_masks[i] = interpolated
+                    if self.to_cpu:
+                        out_masks[i] = out_masks[i].cpu()
+
+        return out_masks
+
+    def _process_boxes_and_labels(
+        self, target_sizes, forced_labels, out_bbox, out_probs
+    ):
+        if out_bbox is None:
+            return None, None, None, None
+        assert len(out_probs) == len(target_sizes)
+        if self.to_cpu:
+            out_probs = out_probs.cpu()
+        scores, labels = out_probs.max(-1)
+        if forced_labels is None:
+            labels = torch.ones_like(labels)
+        else:
+            labels = forced_labels[:, None].expand_as(labels)
+
+        # convert to [x0, y0, x1, y1] format
+        boxes = box_ops.box_cxcywh_to_xyxy(out_bbox)
+
+        img_h, img_w = target_sizes.unbind(1)
+        scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1)
+        boxes = boxes * scale_fct[:, None, :]
+
+        if self.to_cpu:
+            boxes = boxes.cpu()
+
+        keep = None
+        if self.detection_threshold > 0:
+            # Filter out the boxes with scores below the detection threshold
+            keep = scores > self.detection_threshold
+            assert len(keep) == len(boxes) == len(scores) == len(labels)
+
+            boxes = [b[k.to(b.device)] for b, k in zip(boxes, keep)]
+            scores = [s[k.to(s.device)] for s, k in zip(scores, keep)]
+            labels = [l[k.to(l.device)] for l, k in zip(labels, keep)]
+
+        return boxes, scores, labels, keep
+
+    def process_results(
+        self, find_stages, find_metadatas: List[BatchedInferenceMetadata], **kwargs
+    ):
+        if find_stages.loss_stages is not None:
+            find_metadatas = [find_metadatas[i] for i in find_stages.loss_stages]
+        assert len(find_stages) == len(find_metadatas)
+        results = {}
+        for outputs, meta in zip(find_stages, find_metadatas):
+            img_size_for_boxes = (
+                meta.original_size
+                if self.use_original_sizes_box
+                else torch.ones_like(meta.original_size)
+            )
+            img_size_for_masks = (
+                meta.original_size
+                if self.use_original_sizes_mask
+                else torch.ones_like(meta.original_size)
+            )
+            detection_results = self(
+                outputs,
+                img_size_for_boxes,
+                img_size_for_masks,
+                forced_labels=(
+                    meta.original_category_id if self.use_original_ids else None
+                ),
+            )
+            ids = (
+                meta.original_image_id if self.use_original_ids else meta.coco_image_id
+            )
+            assert len(detection_results) == len(ids)
+            for img_id, result in zip(ids, detection_results):
+                if img_id.item() not in results:
+                    results[img_id.item()] = result
+                else:
+                    assert set(results[img_id.item()].keys()) == set(result.keys())
+                    for k in result.keys():
+                        if isinstance(result[k], torch.Tensor):
+                            results[img_id.item()][k] = torch.cat(
+                                [results[img_id.item()][k], result[k]], dim=0
+                            )
+                        elif isinstance(result[k], list):
+                            results[img_id.item()][k] += result[k]
+                        else:
+                            raise NotImplementedError(
+                                f"Unexpected type {type(result[k])} in result."
+                            )
+        # Prune the results to the max number of detections per image.
+        for img_id, result in results.items():
+            if (
+                self.max_dets_per_img > 0
+                and len(result["scores"]) > self.max_dets_per_img
+            ):
+                _, topk_indexes = torch.topk(
+                    result["scores"], self.max_dets_per_img, dim=0
+                )
+                if self.to_cpu:
+                    topk_indexes = topk_indexes.cpu()
+                for k in result.keys():
+                    if isinstance(results[img_id][k], list):
+                        results[img_id][k] = [
+                            results[img_id][k][i] for i in topk_indexes.tolist()
+                        ]
+                    else:
+                        results[img_id][k] = results[img_id][k].to(topk_indexes.device)[
+                            topk_indexes
+                        ]
+
+        return results
+
+
+class PostProcessAPIVideo(PostProcessImage):
+    """This module converts the video model's output into the format expected by the YT-VIS api"""
+
+    def __init__(
+        self,
+        *args,
+        to_cpu: bool = True,
+        convert_mask_to_rle: bool = False,
+        always_interpolate_masks_on_gpu: bool = True,
+        prob_thresh: float = 0.5,
+        use_presence: bool = False,
+        **kwargs,
+    ):
+        super().__init__(
+            *args,
+            # Here we always set `convert_mask_to_rle=False` in the base `PostProcessAPI` class
+            # (so that its `_process_masks` won't return a list of RLEs). If we want to return
+            # RLEs for video masklets, we handle it in this `PostProcessAPIVideo` class instead.
+            convert_mask_to_rle=False,
+            # Here we always set `to_cpu=False` in the base `PostProcessAPI` class (so that
+            # the interpolated masks won't be automatically moved back to CPU). We will handle
+            # it in this `PostProcessAPIVideo` class instead.
+            always_interpolate_masks_on_gpu=always_interpolate_masks_on_gpu,
+            use_presence=use_presence,
+            **kwargs,
+        )
+        # Expected keys in the output dict to postprocess
+        self.EXPECTED_KEYS = [
+            "pred_logits",
+            "pred_boxes",
+            "pred_masks",
+        ]
+        # Whether to post-process video masklets (under packed representation) into RLE format
+        self.convert_mask_to_rle_for_video = convert_mask_to_rle
+        self.to_cpu_for_video = to_cpu
+        self.prob_thresh = prob_thresh
+
+    def process_results(
+        self, find_stages, find_metadatas: List[BatchedInferenceMetadata], **kwargs
+    ):
+        """
+        Tracking Postprocessor for SAM 3 video model.
+        This function takes in the output of the SAM 3 video model and processes it to extract all the tracklet predictions.
+        Args:
+            find_stages: A list of tensors representing the output of the SAM 3 video model.
+            find_metadatas: A list of BatchedInferenceMetadata objects containing metadata about each frame.
+            **kwargs: Additional keyword arguments.
+        Returns:
+            A dictionary of predcitions with video_id as key.
+        """
+
+        # Import tensordict here to avoid global dependency.
+        try:
+            from tensordict import TensorDict
+        except ImportError as e:
+            logging.error(
+                "tensordict is not installed, please install by running `pip install tensordict --no-deps`"
+            )
+            raise e
+        # Notes and assumptions:
+        # 1- This postprocessor assumes results only for a single video.
+        # 2- There are N stage outputs corresponding to N video frames
+        # 3- Each stage outputs contains PxQ preds, where P is number of prompts and Q is number of object queries. The output should also contain the tracking object ids corresponding to each object query.
+        # 4- The tracking object id has a default value of -1, indicating that the object query is not tracking any object in the frame, and hence its predictions can be ingored for a given frame.
+        # 5- Some objects may be tracked in a subset of frames only. So, we first extract the predictions in a packed representation (for efficient postprocessing -- specially memory)
+        # and then we convert the packed representation into a padded one, where we zero pad boxes/masks for objects that are not tracked in some frames.
+        # 6- We refer to objects by an object id, which is a tuple (prompt_idx, obj_id)
+
+        assert len(find_stages) > 0, "There is nothing to postprocess?"
+        PROMPT_AXIS, OBJ_QUERY_AXIS = (0, 1)
+        NO_OBJ_ID = -1
+        # Maps object ID -> [indices in packed tensor]
+        tracked_objects_packed_idx = defaultdict(list)
+        # Maps object ID -> [indices in padded tensor (abs frame index)]
+        tracked_objects_frame_idx = defaultdict(list)
+        total_num_preds = 0
+        # This will hold the packed representation of predictions.
+        vid_preds_packed: List[TensorDict] = []
+        vid_masklets_rle_packed: List[Optional[Dict]] = []
+        video_id = -1  # We assume single video postprocessing, this ID should be unique in the datapoint.
+
+        for frame_idx, (frame_outs, meta) in enumerate(
+            zip(find_stages, find_metadatas)
+        ):
+            # only store keys we need to extract the results
+            frame_outs_td = TensorDict(
+                {k: frame_outs[k] for k in self.EXPECTED_KEYS}
+            ).auto_batch_size_()  # Shape is [P,Q,...]
+            meta_td = TensorDict(
+                dataclasses.asdict(meta)
+            ).auto_batch_size_()  # Shape is [P,...]
+            unique_vid_id = meta.original_image_id.unique()
+            assert unique_vid_id.size(0) == 1
+            if video_id == -1:
+                video_id = unique_vid_id.item()
+            else:
+                assert (
+                    video_id == unique_vid_id.item()
+                ), "We can only postprocess one video per datapoint"
+            # keeping track of which objects appear in the current frame
+            obj_ids_per_frame = frame_outs["pred_object_ids"]
+            assert obj_ids_per_frame.size(-1) == frame_outs["pred_logits"].size(-2)
+            if self.prob_thresh is not None:
+                # only keep the predictions on this frame with probability above the threshold
+                # (remove those predictions during the keep-alive period of a tracking query,
+                # where its "pred_object_ids" is still the tracked object ID rather than -1)
+                pred_probs = frame_outs["pred_logits"].sigmoid().squeeze(-1)
+                obj_ids_per_frame = torch.where(
+                    pred_probs >= self.prob_thresh, obj_ids_per_frame, NO_OBJ_ID
+                )
+            tracked_obj_ids_idx = torch.where(obj_ids_per_frame != NO_OBJ_ID)
+            # Object id is a tuple of (prompt_idx, obj_id). This is because the model can assign same obj_id for two different prompts.
+            tracked_obj_ids = [
+                (p_id.item(), obj_ids_per_frame[p_id, q_id].item())
+                for p_id, q_id in zip(
+                    tracked_obj_ids_idx[PROMPT_AXIS],
+                    tracked_obj_ids_idx[OBJ_QUERY_AXIS],
+                )
+            ]
+            if len(tracked_obj_ids) == 0:
+                continue
+            # For each object, we keep track of the packed and padded (frame index) indices
+            for oid in tracked_obj_ids:
+                tracked_objects_packed_idx[oid].append(total_num_preds)
+                tracked_objects_frame_idx[oid].append(frame_idx)
+                total_num_preds += 1
+
+            # Since we have P*Q masks per frame, mask interpolation is the GPU memory bottleneck or time bottleneck in case of cpu processing.
+            # Instead, we first extract results only for tracked objects, reducing the number of masks to K = sum_i(tracked_objs_per_ith_prompt), hopefully <<< P*Q
+            tracked_objs_outs_td = frame_outs_td[
+                tracked_obj_ids_idx
+            ]  # [P,Q,...] --> [K,...]
+            meta_td = meta_td[tracked_obj_ids_idx[PROMPT_AXIS].cpu()]
+            if self.always_interpolate_masks_on_gpu:
+                gpu_device = meta_td["original_size"].device
+                assert gpu_device.type == "cuda"
+                tracked_objs_outs_td = tracked_objs_outs_td.to(device=gpu_device)
+            frame_results_td = self(
+                tracked_objs_outs_td.unsqueeze(1),
+                (
+                    meta_td["original_size"]
+                    if self.use_original_sizes
+                    else torch.ones_like(meta_td["original_size"])
+                ),
+                forced_labels=(
+                    meta_td["original_category_id"] if self.use_original_ids else None
+                ),
+                consistent=True,
+                ret_tensordict=True,
+            ).squeeze(1)
+            del tracked_objs_outs_td
+
+            # Optionally, remove "masks" from output tensor dict and directly encode them
+            # to RLE format under packed representations
+            if self.convert_mask_to_rle_for_video:
+                interpolated_binary_masks = frame_results_td.pop("masks")
+                rle_list = rle_encode(interpolated_binary_masks, return_areas=True)
+                vid_masklets_rle_packed.extend(rle_list)
+            # Optionally, move output TensorDict to CPU (do this after RLE encoding step above)
+            if self.to_cpu_for_video:
+                frame_results_td = frame_results_td.cpu()
+            vid_preds_packed.append(frame_results_td)
+
+        if len(vid_preds_packed) == 0:
+            logging.debug(f"Video {video_id} has no predictions")
+            return {video_id: []}
+
+        vid_preds_packed = torch.cat(vid_preds_packed, dim=0)
+        ############### Construct a padded representation of the predictions ###############
+        num_preds = len(tracked_objects_packed_idx)
+        num_frames = len(find_stages)
+        # We zero pad any missing prediction
+        # NOTE: here, we also have padded tensors for "scores" and "labels", but we overwrite them later.
+        padded_frames_results = TensorDict(
+            {
+                k: torch.zeros(
+                    num_preds, num_frames, *v.shape[1:], device=v.device, dtype=v.dtype
+                )
+                for k, v in vid_preds_packed.items()
+            },
+            batch_size=[
+                num_preds,
+                num_frames,
+            ],
+        )
+        padded_frames_results["scores"][...] = -1e8  # a very low score for empty object
+        # Track scores and labels of each pred tracklet, only for frames where the model was able to track that object
+        tracklet_scores = []
+        tracklet_labels = []
+        # Optionally, fill the list of RLEs for masklets
+        # note: only frames with actual predicted masks (in packed format) will be
+        # filled with RLEs; the rest will remains None in results["masks_rle"]
+        if self.convert_mask_to_rle_for_video:
+            vid_masklets_rle_padded = [[None] * num_frames for _ in range(num_preds)]
+        for o_idx, oid in enumerate(tracked_objects_packed_idx):
+            oid2packed_idx = tracked_objects_packed_idx[oid]
+            oid2padded_idx = tracked_objects_frame_idx[oid]
+            obj_packed_results = vid_preds_packed[oid2packed_idx]
+            padded_frames_results[o_idx][oid2padded_idx] = obj_packed_results
+            if self.convert_mask_to_rle_for_video:
+                for packed_idx, padded_idx in zip(oid2packed_idx, oid2padded_idx):
+                    vid_masklets_rle_padded[o_idx][padded_idx] = (
+                        vid_masklets_rle_packed[packed_idx]
+                    )
+            # NOTE: We need a single confidence score per tracklet for the mAP metric.
+            # We use the average confidence score across time. (How does this impact AP?)
+            tracklet_scores.append(obj_packed_results["scores"].mean())
+            # We also need to have a unique category Id per tracklet.
+            # This is not a problem for phrase AP, however, for mAP we do majority voting across time.
+            tracklet_labels.append(obj_packed_results["labels"].mode()[0])
+
+        results = padded_frames_results.to_dict()
+        results["scores"] = torch.stack(tracklet_scores, dim=0)
+        results["labels"] = torch.stack(tracklet_labels, dim=0)
+        if self.convert_mask_to_rle_for_video:
+            results["masks_rle"] = vid_masklets_rle_padded
+        # we keep the frame-level scores since it's needed by some evaluation scripts
+        results["per_frame_scores"] = padded_frames_results["scores"]
+
+        return {video_id: results}
+
+
+class PostProcessTracking(PostProcessImage):
+    """This module converts the model's output into the format expected by the coco api"""
+
+    def __init__(
+        self,
+        max_dets_per_img: int,
+        iou_type="bbox",
+        force_single_mask: bool = False,
+        **kwargs,
+    ) -> None:
+        super().__init__(max_dets_per_img=max_dets_per_img, iou_type=iou_type, **kwargs)
+        self.force_single_mask = force_single_mask
+
+    def process_results(
+        self, find_stages, find_metadatas: BatchedInferenceMetadata, **kwargs
+    ):
+        assert len(find_stages) == len(find_metadatas)
+        results = {}
+        for outputs, meta in zip(find_stages, find_metadatas):
+            if self.force_single_mask:
+                scores, labels = outputs["pred_logits"].max(-1)
+                m = []
+                for i in range(len(outputs["pred_masks"])):
+                    score, idx = scores[i].max(0)
+                    m.append(outputs["pred_masks"][i][idx])
+                outputs["pred_masks"] = torch.stack(m, 0).unsqueeze(1)
+            detection_results = self(outputs, meta.original_size, consistent=False)
+            assert len(detection_results) == len(meta.coco_image_id)
+            results.update(
+                {
+                    (media_id.item(), object_id.item(), frame_index.item()): result
+                    for media_id, object_id, frame_index, result in zip(
+                        meta.original_image_id,
+                        meta.object_id,
+                        meta.frame_index,
+                        detection_results,
+                    )
+                }
+            )
+        return results
+
+
+class PostProcessCounting(nn.Module):
+    """This module converts the model's output to be evaluated for counting tasks"""
+
+    def __init__(
+        self,
+        use_original_ids: bool = False,
+        threshold: float = 0.5,
+        use_presence: bool = False,
+    ) -> None:
+        """
+        Args:
+            use_original_ids: whether to use the original image ids or the coco ids
+            threshold: threshold for counting (values above this are counted)
+        """
+        super().__init__()
+        self.use_original_ids = use_original_ids
+        self.threshold = threshold
+        self.use_presence = use_presence
+
+    def forward(self, outputs, target_sizes):
+        """Perform the computation
+        Parameters:
+            outputs: raw outputs of the model
+            target_sizes: tensor of dimension [batch_size x 2] containing the size of each images of the batch
+        """
+        # Extract scores from model outputs and apply sigmoid
+        scores = torch.sigmoid(outputs["pred_logits"]).squeeze(-1)  # [B, N]
+        if self.use_presence:
+            presence_score = outputs["presence_logit_dec"].sigmoid()
+            if presence_score.ndim == 1:
+                presence_score = presence_score.unsqueeze(1)  # [B, 1]
+            scores = scores * presence_score  # [B, N]
+
+        # Calculate counts by summing values above threshold
+        counts = (scores > self.threshold).float().sum(dim=1)
+
+        assert len(counts) == len(target_sizes)
+        results = []
+        for count in counts:
+            results.append({"count": count.item()})
+
+        return results
+
+    @torch.no_grad()
+    def process_results(
+        self, find_stages, find_metadatas: List[BatchedInferenceMetadata], **kwargs
+    ):
+        assert len(find_stages) == len(find_metadatas)
+        results = {}
+        for outputs, meta in zip(find_stages, find_metadatas):
+            detection_results = self(
+                outputs,
+                meta.original_size,
+            )
+            ids = (
+                meta.original_image_id if self.use_original_ids else meta.coco_image_id
+            )
+            assert len(detection_results) == len(ids)
+            for img_id, result in zip(ids, detection_results):
+                results[img_id.item()] = result
+
+        return results

sam3/eval/saco_veval_eval.py

@@ -0,0 +1,155 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+import argparse
+import json
+import os
+from collections import defaultdict
+
+from iopath.common.file_io import g_pathmgr
+from sam3.eval.saco_veval_evaluators import (
+    VideoCGF1Evaluator,
+    VideoPhraseApEvaluator,
+    VideoPhraseHotaEvaluator,
+    VideoTetaEvaluator,
+    YTVISPredFileEvaluator,
+)
+
+
+class VEvalEvaluator:
+    def __init__(self, gt_annot_file: str, eval_res_file: str):
+        self.gt_annot_file = gt_annot_file
+        self.eval_res_file = eval_res_file
+        self.evaluators = [
+            # mAP
+            YTVISPredFileEvaluator(gt_annot_file),
+            # Phrase AP
+            VideoPhraseApEvaluator(gt_annot_file),
+            # TETA
+            VideoTetaEvaluator(gt_annot_file, use_mask=True, is_exhaustive=True),
+            # HOTA
+            VideoPhraseHotaEvaluator(gt_annot_file),
+            # cgF1
+            VideoCGF1Evaluator(gt_annot_file),
+        ]
+
+    def run_eval(self, pred_file: str):
+        dataset_results = {}
+        video_np_results = defaultdict(dict)
+        for evaluator in self.evaluators:
+            d_res, v_np_res = evaluator.evaluate(pred_file)
+            dataset_results.update(d_res)
+            for (video_id, category_id), res in v_np_res.items():
+                video_np_results[(video_id, category_id)].update(res)
+
+        if len(dataset_results) == 0:
+            dataset_results = {"": 0.0}
+
+        formatted_video_np_results = [
+            {"video_id": video_id, "category_id": category_id, **res}
+            for (video_id, category_id), res in video_np_results.items()
+        ]
+        eval_metrics = {
+            "dataset_results": dataset_results,
+            "video_np_results": formatted_video_np_results,
+        }
+
+        with g_pathmgr.open(self.eval_res_file, "w") as f:
+            json.dump(eval_metrics, f)
+
+        return eval_metrics
+
+
+def run_main_all(dataset_name, args):
+    gt_annot_file = os.path.join(args.gt_annot_dir, dataset_name + ".json")
+    pred_file = os.path.join(args.pred_dir, dataset_name + "_preds.json")
+    eval_res_file = os.path.join(args.eval_res_dir, dataset_name + "_eval_res.json")
+    print(f"=== Running evaluation for Pred {pred_file} vs GT {gt_annot_file} ===")
+    veval_evaluator = VEvalEvaluator(
+        gt_annot_file=gt_annot_file, eval_res_file=eval_res_file
+    )
+    _ = veval_evaluator.run_eval(pred_file=pred_file)
+
+    print(f"=== Results saved to {eval_res_file} ===")
+
+
+def main_all(args):
+    saco_veval_dataset_names = [
+        "saco_veval_sav_test",
+        "saco_veval_sav_val",
+        "saco_veval_yt1b_test",
+        "saco_veval_yt1b_val",
+        "saco_veval_smartglasses_test",
+        "saco_veval_smartglasses_val",
+    ]
+
+    # multiprocessing may not really work as inner evaluator also using multiprocessing
+    # so we just for loop
+    for dataset_name in saco_veval_dataset_names:
+        print(f"=== Running evaluation for dataset {dataset_name} ===")
+        run_main_all(dataset_name=dataset_name, args=args)
+
+
+def main_one(args):
+    gt_annot_file = args.gt_annot_file
+    pred_file = args.pred_file
+    eval_res_file = args.eval_res_file
+
+    print(f"=== Running evaluation for Pred {pred_file} vs GT {gt_annot_file} ===")
+    veval_evaluator = VEvalEvaluator(
+        gt_annot_file=gt_annot_file, eval_res_file=eval_res_file
+    )
+    _ = veval_evaluator.run_eval(pred_file=pred_file)
+
+    print(f"=== Results saved to {eval_res_file} ===")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Run video grounding evaluators")
+
+    # Create subparsers for different commands
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    # Run evaluation for all datasets
+    all_parser = subparsers.add_parser("all", help="Run evaluation for all datasets")
+    all_parser.add_argument(
+        "--gt_annot_dir",
+        type=str,
+        help="Directory that contains the ground truth annotation files",
+    )
+    all_parser.add_argument(
+        "--pred_dir",
+        type=str,
+        help="Directory that contains the prediction files",
+    )
+    all_parser.add_argument(
+        "--eval_res_dir",
+        type=str,
+        help="Directory that contains the eval results files",
+    )
+    all_parser.set_defaults(func=main_all)
+
+    # Run evaluation for one dataset
+    one_parser = subparsers.add_parser("one", help="Run evaluation for one dataset")
+    one_parser.add_argument(
+        "--gt_annot_file",
+        type=str,
+        help="Path to the ground truth annotation file",
+    )
+    one_parser.add_argument(
+        "--pred_file",
+        type=str,
+        help="Path to the prediction file",
+    )
+    one_parser.add_argument(
+        "--eval_res_file",
+        type=str,
+        help="Path to the eval results file",
+    )
+    one_parser.set_defaults(func=main_one)
+
+    # Parse and dispatch
+    args = parser.parse_args()
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()

sam3/eval/saco_veval_evaluators.py

@@ -0,0 +1,838 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+import json
+import os
+import tempfile
+from collections import defaultdict
+from typing import Dict, Optional, Sequence, Tuple
+
+import numpy as np
+import pycocotools.mask
+from sam3.eval.cgf1_eval import CGF1_METRICS
+from sam3.eval.conversion_util import (
+    convert_ytbvis_to_cocovid_gt,
+    convert_ytbvis_to_cocovid_pred,
+)
+from sam3.eval.hota_eval_toolkit.run_ytvis_eval import run_ytvis_eval
+from sam3.eval.teta_eval_toolkit import config, Evaluator, metrics
+from sam3.eval.teta_eval_toolkit.datasets import COCO, TAO
+from sam3.eval.ytvis_coco_wrapper import YTVIS
+from sam3.eval.ytvis_eval import VideoDemoF1Eval, YTVISeval
+from sam3.train.nms_helper import process_frame_level_nms, process_track_level_nms
+
+
+def _get_metric_index(metric_name: str, iou_threshold: Optional[float] = None) -> int:
+    """
+    Find the index of a metric in CGF1_METRICS by name and IoU threshold.
+
+    Args:
+        metric_name: Name of the metric (e.g., "cgF1", "precision", "recall")
+        iou_threshold: IoU threshold (None for average over 0.5:0.95, or specific value like 0.5, 0.75)
+
+    Returns:
+        Index of the metric in CGF1_METRICS
+
+    Raises:
+        ValueError: If metric not found
+    """
+    for idx, metric in enumerate(CGF1_METRICS):
+        if metric.name == metric_name and metric.iou_threshold == iou_threshold:
+            return idx
+    raise ValueError(
+        f"Metric '{metric_name}' with IoU threshold {iou_threshold} not found in CGF1_METRICS"
+    )
+
+
+class BasePredFileEvaluator:
+    """A base class for evaluating a prediction file."""
+
+    pass
+
+
+class YTVISPredFileEvaluator(BasePredFileEvaluator):
+    """Evaluate class mAP for YT-VIS prediction files."""
+
+    def __init__(
+        self,
+        gt_ann_file: str,
+        dataset_name: str = "video",
+        iou_types: Optional[Sequence[str]] = None,
+    ):
+        self.gt_ann_file = gt_ann_file
+        self.dataset_name = dataset_name
+        self.iou_types = list(iou_types) if iou_types is not None else ["bbox", "segm"]
+        assert all(iou_type in ["bbox", "segm"] for iou_type in self.iou_types)
+
+    def evaluate(self, pred_file: str) -> Dict[str, float]:
+        # use our internal video evaluation toolkit for YT-VIS pred file
+        # (i.e. the same one we're using for video phrase AP)
+        results = {}
+        use_cats = True  # YT-VIS mAP evaluation uses categories
+        ytvisGT = YTVIS(self.gt_ann_file, ignore_gt_cats=not use_cats)
+        # the original YT-VIS GT annotations have uncompressed RLEs ("counts" is an integer list)
+        # rather than compressed RLEs ("counts" is a string), so we first convert them here.
+        if "segm" in self.iou_types:
+            for ann in ytvisGT.dataset["annotations"]:
+                ann["segmentations"] = [
+                    _compress_rle(rle) for rle in ann["segmentations"]
+                ]
+
+        with open(pred_file) as f:
+            dt = json.load(f)
+        # Our prediction file saves "video_id" and absolute (unnormalized) boxes.
+        # Note that we should use the official (original) YT-VIS annotations (i.e. the one
+        # saved via "scripts/datasets/training/ytvis_split.py", instead of the one saved
+        # via "scripts/api_db_to_ytvis_json.py") in this evaluator, which contain absolute
+        # boxes coordinates in its GT annotations.
+        for d in dt:
+            d["image_id"] = d["video_id"]
+        ytvisDT = ytvisGT.loadRes(dt)
+
+        for iou_type in self.iou_types:
+            ytvisEval = YTVISeval(ytvisGT, ytvisDT, iou_type)
+
+            # set the area ranges for small, medium, and large objects (using
+            # absolute pixel areas) as in the official YT-VIS evaluation toolkit:
+            # https://github.com/achalddave/ytvosapi/blob/eca601117c9f86bad084cb91f1d918e9ab665a75/PythonAPI/ytvostools/ytvoseval.py#L538
+            ytvisEval.params.areaRng = [
+                [0**2, 1e5**2],
+                [0**2, 128**2],
+                [128**2, 256**2],
+                [256**2, 1e5**2],
+            ]
+            ytvisEval.params.areaRngLbl = ["all", "small", "medium", "large"]
+            ytvisEval.params.useCats = use_cats
+
+            ytvisEval.evaluate()
+            ytvisEval.accumulate()
+            ytvisEval.summarize()
+            result_key = f"{self.dataset_name}_{'mask' if iou_type == 'segm' else 'bbox'}_mAP_50_95"
+            results[result_key] = ytvisEval.stats[0]
+
+        # video-NP level results not supported for `YTVISPredFileEvaluator` yet
+        video_np_level_results = {}
+        return results, video_np_level_results
+
+
+class VideoPhraseApEvaluator(BasePredFileEvaluator):
+    """Evaluate Video Phrase AP with YT-VIS format prediction and GT files."""
+
+    def __init__(
+        self,
+        gt_ann_file: str,
+        dataset_name: str = "video",
+        iou_types: Optional[Sequence[str]] = None,
+    ):
+        self.gt_ann_file = gt_ann_file
+        self.dataset_name = dataset_name
+        self.iou_types = list(iou_types) if iou_types is not None else ["bbox", "segm"]
+        assert all(iou_type in ["bbox", "segm"] for iou_type in self.iou_types)
+
+    def evaluate(self, pred_file: str) -> Dict[str, float]:
+        with open(self.gt_ann_file) as f:
+            gt = json.load(f)
+        with open(pred_file) as f:
+            dt = json.load(f)
+        # For phrase AP and demo F1 evaluation, we need to remap each pair of (video_id, category_id) to
+        # a new unique video_id, so that we don't mix detections from different categories under `useCat=False`
+        gt, dt = remap_video_category_pairs_to_unique_video_ids(gt, dt)
+        if "segm" in self.iou_types:
+            for ann in gt["annotations"]:
+                ann["segmentations"] = [
+                    _compress_rle(rle) for rle in ann["segmentations"]
+                ]
+        for d in dt:
+            d["image_id"] = d["video_id"]
+
+        results = {}
+        use_cats = False  # Phrase AP evaluation does not use categories
+        ytvisGT = YTVIS(annotation_file=None, ignore_gt_cats=not use_cats)
+        ytvisGT.dataset = gt
+        ytvisGT.createIndex()
+        ytvisDT = ytvisGT.loadRes(dt)
+
+        for iou_type in self.iou_types:
+            phraseApEval = YTVISeval(ytvisGT, ytvisDT, iou_type)
+
+            # set the area ranges for small, medium, and large objects (using
+            # absolute pixel areas) as in the official YT-VIS evaluation toolkit:
+            # https://github.com/achalddave/ytvosapi/blob/eca601117c9f86bad084cb91f1d918e9ab665a75/PythonAPI/ytvostools/ytvoseval.py#L538
+            phraseApEval.params.areaRng = [
+                [0**2, 1e5**2],
+                [0**2, 128**2],
+                [128**2, 256**2],
+                [256**2, 1e5**2],
+            ]
+            phraseApEval.params.areaRngLbl = ["all", "small", "medium", "large"]
+            phraseApEval.params.useCats = use_cats
+
+            phraseApEval.evaluate()
+            phraseApEval.accumulate()
+            phraseApEval.summarize()
+            result_prefix = f"{self.dataset_name}"
+            result_prefix += f"_{'mask' if iou_type == 'segm' else 'bbox'}_phrase_ap"
+            # fetch Phrase AP results from the corresponding indices in `phraseApEval.stats`
+            # (see `_summarizeDets` in https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocotools/cocoeval.py)
+            results[result_prefix + "_50_95"] = phraseApEval.stats[0]  # IoU=0.5:0.95
+            results[result_prefix + "_50"] = phraseApEval.stats[1]  # IoU=0.5
+            results[result_prefix + "_75"] = phraseApEval.stats[2]  # IoU=0.75
+
+        # video-NP level results not supported for `VideoPhraseApEvaluator` yet
+        video_np_level_results = {}
+        return results, video_np_level_results
+
+
+class VideoCGF1Evaluator(BasePredFileEvaluator):
+    """Evaluate Video Demo F1 with YT-VIS format prediction and GT files."""
+
+    def __init__(
+        self,
+        gt_ann_file: str,
+        dataset_name: str = "video",
+        prob_thresh: float = 0.5,
+        iou_types: Optional[Sequence[str]] = None,
+    ):
+        self.gt_ann_file = gt_ann_file
+        self.dataset_name = dataset_name
+        self.prob_thresh = prob_thresh
+        self.iou_types = list(iou_types) if iou_types is not None else ["bbox", "segm"]
+        assert all(iou_type in ["bbox", "segm"] for iou_type in self.iou_types)
+
+    def evaluate(self, pred_file: str) -> Dict[str, float]:
+        with open(self.gt_ann_file) as f:
+            gt = json.load(f)
+        with open(pred_file) as f:
+            dt = json.load(f)
+        # compute IL_MCC and CG-F1 can only be computed if we have "video_np_pairs" keys in the GT JSON
+        compute_ilmcc_and_cgf1 = "video_np_pairs" in gt
+        if not compute_ilmcc_and_cgf1:
+            print(
+                f"Warning: IL_MCC and CG-F1 are not computed for {pred_file=} as it does not have 'video_np_pairs' keys in the GT JSON"
+            )
+        # For phrase AP and demo F1 evaluation, we need to remap each pair of (video_id, category_id) to
+        # a new unique video_id, so that we don't mix detections from different categories under `useCat=False`
+        gt, dt = remap_video_category_pairs_to_unique_video_ids(
+            gt, dt, add_negative_np_pairs=compute_ilmcc_and_cgf1
+        )
+        if "segm" in self.iou_types:
+            for ann in gt["annotations"]:
+                ann["segmentations"] = [
+                    _compress_rle(rle) for rle in ann["segmentations"]
+                ]
+        for d in dt:
+            d["image_id"] = d["video_id"]
+
+        results = {}
+        use_cats = False  # Demo F1 evaluation does not use categories
+        ytvisGT = YTVIS(annotation_file=None, ignore_gt_cats=not use_cats)
+        ytvisGT.dataset = gt
+        ytvisGT.createIndex()
+        ytvisDT = ytvisGT.loadRes(dt)
+
+        video_np_level_results = {}
+        for iou_type in self.iou_types:
+            demoF1Eval = VideoDemoF1Eval(ytvisGT, ytvisDT, iou_type, self.prob_thresh)
+
+            demoF1Eval.params.useCats = use_cats
+            demoF1Eval.params.areaRng = [[0**2, 1e5**2]]
+            demoF1Eval.params.areaRngLbl = ["all"]
+            demoF1Eval.params.maxDets = [100000]
+
+            demoF1Eval.evaluate()
+            demoF1Eval.accumulate()
+            demoF1Eval.summarize()
+            result_prefix = f"{self.dataset_name}"
+            result_prefix += f"_{'mask' if iou_type == 'segm' else 'bbox'}_demo"
+
+            stats = demoF1Eval.stats
+
+            if compute_ilmcc_and_cgf1:
+                # Average IoU threshold (0.5:0.95)
+                cgf1_micro_avg_idx = _get_metric_index("cgF1", None)
+                positive_micro_f1_avg_idx = _get_metric_index("positive_micro_F1", None)
+                ilmcc_avg_idx = _get_metric_index("IL_MCC", None)
+                results[result_prefix + "_cgf1_micro_50_95"] = stats[cgf1_micro_avg_idx]
+                results[result_prefix + "_ilmcc_50_95"] = stats[ilmcc_avg_idx]
+                results[result_prefix + "_positive_micro_f1_50_95"] = stats[
+                    positive_micro_f1_avg_idx
+                ]
+
+                # IoU = 0.5
+                cgf1_micro_50_idx = _get_metric_index("cgF1", 0.5)
+                positive_micro_f1_50_idx = _get_metric_index("positive_micro_F1", 0.5)
+                results[result_prefix + "_cgf1_micro_50"] = stats[cgf1_micro_50_idx]
+                results[result_prefix + "_ilmcc_50"] = float(
+                    np.array(stats[cgf1_micro_50_idx])
+                    / np.array(stats[positive_micro_f1_50_idx])
+                )
+                results[result_prefix + "_positive_micro_f1_50"] = stats[
+                    positive_micro_f1_50_idx
+                ]
+
+                # IoU = 0.75
+                cgf1_micro_75_idx = _get_metric_index("cgF1", 0.75)
+                positive_micro_f1_75_idx = _get_metric_index("positive_micro_F1", 0.75)
+                results[result_prefix + "_cgf1_micro_75"] = stats[cgf1_micro_75_idx]
+                results[result_prefix + "_ilmcc_75"] = float(
+                    np.array(stats[cgf1_micro_75_idx])
+                    / np.array(stats[positive_micro_f1_75_idx])
+                )
+                results[result_prefix + "_positive_micro_f1_75"] = stats[
+                    positive_micro_f1_75_idx
+                ]
+
+            self.extract_video_np_level_results(demoF1Eval, video_np_level_results)
+
+        return results, video_np_level_results
+
+    def extract_video_np_level_results(self, demoF1Eval, video_np_level_results):
+        """Aggregate statistics for video-level metrics."""
+        num_iou_thrs = len(demoF1Eval.params.iouThrs)
+        iou_50_index = int(np.where(demoF1Eval.params.iouThrs == 0.5)[0])
+        iou_75_index = int(np.where(demoF1Eval.params.iouThrs == 0.75)[0])
+
+        result_prefix = "mask" if demoF1Eval.params.iouType == "segm" else "bbox"
+
+        assert len(demoF1Eval.evalImgs) == len(demoF1Eval.cocoGt.dataset["images"])
+        for i, video in enumerate(demoF1Eval.cocoGt.dataset["images"]):
+            # the original video id and category id before remapping
+            video_id = video["orig_video_id"]
+            category_id = video["orig_category_id"]
+            eval_img_dict = demoF1Eval.evalImgs[i]
+
+            TPs = eval_img_dict.get("TPs", np.zeros(num_iou_thrs, dtype=np.int64))
+            FPs = eval_img_dict.get("FPs", np.zeros(num_iou_thrs, dtype=np.int64))
+            FNs = eval_img_dict.get("FNs", np.zeros(num_iou_thrs, dtype=np.int64))
+            assert len(TPs) == len(FPs) == len(FNs) == num_iou_thrs
+            # F1 = 2*TP / (2*TP + FP + FN), and we set F1 to 1.0 if denominator is 0
+            denominator = 2 * TPs + FPs + FNs
+            F1s = np.where(denominator > 0, 2 * TPs / np.maximum(denominator, 1), 1.0)
+            local_results = {
+                f"{result_prefix}_TP_50_95": float(TPs.mean()),
+                f"{result_prefix}_FP_50_95": float(FPs.mean()),
+                f"{result_prefix}_FN_50_95": float(FNs.mean()),
+                f"{result_prefix}_F1_50_95": float(F1s.mean()),
+                f"{result_prefix}_TP_50": float(TPs[iou_50_index]),
+                f"{result_prefix}_FP_50": float(FPs[iou_50_index]),
+                f"{result_prefix}_FN_50": float(FNs[iou_50_index]),
+                f"{result_prefix}_F1_50": float(F1s[iou_50_index]),
+                f"{result_prefix}_TP_75": float(TPs[iou_75_index]),
+                f"{result_prefix}_FP_75": float(FPs[iou_75_index]),
+                f"{result_prefix}_FN_75": float(FNs[iou_75_index]),
+                f"{result_prefix}_F1_75": float(F1s[iou_75_index]),
+            }
+            if (video_id, category_id) not in video_np_level_results:
+                video_np_level_results[(video_id, category_id)] = {}
+            video_np_level_results[(video_id, category_id)].update(local_results)
+
+
+class VideoTetaEvaluator(BasePredFileEvaluator):
+    """Evaluate TETA metric using YouTubeVIS format prediction and GT files."""
+
+    def __init__(
+        self,
+        gt_ann_file: str,
+        dataset_name: str = "video",
+        tracker_name: str = "Sam3",
+        nms_threshold: float = 0.5,
+        nms_strategy: str = "none",  # "track", "frame", or "none"
+        prob_thresh: float = 0.5,
+        is_exhaustive: bool = False,
+        use_mask: bool = False,
+        num_parallel_cores: int = 8,
+    ):
+        self.gt_ann_file = gt_ann_file
+        self.dataset_name = dataset_name
+        self.tracker_name = tracker_name
+        self.nms_threshold = nms_threshold
+        self.nms_strategy = nms_strategy.lower()  # Convert to lowercase for consistency
+        self.prob_thresh = prob_thresh
+        self.metric_prefix = "TETA"
+        self.is_exhaustive = is_exhaustive
+        self.use_mask = use_mask
+        self.num_parallel_cores = num_parallel_cores
+
+        # Verify NMS strategy is valid
+        valid_strategies = ["track", "frame", "none"]
+        print("current nms_strategy:", self.nms_strategy)
+        if self.nms_strategy not in valid_strategies:
+            raise ValueError(
+                f"Invalid NMS strategy: {self.nms_strategy}. Must be one of {valid_strategies}"
+            )
+
+        print(f"Initialized VideoTetaEvaluator with NMS strategy: {self.nms_strategy}")
+        print(f"Probability threshold set to: {self.prob_thresh}")
+        print(f"Dataset exhaustivity set to: {self.is_exhaustive}")
+        print(f"Tracker name set to: {self.tracker_name}")
+        print(f"Dataset name set to: {self.dataset_name}")
+        print(f"Use mask set to: {self.use_mask}")
+
+    def process_predictions(self, pred_file: str, tmp_dir: str) -> str:
+        """Process predictions with selected NMS strategy"""
+        with open(pred_file, "r") as f:
+            raw_preds = json.load(f)
+        print(f"Processing predictions with {self.nms_strategy} NMS strategy")
+
+        # Filter by score threshold
+        if self.prob_thresh > 0:
+            raw_preds = [d for d in raw_preds if d["score"] >= self.prob_thresh]
+            print(
+                f"Filtered to {len(raw_preds)} predictions with score >= {self.prob_thresh}"
+            )
+        # Group predictions by video_id
+        video_groups = defaultdict(list)
+        for pred in raw_preds:
+            video_groups[pred["video_id"]].append(pred)
+        # Process based on NMS strategy
+        if self.nms_strategy == "track":
+            process_track_level_nms(video_groups, nms_threshold=self.nms_threshold)
+        elif self.nms_strategy == "frame":
+            process_frame_level_nms(video_groups, nms_threshold=self.nms_threshold)
+        elif self.nms_strategy == "none":
+            print("Skipping NMS processing as strategy is set to 'none'")
+            # No processing needed for "none" strategy
+        # Save processed predictions
+        processed_preds = [
+            track for tracks in video_groups.values() for track in tracks
+        ]
+        processed_path = os.path.join(tmp_dir, "processed_preds.json")
+        with open(processed_path, "w") as f:
+            json.dump(processed_preds, f)
+
+        print(f"Saved processed predictions to {processed_path}")
+        return processed_path
+
+    def evaluate(self, pred_file: str) -> Tuple[Dict[str, float], Dict]:
+        """Main evaluation method"""
+
+        print(f"Evaluating TETA Metric with {self.nms_strategy.upper()} NMS strategy")
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            # Process predictions first
+            processed_pred_file = self.process_predictions(pred_file, tmp_dir)
+
+            # Convert GT to COCO-vid format
+            gt_dir = os.path.join(tmp_dir, "gt")
+            os.makedirs(gt_dir, exist_ok=True)
+            gt_coco_path = os.path.join(gt_dir, "annotations.json")
+            convert_ytbvis_to_cocovid_gt(self.gt_ann_file, gt_coco_path)
+
+            # Convert processed predictions to COCO-vid format
+            pred_dir = os.path.join(tmp_dir, "predictions")
+            tracker_dir = os.path.join(pred_dir, self.tracker_name)
+            os.makedirs(tracker_dir, exist_ok=True)
+            pred_coco_path = os.path.join(tracker_dir, "track_results_cocofmt.json")
+            convert_ytbvis_to_cocovid_pred(
+                youtubevis_pred_path=processed_pred_file,
+                converted_dataset_path=gt_coco_path,
+                output_path=pred_coco_path,
+            )
+            # Configure TETA evaluator
+            default_eval_config = config.get_default_eval_config()
+            default_eval_config["PRINT_ONLY_COMBINED"] = True
+            default_eval_config["DISPLAY_LESS_PROGRESS"] = True
+            default_eval_config["OUTPUT_TEMP_RAW_DATA"] = True
+            default_eval_config["NUM_PARALLEL_CORES"] = self.num_parallel_cores
+            default_dataset_config = config.get_default_dataset_config()
+            default_dataset_config["TRACKERS_TO_EVAL"] = [self.tracker_name]
+            default_dataset_config["GT_FOLDER"] = gt_dir
+            default_dataset_config["OUTPUT_FOLDER"] = pred_dir
+            default_dataset_config["TRACKER_SUB_FOLDER"] = tracker_dir
+            default_dataset_config["USE_MASK"] = self.use_mask
+
+            evaluator = Evaluator(default_eval_config)
+            if self.is_exhaustive:
+                dataset_list = [COCO(default_dataset_config)]
+                dataset_parsing_key = "COCO"
+            else:
+                dataset_list = [TAO(default_dataset_config)]
+                dataset_parsing_key = "TAO"
+
+            # Run evaluation
+            eval_results, _ = evaluator.evaluate(
+                dataset_list, [metrics.TETA(exhaustive=self.is_exhaustive)]
+            )
+
+            # Extract and format results
+            results = {
+                f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_teta": float(
+                    eval_results[dataset_parsing_key]["TETA"][0]
+                ),
+                f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_loc_a": float(
+                    eval_results[dataset_parsing_key]["TETA"][1]
+                ),
+                f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_assoc_a": float(
+                    eval_results[dataset_parsing_key]["TETA"][2]
+                ),
+                f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_cls_a": float(
+                    eval_results[dataset_parsing_key]["TETA"][3]
+                ),
+                f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_loc_re": float(
+                    eval_results[dataset_parsing_key]["TETA"][4]
+                ),
+                f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_loc_pr": float(
+                    eval_results[dataset_parsing_key]["TETA"][5]
+                ),
+                f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_assoc_re": float(
+                    eval_results[dataset_parsing_key]["TETA"][6]
+                ),
+                f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_assoc_pr": float(
+                    eval_results[dataset_parsing_key]["TETA"][7]
+                ),
+                f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_cls_re": float(
+                    eval_results[dataset_parsing_key]["TETA"][8]
+                ),
+                f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_cls_pr": float(
+                    eval_results[dataset_parsing_key]["TETA"][9]
+                ),
+            }
+
+        # video-NP level results not supported for `VideoTetaEvaluator` yet
+        video_np_level_results = {}
+        return results, video_np_level_results
+
+
+class VideoPhraseHotaEvaluator(BasePredFileEvaluator):
+    """Evaluate Video Phrase HOTA with YT-VIS format prediction and GT files."""
+
+    def __init__(
+        self,
+        gt_ann_file: str,
+        dataset_name: str = "video",
+        prob_thresh: float = 0.5,
+        iou_types: Optional[Sequence[str]] = None,
+        compute_video_mot_hota: bool = False,
+    ):
+        self.gt_ann_file = gt_ann_file
+        self.dataset_name = dataset_name
+        self.prob_thresh = prob_thresh
+        self.metric_prefix = "phrase"
+        # the list of metrics to collect from the HOTA evaluation results
+        self.metric_to_collect = [
+            "HOTA",
+            "DetA",
+            "AssA",
+            "DetRe",
+            "DetPr",
+            "AssRe",
+            "AssPr",
+            "LocA",
+            "OWTA",
+        ]
+        self.iou_types = list(iou_types) if iou_types is not None else ["bbox", "segm"]
+        assert all(iou_type in ["bbox", "segm"] for iou_type in self.iou_types)
+
+        # If True, compute video MOT HOTA, aggregating predictions/GT from all categories.
+        self.compute_video_mot_hota = compute_video_mot_hota
+
+    def evaluate(self, pred_file: str) -> Dict[str, float]:
+        # use the YT-VIS evaluation toolkit in TrackEval
+
+        with open(self.gt_ann_file) as f:
+            gt = json.load(f)
+        with open(pred_file) as f:
+            dt = json.load(f)
+        # keep only predictions with score above the probability threshold
+        dt = [d for d in dt if d["score"] > self.prob_thresh]
+        for d in dt:
+            assert len(d["areas"]) == len(d["bboxes"])
+            assert len(d["areas"]) == len(d["segmentations"])
+            # remove empty boxes (otherwise they will count as false positives for during
+            # per-frame detection accuracy in HOTA evaluation)
+            for t in range(len(d["bboxes"])):
+                bbox = d["bboxes"][t]
+                if d["areas"][t] == 0 or bbox is None or all(x == 0 for x in bbox):
+                    d["segmentations"][t] = None
+                    d["bboxes"][t] = None
+                    d["areas"][t] = None
+            # check that box occurence and mask occurence are consistent
+            for bbox, mask, area in zip(d["bboxes"], d["segmentations"], d["areas"]):
+                assert (area is None) == (bbox is None)
+                assert (area is None) == (mask is None)
+            # set all scores to 1.0 for HOTA evaluation (just like Demo F1, the exact score
+            # value is not used in HOTA metrics; it will be treated as a detection prediction
+            # as long as its score is above the threshold)
+            d["score"] = 1.0
+
+        # remap the GT and DT annotations for phrase HOTA evaluation
+        gt = _fill_in_ann_height_width(gt)
+        if not self.compute_video_mot_hota:
+            # remap the GT and DT annotations for phrase HOTA evaluation
+            gt, dt = self._remap_gt_dt(gt, dt)
+        else:
+            # Compute video-level MOT HOTA
+            # Apply track-level NMS
+            video_groups = defaultdict(list)
+            for pred in dt:
+                video_groups[pred["video_id"]].append(pred)
+            process_track_level_nms(video_groups, nms_threshold=0.5)
+            dt = [track for tracks in video_groups.values() for track in tracks]
+
+            # Remap GT track ids for class-agnostic HOTA
+            gt, dt = remap_gt_dt_class_agnostic(gt, dt)
+
+        # run the HOTA evaluation using TrackEval on the remapped (video_id, category_id) pairs
+        out_dict = {}
+        video_np_level_results = {}
+        for iou_type in self.iou_types:
+            output_res, _ = run_ytvis_eval(
+                args=[
+                    "--METRICS",
+                    "HOTA",
+                    "--IOU_TYPE",
+                    iou_type,
+                    "--DATASET_NAME",
+                    self.dataset_name,
+                    "--USE_PARALLEL",
+                    "True",
+                    "--NUM_PARALLEL_CORES",
+                    "8",
+                    "--PLOT_CURVES",
+                    "False",
+                    "--LOG_ON_ERROR",
+                    "None",
+                    "--PRINT_ONLY_COMBINED",
+                    "True",
+                    "--OUTPUT_SUMMARY",
+                    "False",
+                    "--OUTPUT_DETAILED",
+                    "False",
+                    "--TIME_PROGRESS",
+                    "False",
+                    "--PRINT_CONFIG",
+                    "False",
+                ],
+                gt_json=gt,
+                dt_json=dt,
+            )
+            self.extract_video_np_level_results(
+                iou_type=iou_type,
+                remapped_gt=gt,
+                raw_results=output_res[self.dataset_name]["tracker"],
+                video_np_level_results=video_np_level_results,
+            )
+
+            def _summarize_results(output_res, iou_type, field, suffix):
+                eval_res = output_res[self.dataset_name]["tracker"][field]
+                result_prefix = f"{self.dataset_name}_{'mask' if iou_type == 'segm' else 'bbox'}_{suffix}"
+                for metric_name in self.metric_to_collect:
+                    eval_res_hota = eval_res["cls_comb_cls_av"]["HOTA"]
+                    result_key = f"{result_prefix}_{self.metric_prefix}_{metric_name}"
+                    result_value = float(np.mean(eval_res_hota[metric_name]))
+                    out_dict[result_key] = result_value
+
+            _summarize_results(output_res, iou_type, "COMBINED_SEQ", "all")
+            if "COMBINED_SEQ_CHALLENGING" in output_res[self.dataset_name]["tracker"]:
+                _summarize_results(
+                    output_res, iou_type, "COMBINED_SEQ_CHALLENGING", "challenging"
+                )
+
+        # video-NP level results not supported for `VideoPhraseHotaEvaluator` yet
+        return out_dict, video_np_level_results
+
+    def _remap_gt_dt(self, gt, dt):
+        # For phrase HOTA evaluation, we need to remap each pair of (video_id, category_id) to
+        # a new unique video_id, so that we don't mix detections from different categories
+        gt, dt = remap_video_category_pairs_to_unique_video_ids(gt, dt)
+        # We further map all the categories to category_id=1 in HOTA evaluation toolkit
+        # for phrase HOTA (similar to "useCat=False" for video phrase AP)
+        remapped_category_id = 1
+        gt["categories"] = [
+            {
+                "supercategory": "object",
+                "id": remapped_category_id,
+                "name": "_REMAPPED_FOR_PHRASE_METRICS_",
+            }
+        ]
+        for ann in gt["annotations"]:
+            ann["category_id"] = remapped_category_id
+        for d in dt:
+            d["category_id"] = remapped_category_id
+        # To be compatible with the TrackEval YT-VIS evaluation toolkit, we need to give
+        # unique filenames to each remapped video, so we add remapped video_id as prefix.
+        for video in gt["videos"]:
+            new_video_id = video["id"]
+            video["file_names"] = [
+                f"remapped_vid_{new_video_id:012d}/{name}"
+                for name in video["file_names"]
+            ]
+        return gt, dt
+
+    def extract_video_np_level_results(
+        self, iou_type, remapped_gt, raw_results, video_np_level_results
+    ):
+        """Aggregate statistics for video-level metrics."""
+        result_prefix = "mask" if iou_type == "segm" else "bbox"
+        for video in remapped_gt["videos"]:
+            # the original video id and category id before remapping
+            video_id = video["orig_video_id"]
+            category_id = video["orig_category_id"]
+            video_key = f"remapped_vid_{video['id']:012d}"
+            results = raw_results[video_key]["_REMAPPED_FOR_PHRASE_METRICS_"]["HOTA"]
+
+            local_results = {}
+            for metric_name in self.metric_to_collect:
+                result_key = f"{result_prefix}_{metric_name}"
+                local_results[result_key] = float(results[metric_name].mean())
+            if (video_id, category_id) not in video_np_level_results:
+                video_np_level_results[(video_id, category_id)] = {}
+            video_np_level_results[(video_id, category_id)].update(local_results)
+
+
+class VideoClassBasedHotaEvaluator(VideoPhraseHotaEvaluator):
+    def __init__(
+        self,
+        gt_ann_file: str,
+        dataset_name: str = "video",
+        prob_thresh: float = 0.5,
+    ):
+        super().__init__(gt_ann_file, dataset_name, prob_thresh)
+        self.metric_prefix = "class"
+
+    def _remap_gt_dt(self, gt, dt):
+        return gt, dt  # no remapping needed for class-based HOTA evaluation
+
+    def extract_video_np_level_results(self, *args, **kwargs):
+        pass  # no video-NP level results for class-based HOTA evaluation
+
+
+def _compress_rle(rle):
+    """Convert RLEs from uncompressed (integer list) to compressed (string) format."""
+    if rle is None:
+        return None
+    if isinstance(rle["counts"], list):
+        rle = pycocotools.mask.frPyObjects(rle, rle["size"][0], rle["size"][1])
+        rle["counts"] = rle["counts"].decode()
+    return rle
+
+
+def remap_video_category_pairs_to_unique_video_ids(
+    gt_json, dt_json, add_negative_np_pairs=False
+):
+    """
+    Remap each pair of (video_id, category_id) to a new unique video_id. This is useful
+    for phrase AP and demo F1 evaluation on videos, where we have `useCat=False` and
+    rely on separating different NPs (from the same video) into different new video ids,
+    so that we don't mix detections from different categories in computeIoU under `useCat=False`.
+
+    This is consistent with how do we phrase AP and demo F1 evaluation on images, where we
+    use a remapped unique coco_image_id for each image-NP pair (based in its query["id"] in
+    CustomCocoDetectionAPI.load_queries in modulated_detection_api.py)
+    """
+    # collect the unique video_id-category_id pairs
+    video_id_to_video = {v["id"]: v for v in gt_json["videos"]}
+    video_id_category_id_pairs = set()
+    for pred in dt_json:
+        video_id_category_id_pairs.add((pred["video_id"], pred["category_id"]))
+    for ann in gt_json["annotations"]:
+        video_id_category_id_pairs.add((ann["video_id"], ann["category_id"]))
+
+    # assign the video_id-category_id pairs to unique video ids
+    video_id_category_id_pairs = sorted(video_id_category_id_pairs)
+    video_id_category_id_to_new_video_id = {
+        pair: (i + 1) for i, pair in enumerate(video_id_category_id_pairs)
+    }
+    # also map the negative NP pairs -- this is needed for IL_MCC and CG-F1 evaluation
+    if add_negative_np_pairs:
+        for vnp in gt_json["video_np_pairs"]:
+            pair = (vnp["video_id"], vnp["category_id"])
+            if pair not in video_id_category_id_to_new_video_id:
+                video_id_category_id_to_new_video_id[pair] = (
+                    len(video_id_category_id_to_new_video_id) + 1
+                )
+
+    # map the "video_id" in predictions
+    for pred in dt_json:
+        pred["video_id"] = video_id_category_id_to_new_video_id[
+            (pred["video_id"], pred["category_id"])
+        ]
+    # map the "video_id" in gt_json["annotations"]
+    for ann in gt_json["annotations"]:
+        ann["video_id"] = video_id_category_id_to_new_video_id[
+            (ann["video_id"], ann["category_id"])
+        ]
+    # map and duplicate gt_json["videos"]
+    new_videos = []
+    for (
+        video_id,
+        category_id,
+    ), new_video_id in video_id_category_id_to_new_video_id.items():
+        video = video_id_to_video[video_id].copy()
+        video["id"] = new_video_id
+        # preserve the original video_id and category_id of each remapped video entry,
+        # so that we can associate sample-level eval metrics with the original video-NP pairs
+        video["orig_video_id"] = video_id
+        video["orig_category_id"] = category_id
+        new_videos.append(video)
+    gt_json["videos"] = new_videos
+
+    return gt_json, dt_json
+
+
+def remap_gt_dt_class_agnostic(gt, dt):
+    """
+    For class-agnostic HOTA, merge all GT tracks for each video (across NPs),
+    ensure unique track_ids, and set all category_id to 1.
+    Also, add orig_video_id and orig_category_id for compatibility.
+    """
+    # 1. Remap all GT track_ids to be unique per video
+    gt_anns_by_video = defaultdict(list)
+    for ann in gt["annotations"]:
+        gt_anns_by_video[ann["video_id"]].append(ann)
+
+    # Ensure unique track ids across tracks of all videos
+    next_tid = 1
+    for _, anns in gt_anns_by_video.items():
+        # Map old track_ids to new unique ones
+        old_to_new_tid = {}
+        for ann in anns:
+            old_tid = ann["id"]
+            if old_tid not in old_to_new_tid:
+                old_to_new_tid[old_tid] = next_tid
+                next_tid += 1
+            ann["id"] = old_to_new_tid[old_tid]
+            # Set category_id to 1 for class-agnostic
+            ann["category_id"] = 1
+
+    # Set all GT categories to a single category
+    gt["categories"] = [
+        {
+            "supercategory": "object",
+            "id": 1,
+            "name": "_REMAPPED_FOR_PHRASE_METRICS_",
+        }
+    ]
+
+    # Add orig_video_id and orig_category_id to each video for compatibility
+    anns_by_video = defaultdict(list)
+    for ann in gt["annotations"]:
+        anns_by_video[ann["video_id"]].append(ann)
+    for video in gt["videos"]:
+        video["orig_video_id"] = video["id"]
+        # Use the first annotation's original category_id if available, else None
+        orig_cat = (
+            anns_by_video[video["id"]][0]["category_id"]
+            if anns_by_video[video["id"]]
+            else None
+        )
+        video["orig_category_id"] = orig_cat
+        video["file_names"] = [
+            f"remapped_vid_{video['id']:012d}/{name}" for name in video["file_names"]
+        ]
+
+    # Set all DT category_id to 1
+    for d in dt:
+        d["category_id"] = 1
+    return gt, dt
+
+
+def _fill_in_ann_height_width(gt_json):
+    """Fill in missing height/width in GT annotations from its video info."""
+    video_id_to_video = {v["id"]: v for v in gt_json["videos"]}
+    for ann in gt_json["annotations"]:
+        if "height" not in ann or "width" not in ann:
+            video = video_id_to_video[ann["video_id"]]
+            if "height" not in ann:
+                ann["height"] = video["height"]
+            if "width" not in ann:
+                ann["width"] = video["width"]
+
+    return gt_json

sam3/eval/teta_eval_toolkit/__init__.py

@@ -0,0 +1,5 @@
+# fmt: off
+# flake8: noqa
+
+from . import config, datasets, metrics, utils
+from .eval import Evaluator

sam3/eval/teta_eval_toolkit/_timing.py

@@ -0,0 +1,69 @@
+# fmt: off
+# flake8: noqa
+
+import inspect
+from functools import wraps
+from time import perf_counter
+
+DO_TIMING = False
+DISPLAY_LESS_PROGRESS = False
+timer_dict = {}
+counter = 0
+
+
+def time(f):
+    @wraps(f)
+    def wrap(*args, **kw):
+        if DO_TIMING:
+            # Run function with timing
+            ts = perf_counter()
+            result = f(*args, **kw)
+            te = perf_counter()
+            tt = te - ts
+
+            # Get function name
+            arg_names = inspect.getfullargspec(f)[0]
+            if arg_names[0] == "self" and DISPLAY_LESS_PROGRESS:
+                return result
+            elif arg_names[0] == "self":
+                method_name = type(args[0]).__name__ + "." + f.__name__
+            else:
+                method_name = f.__name__
+
+            # Record accumulative time in each function for analysis
+            if method_name in timer_dict.keys():
+                timer_dict[method_name] += tt
+            else:
+                timer_dict[method_name] = tt
+
+            # If code is finished, display timing summary
+            if method_name == "Evaluator.evaluate":
+                print("")
+                print("Timing analysis:")
+                for key, value in timer_dict.items():
+                    print("%-70s %2.4f sec" % (key, value))
+            else:
+                # Get function argument values for printing special arguments of interest
+                arg_titles = ["tracker", "seq", "cls"]
+                arg_vals = []
+                for i, a in enumerate(arg_names):
+                    if a in arg_titles:
+                        arg_vals.append(args[i])
+                arg_text = "(" + ", ".join(arg_vals) + ")"
+
+                # Display methods and functions with different indentation.
+                if arg_names[0] == "self":
+                    print("%-74s %2.4f sec" % (" " * 4 + method_name + arg_text, tt))
+                elif arg_names[0] == "test":
+                    pass
+                else:
+                    global counter
+                    counter += 1
+                    print("%i %-70s %2.4f sec" % (counter, method_name + arg_text, tt))
+
+            return result
+        else:
+            # If config["TIME_PROGRESS"] is false, or config["USE_PARALLEL"] is true, run functions normally without timing.
+            return f(*args, **kw)
+
+    return wrap

sam3/eval/teta_eval_toolkit/config.py

@@ -0,0 +1,153 @@
+# fmt: off
+# flake8: noqa
+
+"""Config."""
+import argparse
+import os
+
+
+def parse_configs():
+    """Parse command line."""
+    default_eval_config = get_default_eval_config()
+    default_eval_config["DISPLAY_LESS_PROGRESS"] = True
+    default_dataset_config = get_default_dataset_config()
+    default_metrics_config = {"METRICS": ["TETA"]}
+    config = {
+        **default_eval_config,
+        **default_dataset_config,
+        **default_metrics_config,
+    }
+    parser = argparse.ArgumentParser()
+    for setting in config.keys():
+        if type(config[setting]) == list or type(config[setting]) == type(None):
+            parser.add_argument("--" + setting, nargs="+")
+        else:
+            parser.add_argument("--" + setting)
+    args = parser.parse_args().__dict__
+    for setting in args.keys():
+        if args[setting] is not None:
+            if type(config[setting]) == type(True):
+                if args[setting] == "True":
+                    x = True
+                elif args[setting] == "False":
+                    x = False
+                else:
+                    raise Exception(
+                        f"Command line parameter {setting} must be True/False"
+                    )
+            elif type(config[setting]) == type(1):
+                x = int(args[setting])
+            elif type(args[setting]) == type(None):
+                x = None
+            else:
+                x = args[setting]
+            config[setting] = x
+    eval_config = {k: v for k, v in config.items() if k in default_eval_config.keys()}
+    dataset_config = {
+        k: v for k, v in config.items() if k in default_dataset_config.keys()
+    }
+    metrics_config = {
+        k: v for k, v in config.items() if k in default_metrics_config.keys()
+    }
+
+    return eval_config, dataset_config, metrics_config
+
+
+def get_default_eval_config():
+    """Returns the default config values for evaluation."""
+    code_path = get_code_path()
+    default_config = {
+        "USE_PARALLEL": True,
+        "NUM_PARALLEL_CORES": 8,
+        "BREAK_ON_ERROR": True,
+        "RETURN_ON_ERROR": False,
+        "LOG_ON_ERROR": os.path.join(code_path, "error_log.txt"),
+        "PRINT_RESULTS": True,
+        "PRINT_ONLY_COMBINED": True,
+        "PRINT_CONFIG": True,
+        "TIME_PROGRESS": True,
+        "DISPLAY_LESS_PROGRESS": True,
+        "OUTPUT_SUMMARY": True,
+        "OUTPUT_EMPTY_CLASSES": True,
+        "OUTPUT_TEM_RAW_DATA": True,
+        "OUTPUT_PER_SEQ_RES": True,
+    }
+    return default_config
+
+
+def get_default_dataset_config():
+    """Default class config values"""
+    code_path = get_code_path()
+    default_config = {
+        "GT_FOLDER": os.path.join(
+            code_path, "data/gt/tao/tao_training"
+        ),  # Location of GT data
+        "TRACKERS_FOLDER": os.path.join(
+            code_path, "data/trackers/tao/tao_training"
+        ),  # Trackers location
+        "OUTPUT_FOLDER": None,  # Where to save eval results (if None, same as TRACKERS_FOLDER)
+        "TRACKERS_TO_EVAL": ['TETer'],  # Filenames of trackers to eval (if None, all in folder)
+        "CLASSES_TO_EVAL": None,  # Classes to eval (if None, all classes)
+        "SPLIT_TO_EVAL": "training",  # Valid: 'training', 'val'
+        "PRINT_CONFIG": True,  # Whether to print current config
+        "TRACKER_SUB_FOLDER": "data",  # Tracker files are in TRACKER_FOLDER/tracker_name/TRACKER_SUB_FOLDER
+        "OUTPUT_SUB_FOLDER": "",  # Output files are saved in OUTPUT_FOLDER/tracker_name/OUTPUT_SUB_FOLDER
+        "TRACKER_DISPLAY_NAMES": None,  # Names of trackers to display, if None: TRACKERS_TO_EVAL
+        "MAX_DETECTIONS": 0,  # Number of maximal allowed detections per image (0 for unlimited)
+        "USE_MASK": False,  # Whether to use mask data for evaluation
+    }
+    return default_config
+
+
+def init_config(config, default_config, name=None):
+    """Initialize non-given config values with defaults."""
+    if config is None:
+        config = default_config
+    else:
+        for k in default_config.keys():
+            if k not in config.keys():
+                config[k] = default_config[k]
+    if name and config["PRINT_CONFIG"]:
+        print("\n%s Config:" % name)
+        for c in config.keys():
+            print("%-20s : %-30s" % (c, config[c]))
+    return config
+
+
+def update_config(config):
+    """
+    Parse the arguments of a script and updates the config values for a given value if specified in the arguments.
+    :param config: the config to update
+    :return: the updated config
+    """
+    parser = argparse.ArgumentParser()
+    for setting in config.keys():
+        if type(config[setting]) == list or type(config[setting]) == type(None):
+            parser.add_argument("--" + setting, nargs="+")
+        else:
+            parser.add_argument("--" + setting)
+    args = parser.parse_args().__dict__
+    for setting in args.keys():
+        if args[setting] is not None:
+            if type(config[setting]) == type(True):
+                if args[setting] == "True":
+                    x = True
+                elif args[setting] == "False":
+                    x = False
+                else:
+                    raise Exception(
+                        "Command line parameter " + setting + "must be True or False"
+                    )
+            elif type(config[setting]) == type(1):
+                x = int(args[setting])
+            elif type(args[setting]) == type(None):
+                x = None
+            else:
+                x = args[setting]
+            config[setting] = x
+    return config
+
+
+def get_code_path():
+    """Get base path where code is"""
+    return os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))

sam3/eval/teta_eval_toolkit/datasets/__init__.py

@@ -0,0 +1,5 @@
+# fmt: off
+# flake8: noqa
+"""Datasets."""
+from .coco import COCO
+from .tao import TAO

sam3/eval/teta_eval_toolkit/datasets/_base_dataset.py

@@ -0,0 +1,379 @@
+# fmt: off
+# flake8: noqa
+
+import csv
+import io
+import os
+import traceback
+import zipfile
+from abc import ABC, abstractmethod
+from copy import deepcopy
+
+import numpy as np
+
+from .. import _timing
+from ..utils import TrackEvalException
+
+
+class _BaseDataset(ABC):
+    @abstractmethod
+    def __init__(self):
+        self.tracker_list = None
+        self.seq_list = None
+        self.class_list = None
+        self.output_fol = None
+        self.output_sub_fol = None
+        self.should_classes_combine = True
+        self.use_super_categories = False
+
+    # Functions to implement:
+
+    @abstractmethod
+    def _load_raw_file(self, tracker, seq, is_gt):
+        ...
+
+    @_timing.time
+    @abstractmethod
+    def get_preprocessed_seq_data(self, raw_data, cls):
+        ...
+
+    @abstractmethod
+    def _calculate_similarities(self, gt_dets_t, tracker_dets_t):
+        ...
+
+    # Helper functions for all datasets:
+
+    @classmethod
+    def get_class_name(cls):
+        return cls.__name__
+
+    def get_name(self):
+        return self.get_class_name()
+
+    def get_output_fol(self, tracker):
+        return os.path.join(self.output_fol, tracker, self.output_sub_fol)
+
+    def get_display_name(self, tracker):
+        """Can be overwritten if the trackers name (in files) is different to how it should be displayed.
+        By default this method just returns the trackers name as is.
+        """
+        return tracker
+
+    def get_eval_info(self):
+        """Return info about the dataset needed for the Evaluator"""
+        return self.tracker_list, self.seq_list, self.class_list
+
+    @_timing.time
+    def get_raw_seq_data(self, tracker, seq):
+        """Loads raw data (tracker and ground-truth) for a single tracker on a single sequence.
+        Raw data includes all of the information needed for both preprocessing and evaluation, for all classes.
+        A later function (get_processed_seq_data) will perform such preprocessing and extract relevant information for
+        the evaluation of each class.
+
+        This returns a dict which contains the fields:
+        [num_timesteps]: integer
+        [gt_ids, tracker_ids, gt_classes, tracker_classes, tracker_confidences]:
+                                                                list (for each timestep) of 1D NDArrays (for each det).
+        [gt_dets, tracker_dets, gt_crowd_ignore_regions]: list (for each timestep) of lists of detections.
+        [similarity_scores]: list (for each timestep) of 2D NDArrays.
+        [gt_extras]: dict (for each extra) of lists (for each timestep) of 1D NDArrays (for each det).
+
+        gt_extras contains dataset specific information used for preprocessing such as occlusion and truncation levels.
+
+        Note that similarities are extracted as part of the dataset and not the metric, because almost all metrics are
+        independent of the exact method of calculating the similarity. However datasets are not (e.g. segmentation
+        masks vs 2D boxes vs 3D boxes).
+        We calculate the similarity before preprocessing because often both preprocessing and evaluation require it and
+        we don't wish to calculate this twice.
+        We calculate similarity between all gt and tracker classes (not just each class individually) to allow for
+        calculation of metrics such as class confusion matrices. Typically the impact of this on performance is low.
+        """
+        # Load raw data.
+        raw_gt_data = self._load_raw_file(tracker, seq, is_gt=True)
+        raw_tracker_data = self._load_raw_file(tracker, seq, is_gt=False)
+        raw_data = {**raw_tracker_data, **raw_gt_data}  # Merges dictionaries
+
+        # Calculate similarities for each timestep.
+        similarity_scores = []
+        for _, (gt_dets_t, tracker_dets_t) in enumerate(
+            zip(raw_data["gt_dets"], raw_data["tk_dets"])
+        ):
+            ious = self._calculate_similarities(gt_dets_t, tracker_dets_t)
+            similarity_scores.append(ious)
+        raw_data["similarity_scores"] = similarity_scores
+        return raw_data
+
+    @staticmethod
+    def _load_simple_text_file(
+        file,
+        time_col=0,
+        id_col=None,
+        remove_negative_ids=False,
+        valid_filter=None,
+        crowd_ignore_filter=None,
+        convert_filter=None,
+        is_zipped=False,
+        zip_file=None,
+        force_delimiters=None,
+    ):
+        """Function that loads data which is in a commonly used text file format.
+        Assumes each det is given by one row of a text file.
+        There is no limit to the number or meaning of each column,
+        however one column needs to give the timestep of each det (time_col) which is default col 0.
+
+        The file dialect (deliminator, num cols, etc) is determined automatically.
+        This function automatically separates dets by timestep,
+        and is much faster than alternatives such as np.loadtext or pandas.
+
+        If remove_negative_ids is True and id_col is not None, dets with negative values in id_col are excluded.
+        These are not excluded from ignore data.
+
+        valid_filter can be used to only include certain classes.
+        It is a dict with ints as keys, and lists as values,
+        such that a row is included if "row[key].lower() is in value" for all key/value pairs in the dict.
+        If None, all classes are included.
+
+        crowd_ignore_filter can be used to read crowd_ignore regions separately. It has the same format as valid filter.
+
+        convert_filter can be used to convert value read to another format.
+        This is used most commonly to convert classes given as string to a class id.
+        This is a dict such that the key is the column to convert, and the value is another dict giving the mapping.
+
+        Optionally, input files could be a zip of multiple text files for storage efficiency.
+
+        Returns read_data and ignore_data.
+        Each is a dict (with keys as timesteps as strings) of lists (over dets) of lists (over column values).
+        Note that all data is returned as strings, and must be converted to float/int later if needed.
+        Note that timesteps will not be present in the returned dict keys if there are no dets for them
+        """
+
+        if remove_negative_ids and id_col is None:
+            raise TrackEvalException(
+                "remove_negative_ids is True, but id_col is not given."
+            )
+        if crowd_ignore_filter is None:
+            crowd_ignore_filter = {}
+        if convert_filter is None:
+            convert_filter = {}
+        try:
+            if is_zipped:  # Either open file directly or within a zip.
+                if zip_file is None:
+                    raise TrackEvalException(
+                        "is_zipped set to True, but no zip_file is given."
+                    )
+                archive = zipfile.ZipFile(os.path.join(zip_file), "r")
+                fp = io.TextIOWrapper(archive.open(file, "r"))
+            else:
+                fp = open(file)
+            read_data = {}
+            crowd_ignore_data = {}
+            fp.seek(0, os.SEEK_END)
+            # check if file is empty
+            if fp.tell():
+                fp.seek(0)
+                dialect = csv.Sniffer().sniff(
+                    fp.readline(), delimiters=force_delimiters
+                )  # Auto determine structure.
+                dialect.skipinitialspace = (
+                    True  # Deal with extra spaces between columns
+                )
+                fp.seek(0)
+                reader = csv.reader(fp, dialect)
+                for row in reader:
+                    try:
+                        # Deal with extra trailing spaces at the end of rows
+                        if row[-1] in "":
+                            row = row[:-1]
+                        timestep = str(int(float(row[time_col])))
+                        # Read ignore regions separately.
+                        is_ignored = False
+                        for ignore_key, ignore_value in crowd_ignore_filter.items():
+                            if row[ignore_key].lower() in ignore_value:
+                                # Convert values in one column (e.g. string to id)
+                                for (
+                                    convert_key,
+                                    convert_value,
+                                ) in convert_filter.items():
+                                    row[convert_key] = convert_value[
+                                        row[convert_key].lower()
+                                    ]
+                                # Save data separated by timestep.
+                                if timestep in crowd_ignore_data.keys():
+                                    crowd_ignore_data[timestep].append(row)
+                                else:
+                                    crowd_ignore_data[timestep] = [row]
+                                is_ignored = True
+                        if (
+                            is_ignored
+                        ):  # if det is an ignore region, it cannot be a normal det.
+                            continue
+                        # Exclude some dets if not valid.
+                        if valid_filter is not None:
+                            for key, value in valid_filter.items():
+                                if row[key].lower() not in value:
+                                    continue
+                        if remove_negative_ids:
+                            if int(float(row[id_col])) < 0:
+                                continue
+                        # Convert values in one column (e.g. string to id)
+                        for convert_key, convert_value in convert_filter.items():
+                            row[convert_key] = convert_value[row[convert_key].lower()]
+                        # Save data separated by timestep.
+                        if timestep in read_data.keys():
+                            read_data[timestep].append(row)
+                        else:
+                            read_data[timestep] = [row]
+                    except Exception:
+                        exc_str_init = (
+                            "In file %s the following line cannot be read correctly: \n"
+                            % os.path.basename(file)
+                        )
+                        exc_str = " ".join([exc_str_init] + row)
+                        raise TrackEvalException(exc_str)
+            fp.close()
+        except Exception:
+            print("Error loading file: %s, printing traceback." % file)
+            traceback.print_exc()
+            raise TrackEvalException(
+                "File %s cannot be read because it is either not present or invalidly formatted"
+                % os.path.basename(file)
+            )
+        return read_data, crowd_ignore_data
+
+    @staticmethod
+    def _calculate_mask_ious(masks1, masks2, is_encoded=False, do_ioa=False):
+        """Calculates the IOU (intersection over union) between two arrays of segmentation masks.
+        If is_encoded a run length encoding with pycocotools is assumed as input format, otherwise an input of numpy
+        arrays of the shape (num_masks, height, width) is assumed and the encoding is performed.
+        If do_ioa (intersection over area) , then calculates the intersection over the area of masks1 - this is commonly
+        used to determine if detections are within crowd ignore region.
+        :param masks1:  first set of masks (numpy array of shape (num_masks, height, width) if not encoded,
+                        else pycocotools rle encoded format)
+        :param masks2:  second set of masks (numpy array of shape (num_masks, height, width) if not encoded,
+                        else pycocotools rle encoded format)
+        :param is_encoded: whether the input is in pycocotools rle encoded format
+        :param do_ioa: whether to perform IoA computation
+        :return: the IoU/IoA scores
+        """
+
+        # Only loaded when run to reduce minimum requirements
+        from pycocotools import mask as mask_utils
+
+        # use pycocotools for run length encoding of masks
+        if not is_encoded:
+            masks1 = mask_utils.encode(
+                np.array(np.transpose(masks1, (1, 2, 0)), order="F")
+            )
+            masks2 = mask_utils.encode(
+                np.array(np.transpose(masks2, (1, 2, 0)), order="F")
+            )
+
+        # use pycocotools for iou computation of rle encoded masks
+        ious = mask_utils.iou(masks1, masks2, [do_ioa] * len(masks2))
+        if len(masks1) == 0 or len(masks2) == 0:
+            ious = np.asarray(ious).reshape(len(masks1), len(masks2))
+        assert (ious >= 0 - np.finfo("float").eps).all()
+        assert (ious <= 1 + np.finfo("float").eps).all()
+
+        return ious
+
+    @staticmethod
+    def _calculate_box_ious(bboxes1, bboxes2, box_format="xywh", do_ioa=False):
+        """Calculates the IOU (intersection over union) between two arrays of boxes.
+        Allows variable box formats ('xywh' and 'x0y0x1y1').
+        If do_ioa (intersection over area) , then calculates the intersection over the area of boxes1 - this is commonly
+        used to determine if detections are within crowd ignore region.
+        """
+        if box_format in "xywh":
+            # layout: (x0, y0, w, h)
+            bboxes1 = deepcopy(bboxes1)
+            bboxes2 = deepcopy(bboxes2)
+
+            bboxes1[:, 2] = bboxes1[:, 0] + bboxes1[:, 2]
+            bboxes1[:, 3] = bboxes1[:, 1] + bboxes1[:, 3]
+            bboxes2[:, 2] = bboxes2[:, 0] + bboxes2[:, 2]
+            bboxes2[:, 3] = bboxes2[:, 1] + bboxes2[:, 3]
+        elif box_format not in "x0y0x1y1":
+            raise (TrackEvalException("box_format %s is not implemented" % box_format))
+
+        # layout: (x0, y0, x1, y1)
+        min_ = np.minimum(bboxes1[:, np.newaxis, :], bboxes2[np.newaxis, :, :])
+        max_ = np.maximum(bboxes1[:, np.newaxis, :], bboxes2[np.newaxis, :, :])
+        intersection = np.maximum(min_[..., 2] - max_[..., 0], 0) * np.maximum(
+            min_[..., 3] - max_[..., 1], 0
+        )
+        area1 = (bboxes1[..., 2] - bboxes1[..., 0]) * (
+            bboxes1[..., 3] - bboxes1[..., 1]
+        )
+
+        if do_ioa:
+            ioas = np.zeros_like(intersection)
+            valid_mask = area1 > 0 + np.finfo("float").eps
+            ioas[valid_mask, :] = (
+                intersection[valid_mask, :] / area1[valid_mask][:, np.newaxis]
+            )
+
+            return ioas
+        else:
+            area2 = (bboxes2[..., 2] - bboxes2[..., 0]) * (
+                bboxes2[..., 3] - bboxes2[..., 1]
+            )
+            union = area1[:, np.newaxis] + area2[np.newaxis, :] - intersection
+            intersection[area1 <= 0 + np.finfo("float").eps, :] = 0
+            intersection[:, area2 <= 0 + np.finfo("float").eps] = 0
+            intersection[union <= 0 + np.finfo("float").eps] = 0
+            union[union <= 0 + np.finfo("float").eps] = 1
+            ious = intersection / union
+            return ious
+
+    @staticmethod
+    def _calculate_euclidean_similarity(dets1, dets2, zero_distance=2.0):
+        """Calculates the euclidean distance between two sets of detections, and then converts this into a similarity
+        measure with values between 0 and 1 using the following formula: sim = max(0, 1 - dist/zero_distance).
+        The default zero_distance of 2.0, corresponds to the default used in MOT15_3D, such that a 0.5 similarity
+        threshold corresponds to a 1m distance threshold for TPs.
+        """
+        dist = np.linalg.norm(dets1[:, np.newaxis] - dets2[np.newaxis, :], axis=2)
+        sim = np.maximum(0, 1 - dist / zero_distance)
+        return sim
+
+    @staticmethod
+    def _check_unique_ids(data, after_preproc=False):
+        """Check the requirement that the tracker_ids and gt_ids are unique per timestep"""
+        gt_ids = data["gt_ids"]
+        tracker_ids = data["tk_ids"]
+        for t, (gt_ids_t, tracker_ids_t) in enumerate(zip(gt_ids, tracker_ids)):
+            if len(tracker_ids_t) > 0:
+                unique_ids, counts = np.unique(tracker_ids_t, return_counts=True)
+                if np.max(counts) != 1:
+                    duplicate_ids = unique_ids[counts > 1]
+                    exc_str_init = (
+                        "Tracker predicts the same ID more than once in a single timestep "
+                        "(seq: %s, frame: %i, ids:" % (data["seq"], t + 1)
+                    )
+                    exc_str = (
+                        " ".join([exc_str_init] + [str(d) for d in duplicate_ids]) + ")"
+                    )
+                    if after_preproc:
+                        exc_str_init += (
+                            "\n Note that this error occurred after preprocessing (but not before), "
+                            "so ids may not be as in file, and something seems wrong with preproc."
+                        )
+                    raise TrackEvalException(exc_str)
+            if len(gt_ids_t) > 0:
+                unique_ids, counts = np.unique(gt_ids_t, return_counts=True)
+                if np.max(counts) != 1:
+                    duplicate_ids = unique_ids[counts > 1]
+                    exc_str_init = (
+                        "Ground-truth has the same ID more than once in a single timestep "
+                        "(seq: %s, frame: %i, ids:" % (data["seq"], t + 1)
+                    )
+                    exc_str = (
+                        " ".join([exc_str_init] + [str(d) for d in duplicate_ids]) + ")"
+                    )
+                    if after_preproc:
+                        exc_str_init += (
+                            "\n Note that this error occurred after preprocessing (but not before), "
+                            "so ids may not be as in file, and something seems wrong with preproc."
+                        )
+                    raise TrackEvalException(exc_str)

sam3/eval/teta_eval_toolkit/datasets/coco.py

@@ -0,0 +1,637 @@
+# fmt: off
+# flake8: noqa
+
+"""COCO Dataset."""
+import copy
+import itertools
+import json
+import os
+from collections import defaultdict
+
+import numpy as np
+from scipy.optimize import linear_sum_assignment
+
+from .. import _timing, utils
+from ..config import get_default_dataset_config, init_config
+from ..utils import TrackEvalException
+from ._base_dataset import _BaseDataset
+
+
+class COCO(_BaseDataset):
+    """Tracking datasets in COCO format."""
+
+    def __init__(self, config=None):
+        """Initialize dataset, checking that all required files are present."""
+        super().__init__()
+        # Fill non-given config values with defaults
+        self.config = init_config(config, get_default_dataset_config(), self.get_name())
+        self.gt_fol = self.config["GT_FOLDER"]
+        self.tracker_fol = self.config["TRACKERS_FOLDER"]
+        self.should_classes_combine = True
+        self.use_super_categories = False
+        self.use_mask = self.config["USE_MASK"]
+
+        self.tracker_sub_fol = self.config["TRACKER_SUB_FOLDER"]
+        self.output_fol = self.config["OUTPUT_FOLDER"]
+        if self.output_fol is None:
+            self.output_fol = self.tracker_fol
+        self.output_sub_fol = self.config["OUTPUT_SUB_FOLDER"]
+
+        if self.gt_fol.endswith(".json"):
+            self.gt_data = json.load(open(self.gt_fol, "r"))
+        else:
+            gt_dir_files = [
+                file for file in os.listdir(self.gt_fol) if file.endswith(".json")
+            ]
+            if len(gt_dir_files) != 1:
+                raise TrackEvalException(
+                    f"{self.gt_fol} does not contain exactly one json file."
+                )
+
+            with open(os.path.join(self.gt_fol, gt_dir_files[0])) as f:
+                self.gt_data = json.load(f)
+
+        # fill missing video ids
+        self._fill_video_ids_inplace(self.gt_data["annotations"])
+
+        # get sequences to eval and sequence information
+        self.seq_list = [
+            vid["name"].replace("/", "-") for vid in self.gt_data["videos"]
+        ]
+        self.seq_name2seqid = {
+            vid["name"].replace("/", "-"): vid["id"] for vid in self.gt_data["videos"]
+        }
+        # compute mappings from videos to annotation data
+        self.video2gt_track, self.video2gt_image = self._compute_vid_mappings(
+            self.gt_data["annotations"]
+        )
+        # compute sequence lengths
+        self.seq_lengths = {vid["id"]: 0 for vid in self.gt_data["videos"]}
+        for img in self.gt_data["images"]:
+            self.seq_lengths[img["video_id"]] += 1
+        self.seq2images2timestep = self._compute_image_to_timestep_mappings()
+        self.seq2cls = {
+            vid["id"]: {
+                "pos_cat_ids": list(
+                    {track["category_id"] for track in self.video2gt_track[vid["id"]]}
+                ),
+            }
+            for vid in self.gt_data["videos"]
+        }
+
+        # Get classes to eval
+        considered_vid_ids = [self.seq_name2seqid[vid] for vid in self.seq_list]
+        seen_cats = set(
+            [
+                cat_id
+                for vid_id in considered_vid_ids
+                for cat_id in self.seq2cls[vid_id]["pos_cat_ids"]
+            ]
+        )
+        # only classes with ground truth are evaluated in TAO
+        self.valid_classes = [
+            cls["name"] for cls in self.gt_data["categories"] if cls["id"] in seen_cats
+        ]
+        cls_name2clsid_map = {
+            cls["name"]: cls["id"] for cls in self.gt_data["categories"]
+        }
+
+        if self.config["CLASSES_TO_EVAL"]:
+            self.class_list = [
+                cls.lower() if cls.lower() in self.valid_classes else None
+                for cls in self.config["CLASSES_TO_EVAL"]
+            ]
+            if not all(self.class_list):
+                valid_cls = ", ".join(self.valid_classes)
+                raise TrackEvalException(
+                    "Attempted to evaluate an invalid class. Only classes "
+                    f"{valid_cls} are valid (classes present in ground truth"
+                    " data)."
+                )
+        else:
+            self.class_list = [cls for cls in self.valid_classes]
+        self.cls_name2clsid = {
+            k: v for k, v in cls_name2clsid_map.items() if k in self.class_list
+        }
+        self.clsid2cls_name = {
+            v: k for k, v in cls_name2clsid_map.items() if k in self.class_list
+        }
+        # get trackers to eval
+        if self.config["TRACKERS_TO_EVAL"] is None:
+            self.tracker_list = os.listdir(self.tracker_fol)
+        else:
+            self.tracker_list = self.config["TRACKERS_TO_EVAL"]
+
+        if self.config["TRACKER_DISPLAY_NAMES"] is None:
+            self.tracker_to_disp = dict(zip(self.tracker_list, self.tracker_list))
+        elif (self.config["TRACKERS_TO_EVAL"] is not None) and (
+            len(self.config["TK_DISPLAY_NAMES"]) == len(self.tracker_list)
+        ):
+            self.tracker_to_disp = dict(
+                zip(self.tracker_list, self.config["TK_DISPLAY_NAMES"])
+            )
+        else:
+            raise TrackEvalException(
+                "List of tracker files and tracker display names do not match."
+            )
+
+        self.tracker_data = {tracker: dict() for tracker in self.tracker_list}
+
+        for tracker in self.tracker_list:
+            if self.tracker_sub_fol.endswith(".json"):
+                with open(os.path.join(self.tracker_sub_fol)) as f:
+                    curr_data = json.load(f)
+            else:
+                tr_dir = os.path.join(self.tracker_fol, tracker, self.tracker_sub_fol)
+                tr_dir_files = [
+                    file for file in os.listdir(tr_dir) if file.endswith(".json")
+                ]
+                if len(tr_dir_files) != 1:
+                    raise TrackEvalException(
+                        f"{tr_dir} does not contain exactly one json file."
+                    )
+                with open(os.path.join(tr_dir, tr_dir_files[0])) as f:
+                    curr_data = json.load(f)
+
+            # limit detections if MAX_DETECTIONS > 0
+            if self.config["MAX_DETECTIONS"]:
+                curr_data = self._limit_dets_per_image(curr_data)
+
+            # fill missing video ids
+            self._fill_video_ids_inplace(curr_data)
+
+            # make track ids unique over whole evaluation set
+            self._make_tk_ids_unique(curr_data)
+
+            # get tracker sequence information
+            curr_vids2tracks, curr_vids2images = self._compute_vid_mappings(curr_data)
+            self.tracker_data[tracker]["vids_to_tracks"] = curr_vids2tracks
+            self.tracker_data[tracker]["vids_to_images"] = curr_vids2images
+
+    def get_display_name(self, tracker):
+        return self.tracker_to_disp[tracker]
+
+    def _load_raw_file(self, tracker, seq, is_gt):
+        """Load a file (gt or tracker) in the TAO format
+
+        If is_gt, this returns a dict which contains the fields:
+        [gt_ids, gt_classes]:
+            list (for each timestep) of 1D NDArrays (for each det).
+        [gt_dets]: list (for each timestep) of lists of detections.
+
+        if not is_gt, this returns a dict which contains the fields:
+        [tk_ids, tk_classes]:
+            list (for each timestep) of 1D NDArrays (for each det).
+        [tk_dets]: list (for each timestep) of lists of detections.
+        """
+        seq_id = self.seq_name2seqid[seq]
+        # file location
+        if is_gt:
+            imgs = self.video2gt_image[seq_id]
+        else:
+            imgs = self.tracker_data[tracker]["vids_to_images"][seq_id]
+
+        # convert data to required format
+        num_timesteps = self.seq_lengths[seq_id]
+        img_to_timestep = self.seq2images2timestep[seq_id]
+        data_keys = ["ids", "classes", "dets"]
+        # if not is_gt:
+        #     data_keys += ["tk_confidences"]
+        raw_data = {key: [None] * num_timesteps for key in data_keys}
+        for img in imgs:
+            # some tracker data contains images without any ground truth info,
+            # these are ignored
+            if img["id"] not in img_to_timestep:
+                continue
+            t = img_to_timestep[img["id"]]
+            anns = img["annotations"]
+            tk_str = utils.get_track_id_str(anns[0])
+
+            if self.use_mask:
+                # When using mask, extract segmentation data
+                raw_data["dets"][t] = [ann.get("segmentation") for ann in anns]
+            else:
+                # When using bbox, extract bbox data
+                raw_data["dets"][t] = np.atleast_2d([ann["bbox"] for ann in anns]).astype(
+                    float
+                )
+            raw_data["ids"][t] = np.atleast_1d([ann[tk_str] for ann in anns]).astype(
+                int
+            )
+            raw_data["classes"][t] = np.atleast_1d(
+                [ann["category_id"] for ann in anns]
+            ).astype(int)
+            # if not is_gt:
+            #     raw_data["tk_confidences"][t] = np.atleast_1d(
+            #         [ann["score"] for ann in anns]
+            #     ).astype(float)
+
+        for t, d in enumerate(raw_data["dets"]):
+            if d is None:
+                raw_data["dets"][t] = np.empty((0, 4)).astype(float)
+                raw_data["ids"][t] = np.empty(0).astype(int)
+                raw_data["classes"][t] = np.empty(0).astype(int)
+                # if not is_gt:
+                #     raw_data["tk_confidences"][t] = np.empty(0)
+
+        if is_gt:
+            key_map = {"ids": "gt_ids", "classes": "gt_classes", "dets": "gt_dets"}
+        else:
+            key_map = {"ids": "tk_ids", "classes": "tk_classes", "dets": "tk_dets"}
+        for k, v in key_map.items():
+            raw_data[v] = raw_data.pop(k)
+
+        raw_data["num_timesteps"] = num_timesteps
+        raw_data["seq"] = seq
+        return raw_data
+
+    def get_preprocessed_seq_data_thr(self, raw_data, cls, assignment=None):
+        """Preprocess data for a single sequence for a single class.
+
+        Inputs:
+            raw_data: dict containing the data for the sequence already
+                read in by get_raw_seq_data().
+            cls: class to be evaluated.
+        Outputs:
+            gt_ids:
+                list (for each timestep) of ids of GT tracks
+            tk_ids:
+                list (for each timestep) of ids of predicted tracks (all for TP
+                matching (Det + AssocA))
+            tk_overlap_ids:
+                list (for each timestep) of ids of predicted tracks that overlap
+                with GTs
+            tk_dets:
+                list (for each timestep) of lists of detections that
+                corresponding to the tk_ids
+            tk_classes:
+                list (for each timestep) of lists of classes that corresponding
+                to the tk_ids
+            tk_confidences:
+                list (for each timestep) of lists of classes that corresponding
+                to the tk_ids
+            sim_scores:
+                similarity score between gt_ids and tk_ids.
+        """
+        if cls != "all":
+            cls_id = self.cls_name2clsid[cls]
+
+        data_keys = [
+            "gt_ids",
+            "tk_ids",
+            "gt_id_map",
+            "tk_id_map",
+            "gt_dets",
+            "gt_classes",
+            "gt_class_name",
+            "tk_overlap_classes",
+            "tk_overlap_ids",
+            "tk_class_eval_tk_ids",
+            "tk_dets",
+            "tk_classes",
+            # "tk_confidences",
+            "tk_exh_ids",
+            "sim_scores",
+        ]
+        data = {key: [None] * raw_data["num_timesteps"] for key in data_keys}
+        unique_gt_ids = []
+        unique_tk_ids = []
+        num_gt_dets = 0
+        num_tk_cls_dets = 0
+        num_tk_overlap_dets = 0
+        overlap_ious_thr = 0.5
+        loc_and_asso_tk_ids = []
+        exh_class_tk_ids = []
+
+        for t in range(raw_data["num_timesteps"]):
+            # only extract relevant dets for this class for preproc and eval
+            if cls == "all":
+                gt_class_mask = np.ones_like(raw_data["gt_classes"][t]).astype(bool)
+            else:
+                gt_class_mask = np.atleast_1d(
+                    raw_data["gt_classes"][t] == cls_id
+                ).astype(bool)
+
+            # select GT that is not in the evaluating classes
+            if assignment is not None and assignment:
+                all_gt_ids = list(assignment[t].keys())
+                gt_ids_in = raw_data["gt_ids"][t][gt_class_mask]
+                gt_ids_out = set(all_gt_ids) - set(gt_ids_in)
+                tk_ids_out = set([assignment[t][key] for key in list(gt_ids_out)])
+
+            # compute overlapped tracks and add their ids to overlap_tk_ids
+            sim_scores = raw_data["similarity_scores"]
+            overlap_ids_masks = (sim_scores[t][gt_class_mask] >= overlap_ious_thr).any(
+                axis=0
+            )
+            overlap_tk_ids_t = raw_data["tk_ids"][t][overlap_ids_masks]
+            if assignment is not None and assignment:
+                data["tk_overlap_ids"][t] = list(set(overlap_tk_ids_t) - tk_ids_out)
+            else:
+                data["tk_overlap_ids"][t] = list(set(overlap_tk_ids_t))
+
+            loc_and_asso_tk_ids += data["tk_overlap_ids"][t]
+
+            data["tk_exh_ids"][t] = []
+            if cls == "all":
+                continue
+
+            # add the track ids of exclusive annotated class to exh_class_tk_ids
+            tk_exh_mask = np.atleast_1d(raw_data["tk_classes"][t] == cls_id)
+            tk_exh_mask = tk_exh_mask.astype(bool)
+            exh_class_tk_ids_t = raw_data["tk_ids"][t][tk_exh_mask]
+            exh_class_tk_ids.append(exh_class_tk_ids_t)
+            data["tk_exh_ids"][t] = exh_class_tk_ids_t
+
+        # remove tk_ids that has been assigned to GT belongs to other classes.
+        loc_and_asso_tk_ids = list(set(loc_and_asso_tk_ids))
+
+        # remove all unwanted unmatched tracker detections
+        for t in range(raw_data["num_timesteps"]):
+            # add gt to the data
+            if cls == "all":
+                gt_class_mask = np.ones_like(raw_data["gt_classes"][t]).astype(bool)
+            else:
+                gt_class_mask = np.atleast_1d(
+                    raw_data["gt_classes"][t] == cls_id
+                ).astype(bool)
+                data["gt_classes"][t] = cls_id
+                data["gt_class_name"][t] = cls
+
+            gt_ids = raw_data["gt_ids"][t][gt_class_mask]
+            if self.use_mask:
+                gt_dets = [raw_data['gt_dets'][t][ind] for ind in range(len(gt_class_mask)) if gt_class_mask[ind]]
+            else:
+                gt_dets = raw_data["gt_dets"][t][gt_class_mask]
+            data["gt_ids"][t] = gt_ids
+            data["gt_dets"][t] = gt_dets
+
+            # filter pred and only keep those that highly overlap with GTs
+            tk_mask = np.isin(
+                raw_data["tk_ids"][t], np.array(loc_and_asso_tk_ids), assume_unique=True
+            )
+            tk_overlap_mask = np.isin(
+                raw_data["tk_ids"][t],
+                np.array(data["tk_overlap_ids"][t]),
+                assume_unique=True,
+            )
+
+            tk_ids = raw_data["tk_ids"][t][tk_mask]
+            if self.use_mask:
+                tk_dets = [raw_data['tk_dets'][t][ind] for ind in range(len(tk_mask)) if
+                            tk_mask[ind]]
+            else:
+                tk_dets = raw_data["tk_dets"][t][tk_mask]
+
+            tracker_classes = raw_data["tk_classes"][t][tk_mask]
+
+            # add overlap classes for computing the FP for Cls term
+            tracker_overlap_classes = raw_data["tk_classes"][t][tk_overlap_mask]
+            # tracker_confidences = raw_data["tk_confidences"][t][tk_mask]
+            sim_scores_masked = sim_scores[t][gt_class_mask, :][:, tk_mask]
+
+            # add filtered prediction to the data
+            data["tk_classes"][t] = tracker_classes
+            data["tk_overlap_classes"][t] = tracker_overlap_classes
+            data["tk_ids"][t] = tk_ids
+            data["tk_dets"][t] = tk_dets
+            # data["tk_confidences"][t] = tracker_confidences
+            data["sim_scores"][t] = sim_scores_masked
+            data["tk_class_eval_tk_ids"][t] = set(
+                list(data["tk_overlap_ids"][t]) + list(data["tk_exh_ids"][t])
+            )
+
+            # count total number of detections
+            unique_gt_ids += list(np.unique(data["gt_ids"][t]))
+            # the unique track ids are for association.
+            unique_tk_ids += list(np.unique(data["tk_ids"][t]))
+
+            num_tk_overlap_dets += len(data["tk_overlap_ids"][t])
+            num_tk_cls_dets += len(data["tk_class_eval_tk_ids"][t])
+            num_gt_dets += len(data["gt_ids"][t])
+
+        # re-label IDs such that there are no empty IDs
+        if len(unique_gt_ids) > 0:
+            unique_gt_ids = np.unique(unique_gt_ids)
+            gt_id_map = np.nan * np.ones((np.max(unique_gt_ids) + 1))
+            gt_id_map[unique_gt_ids] = np.arange(len(unique_gt_ids))
+            data["gt_id_map"] = {}
+            for gt_id in unique_gt_ids:
+                new_gt_id = gt_id_map[gt_id].astype(int)
+                data["gt_id_map"][new_gt_id] = gt_id
+
+            for t in range(raw_data["num_timesteps"]):
+                if len(data["gt_ids"][t]) > 0:
+                    data["gt_ids"][t] = gt_id_map[data["gt_ids"][t]].astype(int)
+
+        if len(unique_tk_ids) > 0:
+            unique_tk_ids = np.unique(unique_tk_ids)
+            tk_id_map = np.nan * np.ones((np.max(unique_tk_ids) + 1))
+            tk_id_map[unique_tk_ids] = np.arange(len(unique_tk_ids))
+
+            data["tk_id_map"] = {}
+            for track_id in unique_tk_ids:
+                new_track_id = tk_id_map[track_id].astype(int)
+                data["tk_id_map"][new_track_id] = track_id
+
+            for t in range(raw_data["num_timesteps"]):
+                if len(data["tk_ids"][t]) > 0:
+                    data["tk_ids"][t] = tk_id_map[data["tk_ids"][t]].astype(int)
+                if len(data["tk_overlap_ids"][t]) > 0:
+                    data["tk_overlap_ids"][t] = tk_id_map[
+                        data["tk_overlap_ids"][t]
+                    ].astype(int)
+
+        # record overview statistics.
+        data["num_tk_cls_dets"] = num_tk_cls_dets
+        data["num_tk_overlap_dets"] = num_tk_overlap_dets
+        data["num_gt_dets"] = num_gt_dets
+        data["num_tk_ids"] = len(unique_tk_ids)
+        data["num_gt_ids"] = len(unique_gt_ids)
+        data["num_timesteps"] = raw_data["num_timesteps"]
+        data["seq"] = raw_data["seq"]
+
+        self._check_unique_ids(data)
+
+        return data
+
+    @_timing.time
+    def get_preprocessed_seq_data(
+        self, raw_data, cls, assignment=None, thresholds=[50, 75]
+    ):
+        """Preprocess data for a single sequence for a single class."""
+        data = {}
+        if thresholds is None:
+            thresholds = [50, 75]
+        elif isinstance(thresholds, int):
+            thresholds = [thresholds]
+
+        for thr in thresholds:
+            assignment_thr = None
+            if assignment is not None:
+                assignment_thr = assignment[thr]
+            data[thr] = self.get_preprocessed_seq_data_thr(
+                raw_data, cls, assignment_thr
+            )
+
+        return data
+
+    def _calculate_similarities(self, gt_dets_t, tk_dets_t):
+        """Compute similarity scores."""
+        if self.use_mask:
+            similarity_scores = self._calculate_mask_ious(gt_dets_t, tk_dets_t, is_encoded=True, do_ioa=False)
+        else:
+            similarity_scores = self._calculate_box_ious(gt_dets_t, tk_dets_t)
+        return similarity_scores
+
+    def _compute_vid_mappings(self, annotations):
+        """Computes mappings from videos to corresponding tracks and images."""
+        vids_to_tracks = {}
+        vids_to_imgs = {}
+        vid_ids = [vid["id"] for vid in self.gt_data["videos"]]
+
+        # compute an mapping from image IDs to images
+        images = {}
+        for image in self.gt_data["images"]:
+            images[image["id"]] = image
+
+        tk_str = utils.get_track_id_str(annotations[0])
+        for ann in annotations:
+            ann["area"] = ann["bbox"][2] * ann["bbox"][3]
+
+            vid = ann["video_id"]
+            if ann["video_id"] not in vids_to_tracks.keys():
+                vids_to_tracks[ann["video_id"]] = list()
+            if ann["video_id"] not in vids_to_imgs.keys():
+                vids_to_imgs[ann["video_id"]] = list()
+
+            # fill in vids_to_tracks
+            tid = ann[tk_str]
+            exist_tids = [track["id"] for track in vids_to_tracks[vid]]
+            try:
+                index1 = exist_tids.index(tid)
+            except ValueError:
+                index1 = -1
+            if tid not in exist_tids:
+                curr_track = {
+                    "id": tid,
+                    "category_id": ann["category_id"],
+                    "video_id": vid,
+                    "annotations": [ann],
+                }
+                vids_to_tracks[vid].append(curr_track)
+            else:
+                vids_to_tracks[vid][index1]["annotations"].append(ann)
+
+            # fill in vids_to_imgs
+            img_id = ann["image_id"]
+            exist_img_ids = [img["id"] for img in vids_to_imgs[vid]]
+            try:
+                index2 = exist_img_ids.index(img_id)
+            except ValueError:
+                index2 = -1
+            if index2 == -1:
+                curr_img = {"id": img_id, "annotations": [ann]}
+                vids_to_imgs[vid].append(curr_img)
+            else:
+                vids_to_imgs[vid][index2]["annotations"].append(ann)
+
+        # sort annotations by frame index and compute track area
+        for vid, tracks in vids_to_tracks.items():
+            for track in tracks:
+                track["annotations"] = sorted(
+                    track["annotations"],
+                    key=lambda x: images[x["image_id"]]["frame_id"],
+                )
+                # compute average area
+                track["area"] = sum(x["area"] for x in track["annotations"]) / len(
+                    track["annotations"]
+                )
+
+        # ensure all videos are present
+        for vid_id in vid_ids:
+            if vid_id not in vids_to_tracks.keys():
+                vids_to_tracks[vid_id] = []
+            if vid_id not in vids_to_imgs.keys():
+                vids_to_imgs[vid_id] = []
+
+        return vids_to_tracks, vids_to_imgs
+
+    def _compute_image_to_timestep_mappings(self):
+        """Computes a mapping from images to timestep in sequence."""
+        images = {}
+        for image in self.gt_data["images"]:
+            images[image["id"]] = image
+
+        seq_to_imgs_to_timestep = {vid["id"]: dict() for vid in self.gt_data["videos"]}
+        for vid in seq_to_imgs_to_timestep:
+            curr_imgs = [img["id"] for img in self.video2gt_image[vid]]
+            curr_imgs = sorted(curr_imgs, key=lambda x: images[x]["frame_id"])
+            seq_to_imgs_to_timestep[vid] = {
+                curr_imgs[i]: i for i in range(len(curr_imgs))
+            }
+
+        return seq_to_imgs_to_timestep
+
+    def _limit_dets_per_image(self, annotations):
+        """Limits the number of detections for each image.
+
+        Adapted from https://github.com/TAO-Dataset/.
+        """
+        max_dets = self.config["MAX_DETECTIONS"]
+        img_ann = defaultdict(list)
+        for ann in annotations:
+            img_ann[ann["image_id"]].append(ann)
+
+        for img_id, _anns in img_ann.items():
+            if len(_anns) <= max_dets:
+                continue
+            _anns = sorted(_anns, key=lambda x: x["score"], reverse=True)
+            img_ann[img_id] = _anns[:max_dets]
+
+        return [ann for anns in img_ann.values() for ann in anns]
+
+    def _fill_video_ids_inplace(self, annotations):
+        """Fills in missing video IDs inplace.
+
+        Adapted from https://github.com/TAO-Dataset/.
+        """
+        missing_video_id = [x for x in annotations if "video_id" not in x]
+        if missing_video_id:
+            image_id_to_video_id = {
+                x["id"]: x["video_id"] for x in self.gt_data["images"]
+            }
+            for x in missing_video_id:
+                x["video_id"] = image_id_to_video_id[x["image_id"]]
+
+    @staticmethod
+    def _make_tk_ids_unique(annotations):
+        """Makes track IDs unqiue over the whole annotation set.
+
+        Adapted from https://github.com/TAO-Dataset/.
+        """
+        track_id_videos = {}
+        track_ids_to_update = set()
+        max_track_id = 0
+
+        tk_str = utils.get_track_id_str(annotations[0])
+        for ann in annotations:
+            t = int(ann[tk_str])
+            if t not in track_id_videos:
+                track_id_videos[t] = ann["video_id"]
+
+            if ann["video_id"] != track_id_videos[t]:
+                # track id is assigned to multiple videos
+                track_ids_to_update.add(t)
+            max_track_id = max(max_track_id, t)
+
+        if track_ids_to_update:
+            print("true")
+            next_id = itertools.count(max_track_id + 1)
+            new_tk_ids = defaultdict(lambda: next(next_id))
+            for ann in annotations:
+                t = ann[tk_str]
+                v = ann["video_id"]
+                if t in track_ids_to_update:
+                    ann[tk_str] = new_tk_ids[t, v]
+        return len(track_ids_to_update)

sam3/eval/teta_eval_toolkit/datasets/tao.py

@@ -0,0 +1,659 @@
+# fmt: off
+# flake8: noqa
+
+"""TAO Dataset."""
+import copy
+import itertools
+import json
+import os
+from collections import defaultdict
+
+import numpy as np
+
+from .. import _timing
+from ..config import get_default_dataset_config, init_config
+from ..utils import TrackEvalException
+from ._base_dataset import _BaseDataset
+
+
+class TAO(_BaseDataset):
+    """Dataset class for TAO tracking"""
+
+    def __init__(self, config=None):
+        """Initialize dataset, checking that all required files are present."""
+        super().__init__()
+        # Fill non-given config values with defaults
+        self.config = init_config(config, get_default_dataset_config(), self.get_name())
+        self.gt_fol = self.config["GT_FOLDER"]
+        self.tracker_fol = self.config["TRACKERS_FOLDER"]
+        self.should_classes_combine = True
+        self.use_super_categories = False
+        self.use_mask = self.config["USE_MASK"]
+
+
+        self.tracker_sub_fol = self.config["TRACKER_SUB_FOLDER"]
+        self.output_fol = self.config["OUTPUT_FOLDER"]
+        if self.output_fol is None:
+            self.output_fol = self.tracker_fol
+        self.output_sub_fol = self.config["OUTPUT_SUB_FOLDER"]
+
+        if self.gt_fol.endswith(".json"):
+            self.gt_data = json.load(open(self.gt_fol, "r"))
+        else:
+            gt_dir_files = [
+                file for file in os.listdir(self.gt_fol) if file.endswith(".json")
+            ]
+            if len(gt_dir_files) != 1:
+                raise TrackEvalException(
+                    f"{self.gt_fol} does not contain exactly one json file."
+                )
+
+            with open(os.path.join(self.gt_fol, gt_dir_files[0])) as f:
+                self.gt_data = json.load(f)
+
+        # merge categories marked with a merged tag in TAO dataset
+        self._merge_categories(self.gt_data["annotations"] + self.gt_data["tracks"])
+
+        # get sequences to eval and sequence information
+        self.seq_list = [
+            vid["name"].replace("/", "-") for vid in self.gt_data["videos"]
+        ]
+        self.seq_name2seqid = {
+            vid["name"].replace("/", "-"): vid["id"] for vid in self.gt_data["videos"]
+        }
+        # compute mappings from videos to annotation data
+        self.video2gt_track, self.video2gt_image = self._compute_vid_mappings(
+            self.gt_data["annotations"]
+        )
+        # compute sequence lengths
+        self.seq_lengths = {vid["id"]: 0 for vid in self.gt_data["videos"]}
+        for img in self.gt_data["images"]:
+            self.seq_lengths[img["video_id"]] += 1
+        self.seq2images2timestep = self._compute_image_to_timestep_mappings()
+        self.seq2cls = {
+            vid["id"]: {
+                "pos_cat_ids": list(
+                    {track["category_id"] for track in self.video2gt_track[vid["id"]]}
+                ),
+                "neg_cat_ids": vid["neg_category_ids"],
+                "not_exh_labeled_cat_ids": vid["not_exhaustive_category_ids"],
+            }
+            for vid in self.gt_data["videos"]
+        }
+
+        # Get classes to eval
+        considered_vid_ids = [self.seq_name2seqid[vid] for vid in self.seq_list]
+        seen_cats = set(
+            [
+                cat_id
+                for vid_id in considered_vid_ids
+                for cat_id in self.seq2cls[vid_id]["pos_cat_ids"]
+            ]
+        )
+        # only classes with ground truth are evaluated in TAO
+        self.valid_classes = [
+            cls["name"] for cls in self.gt_data["categories"] if cls["id"] in seen_cats
+        ]
+        cls_name2clsid_map = {
+            cls["name"]: cls["id"] for cls in self.gt_data["categories"]
+        }
+
+        if self.config["CLASSES_TO_EVAL"]:
+            self.class_list = [
+                cls.lower() if cls.lower() in self.valid_classes else None
+                for cls in self.config["CLASSES_TO_EVAL"]
+            ]
+            if not all(self.class_list):
+                valid_cls = ", ".join(self.valid_classes)
+                raise TrackEvalException(
+                    "Attempted to evaluate an invalid class. Only classes "
+                    f"{valid_cls} are valid (classes present in ground truth"
+                    " data)."
+                )
+        else:
+            self.class_list = [cls for cls in self.valid_classes]
+        self.cls_name2clsid = {
+            k: v for k, v in cls_name2clsid_map.items() if k in self.class_list
+        }
+        self.clsid2cls_name = {
+            v: k for k, v in cls_name2clsid_map.items() if k in self.class_list
+        }
+        # get trackers to eval
+        print(self.config["TRACKERS_TO_EVAL"] )
+        if self.config["TRACKERS_TO_EVAL"] is None:
+            self.tracker_list = os.listdir(self.tracker_fol)
+        else:
+            self.tracker_list = self.config["TRACKERS_TO_EVAL"]
+
+        if self.config["TRACKER_DISPLAY_NAMES"] is None:
+            self.tracker_to_disp = dict(zip(self.tracker_list, self.tracker_list))
+        elif (self.config["TRACKERS_TO_EVAL"] is not None) and (
+            len(self.config["TK_DISPLAY_NAMES"]) == len(self.tracker_list)
+        ):
+            self.tracker_to_disp = dict(
+                zip(self.tracker_list, self.config["TK_DISPLAY_NAMES"])
+            )
+        else:
+            raise TrackEvalException(
+                "List of tracker files and tracker display names do not match."
+            )
+
+        self.tracker_data = {tracker: dict() for tracker in self.tracker_list}
+
+        for tracker in self.tracker_list:
+            if self.tracker_sub_fol.endswith(".json"):
+                with open(os.path.join(self.tracker_sub_fol)) as f:
+                    curr_data = json.load(f)
+            else:
+                tr_dir = os.path.join(self.tracker_fol, tracker, self.tracker_sub_fol)
+                tr_dir_files = [
+                    file for file in os.listdir(tr_dir) if file.endswith(".json")
+                ]
+                if len(tr_dir_files) != 1:
+                    raise TrackEvalException(
+                        f"{tr_dir} does not contain exactly one json file."
+                    )
+                with open(os.path.join(tr_dir, tr_dir_files[0])) as f:
+                    curr_data = json.load(f)
+
+            # limit detections if MAX_DETECTIONS > 0
+            if self.config["MAX_DETECTIONS"]:
+                curr_data = self._limit_dets_per_image(curr_data)
+
+            # fill missing video ids
+            self._fill_video_ids_inplace(curr_data)
+
+            # make track ids unique over whole evaluation set
+            self._make_tk_ids_unique(curr_data)
+
+            # merge categories marked with a merged tag in TAO dataset
+            self._merge_categories(curr_data)
+
+            # get tracker sequence information
+            curr_vids2tracks, curr_vids2images = self._compute_vid_mappings(curr_data)
+            self.tracker_data[tracker]["vids_to_tracks"] = curr_vids2tracks
+            self.tracker_data[tracker]["vids_to_images"] = curr_vids2images
+
+    def get_display_name(self, tracker):
+        return self.tracker_to_disp[tracker]
+
+    def _load_raw_file(self, tracker, seq, is_gt):
+        """Load a file (gt or tracker) in the TAO format
+
+        If is_gt, this returns a dict which contains the fields:
+        [gt_ids, gt_classes]:
+            list (for each timestep) of 1D NDArrays (for each det).
+        [gt_dets]: list (for each timestep) of lists of detections.
+
+        if not is_gt, this returns a dict which contains the fields:
+        [tk_ids, tk_classes, tk_confidences]:
+            list (for each timestep) of 1D NDArrays (for each det).
+        [tk_dets]: list (for each timestep) of lists of detections.
+        """
+        seq_id = self.seq_name2seqid[seq]
+        # file location
+        if is_gt:
+            imgs = self.video2gt_image[seq_id]
+        else:
+            imgs = self.tracker_data[tracker]["vids_to_images"][seq_id]
+
+        # convert data to required format
+        num_timesteps = self.seq_lengths[seq_id]
+        img_to_timestep = self.seq2images2timestep[seq_id]
+        data_keys = ["ids", "classes", "dets"]
+        if not is_gt:
+            data_keys += ["tk_confidences"]
+        raw_data = {key: [None] * num_timesteps for key in data_keys}
+        for img in imgs:
+            # some tracker data contains images without any ground truth info,
+            # these are ignored
+            if img["id"] not in img_to_timestep:
+                continue
+            t = img_to_timestep[img["id"]]
+            anns = img["annotations"]
+            if self.use_mask:
+                # When using mask, extract segmentation data
+                raw_data["dets"][t] = [ann.get("segmentation") for ann in anns]
+            else:
+                # When using bbox, extract bbox data
+                raw_data["dets"][t] = np.atleast_2d([ann["bbox"] for ann in anns]).astype(
+                    float
+                )
+            raw_data["ids"][t] = np.atleast_1d(
+                [ann["track_id"] for ann in anns]
+            ).astype(int)
+            raw_data["classes"][t] = np.atleast_1d(
+                [ann["category_id"] for ann in anns]
+            ).astype(int)
+            if not is_gt:
+                raw_data["tk_confidences"][t] = np.atleast_1d(
+                    [ann["score"] for ann in anns]
+                ).astype(float)
+
+        for t, d in enumerate(raw_data["dets"]):
+            if d is None:
+                raw_data["dets"][t] = np.empty((0, 4)).astype(float)
+                raw_data["ids"][t] = np.empty(0).astype(int)
+                raw_data["classes"][t] = np.empty(0).astype(int)
+                if not is_gt:
+                    raw_data["tk_confidences"][t] = np.empty(0)
+
+        if is_gt:
+            key_map = {"ids": "gt_ids", "classes": "gt_classes", "dets": "gt_dets"}
+        else:
+            key_map = {"ids": "tk_ids", "classes": "tk_classes", "dets": "tk_dets"}
+        for k, v in key_map.items():
+            raw_data[v] = raw_data.pop(k)
+
+        raw_data["num_timesteps"] = num_timesteps
+        raw_data["neg_cat_ids"] = self.seq2cls[seq_id]["neg_cat_ids"]
+        raw_data["not_exh_labeled_cls"] = self.seq2cls[seq_id][
+            "not_exh_labeled_cat_ids"
+        ]
+        raw_data["seq"] = seq
+        return raw_data
+
+    def get_preprocessed_seq_data_thr(self, raw_data, cls, assignment=None):
+        """Preprocess data for a single sequence for a single class.
+
+        Inputs:
+            raw_data: dict containing the data for the sequence already
+                read in by get_raw_seq_data().
+            cls: class to be evaluated.
+        Outputs:
+            gt_ids:
+                list (for each timestep) of ids of GT tracks
+            tk_ids:
+                list (for each timestep) of ids of predicted tracks (all for TP
+                matching (Det + AssocA))
+            tk_overlap_ids:
+                list (for each timestep) of ids of predicted tracks that overlap
+                with GTs
+            tk_neg_ids:
+                list (for each timestep) of ids of predicted tracks that with
+                the class id on the negative list for the current sequence.
+            tk_exh_ids:
+                list (for each timestep) of ids of predicted tracks that do not
+                overlap with existing GTs but have the class id on the
+                exhaustive annotated class list for the current sequence.
+            tk_dets:
+                list (for each timestep) of lists of detections that
+                corresponding to the tk_ids
+            tk_classes:
+                list (for each timestep) of lists of classes that corresponding
+                to the tk_ids
+            tk_confidences:
+                list (for each timestep) of lists of classes that corresponding
+                to the tk_ids
+            sim_scores:
+                similarity score between gt_ids and tk_ids.
+        """
+        if cls != "all":
+            cls_id = self.cls_name2clsid[cls]
+
+        data_keys = [
+            "gt_ids",
+            "tk_ids",
+            "gt_id_map",
+            "tk_id_map",
+            "gt_dets",
+            "gt_classes",
+            "gt_class_name",
+            "tk_overlap_classes",
+            "tk_overlap_ids",
+            "tk_neg_ids",
+            "tk_exh_ids",
+            "tk_class_eval_tk_ids",
+            "tk_dets",
+            "tk_classes",
+            "tk_confidences",
+            "sim_scores",
+        ]
+        data = {key: [None] * raw_data["num_timesteps"] for key in data_keys}
+        unique_gt_ids = []
+        unique_tk_ids = []
+        num_gt_dets = 0
+        num_tk_cls_dets = 0
+        num_tk_overlap_dets = 0
+        overlap_ious_thr = 0.5
+        loc_and_asso_tk_ids = []
+
+        for t in range(raw_data["num_timesteps"]):
+            # only extract relevant dets for this class for preproc and eval
+            if cls == "all":
+                gt_class_mask = np.ones_like(raw_data["gt_classes"][t]).astype(bool)
+            else:
+                gt_class_mask = np.atleast_1d(
+                    raw_data["gt_classes"][t] == cls_id
+                ).astype(bool)
+
+            # select GT that is not in the evaluating classes
+            if assignment is not None and assignment:
+                all_gt_ids = list(assignment[t].keys())
+                gt_ids_in = raw_data["gt_ids"][t][gt_class_mask]
+                gt_ids_out = set(all_gt_ids) - set(gt_ids_in)
+                tk_ids_out = set([assignment[t][key] for key in list(gt_ids_out)])
+
+            # compute overlapped tracks and add their ids to overlap_tk_ids
+            sim_scores = raw_data["similarity_scores"]
+            overlap_ids_masks = (sim_scores[t][gt_class_mask] >= overlap_ious_thr).any(
+                axis=0
+            )
+            overlap_tk_ids_t = raw_data["tk_ids"][t][overlap_ids_masks]
+            if assignment is not None and assignment:
+                data["tk_overlap_ids"][t] = list(set(overlap_tk_ids_t) - tk_ids_out)
+            else:
+                data["tk_overlap_ids"][t] = list(set(overlap_tk_ids_t))
+
+            loc_and_asso_tk_ids += data["tk_overlap_ids"][t]
+
+            data["tk_exh_ids"][t] = []
+            data["tk_neg_ids"][t] = []
+
+            if cls == "all":
+                continue
+
+        # remove tk_ids that has been assigned to GT belongs to other classes.
+        loc_and_asso_tk_ids = list(set(loc_and_asso_tk_ids))
+
+        # remove all unwanted unmatched tracker detections
+        for t in range(raw_data["num_timesteps"]):
+            # add gt to the data
+            if cls == "all":
+                gt_class_mask = np.ones_like(raw_data["gt_classes"][t]).astype(bool)
+            else:
+                gt_class_mask = np.atleast_1d(
+                    raw_data["gt_classes"][t] == cls_id
+                ).astype(bool)
+                data["gt_classes"][t] = cls_id
+                data["gt_class_name"][t] = cls
+
+            gt_ids = raw_data["gt_ids"][t][gt_class_mask]
+            if self.use_mask:
+                gt_dets = [raw_data['gt_dets'][t][ind] for ind in range(len(gt_class_mask)) if gt_class_mask[ind]]
+            else:
+                gt_dets = raw_data["gt_dets"][t][gt_class_mask]
+            data["gt_ids"][t] = gt_ids
+            data["gt_dets"][t] = gt_dets
+
+            # filter pred and only keep those that highly overlap with GTs
+            tk_mask = np.isin(
+                raw_data["tk_ids"][t], np.array(loc_and_asso_tk_ids), assume_unique=True
+            )
+            tk_overlap_mask = np.isin(
+                raw_data["tk_ids"][t],
+                np.array(data["tk_overlap_ids"][t]),
+                assume_unique=True,
+            )
+
+            tk_ids = raw_data["tk_ids"][t][tk_mask]
+            if self.use_mask:
+                tk_dets = [raw_data['tk_dets'][t][ind] for ind in range(len(tk_mask)) if
+                            tk_mask[ind]]
+            else:
+                tk_dets = raw_data["tk_dets"][t][tk_mask]
+            tracker_classes = raw_data["tk_classes"][t][tk_mask]
+
+            # add overlap classes for computing the FP for Cls term
+            tracker_overlap_classes = raw_data["tk_classes"][t][tk_overlap_mask]
+            tracker_confidences = raw_data["tk_confidences"][t][tk_mask]
+            sim_scores_masked = sim_scores[t][gt_class_mask, :][:, tk_mask]
+
+            # add filtered prediction to the data
+            data["tk_classes"][t] = tracker_classes
+            data["tk_overlap_classes"][t] = tracker_overlap_classes
+            data["tk_ids"][t] = tk_ids
+            data["tk_dets"][t] = tk_dets
+            data["tk_confidences"][t] = tracker_confidences
+            data["sim_scores"][t] = sim_scores_masked
+            data["tk_class_eval_tk_ids"][t] = set(
+                list(data["tk_overlap_ids"][t])
+                + list(data["tk_neg_ids"][t])
+                + list(data["tk_exh_ids"][t])
+            )
+
+            # count total number of detections
+            unique_gt_ids += list(np.unique(data["gt_ids"][t]))
+            # the unique track ids are for association.
+            unique_tk_ids += list(np.unique(data["tk_ids"][t]))
+
+            num_tk_overlap_dets += len(data["tk_overlap_ids"][t])
+            num_tk_cls_dets += len(data["tk_class_eval_tk_ids"][t])
+            num_gt_dets += len(data["gt_ids"][t])
+
+        # re-label IDs such that there are no empty IDs
+        if len(unique_gt_ids) > 0:
+            unique_gt_ids = np.unique(unique_gt_ids)
+            gt_id_map = np.nan * np.ones((np.max(unique_gt_ids) + 1))
+            gt_id_map[unique_gt_ids] = np.arange(len(unique_gt_ids))
+            data["gt_id_map"] = {}
+            for gt_id in unique_gt_ids:
+                new_gt_id = gt_id_map[gt_id].astype(int)
+                data["gt_id_map"][new_gt_id] = gt_id
+
+            for t in range(raw_data["num_timesteps"]):
+                if len(data["gt_ids"][t]) > 0:
+                    data["gt_ids"][t] = gt_id_map[data["gt_ids"][t]].astype(int)
+
+        if len(unique_tk_ids) > 0:
+            unique_tk_ids = np.unique(unique_tk_ids)
+            tk_id_map = np.nan * np.ones((np.max(unique_tk_ids) + 1))
+            tk_id_map[unique_tk_ids] = np.arange(len(unique_tk_ids))
+
+            data["tk_id_map"] = {}
+            for track_id in unique_tk_ids:
+                new_track_id = tk_id_map[track_id].astype(int)
+                data["tk_id_map"][new_track_id] = track_id
+
+            for t in range(raw_data["num_timesteps"]):
+                if len(data["tk_ids"][t]) > 0:
+                    data["tk_ids"][t] = tk_id_map[data["tk_ids"][t]].astype(int)
+                if len(data["tk_overlap_ids"][t]) > 0:
+                    data["tk_overlap_ids"][t] = tk_id_map[
+                        data["tk_overlap_ids"][t]
+                    ].astype(int)
+
+        # record overview statistics.
+        data["num_tk_cls_dets"] = num_tk_cls_dets
+        data["num_tk_overlap_dets"] = num_tk_overlap_dets
+        data["num_gt_dets"] = num_gt_dets
+        data["num_tk_ids"] = len(unique_tk_ids)
+        data["num_gt_ids"] = len(unique_gt_ids)
+        data["num_timesteps"] = raw_data["num_timesteps"]
+        data["seq"] = raw_data["seq"]
+
+        self._check_unique_ids(data)
+
+        return data
+
+    @_timing.time
+    def get_preprocessed_seq_data(
+        self, raw_data, cls, assignment=None, thresholds=[50, 75]
+    ):
+        """Preprocess data for a single sequence for a single class."""
+        data = {}
+        if thresholds is None:
+            thresholds = [50]
+        elif isinstance(thresholds, int):
+            thresholds = [thresholds]
+
+        for thr in thresholds:
+            assignment_thr = None
+            if assignment is not None:
+                assignment_thr = assignment[thr]
+            data[thr] = self.get_preprocessed_seq_data_thr(
+                raw_data, cls, assignment_thr
+            )
+
+        return data
+
+    def _calculate_similarities(self, gt_dets_t, tk_dets_t):
+        """Compute similarity scores."""
+        if self.use_mask:
+            similarity_scores = self._calculate_mask_ious(gt_dets_t, tk_dets_t, is_encoded=True, do_ioa=False)
+        else:
+            similarity_scores = self._calculate_box_ious(gt_dets_t, tk_dets_t)
+        return similarity_scores
+
+    def _merge_categories(self, annotations):
+        """Merges categories with a merged tag.
+
+        Adapted from https://github.com/TAO-Dataset.
+        """
+        merge_map = {}
+        for category in self.gt_data["categories"]:
+            if "merged" in category:
+                for to_merge in category["merged"]:
+                    merge_map[to_merge["id"]] = category["id"]
+
+        for ann in annotations:
+            ann["category_id"] = merge_map.get(ann["category_id"], ann["category_id"])
+
+    def _compute_vid_mappings(self, annotations):
+        """Computes mappings from videos to corresponding tracks and images."""
+        vids_to_tracks = {}
+        vids_to_imgs = {}
+        vid_ids = [vid["id"] for vid in self.gt_data["videos"]]
+
+        # compute an mapping from image IDs to images
+        images = {}
+        for image in self.gt_data["images"]:
+            images[image["id"]] = image
+
+        for ann in annotations:
+            ann["area"] = ann["bbox"][2] * ann["bbox"][3]
+
+            vid = ann["video_id"]
+            if ann["video_id"] not in vids_to_tracks.keys():
+                vids_to_tracks[ann["video_id"]] = list()
+            if ann["video_id"] not in vids_to_imgs.keys():
+                vids_to_imgs[ann["video_id"]] = list()
+
+            # fill in vids_to_tracks
+            tid = ann["track_id"]
+            exist_tids = [track["id"] for track in vids_to_tracks[vid]]
+            try:
+                index1 = exist_tids.index(tid)
+            except ValueError:
+                index1 = -1
+            if tid not in exist_tids:
+                curr_track = {
+                    "id": tid,
+                    "category_id": ann["category_id"],
+                    "video_id": vid,
+                    "annotations": [ann],
+                }
+                vids_to_tracks[vid].append(curr_track)
+            else:
+                vids_to_tracks[vid][index1]["annotations"].append(ann)
+
+            # fill in vids_to_imgs
+            img_id = ann["image_id"]
+            exist_img_ids = [img["id"] for img in vids_to_imgs[vid]]
+            try:
+                index2 = exist_img_ids.index(img_id)
+            except ValueError:
+                index2 = -1
+            if index2 == -1:
+                curr_img = {"id": img_id, "annotations": [ann]}
+                vids_to_imgs[vid].append(curr_img)
+            else:
+                vids_to_imgs[vid][index2]["annotations"].append(ann)
+
+        # sort annotations by frame index and compute track area
+        for vid, tracks in vids_to_tracks.items():
+            for track in tracks:
+                track["annotations"] = sorted(
+                    track["annotations"],
+                    key=lambda x: images[x["image_id"]]["frame_index"],
+                )
+                # compute average area
+                track["area"] = sum(x["area"] for x in track["annotations"]) / len(
+                    track["annotations"]
+                )
+
+        # ensure all videos are present
+        for vid_id in vid_ids:
+            if vid_id not in vids_to_tracks.keys():
+                vids_to_tracks[vid_id] = []
+            if vid_id not in vids_to_imgs.keys():
+                vids_to_imgs[vid_id] = []
+
+        return vids_to_tracks, vids_to_imgs
+
+    def _compute_image_to_timestep_mappings(self):
+        """Computes a mapping from images to timestep in sequence."""
+        images = {}
+        for image in self.gt_data["images"]:
+            images[image["id"]] = image
+
+        seq_to_imgs_to_timestep = {vid["id"]: dict() for vid in self.gt_data["videos"]}
+        for vid in seq_to_imgs_to_timestep:
+            curr_imgs = [img["id"] for img in self.video2gt_image[vid]]
+            curr_imgs = sorted(curr_imgs, key=lambda x: images[x]["frame_index"])
+            seq_to_imgs_to_timestep[vid] = {
+                curr_imgs[i]: i for i in range(len(curr_imgs))
+            }
+
+        return seq_to_imgs_to_timestep
+
+    def _limit_dets_per_image(self, annotations):
+        """Limits the number of detections for each image.
+
+        Adapted from https://github.com/TAO-Dataset/.
+        """
+        max_dets = self.config["MAX_DETECTIONS"]
+        img_ann = defaultdict(list)
+        for ann in annotations:
+            img_ann[ann["image_id"]].append(ann)
+
+        for img_id, _anns in img_ann.items():
+            if len(_anns) <= max_dets:
+                continue
+            _anns = sorted(_anns, key=lambda x: x["score"], reverse=True)
+            img_ann[img_id] = _anns[:max_dets]
+
+        return [ann for anns in img_ann.values() for ann in anns]
+
+    def _fill_video_ids_inplace(self, annotations):
+        """Fills in missing video IDs inplace.
+
+        Adapted from https://github.com/TAO-Dataset/.
+        """
+        missing_video_id = [x for x in annotations if "video_id" not in x]
+        if missing_video_id:
+            image_id_to_video_id = {
+                x["id"]: x["video_id"] for x in self.gt_data["images"]
+            }
+            for x in missing_video_id:
+                x["video_id"] = image_id_to_video_id[x["image_id"]]
+
+    @staticmethod
+    def _make_tk_ids_unique(annotations):
+        """Makes track IDs unqiue over the whole annotation set.
+
+        Adapted from https://github.com/TAO-Dataset/.
+        """
+        track_id_videos = {}
+        track_ids_to_update = set()
+        max_track_id = 0
+        for ann in annotations:
+            t = ann["track_id"]
+            if t not in track_id_videos:
+                track_id_videos[t] = ann["video_id"]
+
+            if ann["video_id"] != track_id_videos[t]:
+                # track id is assigned to multiple videos
+                track_ids_to_update.add(t)
+            max_track_id = max(max_track_id, t)
+
+        if track_ids_to_update:
+            print("true")
+            next_id = itertools.count(max_track_id + 1)
+            new_tk_ids = defaultdict(lambda: next(next_id))
+            for ann in annotations:
+                t = ann["track_id"]
+                v = ann["video_id"]
+                if t in track_ids_to_update:
+                    ann["track_id"] = new_tk_ids[t, v]
+        return len(track_ids_to_update)

sam3/eval/teta_eval_toolkit/eval.py

@@ -0,0 +1,275 @@
+# fmt: off
+# flake8: noqa
+
+import copy
+import os
+import pickle
+import time
+import traceback
+from functools import partial
+from multiprocessing.pool import Pool
+
+import numpy as np
+
+from . import _timing, utils
+from .config import get_default_eval_config, init_config
+from .utils import TrackEvalException
+
+
+class Evaluator:
+    """Evaluator class for evaluating different metrics for each datasets."""
+
+    def __init__(self, config=None):
+        """Initialize the evaluator with a config file."""
+        self.config = init_config(config, get_default_eval_config(), "Eval")
+        # Only run timing analysis if not run in parallel.
+        if self.config["TIME_PROGRESS"] and not self.config["USE_PARALLEL"]:
+            _timing.DO_TIMING = True
+            if self.config["DISPLAY_LESS_PROGRESS"]:
+                _timing.DISPLAY_LESS_PROGRESS = True
+
+    @_timing.time
+    def evaluate(self, dataset_list, metrics_list):
+        """Evaluate a set of metrics on a set of datasets."""
+        config = self.config
+        metrics_list = metrics_list
+        metric_names = utils.validate_metrics_list(metrics_list)
+        dataset_names = [dataset.get_name() for dataset in dataset_list]
+        output_res = {}
+        output_msg = {}
+
+        for dataset, dname in zip(dataset_list, dataset_names):
+            # Get dataset info about what to evaluate
+            output_res[dname] = {}
+            output_msg[dname] = {}
+            tracker_list, seq_list, class_list = dataset.get_eval_info()
+            print(
+                f"\nEvaluating {len(tracker_list)} tracker(s) on "
+                f"{len(seq_list)} sequence(s) for {len(class_list)} class(es)"
+                f" on {dname} dataset using the following "
+                f'metrics: {", ".join(metric_names)}\n'
+            )
+
+            # Evaluate each tracker
+            for tracker in tracker_list:
+                try:
+                    output_res, output_msg = self.evaluate_tracker(
+                        tracker,
+                        dataset,
+                        dname,
+                        class_list,
+                        metrics_list,
+                        metric_names,
+                        seq_list,
+                        output_res,
+                        output_msg,
+                    )
+                except Exception as err:
+                    output_res[dname][tracker] = None
+                    if type(err) == TrackEvalException:
+                        output_msg[dname][tracker] = str(err)
+                    else:
+                        output_msg[dname][tracker] = "Unknown error occurred."
+                    print("Tracker %s was unable to be evaluated." % tracker)
+                    print(err)
+                    traceback.print_exc()
+                    if config["LOG_ON_ERROR"] is not None:
+                        with open(config["LOG_ON_ERROR"], "a") as f:
+                            print(dname, file=f)
+                            print(tracker, file=f)
+                            print(traceback.format_exc(), file=f)
+                            print("\n\n\n", file=f)
+                    if config["BREAK_ON_ERROR"]:
+                        raise err
+                    elif config["RETURN_ON_ERROR"]:
+                        return output_res, output_msg
+
+        return output_res, output_msg
+
+    def evaluate_tracker(
+        self,
+        tracker,
+        dataset,
+        dname,
+        class_list,
+        metrics_list,
+        metric_names,
+        seq_list,
+        output_res,
+        output_msg,
+    ):
+        """Evaluate each sequence in parallel or in series."""
+        print("\nEvaluating %s\n" % tracker)
+        time_start = time.time()
+        config = self.config
+        if config["USE_PARALLEL"]:
+            with Pool(config["NUM_PARALLEL_CORES"]) as pool:
+                _eval_sequence = partial(
+                    eval_sequence,
+                    dataset=dataset,
+                    tracker=tracker,
+                    class_list=class_list,
+                    metrics_list=metrics_list,
+                    metric_names=metric_names,
+                )
+                results = pool.map(_eval_sequence, seq_list)
+                res = dict(zip(seq_list, results))
+        else:
+            res = {}
+            for curr_seq in sorted(seq_list):
+                res[curr_seq] = eval_sequence(
+                    curr_seq, dataset, tracker, class_list, metrics_list, metric_names
+                )
+
+
+        # collecting combined cls keys (cls averaged, det averaged, super classes)
+        cls_keys = []
+        res["COMBINED_SEQ"] = {}
+        # combine sequences for each class
+        for c_cls in class_list:
+            res["COMBINED_SEQ"][c_cls] = {}
+            for metric, mname in zip(metrics_list, metric_names):
+                curr_res = {
+                    seq_key: seq_value[c_cls][mname]
+                    for seq_key, seq_value in res.items()
+                    if seq_key != "COMBINED_SEQ"
+                }
+                # combine results over all sequences and then over all classes
+                res["COMBINED_SEQ"][c_cls][mname] = metric.combine_sequences(curr_res)
+
+        # combine classes
+        if dataset.should_classes_combine:
+            if config["OUTPUT_PER_SEQ_RES"]:
+                video_keys = res.keys()
+            else:
+                video_keys = ["COMBINED_SEQ"]
+            for v_key in video_keys:
+                cls_keys += ["average"]
+                res[v_key]["average"] = {}
+                for metric, mname in zip(metrics_list, metric_names):
+                    cls_res = {
+                        cls_key: cls_value[mname]
+                        for cls_key, cls_value in res[v_key].items()
+                        if cls_key not in cls_keys
+                    }
+                    res[v_key]["average"][
+                        mname
+                    ] = metric.combine_classes_class_averaged(
+                        cls_res, ignore_empty=True
+                    )
+
+        # combine classes to super classes
+        if dataset.use_super_categories:
+            for cat, sub_cats in dataset.super_categories.items():
+                cls_keys.append(cat)
+                res["COMBINED_SEQ"][cat] = {}
+                for metric, mname in zip(metrics_list, metric_names):
+                    cat_res = {
+                        cls_key: cls_value[mname]
+                        for cls_key, cls_value in res["COMBINED_SEQ"].items()
+                        if cls_key in sub_cats
+                    }
+                    res["COMBINED_SEQ"][cat][
+                        mname
+                    ] = metric.combine_classes_det_averaged(cat_res)
+        # Print and output results in various formats
+        if config["TIME_PROGRESS"]:
+            print(
+                f"\nAll sequences for {tracker} finished in"
+                f" {time.time() - time_start} seconds"
+            )
+        output_fol = dataset.get_output_fol(tracker)
+        os.makedirs(output_fol, exist_ok=True)
+
+        # take a mean of each field of each thr
+        if config["OUTPUT_PER_SEQ_RES"]:
+            all_res = copy.deepcopy(res)
+            summary_keys = res.keys()
+        else:
+            all_res = copy.deepcopy(res["COMBINED_SEQ"])
+            summary_keys = ["COMBINED_SEQ"]
+        thr_key_list = [50]
+        for s_key in summary_keys:
+            for metric, mname in zip(metrics_list, metric_names):
+                if mname != "TETA":
+                    if s_key == "COMBINED_SEQ":
+                        metric.print_table(
+                            {"COMBINED_SEQ": res["COMBINED_SEQ"][cls_keys[0]][mname]},
+                            tracker,
+                            cls_keys[0],
+                        )
+                    continue
+
+                for c_cls in res[s_key].keys():
+                    for thr in thr_key_list:
+                        all_res[s_key][c_cls][mname][thr] = metric._summary_row(
+                            res[s_key][c_cls][mname][thr]
+                        )
+                    x = (
+                        np.array(list(all_res[s_key][c_cls]["TETA"].values()))
+                        .astype("float")
+                        .mean(axis=0)
+                    )
+                    all_res_summary = list(x.round(decimals=2).astype("str"))
+                    all_res[s_key][c_cls][mname]["ALL"] = all_res_summary
+                if config["OUTPUT_SUMMARY"] and s_key == "COMBINED_SEQ":
+                    for t in thr_key_list:
+                        metric.print_summary_table(
+                            all_res[s_key][cls_keys[0]][mname][t],
+                            t,
+                            tracker,
+                            cls_keys[0],
+                        )
+
+        if config["OUTPUT_TEM_RAW_DATA"]:
+            out_file = os.path.join(output_fol, "teta_summary_results.pth")
+            pickle.dump(all_res, open(out_file, "wb"))
+            print("Saved the TETA summary results.")
+
+        # output
+        output_res[dname][mname] = all_res[s_key][cls_keys[0]][mname][t]
+        output_msg[dname][tracker] = "Success"
+
+        return output_res, output_msg
+
+
+@_timing.time
+def eval_sequence(seq, dataset, tracker, class_list, metrics_list, metric_names):
+    """Function for evaluating a single sequence."""
+    raw_data = dataset.get_raw_seq_data(tracker, seq)
+    seq_res = {}
+
+    if "TETA" in metric_names:
+        thresholds = [50]
+        data_all_class = dataset.get_preprocessed_seq_data(
+            raw_data, "all", thresholds=thresholds
+        )
+        teta = metrics_list[metric_names.index("TETA")]
+        assignment = teta.compute_global_assignment(data_all_class)
+
+        # create a dict to save Cls_FP for each class in different thr.
+        cls_fp = {
+            key: {
+                cls: np.zeros((len(np.arange(0.5, 0.99, 0.05)))) for cls in class_list
+            }
+            for key in thresholds
+        }
+
+    for cls in class_list:
+        seq_res[cls] = {}
+        data = dataset.get_preprocessed_seq_data(raw_data, cls, assignment, thresholds)
+
+        for metric, mname in zip(metrics_list, metric_names):
+            if mname == "TETA":
+                seq_res[cls][mname], cls_fp, _ = metric.eval_sequence(
+                    data, cls, dataset.clsid2cls_name, cls_fp
+                )
+            else:
+                seq_res[cls][mname] = metric.eval_sequence(data)
+
+    if "TETA" in metric_names:
+        for thr in thresholds:
+            for cls in class_list:
+                seq_res[cls]["TETA"][thr]["Cls_FP"] += cls_fp[thr][cls]
+
+    return seq_res

sam3/eval/teta_eval_toolkit/metrics/__init__.py

@@ -0,0 +1,4 @@
+# fmt: off
+# flake8: noqa
+
+from .teta import TETA

sam3/eval/teta_eval_toolkit/metrics/_base_metric.py

@@ -0,0 +1,148 @@
+# fmt: off
+# flake8: noqa
+
+from abc import ABC, abstractmethod
+
+import numpy as np
+
+from .. import _timing
+from ..utils import TrackEvalException
+
+
+class _BaseMetric(ABC):
+    @abstractmethod
+    def __init__(self):
+        self.plottable = False
+        self.integer_fields = []
+        self.float_fields = []
+        self.array_labels = []
+        self.integer_array_fields = []
+        self.float_array_fields = []
+        self.fields = []
+        self.summary_fields = []
+        self.registered = False
+
+    #####################################################################
+    # Abstract functions for subclasses to implement
+
+    @_timing.time
+    @abstractmethod
+    def eval_sequence(self, data):
+        ...
+
+    @abstractmethod
+    def combine_sequences(self, all_res):
+        ...
+
+    @abstractmethod
+    def combine_classes_class_averaged(self, all_res, ignore_empty=False):
+        ...
+
+    @abstractmethod
+    def combine_classes_det_averaged(self, all_res):
+        ...
+
+    def plot_single_tracker_results(self, all_res, tracker, output_folder, cls):
+        """Plot results, only valid for metrics with self.plottable."""
+        if self.plottable:
+            raise NotImplementedError(
+                f"plot_results is not implemented for metric {self.get_name()}"
+            )
+        else:
+            pass
+
+    #####################################################################
+    # Helper functions which are useful for all metrics:
+
+    @classmethod
+    def get_name(cls):
+        return cls.__name__
+
+    @staticmethod
+    def _combine_sum(all_res, field):
+        """Combine sequence results via sum"""
+        return sum([all_res[k][field] for k in all_res.keys()])
+
+    @staticmethod
+    def _combine_weighted_av(all_res, field, comb_res, weight_field):
+        """Combine sequence results via weighted average."""
+        return sum(
+            [all_res[k][field] * all_res[k][weight_field] for k in all_res.keys()]
+        ) / np.maximum(1.0, comb_res[weight_field])
+
+    def print_table(self, table_res, tracker, cls):
+        """Print table of results for all sequences."""
+        print("")
+        metric_name = self.get_name()
+        self._row_print(
+            [metric_name + ": " + tracker + "-" + cls] + self.summary_fields
+        )
+        for seq, results in sorted(table_res.items()):
+            if seq == "COMBINED_SEQ":
+                continue
+            summary_res = self._summary_row(results)
+            self._row_print([seq] + summary_res)
+        summary_res = self._summary_row(table_res["COMBINED_SEQ"])
+        self._row_print(["COMBINED"] + summary_res)
+
+    def _summary_row(self, results_):
+        vals = []
+        for h in self.summary_fields:
+            if h in self.float_array_fields:
+                vals.append("{0:1.5g}".format(100 * np.mean(results_[h])))
+            elif h in self.float_fields:
+                vals.append("{0:1.5g}".format(100 * float(results_[h])))
+            elif h in self.integer_fields:
+                vals.append("{0:d}".format(int(results_[h])))
+            else:
+                raise NotImplementedError(
+                    "Summary function not implemented for this field type."
+                )
+        return vals
+
+    @staticmethod
+    def _row_print(*argv):
+        """Print results in evenly spaced rows, with more space in first row."""
+        if len(argv) == 1:
+            argv = argv[0]
+        to_print = "%-35s" % argv[0]
+        for v in argv[1:]:
+            to_print += "%-10s" % str(v)
+        print(to_print)
+
+    def summary_results(self, table_res):
+        """Return a simple summary of final results for a tracker."""
+        return dict(
+            zip(self.summary_fields, self._summary_row(table_res["COMBINED_SEQ"]),)
+        )
+
+    def detailed_results(self, table_res):
+        """Return detailed final results for a tracker."""
+        # Get detailed field information
+        detailed_fields = self.float_fields + self.integer_fields
+        for h in self.float_array_fields + self.integer_array_fields:
+            for alpha in [int(100 * x) for x in self.array_labels]:
+                detailed_fields.append(h + "___" + str(alpha))
+            detailed_fields.append(h + "___AUC")
+
+        # Get detailed results
+        detailed_results = {}
+        for seq, res in table_res.items():
+            detailed_row = self._detailed_row(res)
+            if len(detailed_row) != len(detailed_fields):
+                raise TrackEvalException(
+                    f"Field names and data have different sizes "
+                    f"({len(detailed_row)} and {len(detailed_fields)})"
+                )
+            detailed_results[seq] = dict(zip(detailed_fields, detailed_row))
+        return detailed_results
+
+    def _detailed_row(self, res):
+        detailed_row = []
+        for h in self.float_fields + self.integer_fields:
+            detailed_row.append(res[h])
+        for h in self.float_array_fields + self.integer_array_fields:
+            for i, _ in enumerate([int(100 * x) for x in self.array_labels]):
+                detailed_row.append(res[h][i])
+            detailed_row.append(np.mean(res[h]))
+        return detailed_row

sam3/eval/teta_eval_toolkit/metrics/teta.py

@@ -0,0 +1,399 @@
+# fmt: off
+# flake8: noqa
+
+"""Track Every Thing Accuracy metric."""
+
+import numpy as np
+from scipy.optimize import linear_sum_assignment
+
+from .. import _timing
+from ._base_metric import _BaseMetric
+
+EPS = np.finfo("float").eps  # epsilon
+
+
+class TETA(_BaseMetric):
+    """TETA metric."""
+
+    def __init__(self, exhaustive=False, config=None):
+        """Initialize metric."""
+        super().__init__()
+        self.plottable = True
+        self.array_labels = np.arange(0.0, 0.99, 0.05)
+        self.cls_array_labels = np.arange(0.5, 0.99, 0.05)
+
+        self.integer_array_fields = [
+            "Loc_TP",
+            "Loc_FN",
+            "Loc_FP",
+            "Cls_TP",
+            "Cls_FN",
+            "Cls_FP",
+        ]
+        self.float_array_fields = (
+            ["TETA", "LocA", "AssocA", "ClsA"]
+            + ["LocRe", "LocPr"]
+            + ["AssocRe", "AssocPr"]
+            + ["ClsRe", "ClsPr"]
+        )
+        self.fields = self.float_array_fields + self.integer_array_fields
+        self.summary_fields = self.float_array_fields
+        self.exhaustive = exhaustive
+
+    def compute_global_assignment(self, data_thr, alpha=0.5):
+        """Compute global assignment of TP."""
+        res = {
+            thr: {t: {} for t in range(data_thr[thr]["num_timesteps"])}
+            for thr in data_thr
+        }
+
+        for thr in data_thr:
+            data = data_thr[thr]
+            # return empty result if tracker or gt sequence is empty
+            if data["num_tk_overlap_dets"] == 0 or data["num_gt_dets"] == 0:
+                return res
+
+            # global alignment score
+            ga_score, _, _ = self.compute_global_alignment_score(data)
+
+            # calculate scores for each timestep
+            for t, (gt_ids_t, tk_ids_t) in enumerate(
+                zip(data["gt_ids"], data["tk_ids"])
+            ):
+                # get matches optimizing for TETA
+                amatch_rows, amatch_cols = self.compute_matches(
+                    data, t, ga_score, gt_ids_t, tk_ids_t, alpha=alpha
+                )
+                gt_ids = [data["gt_id_map"][tid] for tid in gt_ids_t[amatch_rows[0]]]
+                matched_ids = [
+                    data["tk_id_map"][tid] for tid in tk_ids_t[amatch_cols[0]]
+                ]
+                res[thr][t] = dict(zip(gt_ids, matched_ids))
+
+        return res
+
+    def eval_sequence_single_thr(self, data, cls, cid2clsname, cls_fp_thr, thr):
+        """Computes TETA metric for one threshold for one sequence."""
+        res = {}
+        class_info_list = []
+        for field in self.float_array_fields + self.integer_array_fields:
+            if field.startswith("Cls"):
+                res[field] = np.zeros(len(self.cls_array_labels), dtype=float)
+            else:
+                res[field] = np.zeros((len(self.array_labels)), dtype=float)
+
+        # return empty result if tracker or gt sequence is empty
+        if data["num_tk_overlap_dets"] == 0:
+            res["Loc_FN"] = data["num_gt_dets"] * np.ones(
+                (len(self.array_labels)), dtype=float
+            )
+            if self.exhaustive:
+                cls_fp_thr[cls] = data["num_tk_cls_dets"] * np.ones(
+                    (len(self.cls_array_labels)), dtype=float
+                )
+            res = self._compute_final_fields(res)
+            return res, cls_fp_thr, class_info_list
+
+        if data["num_gt_dets"] == 0:
+            if self.exhaustive:
+                cls_fp_thr[cls] = data["num_tk_cls_dets"] * np.ones(
+                    (len(self.cls_array_labels)), dtype=float
+                )
+            res = self._compute_final_fields(res)
+            return res, cls_fp_thr, class_info_list
+
+        # global alignment score
+        ga_score, gt_id_count, tk_id_count = self.compute_global_alignment_score(data)
+        matches_counts = [np.zeros_like(ga_score) for _ in self.array_labels]
+
+        # calculate scores for each timestep
+        for t, (gt_ids_t, tk_ids_t, tk_overlap_ids_t, tk_cls_ids_t) in enumerate(
+            zip(
+                data["gt_ids"],
+                data["tk_ids"],
+                data["tk_overlap_ids"],
+                data["tk_class_eval_tk_ids"],
+            )
+        ):
+            # deal with the case that there are no gt_det/tk_det in a timestep
+            if len(gt_ids_t) == 0:
+                if self.exhaustive:
+                    cls_fp_thr[cls] += len(tk_cls_ids_t)
+                continue
+
+            # get matches optimizing for TETA
+            amatch_rows, amatch_cols = self.compute_matches(
+                data, t, ga_score, gt_ids_t, tk_ids_t, list(self.array_labels)
+            )
+
+            # map overlap_ids to original ids.
+            if len(tk_overlap_ids_t) != 0:
+                sorter = np.argsort(tk_ids_t)
+                indexes = sorter[
+                    np.searchsorted(tk_ids_t, tk_overlap_ids_t, sorter=sorter)
+                ]
+                sim_t = data["sim_scores"][t][:, indexes]
+                fpl_candidates = tk_overlap_ids_t[(sim_t >= (thr / 100)).any(axis=0)]
+                fpl_candidates_ori_ids_t = np.array(
+                    [data["tk_id_map"][tid] for tid in fpl_candidates]
+                )
+            else:
+                fpl_candidates_ori_ids_t = []
+
+            if self.exhaustive:
+                cls_fp_thr[cls] += len(tk_cls_ids_t) - len(tk_overlap_ids_t)
+
+            # calculate and accumulate basic statistics
+            for a, alpha in enumerate(self.array_labels):
+                match_row, match_col = amatch_rows[a], amatch_cols[a]
+                num_matches = len(match_row)
+                matched_ori_ids = set(
+                    [data["tk_id_map"][tid] for tid in tk_ids_t[match_col]]
+                )
+                match_tk_cls = data["tk_classes"][t][match_col]
+                wrong_tk_cls = match_tk_cls[match_tk_cls != data["gt_classes"][t]]
+
+                num_class_and_det_matches = np.sum(
+                    match_tk_cls == data["gt_classes"][t]
+                )
+
+                if alpha >= 0.5:
+                    for cid in wrong_tk_cls:
+                        if cid in cid2clsname:
+                            cname = cid2clsname[cid]
+                            cls_fp_thr[cname][a - 10] += 1
+                    res["Cls_TP"][a - 10] += num_class_and_det_matches
+                    res["Cls_FN"][a - 10] += num_matches - num_class_and_det_matches
+
+                res["Loc_TP"][a] += num_matches
+                res["Loc_FN"][a] += len(gt_ids_t) - num_matches
+                res["Loc_FP"][a] += len(set(fpl_candidates_ori_ids_t) - matched_ori_ids)
+
+                if num_matches > 0:
+                    matches_counts[a][gt_ids_t[match_row], tk_ids_t[match_col]] += 1
+
+        # calculate AssocA, AssocRe, AssocPr
+        self.compute_association_scores(res, matches_counts, gt_id_count, tk_id_count)
+
+        # calculate final scores
+        res = self._compute_final_fields(res)
+        return res, cls_fp_thr, class_info_list
+
+    def compute_global_alignment_score(self, data):
+        """Computes global alignment score."""
+        num_matches = np.zeros((data["num_gt_ids"], data["num_tk_ids"]))
+        gt_id_count = np.zeros((data["num_gt_ids"], 1))
+        tk_id_count = np.zeros((1, data["num_tk_ids"]))
+
+        # loop through each timestep and accumulate global track info.
+        for t, (gt_ids_t, tk_ids_t) in enumerate(zip(data["gt_ids"], data["tk_ids"])):
+            # count potential matches between ids in each time step
+            # these are normalized, weighted by match similarity
+            sim = data["sim_scores"][t]
+            sim_iou_denom = sim.sum(0, keepdims=True) + sim.sum(1, keepdims=True) - sim
+            sim_iou = np.zeros_like(sim)
+            mask = sim_iou_denom > (0 + EPS)
+            sim_iou[mask] = sim[mask] / sim_iou_denom[mask]
+            num_matches[gt_ids_t[:, None], tk_ids_t[None, :]] += sim_iou
+
+            # calculate total number of dets for each gt_id and tk_id.
+            gt_id_count[gt_ids_t] += 1
+            tk_id_count[0, tk_ids_t] += 1
+
+        # Calculate overall Jaccard alignment score between IDs
+        ga_score = num_matches / (gt_id_count + tk_id_count - num_matches)
+        return ga_score, gt_id_count, tk_id_count
+
+    def compute_matches(self, data, t, ga_score, gt_ids, tk_ids, alpha):
+        """Compute matches based on alignment score."""
+        sim = data["sim_scores"][t]
+        score_mat = ga_score[gt_ids[:, None], tk_ids[None, :]] * sim
+        # Hungarian algorithm to find best matches
+        match_rows, match_cols = linear_sum_assignment(-score_mat)
+
+        if not isinstance(alpha, list):
+            alpha = [alpha]
+        alpha_match_rows, alpha_match_cols = [], []
+        for a in alpha:
+            matched_mask = sim[match_rows, match_cols] >= a - EPS
+            alpha_match_rows.append(match_rows[matched_mask])
+            alpha_match_cols.append(match_cols[matched_mask])
+        return alpha_match_rows, alpha_match_cols
+
+    def compute_association_scores(self, res, matches_counts, gt_id_count, tk_id_count):
+        """Calculate association scores for each alpha.
+
+        First calculate scores per gt_id/tk_id combo,
+        and then average over the number of detections.
+        """
+        for a, _ in enumerate(self.array_labels):
+            matches_count = matches_counts[a]
+            ass_a = matches_count / np.maximum(
+                1, gt_id_count + tk_id_count - matches_count
+            )
+            res["AssocA"][a] = np.sum(matches_count * ass_a) / np.maximum(
+                1, res["Loc_TP"][a]
+            )
+            ass_re = matches_count / np.maximum(1, gt_id_count)
+            res["AssocRe"][a] = np.sum(matches_count * ass_re) / np.maximum(
+                1, res["Loc_TP"][a]
+            )
+            ass_pr = matches_count / np.maximum(1, tk_id_count)
+            res["AssocPr"][a] = np.sum(matches_count * ass_pr) / np.maximum(
+                1, res["Loc_TP"][a]
+            )
+
+    @_timing.time
+    def eval_sequence(self, data, cls, cls_id_name_mapping, cls_fp):
+        """Evaluate a single sequence across all thresholds."""
+        res = {}
+        class_info_dict = {}
+
+        for thr in data:
+            res[thr], cls_fp[thr], cls_info = self.eval_sequence_single_thr(
+                data[thr], cls, cls_id_name_mapping, cls_fp[thr], thr
+            )
+            class_info_dict[thr] = cls_info
+
+        return res, cls_fp, class_info_dict
+
+    def combine_sequences(self, all_res):
+        """Combines metrics across all sequences."""
+        data = {}
+        res = {}
+
+        if all_res:
+            thresholds = list(list(all_res.values())[0].keys())
+        else:
+            thresholds = [50]
+        for thr in thresholds:
+            data[thr] = {}
+            for seq_key in all_res:
+                data[thr][seq_key] = all_res[seq_key][thr]
+        for thr in thresholds:
+            res[thr] = self._combine_sequences_thr(data[thr])
+
+        return res
+
+    def _combine_sequences_thr(self, all_res):
+        """Combines sequences over each threshold."""
+        res = {}
+        for field in self.integer_array_fields:
+            res[field] = self._combine_sum(all_res, field)
+        for field in ["AssocRe", "AssocPr", "AssocA"]:
+            res[field] = self._combine_weighted_av(
+                all_res, field, res, weight_field="Loc_TP"
+            )
+        res = self._compute_final_fields(res)
+        return res
+
+    def combine_classes_class_averaged(self, all_res, ignore_empty=False):
+        """Combines metrics across all classes by averaging over classes.
+
+        If 'ignore_empty' is True, then it only sums over classes
+        with at least one gt or predicted detection.
+        """
+        data = {}
+        res = {}
+        if all_res:
+            thresholds = list(list(all_res.values())[0].keys())
+        else:
+            thresholds = [50]
+        for thr in thresholds:
+            data[thr] = {}
+            for cls_key in all_res:
+                data[thr][cls_key] = all_res[cls_key][thr]
+        for thr in data:
+            res[thr] = self._combine_classes_class_averaged_thr(
+                data[thr], ignore_empty=ignore_empty
+            )
+        return res
+
+    def _combine_classes_class_averaged_thr(self, all_res, ignore_empty=False):
+        """Combines classes over each threshold."""
+        res = {}
+
+        def check_empty(val):
+            """Returns True if empty."""
+            return not (val["Loc_TP"] + val["Loc_FN"] + val["Loc_FP"] > 0 + EPS).any()
+
+        for field in self.integer_array_fields:
+            if ignore_empty:
+                res_field = {k: v for k, v in all_res.items() if not check_empty(v)}
+            else:
+                res_field = {k: v for k, v in all_res.items()}
+            res[field] = self._combine_sum(res_field, field)
+
+        for field in self.float_array_fields:
+            if ignore_empty:
+                res_field = [v[field] for v in all_res.values() if not check_empty(v)]
+            else:
+                res_field = [v[field] for v in all_res.values()]
+            res[field] = np.mean(res_field, axis=0)
+        return res
+
+    def combine_classes_det_averaged(self, all_res):
+        """Combines metrics across all classes by averaging over detections."""
+        data = {}
+        res = {}
+        if all_res:
+            thresholds = list(list(all_res.values())[0].keys())
+        else:
+            thresholds = [50]
+        for thr in thresholds:
+            data[thr] = {}
+            for cls_key in all_res:
+                data[thr][cls_key] = all_res[cls_key][thr]
+        for thr in data:
+            res[thr] = self._combine_classes_det_averaged_thr(data[thr])
+        return res
+
+    def _combine_classes_det_averaged_thr(self, all_res):
+        """Combines detections over each threshold."""
+        res = {}
+        for field in self.integer_array_fields:
+            res[field] = self._combine_sum(all_res, field)
+        for field in ["AssocRe", "AssocPr", "AssocA"]:
+            res[field] = self._combine_weighted_av(
+                all_res, field, res, weight_field="Loc_TP"
+            )
+        res = self._compute_final_fields(res)
+        return res
+
+    @staticmethod
+    def _compute_final_fields(res):
+        """Calculate final metric values.
+
+        This function is used both for both per-sequence calculation,
+        and in combining values across sequences.
+        """
+        # LocA
+        res["LocRe"] = res["Loc_TP"] / np.maximum(1, res["Loc_TP"] + res["Loc_FN"])
+        res["LocPr"] = res["Loc_TP"] / np.maximum(1, res["Loc_TP"] + res["Loc_FP"])
+        res["LocA"] = res["Loc_TP"] / np.maximum(
+            1, res["Loc_TP"] + res["Loc_FN"] + res["Loc_FP"]
+        )
+
+        # ClsA
+        res["ClsRe"] = res["Cls_TP"] / np.maximum(1, res["Cls_TP"] + res["Cls_FN"])
+        res["ClsPr"] = res["Cls_TP"] / np.maximum(1, res["Cls_TP"] + res["Cls_FP"])
+        res["ClsA"] = res["Cls_TP"] / np.maximum(
+            1, res["Cls_TP"] + res["Cls_FN"] + res["Cls_FP"]
+        )
+
+        res["ClsRe"] = np.mean(res["ClsRe"])
+        res["ClsPr"] = np.mean(res["ClsPr"])
+        res["ClsA"] = np.mean(res["ClsA"])
+
+        res["TETA"] = (res["LocA"] + res["AssocA"] + res["ClsA"]) / 3
+
+        return res
+
+    def print_summary_table(self, thr_res, thr, tracker, cls):
+        """Prints summary table of results."""
+        print("")
+        metric_name = self.get_name()
+        self._row_print(
+            [f"{metric_name}{str(thr)}: {tracker}-{cls}"] + self.summary_fields
+        )
+        self._row_print(["COMBINED"] + thr_res)

sam3/eval/teta_eval_toolkit/utils.py

@@ -0,0 +1,46 @@
+# fmt: off
+# flake8: noqa
+
+import csv
+import os
+from collections import OrderedDict
+
+
+def validate_metrics_list(metrics_list):
+    """Get names of metric class and ensures they are unique, further checks that the fields within each metric class
+    do not have overlapping names.
+    """
+    metric_names = [metric.get_name() for metric in metrics_list]
+    # check metric names are unique
+    if len(metric_names) != len(set(metric_names)):
+        raise TrackEvalException(
+            "Code being run with multiple metrics of the same name"
+        )
+    fields = []
+    for m in metrics_list:
+        fields += m.fields
+    # check metric fields are unique
+    if len(fields) != len(set(fields)):
+        raise TrackEvalException(
+            "Code being run with multiple metrics with fields of the same name"
+        )
+    return metric_names
+
+
+def get_track_id_str(ann):
+    """Get name of track ID in annotation."""
+    if "track_id" in ann:
+        tk_str = "track_id"
+    elif "instance_id" in ann:
+        tk_str = "instance_id"
+    elif "scalabel_id" in ann:
+        tk_str = "scalabel_id"
+    else:
+        assert False, "No track/instance ID."
+    return tk_str
+
+
+class TrackEvalException(Exception):
+    """Custom exception for catching expected errors."""
+
+    ...

sam3/eval/ytvis_coco_wrapper.py

@@ -0,0 +1,146 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+
+import copy
+import json
+import logging
+
+import numpy as np
+import pycocotools.mask as mask_util
+from pycocotools.coco import COCO
+from typing_extensions import override
+
+
+class YTVIS(COCO):
+    """
+    Helper class for reading YT-VIS annotations
+    """
+
+    @override
+    def __init__(self, annotation_file: str = None, ignore_gt_cats: bool = True):
+        """
+        Args:
+            annotation_file: Path to the annotation file
+            ignore_gt_cats: If True, we ignore the ground truth categories and replace them with a dummy "object" category. This is useful for Phrase AP evaluation.
+        """
+        self.ignore_gt_cats = ignore_gt_cats
+        super().__init__(annotation_file=annotation_file)
+
+    @override
+    def createIndex(self):
+        # We rename some keys to match the COCO format before creating the index.
+        if "annotations" in self.dataset:
+            for ann in self.dataset["annotations"]:
+                if "video_id" in ann:
+                    ann["image_id"] = int(ann.pop("video_id"))
+                if self.ignore_gt_cats:
+                    ann["category_id"] = -1
+                else:
+                    ann["category_id"] = int(ann["category_id"])
+                if "bboxes" in ann:
+                    # note that in some datasets we load under this YTVIS class,
+                    # some "bboxes" could be None for when the GT object is invisible,
+                    # so we replace them with [0, 0, 0, 0]
+                    ann["bboxes"] = [
+                        bbox if bbox is not None else [0, 0, 0, 0]
+                        for bbox in ann["bboxes"]
+                    ]
+                if "areas" in ann:
+                    # similar to "bboxes", some areas could be None for when the GT
+                    # object is invisible, so we replace them with 0
+                    areas = [a if a is not None else 0 for a in ann["areas"]]
+                    # Compute average area of tracklet
+                    ann["area"] = np.mean(areas)
+        if "videos" in self.dataset:
+            for vid in self.dataset["videos"]:
+                vid["id"] = int(vid["id"])
+            self.dataset["images"] = self.dataset.pop("videos")
+
+        if self.ignore_gt_cats:
+            self.dataset["categories"] = [
+                {"supercategory": "object", "id": -1, "name": "object"}
+            ]
+        else:
+            for cat in self.dataset["categories"]:
+                cat["id"] = int(cat["id"])
+        super().createIndex()
+
+    @override
+    def getAnnIds(self, imgIds=[], catIds=[], areaRng=[], iscrowd=None):
+        if len(areaRng) > 0:
+            logging.warning(
+                "Note that we filter out objects based on their *average* area across the video, not per frame area"
+            )
+
+        return super().getAnnIds(imgIds=imgIds, catIds=catIds, iscrowd=iscrowd)
+
+    @override
+    def showAnns(self, anns, draw_bbox=False):
+        raise NotImplementedError("Showing annotations is not supported")
+
+    @override
+    def loadRes(self, resFile):
+        # Adapted from COCO.loadRes to support tracklets/masklets
+        res = YTVIS(ignore_gt_cats=self.ignore_gt_cats)
+        res.dataset["images"] = [img for img in self.dataset["images"]]
+
+        if type(resFile) == str:
+            with open(resFile) as f:
+                anns = json.load(f)
+        elif type(resFile) == np.ndarray:
+            anns = self.loadNumpyAnnotations(resFile)
+        else:
+            anns = resFile
+        assert type(anns) == list, "results is not an array of objects"
+        annsImgIds = [ann["image_id"] for ann in anns]
+        assert set(annsImgIds) == (
+            set(annsImgIds) & set(self.getImgIds())
+        ), "Results do not correspond to current coco set"
+        if "bboxes" in anns[0] and not anns[0]["bboxes"] == []:
+            res.dataset["categories"] = copy.deepcopy(self.dataset["categories"])
+            for id, ann in enumerate(anns):
+                bbs = [(bb if bb is not None else [0, 0, 0, 0]) for bb in ann["bboxes"]]
+                xxyy = [[bb[0], bb[0] + bb[2], bb[1], bb[1] + bb[3]] for bb in bbs]
+                if not "segmentations" in ann:
+                    ann["segmentations"] = [
+                        [[x1, y1, x1, y2, x2, y2, x2, y1]] for (x1, x2, y1, y2) in xxyy
+                    ]
+                ann["areas"] = [bb[2] * bb[3] for bb in bbs]
+                # NOTE: We also compute average area of a tracklet across video, allowing us to compute area based mAP.
+                ann["area"] = np.mean(ann["areas"])
+                ann["id"] = id + 1
+                ann["iscrowd"] = 0
+        elif "segmentations" in anns[0]:
+            res.dataset["categories"] = copy.deepcopy(self.dataset["categories"])
+            for id, ann in enumerate(anns):
+                ann["bboxes"] = [
+                    mask_util.toBbox(segm) for segm in ann["segmentations"]
+                ]
+                if "areas" not in ann:
+                    ann["areas"] = [
+                        mask_util.area(segm) for segm in ann["segmentations"]
+                    ]
+                # NOTE: We also compute average area of a tracklet across video, allowing us to compute area based mAP.
+                ann["area"] = np.mean(ann["areas"])
+                ann["id"] = id + 1
+                ann["iscrowd"] = 0
+
+        res.dataset["annotations"] = anns
+        res.createIndex()
+        return res
+
+    @override
+    def download(self, tarDir=None, imgIds=[]):
+        raise NotImplementedError
+
+    @override
+    def loadNumpyAnnotations(self, data):
+        raise NotImplementedError("We don't support numpy annotations for now")
+
+    @override
+    def annToRLE(self, ann):
+        raise NotImplementedError("We expect masks to be already in RLE format")
+
+    @override
+    def annToMask(self, ann):
+        raise NotImplementedError("We expect masks to be already in RLE format")

sam3/eval/ytvis_eval.py

@@ -0,0 +1,411 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+import copy
+import gc
+import logging
+import os
+from collections import defaultdict
+from operator import xor
+from pathlib import Path
+from typing import List, Optional
+
+import numpy as np
+import pycocotools.mask as mask_util
+import torch
+from pycocotools.cocoeval import COCOeval
+from sam3.eval.cgf1_eval import CGF1Eval
+from sam3.eval.coco_eval_offline import convert_to_xywh
+from sam3.model.box_ops import box_xywh_inter_union
+from sam3.train.masks_ops import rle_encode
+from sam3.train.utils import distributed as dist
+from typing_extensions import override
+
+try:
+    import rapidjson as json
+except ModuleNotFoundError:
+    import json
+
+from iopath.common.file_io import g_pathmgr
+
+
+class YTVISevalMixin:
+    """
+    Identical to COCOeval but adapts computeIoU to compute IoU between tracklets/masklets.
+    """
+
+    @override
+    def _prepare(self):
+        """
+        Copied from cocoeval.py but doesn't convert masks to RLEs (we assume they already are RLEs)
+        """
+        p = self.params
+        if p.useCats:
+            gts = self.cocoGt.loadAnns(
+                self.cocoGt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds)
+            )
+            dts = self.cocoDt.loadAnns(
+                self.cocoDt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds)
+            )
+        else:
+            gts = self.cocoGt.loadAnns(self.cocoGt.getAnnIds(imgIds=p.imgIds))
+            dts = self.cocoDt.loadAnns(self.cocoDt.getAnnIds(imgIds=p.imgIds))
+
+        # set ignore flag
+        for gt in gts:
+            gt["ignore"] = gt["ignore"] if "ignore" in gt else 0
+            gt["ignore"] = "iscrowd" in gt and gt["iscrowd"]
+            if p.iouType == "keypoints":
+                gt["ignore"] = (gt["num_keypoints"] == 0) or gt["ignore"]
+        self._gts = defaultdict(list)  # gt for evaluation
+        self._dts = defaultdict(list)  # dt for evaluation
+        for gt in gts:
+            self._gts[gt["image_id"], gt["category_id"]].append(gt)
+        for dt in dts:
+            self._dts[dt["image_id"], dt["category_id"]].append(dt)
+        self.evalImgs = defaultdict(list)  # per-image per-category evaluation results
+        self.eval = {}  # accumulated evaluation results
+
+    def computeIoU(self, imgId, catId):
+        """
+        Compute IoU between tracklets. Copied from cocoeval.py but adapted for videos (in YT-VIS format)
+        """
+        p = self.params
+        if p.useCats:
+            gt = self._gts[imgId, catId]
+            dt = self._dts[imgId, catId]
+        else:
+            gt = [_ for cId in p.catIds for _ in self._gts[imgId, cId]]
+            dt = [_ for cId in p.catIds for _ in self._dts[imgId, cId]]
+        if len(gt) == 0 or len(dt) == 0:
+            return []
+
+        # For class mAP and phrase AP evaluation, we sort the detections in descending order of scores (as in COCOeval).
+        # For demo F1 evaluation, we DO NOT sort the detections (but match them with GTs via Hungarian matching).
+        assert hasattr(self, "sort_inds_by_scores_in_iou"), (
+            "subclasses that inherits YTVISevalMixin should set `self.sort_inds_by_scores_in_iou` "
+            "(True for class mAP and phrase AP, False for demo F1)"
+        )
+        if self.sort_inds_by_scores_in_iou:
+            inds = np.argsort([-d["score"] for d in dt], kind="mergesort")
+            dt = [dt[i] for i in inds]
+            if len(dt) > p.maxDets[-1]:
+                dt = dt[0 : p.maxDets[-1]]
+
+        if p.iouType == "segm":
+            g = [g["segmentations"] for g in gt]
+            d = [d["segmentations"] for d in dt]
+        elif p.iouType == "bbox":
+            g = [g["bboxes"] for g in gt]
+            d = [d["bboxes"] for d in dt]
+        else:
+            raise Exception("unknown iouType for iou computation")
+
+        def iou_tracklets(preds, gts):
+            preds = torch.tensor(preds)
+            gts = torch.tensor(gts)
+            inter, union = box_xywh_inter_union(
+                preds.unsqueeze(1), gts.unsqueeze(0)
+            )  # Num preds x Num GTS x Num frames
+            inter = inter.sum(-1)
+            union = union.sum(-1)
+            assert (
+                union > 0
+            ).all(), (
+                "There exists a tracklet with zero GTs across time. This is suspicious"
+            )
+            return inter / union
+
+        def iou_masklets(preds, gts):
+            inter = 0
+            union = 0
+            for p_i, gt_i in zip(preds, gts):
+                if p_i and gt_i:
+                    # Compute areas of intersection and union
+                    inter += mask_util.area(
+                        mask_util.merge([p_i, gt_i], intersect=True)
+                    )
+                    union += mask_util.area(
+                        mask_util.merge([p_i, gt_i], intersect=False)
+                    )
+                elif gt_i:
+                    union += mask_util.area(gt_i)
+                elif p_i:
+                    union += mask_util.area(p_i)
+            if union > 0:
+                iou = inter / union
+                assert iou >= 0 and iou <= 1, "Encountered an error in IoU computation"
+            else:
+                assert np.isclose(inter, 0) and np.isclose(
+                    union, 0
+                ), "Encountered an error in IoU computation"
+                iou = 1
+            return iou
+
+        if p.iouType == "segm":
+            ious = [[iou_masklets(d_i, g_i) for g_i in g] for d_i in d]
+        else:
+            ious = iou_tracklets(d, g)
+        return np.array(ious)
+
+
+class YTVISeval(YTVISevalMixin, COCOeval):
+    # For class mAP and phrase AP evaluation, we sort the detections in descending order of scores (as in COCOeval).
+    sort_inds_by_scores_in_iou = True
+
+
+class VideoDemoF1Eval(YTVISevalMixin, CGF1Eval):
+    # For demo F1 evaluation, we DO NOT sort the detections (but match them with GTs via Hungarian matching).
+    sort_inds_by_scores_in_iou = False
+
+
+class YTVISResultsWriter:
+    """
+    Gather and dumps predictions in YT-VIS format.
+    Expected flow of API calls: reset() -> N * update() -> compute_synced()
+    """
+
+    def __init__(
+        self,
+        dump_file: str,
+        postprocessor,
+        gather_pred_via_filesys=False,
+        pred_file_evaluators: Optional[List] = None,
+        save_per_frame_scores: bool = False,
+        write_eval_metrics_file: bool = True,
+        eval_metrics_file_suffix: str = ".sam3_eval_metrics",
+    ):
+        self.dump_file = dump_file
+        self.dump = []
+        self.postprocessor = postprocessor
+        self.gather_pred_via_filesys = gather_pred_via_filesys
+        if dist.is_main_process():
+            dirname = os.path.dirname(self.dump_file)
+            if not os.path.exists(dirname):
+                os.makedirs(dirname, exist_ok=True)
+                logging.info(f"Creating folder: {dirname}")
+
+        # the evaluation hooks to be applied to the prediction files
+        self.pred_file_evaluators = pred_file_evaluators or []
+        self.save_per_frame_scores = save_per_frame_scores
+        # in addition to the prediction file, we also write the evaluation metrics
+        # for easier debugging and analysis (stored in another eval_metrics_file
+        # so that we can keep the dumped prediction file under YT-VIS format)
+        self.write_eval_metrics_file = write_eval_metrics_file
+        if self.write_eval_metrics_file:
+            self.eval_metrics_file = self.dump_file + eval_metrics_file_suffix
+            os.makedirs(os.path.dirname(self.eval_metrics_file), exist_ok=True)
+
+    def _dump_vid_preds(self, results):
+        dumped_results = copy.deepcopy(results)
+        self.dump.extend(dumped_results)
+
+    def prepare(self, predictions):
+        ytvis_results = []
+        for video_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+            for k in ["boxes", "scores", "labels"]:
+                assert (
+                    k in prediction
+                ), f"Expected predictions to have `{k}` key, available keys are {prediction.keys()}"
+            if self.save_per_frame_scores:
+                assert (
+                    "per_frame_scores" in prediction
+                ), f"Expected predictions to have `per_frame_scores` key, available keys are {prediction.keys()}"
+            assert xor(
+                "masks" in prediction, "masks_rle" in prediction
+            ), f"Expected predictions to have either `masks` key or `masks_rle` key, available keys are {prediction.keys()}"
+
+            boxes = prediction["boxes"]
+            boxes = convert_to_xywh(boxes).tolist()
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+            if "masks" in prediction:
+                masks = prediction["masks"].squeeze(2)
+                assert (
+                    masks.ndim == 4
+                ), "Expected masks to be of shape(N_preds,T_frames,H,W)"
+
+                areas = [mask.flatten(1).sum(1).tolist() for mask in masks]
+                rles = [rle_encode(masklet) for masklet in masks]
+
+                # memory clean
+                del masks
+                del prediction["masks"]
+            elif "masks_rle" in prediction:
+                rles = prediction.pop("masks_rle")
+                areas = [
+                    [0 if rle is None else rle.pop("area") for rle in rles_per_obj]
+                    for rles_per_obj in rles
+                ]
+            else:
+                raise ValueError(
+                    "Expected either `masks` or `masks_rle` key in the predictions."
+                )
+
+            new_results = [
+                {
+                    "video_id": video_id,
+                    "category_id": track_label,
+                    "bboxes": track_boxes,
+                    "score": track_score,
+                    "segmentations": track_masks,
+                    "areas": track_areas,
+                }
+                for (
+                    track_boxes,
+                    track_masks,
+                    track_areas,
+                    track_score,
+                    track_label,
+                ) in zip(boxes, rles, areas, scores, labels)
+            ]
+            # Optionally, save per-frame scores
+            if self.save_per_frame_scores:
+                per_frame_scores = prediction["per_frame_scores"].tolist()
+                for res, track_per_frame_scores in zip(new_results, per_frame_scores):
+                    res["per_frame_scores"] = track_per_frame_scores
+
+            ytvis_results.extend(new_results)
+
+        return ytvis_results
+
+    def set_sync_device(self, device: torch.device):
+        self._sync_device = device
+
+    def update(self, *args, **kwargs):
+        predictions = self.postprocessor.process_results(*args, **kwargs)
+        results = self.prepare(predictions)
+        self._dump_vid_preds(results)
+
+    def _dump_preds(self):
+        if not dist.is_main_process():
+            self.dump = []
+            gc.collect()
+            return
+        dumped_file = Path(self.dump_file)
+        logging.info(f"YTVIS evaluator: Dumping predictions to {dumped_file}")
+        with g_pathmgr.open(str(dumped_file), "w") as f:
+            json.dump(self.dump, f)
+        self.dump = []
+        gc.collect()
+        return str(dumped_file)
+
+    def synchronize_between_processes(self):
+        logging.info("YT-VIS evaluator: Synchronizing between processes")
+        dump_dict = self._dedup_pre_gather(self.dump)
+        if self.gather_pred_via_filesys:
+            dump_dict_all_gpus = dist.gather_to_rank_0_via_filesys(dump_dict)
+        else:
+            dump_dict_all_gpus = dist.all_gather(dump_dict, force_cpu=True)
+        self.dump = self._dedup_post_gather(dump_dict_all_gpus)
+        logging.info(f"Gathered all {len(self.dump)} predictions")
+
+    def _dedup_pre_gather(self, predictions):
+        """
+        Organize the predictions as a dict-of-list using (video_id, category_id) as keys
+        for deduplication after gathering them across GPUs.
+
+        During evaluation, PyTorch data loader under `drop_last: False` would wrap
+        around the dataset length to be a multiple of world size (GPU num) and duplicate
+        the remaining batches. This causes the same test sample to appear simultaneously
+        in multiple GPUs, resulting in duplicated predictions being saved into prediction
+        files. These duplicates are then counted as false positives under detection mAP
+        metrics (since a ground truth can be matched with only one prediction).
+
+        For example, if there are 4 GPUs and 6 samples [A1, A2, B1, B2, C1, C2], the data
+        loader (under `drop_last: False`) would load it by wrapping it around like
+        `[A1, A2, B1, B2, C1, C2, *A1*, *A2*]` to make a multiple of 4 and then split it as
+
+        - GPU 0: A1, C1
+        - GPU 1: A2, C2
+        - GPU 3: B1, **A1**
+        - GPU 4: B2, **A2**
+        (as in DistributedSampler in https://github.com/pytorch/pytorch/blob/521588519da9f4876d90ddd7a17c10d0eca89dc6/torch/utils/data/distributed.py#L116-L124)
+
+        so the predictions on A1 and A2 will occur twice in the final gathered outputs
+        in the prediction file (and counted as false positives). This also affects our
+        YT-VIS official val evaluation, but to a lesser extent than YT-VIS dev since
+        the latter is much smaller and more susceptible to false positives.
+
+        So we to deduplicate this. The tricky part is that we cannot deduplicate them
+        simply using video id, given that we are sharding the classes in each video
+        across multiple batches (with 20 prompts per batch) in our "orig_cats" eval dbs.
+
+        The solution is to deduplicate based on (video_id, category_id) tuple as keys.
+        We organize the predictions as a dict-of-list using (video_id, category_id) as
+        keys on each GPU, with the list of masklets under this (video_id, category_id)
+        on this GPU as values. Then, we all-gather this dict-of-list across GPUs and
+        if a key (video_id, category_id) appears in multiple GPUs, we only take the
+        prediction masklet list from one GPU.
+        """
+        prediction_dict = defaultdict(list)
+        for p in predictions:
+            prediction_dict[(p["video_id"], p["category_id"])].append(p)
+        return prediction_dict
+
+    def _dedup_post_gather(self, list_of_prediction_dict):
+        """
+        Deduplicate the predictions from all GPUs. See `_dedup_pre_gather` for details.
+        """
+        dedup_prediction_dict = {}
+        duplication_keys = []
+        for prediction_dict in list_of_prediction_dict:
+            for k, v in prediction_dict.items():
+                if k not in dedup_prediction_dict:
+                    dedup_prediction_dict[k] = v
+                else:
+                    duplication_keys.append(k)
+
+        logging.info(
+            f"skipped {len(duplication_keys)} duplicated predictions in YTVISResultsWriter "
+            f"with the following (video_id, category_id) tuples: {duplication_keys}"
+        )
+        dedup_predictions = sum(dedup_prediction_dict.values(), [])
+        return dedup_predictions
+
+    def compute_synced(
+        self,
+    ):
+        self.synchronize_between_processes()
+        dumped_file = self._dump_preds()
+        if not dist.is_main_process():
+            return {"": 0.0}
+
+        # run evaluation hooks on the prediction file
+        meters = {}
+        all_video_np_level_results = defaultdict(dict)
+        for evaluator in self.pred_file_evaluators:
+            gc.collect()
+            results, video_np_level_results = evaluator.evaluate(dumped_file)
+            meters.update(results)
+            for (video_id, category_id), res in video_np_level_results.items():
+                all_video_np_level_results[(video_id, category_id)].update(res)
+
+        gc.collect()
+        if self.write_eval_metrics_file:
+            # convert the nested dict of {(video_id, category_id): per_sample_metric_dict}
+            # to a list of per-sample metric dicts (with video_id and category_id) for JSON,
+            # as JSON doesn't allow using tuples like (video_id, category_id) as dict keys
+            video_np_level_metrics = [
+                {"video_id": video_id, "category_id": category_id, **res}
+                for (video_id, category_id), res in all_video_np_level_results.items()
+            ]
+            eval_metrics = {
+                "dataset_level_metrics": meters,
+                "video_np_level_metrics": video_np_level_metrics,
+            }
+            with g_pathmgr.open(self.eval_metrics_file, "w") as f:
+                json.dump(eval_metrics, f)
+            logging.info(
+                f"YTVIS evaluator: Dumped evaluation metrics to {self.eval_metrics_file}"
+            )
+
+        if len(meters) == 0:
+            meters = {"": 0.0}
+        return meters
+
+    def compute(self):
+        return {"": 0.0}
+
+    def reset(self, *args, **kwargs):
+        self.dump = []

sam3/logger.py

@@ -0,0 +1,54 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+import logging
+import os
+
+LOG_LEVELS = {
+    "DEBUG": logging.DEBUG,
+    "INFO": logging.INFO,
+    "WARNING": logging.WARNING,
+    "ERROR": logging.ERROR,
+    "CRITICAL": logging.CRITICAL,
+}
+
+
+class ColoredFormatter(logging.Formatter):
+    """A command line formatter with different colors for each level."""
+
+    def __init__(self):
+        super().__init__()
+        reset = "\033[0m"
+        colors = {
+            logging.DEBUG: f"{reset}\033[36m",  # cyan,
+            logging.INFO: f"{reset}\033[32m",  # green
+            logging.WARNING: f"{reset}\033[33m",  # yellow
+            logging.ERROR: f"{reset}\033[31m",  # red
+            logging.CRITICAL: f"{reset}\033[35m",  # magenta
+        }
+        fmt_str = "{color}%(levelname)s %(asctime)s %(process)d %(filename)s:%(lineno)4d:{reset} %(message)s"
+        self.formatters = {
+            level: logging.Formatter(fmt_str.format(color=color, reset=reset))
+            for level, color in colors.items()
+        }
+        self.default_formatter = self.formatters[logging.INFO]
+
+    def format(self, record):
+        formatter = self.formatters.get(record.levelno, self.default_formatter)
+        return formatter.format(record)
+
+
+def get_logger(name, level=logging.INFO):
+    """A command line logger."""
+    if "LOG_LEVEL" in os.environ:
+        level = os.environ["LOG_LEVEL"].upper()
+        assert (
+            level in LOG_LEVELS
+        ), f"Invalid LOG_LEVEL: {level}, must be one of {list(LOG_LEVELS.keys())}"
+        level = LOG_LEVELS[level]
+    logger = logging.getLogger(name)
+    logger.setLevel(level)
+    logger.propagate = False
+    ch = logging.StreamHandler()
+    ch.setLevel(level)
+    ch.setFormatter(ColoredFormatter())
+    logger.addHandler(ch)
+    return logger

sam3/model/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved

sam3/model/act_ckpt_utils.py

@@ -0,0 +1,114 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import inspect
+from functools import wraps
+from typing import Callable, TypeVar, Union
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as checkpoint
+from torch.utils._pytree import tree_map_only
+
+# Type variables for better type hinting
+T = TypeVar("T")
+Module = TypeVar("Module", bound=nn.Module)
+
+
+def activation_ckpt_wrapper(module: Union[nn.Module, Callable]) -> Callable:
+    """
+    Wraps a given module to enable or disable activation checkpointing.
+
+    Activation checkpointing (gradient checkpointing) trades compute for memory by
+    recomputing intermediate activations during the backward pass instead of storing
+    them in memory during the forward pass.
+
+    When activation checkpointing is enabled, the wrapper expects only keyword arguments,
+    and it maps these to positional arguments based on the module's signature.
+
+    Args:
+        module: The module or function to wrap with activation checkpointing
+
+    Returns:
+        A wrapped callable that supports activation checkpointing
+
+    Usage:
+        The returned wrapper function can be called with the same arguments as the
+        original module, with an additional `act_ckpt_enable` keyword argument to control
+        activation checkpointing and optional `use_reentrant` parameter.
+
+    Example:
+        ```python
+        wrapped_module = activation_ckpt_wrapper(my_module)
+        output = wrapped_module(x=input_tensor, y=another_tensor, act_ckpt_enable=True)
+        ```
+    """
+
+    @wraps(module)
+    def act_ckpt_wrapper(
+        *args, act_ckpt_enable: bool = True, use_reentrant: bool = False, **kwargs
+    ):
+        if act_ckpt_enable:
+            if len(args) > 0:
+                raise ValueError(
+                    "This wrapper expects keyword arguments only when `act_ckpt_enable=True`"
+                )
+            # Get the signature of the target function/module
+            callable_fn = module.forward if isinstance(module, nn.Module) else module
+            sig = inspect.signature(callable_fn)
+            # Create a mapping of parameter names to their default values
+            param_defaults = {
+                name: param.default for name, param in sig.parameters.items()
+            }
+            args = []
+            for p_name in param_defaults.keys():
+                if p_name in kwargs:
+                    args.append(kwargs.pop(p_name))
+                elif param_defaults[p_name] is not inspect.Parameter.empty:
+                    # Set arg to default value if it's not in kwargs. Useful for primitive types or args that default to None
+                    args.append(param_defaults[p_name])
+                elif (
+                    sig.parameters[p_name].kind is not inspect.Parameter.VAR_KEYWORD
+                ):  # Skip **kwargs parameter
+                    raise ValueError(f"Missing positional argument: {p_name}")
+
+            # Scan remaining kwargs for torch.Tensor
+            remaining_keys = list(kwargs.keys())
+            for key in remaining_keys:
+                if isinstance(kwargs[key], torch.Tensor):
+                    # Remove the tensor from kwargs, assuming it's not required by the module.
+                    # If it is required, the module's signature should be modified to accept it as a positional or keyword argument.
+                    kwargs[key] = "_REMOVED_BY_ACT_CKPT_WRAPPER_"
+
+            ret = checkpoint.checkpoint(
+                module, *args, use_reentrant=use_reentrant, **kwargs
+            )
+        else:
+            ret = module(*args, **kwargs)
+
+        return ret
+
+    return act_ckpt_wrapper
+
+
+def clone_output_wrapper(f: Callable[..., T]) -> Callable[..., T]:
+    """
+    Clone the CUDA output tensors of a function to avoid in-place operations.
+
+    This wrapper is useful when working with torch.compile to prevent errors
+    related to in-place operations on tensors.
+
+    Args:
+        f: The function whose CUDA tensor outputs should be cloned
+
+    Returns:
+        A wrapped function that clones any CUDA tensor outputs
+    """
+
+    @wraps(f)
+    def wrapped(*args, **kwargs):
+        outputs = f(*args, **kwargs)
+        return tree_map_only(
+            torch.Tensor, lambda t: t.clone() if t.is_cuda else t, outputs
+        )
+
+    return wrapped

sam3/model/box_ops.py

@@ -0,0 +1,217 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+"""
+Utilities for bounding box manipulation and GIoU.
+"""
+
+from typing import Tuple
+
+import torch
+
+
+def box_cxcywh_to_xyxy(x):
+    x_c, y_c, w, h = x.unbind(-1)
+    b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)]
+    return torch.stack(b, dim=-1)
+
+
+def box_cxcywh_to_xywh(x):
+    x_c, y_c, w, h = x.unbind(-1)
+    b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (w), (h)]
+    return torch.stack(b, dim=-1)
+
+
+def box_xywh_to_xyxy(x):
+    x, y, w, h = x.unbind(-1)
+    b = [(x), (y), (x + w), (y + h)]
+    return torch.stack(b, dim=-1)
+
+
+def box_xywh_to_cxcywh(x):
+    x, y, w, h = x.unbind(-1)
+    b = [(x + 0.5 * w), (y + 0.5 * h), (w), (h)]
+    return torch.stack(b, dim=-1)
+
+
+def box_xyxy_to_xywh(x):
+    x, y, X, Y = x.unbind(-1)
+    b = [(x), (y), (X - x), (Y - y)]
+    return torch.stack(b, dim=-1)
+
+
+def box_xyxy_to_cxcywh(x):
+    x0, y0, x1, y1 = x.unbind(-1)
+    b = [(x0 + x1) / 2, (y0 + y1) / 2, (x1 - x0), (y1 - y0)]
+    return torch.stack(b, dim=-1)
+
+
+def box_area(boxes):
+    """
+    Batched version of box area. Boxes should be in [x0, y0, x1, y1] format.
+
+    Inputs:
+    - boxes: Tensor of shape (..., 4)
+
+    Returns:
+    - areas: Tensor of shape (...,)
+    """
+    x0, y0, x1, y1 = boxes.unbind(-1)
+    return (x1 - x0) * (y1 - y0)
+
+
+def masks_to_boxes(masks):
+    """Compute the bounding boxes around the provided masks
+
+    The masks should be in format [N, H, W] where N is the number of masks, (H, W) are the spatial dimensions.
+
+    Returns a [N, 4] tensors, with the boxes in xyxy format
+    """
+    if masks.numel() == 0:
+        return torch.zeros((0, 4), device=masks.device)
+
+    h, w = masks.shape[-2:]
+
+    y = torch.arange(0, h, dtype=torch.float, device=masks.device)
+    x = torch.arange(0, w, dtype=torch.float, device=masks.device)
+    y, x = torch.meshgrid(y, x)
+
+    x_mask = masks * x.unsqueeze(0)
+    x_max = x_mask.flatten(1).max(-1)[0] + 1
+    x_min = x_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]
+
+    y_mask = masks * y.unsqueeze(0)
+    y_max = y_mask.flatten(1).max(-1)[0] + 1
+    y_min = y_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]
+
+    boxes = torch.stack([x_min, y_min, x_max, y_max], 1)
+    # Invalidate boxes corresponding to empty masks.
+    boxes = boxes * masks.flatten(-2).any(-1)
+    return boxes
+
+
+def box_iou(boxes1, boxes2):
+    """
+    Batched version of box_iou. Boxes should be in [x0, y0, x1, y1] format.
+
+    Inputs:
+    - boxes1: Tensor of shape (..., N, 4)
+    - boxes2: Tensor of shape (..., M, 4)
+
+    Returns:
+    - iou, union: Tensors of shape (..., N, M)
+    """
+    area1 = box_area(boxes1)
+    area2 = box_area(boxes2)
+
+    # boxes1: (..., N, 4) -> (..., N, 1, 2)
+    # boxes2: (..., M, 4) -> (..., 1, M, 2)
+    lt = torch.max(boxes1[..., :, None, :2], boxes2[..., None, :, :2])
+    rb = torch.min(boxes1[..., :, None, 2:], boxes2[..., None, :, 2:])
+
+    wh = (rb - lt).clamp(min=0)  # (..., N, M, 2)
+    inter = wh[..., 0] * wh[..., 1]  # (..., N, M)
+
+    union = area1[..., None] + area2[..., None, :] - inter
+
+    iou = inter / union
+    return iou, union
+
+
+def generalized_box_iou(boxes1, boxes2):
+    """
+    Batched version of Generalized IoU from https://giou.stanford.edu/
+
+    Boxes should be in [x0, y0, x1, y1] format
+
+    Inputs:
+    - boxes1: Tensor of shape (..., N, 4)
+    - boxes2: Tensor of shape (..., M, 4)
+
+    Returns:
+    - giou: Tensor of shape (..., N, M)
+    """
+    iou, union = box_iou(boxes1, boxes2)
+
+    # boxes1: (..., N, 4) -> (..., N, 1, 2)
+    # boxes2: (..., M, 4) -> (..., 1, M, 2)
+    lt = torch.min(boxes1[..., :, None, :2], boxes2[..., None, :, :2])
+    rb = torch.max(boxes1[..., :, None, 2:], boxes2[..., None, :, 2:])
+
+    wh = (rb - lt).clamp(min=0)  # (..., N, M, 2)
+    area = wh[..., 0] * wh[..., 1]  # (..., N, M)
+
+    return iou - (area - union) / area
+
+
+@torch.jit.script
+def fast_diag_generalized_box_iou(boxes1, boxes2):
+    assert len(boxes1) == len(boxes2)
+    box1_xy = boxes1[:, 2:]
+    box1_XY = boxes1[:, :2]
+    box2_xy = boxes2[:, 2:]
+    box2_XY = boxes2[:, :2]
+    # assert (box1_xy >= box1_XY).all()
+    # assert (box2_xy >= box2_XY).all()
+    area1 = (box1_xy - box1_XY).prod(-1)
+    area2 = (box2_xy - box2_XY).prod(-1)
+
+    lt = torch.max(box1_XY, box2_XY)  # [N,2]
+    lt2 = torch.min(box1_XY, box2_XY)
+    rb = torch.min(box1_xy, box2_xy)  # [N,2]
+    rb2 = torch.max(box1_xy, box2_xy)
+
+    inter = (rb - lt).clamp(min=0).prod(-1)
+    tot_area = (rb2 - lt2).clamp(min=0).prod(-1)
+
+    union = area1 + area2 - inter
+
+    iou = inter / union
+
+    return iou - (tot_area - union) / tot_area
+
+
+@torch.jit.script
+def fast_diag_box_iou(boxes1, boxes2):
+    assert len(boxes1) == len(boxes2)
+    box1_xy = boxes1[:, 2:]
+    box1_XY = boxes1[:, :2]
+    box2_xy = boxes2[:, 2:]
+    box2_XY = boxes2[:, :2]
+    # assert (box1_xy >= box1_XY).all()
+    # assert (box2_xy >= box2_XY).all()
+    area1 = (box1_xy - box1_XY).prod(-1)
+    area2 = (box2_xy - box2_XY).prod(-1)
+
+    lt = torch.max(box1_XY, box2_XY)  # [N,2]
+    rb = torch.min(box1_xy, box2_xy)  # [N,2]
+
+    inter = (rb - lt).clamp(min=0).prod(-1)
+
+    union = area1 + area2 - inter
+
+    iou = inter / union
+
+    return iou
+
+
+def box_xywh_inter_union(
+    boxes1: torch.Tensor, boxes2: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    # Asuumes boxes in xywh format
+    assert boxes1.size(-1) == 4 and boxes2.size(-1) == 4
+    boxes1 = box_xywh_to_xyxy(boxes1)
+    boxes2 = box_xywh_to_xyxy(boxes2)
+    box1_tl_xy = boxes1[..., :2]
+    box1_br_xy = boxes1[..., 2:]
+    box2_tl_xy = boxes2[..., :2]
+    box2_br_xy = boxes2[..., 2:]
+    area1 = (box1_br_xy - box1_tl_xy).prod(-1)
+    area2 = (box2_br_xy - box2_tl_xy).prod(-1)
+
+    assert (area1 >= 0).all() and (area2 >= 0).all()
+    tl = torch.max(box1_tl_xy, box2_tl_xy)
+    br = torch.min(box1_br_xy, box2_br_xy)
+
+    inter = (br - tl).clamp(min=0).prod(-1)
+    union = area1 + area2 - inter
+
+    return inter, union

sam3/model/data_misc.py

@@ -0,0 +1,209 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+"""
+Misc functions, including distributed helpers.
+"""
+
+import collections
+import re
+
+from dataclasses import dataclass, field as field_ptr_behaviour, fields, is_dataclass
+from typing import Any, get_args, get_origin, List, Mapping, Optional, Sequence, Union
+
+import torch
+
+
+MyTensor = Union[torch.Tensor, List[Any]]
+
+
+def interpolate(
+    input, size=None, scale_factor=None, mode="nearest", align_corners=None
+):
+    # type: (Tensor, Optional[List[int]], Optional[float], str, Optional[bool]) -> Tensor
+    """
+    Equivalent to nn.functional.interpolate, but with support for empty channel sizes.
+    """
+    if input.numel() > 0:
+        return torch.nn.functional.interpolate(
+            input, size, scale_factor, mode, align_corners
+        )
+
+    assert (
+        input.shape[0] != 0 or input.shape[1] != 0
+    ), "At least one of the two first dimensions must be non zero"
+
+    if input.shape[1] == 0:
+        # Pytorch doesn't support null dimension on the channel dimension, so we transpose to fake a null batch dim
+        return torch.nn.functional.interpolate(
+            input.transpose(0, 1), size, scale_factor, mode, align_corners
+        ).transpose(0, 1)
+
+    # empty batch dimension is now supported in pytorch
+    return torch.nn.functional.interpolate(
+        input, size, scale_factor, mode, align_corners
+    )
+
+
+@dataclass
+class BatchedPointer:
+    stage_ids: MyTensor
+    stage_ids__type = torch.long
+    query_ids: MyTensor
+    query_ids__type = torch.long
+    object_ids: MyTensor
+    object_ids__type = torch.long
+    ptr_mask: MyTensor
+    ptr_mask__type = torch.bool
+    ptr_types: MyTensor
+    ptr_types__type = torch.long
+
+
+@dataclass
+class FindStage:
+    img_ids: MyTensor
+    img_ids__type = torch.long
+    text_ids: MyTensor
+    text_ids__type = torch.long
+
+    input_boxes: MyTensor
+    input_boxes__type = torch.float
+    input_boxes_mask: MyTensor
+    input_boxes_mask__type = torch.bool
+    input_boxes_label: MyTensor
+    input_boxes_label__type = torch.long
+
+    input_points: MyTensor
+    input_points__type = torch.float
+    input_points_mask: MyTensor
+    input_points_mask__type = torch.bool
+
+    # We track the object ids referred to by this query.
+    # This is beneficial for tracking in videos without the need for pointers.
+    object_ids: Optional[List[List]] = None  # List of objects per query
+
+
+@dataclass
+class BatchedFindTarget:
+    # The number of boxes in each find query
+    num_boxes: MyTensor
+    num_boxes__type = torch.long
+
+    # Target boxes in normalized CxCywh format
+    boxes: MyTensor
+    boxes__type = torch.float
+    # Target boxes in normalized CxCywh format but in padded representation
+    # as used in BinaryHungarianMatcherV2 (unlike the packed ones in `boxes`)
+    boxes_padded: MyTensor
+    boxes_padded__type = torch.float
+
+    # For hybrid matching, we repeat the boxes
+    repeated_boxes: MyTensor
+    repeated_boxes__type = torch.float
+
+    # Target Segmentation masks
+    segments: Optional[MyTensor]
+    segments__type = torch.bool
+
+    # Target Semantic Segmentation masks
+    semantic_segments: Optional[MyTensor]
+    semantic_segments__type = torch.bool
+
+    is_valid_segment: Optional[MyTensor]
+    is_valid_segment__type = torch.bool
+
+    # Whether annotations are exhaustive for each query
+    is_exhaustive: MyTensor
+    is_exhaustive__type = torch.bool
+
+    # The object id for each ground-truth box, in both packed and padded representations
+    object_ids: MyTensor
+    object_ids__type = torch.long
+    object_ids_padded: MyTensor
+    object_ids_padded__type = torch.long
+
+
+@dataclass
+class BatchedInferenceMetadata:
+    """All metadata required to post-process a find stage"""
+
+    # Coco id that corresponds to the "image" for evaluation by the coco evaluator
+    coco_image_id: MyTensor
+    coco_image_id__type = torch.long
+
+    # id in the original dataset, such that we can use the original evaluator
+    original_image_id: MyTensor
+    original_image_id__type = torch.long
+
+    # Original category id (if we want to use the original evaluator)
+    original_category_id: MyTensor
+    original_category_id__type = torch.int
+
+    # Size of the raw image (height, width)
+    original_size: MyTensor
+    original_size__type = torch.long
+
+    # id of the object in the media (track_id for a video)
+    object_id: MyTensor
+    object_id__type = torch.long
+
+    # index of the frame in the media (0 in the case of a single-frame media)
+    frame_index: MyTensor
+    frame_index__type = torch.long
+
+    # Adding for relations inference
+    # get_text_input: List[Optional[str]]
+
+    # Adding for TA conditional inference
+    is_conditioning_only: List[Optional[bool]]
+
+
+@dataclass
+class BatchedDatapoint:
+    img_batch: torch.Tensor
+    find_text_batch: List[str]
+    find_inputs: List[FindStage]
+    find_targets: List[BatchedFindTarget]
+    find_metadatas: List[BatchedInferenceMetadata]
+    raw_images: Optional[List[Any]] = None
+
+
+def convert_my_tensors(obj):
+    def is_optional_field(field) -> bool:
+        return get_origin(field) is Union and type(None) in get_args(field)
+
+    for field in fields(obj):
+        if is_dataclass(getattr(obj, field.name)):
+            convert_my_tensors(getattr(obj, field.name))
+            continue
+
+        field_type = field.type
+        if is_optional_field(field.type):
+            field_type = Union[get_args(field.type)[:-1]]  # Get the Optional field type
+
+        if field_type != MyTensor or getattr(obj, field.name) is None:
+            continue
+
+        elif len(getattr(obj, field.name)) and isinstance(
+            getattr(obj, field.name)[0], torch.Tensor
+        ):
+            stack_dim = 0
+            if field.name in [
+                "input_boxes",
+                "input_boxes_label",
+            ]:
+                stack_dim = 1
+            setattr(
+                obj,
+                field.name,
+                torch.stack(getattr(obj, field.name), dim=stack_dim).to(
+                    getattr(obj, field.name + "__type")
+                ),
+            )
+        else:
+            setattr(
+                obj,
+                field.name,
+                torch.as_tensor(
+                    getattr(obj, field.name), dtype=getattr(obj, field.name + "__type")
+                ),
+            )
+    return obj

sam3/model/decoder.py

@@ -0,0 +1,956 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+"""
+Transformer decoder.
+Inspired from Pytorch's version, adds the pre-norm variant
+"""
+
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+
+import torch
+
+from sam3.sam.transformer import RoPEAttention
+
+from torch import nn, Tensor
+from torchvision.ops.roi_align import RoIAlign
+
+from .act_ckpt_utils import activation_ckpt_wrapper
+
+from .box_ops import box_cxcywh_to_xyxy
+
+from .model_misc import (
+    gen_sineembed_for_position,
+    get_activation_fn,
+    get_clones,
+    inverse_sigmoid,
+    MLP,
+)
+
+
+class TransformerDecoderLayer(nn.Module):
+    def __init__(
+        self,
+        activation: str,
+        d_model: int,
+        dim_feedforward: int,
+        dropout: float,
+        cross_attention: nn.Module,
+        n_heads: int,
+        use_text_cross_attention: bool = False,
+    ):
+        super().__init__()
+
+        # cross attention
+        self.cross_attn = cross_attention
+        self.dropout1 = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+        self.norm1 = nn.LayerNorm(d_model)
+
+        # cross attention text
+        self.use_text_cross_attention = use_text_cross_attention
+        if use_text_cross_attention:
+            self.ca_text = nn.MultiheadAttention(d_model, n_heads, dropout=dropout)
+            self.catext_dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+            self.catext_norm = nn.LayerNorm(d_model)
+
+        # self attention
+        self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout)
+        self.dropout2 = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+        self.norm2 = nn.LayerNorm(d_model)
+
+        # ffn
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.activation = get_activation_fn(activation)
+        self.dropout3 = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+        self.dropout4 = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+        self.norm3 = nn.LayerNorm(d_model)
+
+    @staticmethod
+    def with_pos_embed(tensor, pos):
+        return tensor if pos is None else tensor + pos
+
+    def forward_ffn(self, tgt):
+        with torch.amp.autocast(device_type="cuda", enabled=False):
+            tgt2 = self.linear2(self.dropout3(self.activation(self.linear1(tgt))))
+        tgt = tgt + self.dropout4(tgt2)
+        tgt = self.norm3(tgt)
+        return tgt
+
+    def forward(
+        self,
+        # for tgt
+        tgt: Optional[Tensor],  # nq, bs, d_model
+        tgt_query_pos: Optional[Tensor] = None,  # pos for query. MLP(Sine(pos))
+        tgt_query_sine_embed: Optional[Tensor] = None,  # pos for query. Sine(pos)
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        tgt_reference_points: Optional[Tensor] = None,  # nq, bs, 4
+        memory_text: Optional[Tensor] = None,  # num_token, bs, d_model
+        text_attention_mask: Optional[Tensor] = None,  # bs, num_token
+        # for memory
+        memory: Optional[Tensor] = None,  # hw, bs, d_model
+        memory_key_padding_mask: Optional[Tensor] = None,
+        memory_level_start_index: Optional[Tensor] = None,  # num_levels
+        memory_spatial_shapes: Optional[Tensor] = None,  # bs, num_levels, 2
+        memory_pos: Optional[Tensor] = None,  # pos for memory
+        # sa
+        self_attn_mask: Optional[Tensor] = None,  # mask used for self-attention
+        cross_attn_mask: Optional[Tensor] = None,  # mask used for cross-attention
+        # dac
+        dac=False,
+        dac_use_selfatt_ln=True,
+        presence_token=None,
+        # skip inside deformable attn
+        identity=0.0,
+        **kwargs,  # additional kwargs for compatibility
+    ):
+        """
+        Input:
+            - tgt/tgt_query_pos: nq, bs, d_model
+            -
+        """
+        # self attention
+        if self.self_attn is not None:
+            if dac:
+                # we only apply self attention to the first half of the queries
+                assert tgt.shape[0] % 2 == 0
+                num_o2o_queries = tgt.shape[0] // 2
+                tgt_o2o = tgt[:num_o2o_queries]
+                tgt_query_pos_o2o = tgt_query_pos[:num_o2o_queries]
+                tgt_o2m = tgt[num_o2o_queries:]
+            else:
+                tgt_o2o = tgt
+                tgt_query_pos_o2o = tgt_query_pos
+
+            if presence_token is not None:
+                tgt_o2o = torch.cat([presence_token, tgt_o2o], dim=0)
+                tgt_query_pos_o2o = torch.cat(
+                    [torch.zeros_like(presence_token), tgt_query_pos_o2o], dim=0
+                )
+                tgt_query_pos = torch.cat(
+                    [torch.zeros_like(presence_token), tgt_query_pos], dim=0
+                )
+
+            q = k = self.with_pos_embed(tgt_o2o, tgt_query_pos_o2o)
+            tgt2 = self.self_attn(q, k, tgt_o2o, attn_mask=self_attn_mask)[0]
+            tgt_o2o = tgt_o2o + self.dropout2(tgt2)
+            if dac:
+                if not dac_use_selfatt_ln:
+                    tgt_o2o = self.norm2(tgt_o2o)
+                tgt = torch.cat((tgt_o2o, tgt_o2m), dim=0)  # Recombine
+                if dac_use_selfatt_ln:
+                    tgt = self.norm2(tgt)
+            else:
+                tgt = tgt_o2o
+                tgt = self.norm2(tgt)
+
+        if self.use_text_cross_attention:
+            tgt2 = self.ca_text(
+                self.with_pos_embed(tgt, tgt_query_pos),
+                memory_text,
+                memory_text,
+                key_padding_mask=text_attention_mask,
+            )[0]
+            tgt = tgt + self.catext_dropout(tgt2)
+            tgt = self.catext_norm(tgt)
+
+        if presence_token is not None:
+            presence_token_mask = torch.zeros_like(cross_attn_mask[:, :1, :])
+            cross_attn_mask = torch.cat(
+                [presence_token_mask, cross_attn_mask], dim=1
+            )  # (bs*nheads, 1+nq, hw)
+
+        # Cross attention to image
+        tgt2 = self.cross_attn(
+            query=self.with_pos_embed(tgt, tgt_query_pos),
+            key=self.with_pos_embed(memory, memory_pos),
+            value=memory,
+            attn_mask=cross_attn_mask,
+            key_padding_mask=(
+                memory_key_padding_mask.transpose(0, 1)
+                if memory_key_padding_mask is not None
+                else None
+            ),
+        )[0]
+
+        tgt = tgt + self.dropout1(tgt2)
+        tgt = self.norm1(tgt)
+
+        # ffn
+        tgt = self.forward_ffn(tgt)
+
+        presence_token_out = None
+        if presence_token is not None:
+            presence_token_out = tgt[:1]
+            tgt = tgt[1:]
+
+        return tgt, presence_token_out
+
+
+class TransformerDecoder(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        frozen: bool,
+        interaction_layer,
+        layer,
+        num_layers: int,
+        num_queries: int,
+        return_intermediate: bool,
+        box_refine: bool = False,
+        num_o2m_queries: int = 0,
+        dac: bool = False,
+        boxRPB: str = "none",
+        # Experimental: An object query for SAM 2 tasks
+        instance_query: bool = False,
+        # Defines the number of additional instance queries,
+        # 1 or 4 are the most likely for single vs multi mask support
+        num_instances: int = 1,  # Irrelevant if instance_query is False
+        dac_use_selfatt_ln: bool = True,
+        use_act_checkpoint: bool = False,
+        compile_mode=None,
+        presence_token: bool = False,
+        clamp_presence_logits: bool = True,
+        clamp_presence_logit_max_val: float = 10.0,
+        use_normed_output_consistently: bool = True,
+        separate_box_head_instance: bool = False,
+        separate_norm_instance: bool = False,
+        resolution: Optional[int] = None,
+        stride: Optional[int] = None,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.layers = get_clones(layer, num_layers)
+        self.fine_layers = (
+            get_clones(interaction_layer, num_layers)
+            if interaction_layer is not None
+            else [None] * num_layers
+        )
+        self.num_layers = num_layers
+        self.num_queries = num_queries
+        self.dac = dac
+        if dac:
+            self.num_o2m_queries = num_queries
+            tot_num_queries = num_queries
+        else:
+            self.num_o2m_queries = num_o2m_queries
+            tot_num_queries = num_queries + num_o2m_queries
+        self.norm = nn.LayerNorm(d_model)
+        self.return_intermediate = return_intermediate
+        self.bbox_embed = MLP(d_model, d_model, 4, 3)
+        self.query_embed = nn.Embedding(tot_num_queries, d_model)
+        self.instance_query_embed = None
+        self.instance_query_reference_points = None
+        self.use_instance_query = instance_query
+        self.num_instances = num_instances
+        self.use_normed_output_consistently = use_normed_output_consistently
+
+        self.instance_norm = nn.LayerNorm(d_model) if separate_norm_instance else None
+        self.instance_bbox_embed = None
+        if separate_box_head_instance:
+            self.instance_bbox_embed = MLP(d_model, d_model, 4, 3)
+        if instance_query:
+            self.instance_query_embed = nn.Embedding(num_instances, d_model)
+        self.box_refine = box_refine
+        if box_refine:
+            nn.init.constant_(self.bbox_embed.layers[-1].weight.data, 0)
+            nn.init.constant_(self.bbox_embed.layers[-1].bias.data, 0)
+
+            self.reference_points = nn.Embedding(num_queries, 4)
+            if instance_query:
+                self.instance_reference_points = nn.Embedding(num_instances, 4)
+
+        assert boxRPB in ["none", "log", "linear", "both"]
+        self.boxRPB = boxRPB
+        if boxRPB != "none":
+            try:
+                nheads = self.layers[0].cross_attn_image.num_heads
+            except AttributeError:
+                nheads = self.layers[0].cross_attn.num_heads
+
+            n_input = 4 if boxRPB == "both" else 2
+            self.boxRPB_embed_x = MLP(n_input, d_model, nheads, 2)
+            self.boxRPB_embed_y = MLP(n_input, d_model, nheads, 2)
+            self.compilable_cord_cache = None
+            self.compilable_stored_size = None
+            self.coord_cache = {}
+
+            if resolution is not None and stride is not None:
+                feat_size = resolution // stride
+                coords_h, coords_w = self._get_coords(
+                    feat_size, feat_size, device="cuda"
+                )
+                self.compilable_cord_cache = (coords_h, coords_w)
+                self.compilable_stored_size = (feat_size, feat_size)
+
+        self.roi_pooler = (
+            RoIAlign(output_size=7, spatial_scale=1, sampling_ratio=-1, aligned=True)
+            if interaction_layer is not None
+            else None
+        )
+        if frozen:
+            for p in self.parameters():
+                p.requires_grad_(False)
+
+        self.presence_token = None
+        self.clamp_presence_logits = clamp_presence_logits
+        self.clamp_presence_logit_max_val = clamp_presence_logit_max_val
+        if presence_token:
+            self.presence_token = nn.Embedding(1, d_model)
+            self.presence_token_head = MLP(d_model, d_model, 1, 3)
+            self.presence_token_out_norm = nn.LayerNorm(d_model)
+
+        self.ref_point_head = MLP(2 * self.d_model, self.d_model, self.d_model, 2)
+        self.dac_use_selfatt_ln = dac_use_selfatt_ln
+        self.use_act_checkpoint = use_act_checkpoint
+
+        nn.init.normal_(self.query_embed.weight.data)
+        if self.instance_query_embed is not None:
+            nn.init.normal_(self.instance_query_embed.weight.data)
+
+        assert self.roi_pooler is None
+        assert self.return_intermediate, "support return_intermediate only"
+        assert self.box_refine, "support box refine only"
+
+        self.compile_mode = compile_mode
+        self.compiled = False
+        # We defer compilation till after the first forward, to first warm-up the boxRPB cache
+
+        # assign layer index to each layer so that some layers can decide what to do
+        # based on which layer index they are (e.g. cross attention to memory bank only
+        # in selected layers)
+        for layer_idx, layer in enumerate(self.layers):
+            layer.layer_idx = layer_idx
+
+    @staticmethod
+    def _get_coords(H, W, device):
+        coords_h = torch.arange(0, H, device=device, dtype=torch.float32) / H
+        coords_w = torch.arange(0, W, device=device, dtype=torch.float32) / W
+        return coords_h, coords_w
+
+    def _get_rpb_matrix(self, reference_boxes, feat_size):
+        H, W = feat_size
+        boxes_xyxy = box_cxcywh_to_xyxy(reference_boxes).transpose(0, 1)
+        bs, num_queries, _ = boxes_xyxy.shape
+        if self.compilable_cord_cache is None:
+            self.compilable_cord_cache = self._get_coords(H, W, reference_boxes.device)
+            self.compilable_stored_size = (H, W)
+
+        if torch.compiler.is_dynamo_compiling() or self.compilable_stored_size == (
+            H,
+            W,
+        ):
+            # good, hitting the cache, will be compilable
+            coords_h, coords_w = self.compilable_cord_cache
+        else:
+            # cache miss, will create compilation issue
+            # In case we're not compiling, we'll still rely on the dict-based cache
+            if feat_size not in self.coord_cache:
+                self.coord_cache[feat_size] = self._get_coords(
+                    H, W, reference_boxes.device
+                )
+            coords_h, coords_w = self.coord_cache[feat_size]
+
+            assert coords_h.shape == (H,)
+            assert coords_w.shape == (W,)
+
+        deltas_y = coords_h.view(1, -1, 1) - boxes_xyxy.reshape(-1, 1, 4)[:, :, 1:4:2]
+        deltas_y = deltas_y.view(bs, num_queries, -1, 2)
+        deltas_x = coords_w.view(1, -1, 1) - boxes_xyxy.reshape(-1, 1, 4)[:, :, 0:3:2]
+        deltas_x = deltas_x.view(bs, num_queries, -1, 2)
+
+        if self.boxRPB in ["log", "both"]:
+            deltas_x_log = deltas_x * 8  # normalize to -8, 8
+            deltas_x_log = (
+                torch.sign(deltas_x_log)
+                * torch.log2(torch.abs(deltas_x_log) + 1.0)
+                / np.log2(8)
+            )
+
+            deltas_y_log = deltas_y * 8  # normalize to -8, 8
+            deltas_y_log = (
+                torch.sign(deltas_y_log)
+                * torch.log2(torch.abs(deltas_y_log) + 1.0)
+                / np.log2(8)
+            )
+            if self.boxRPB == "log":
+                deltas_x = deltas_x_log
+                deltas_y = deltas_y_log
+            else:
+                deltas_x = torch.cat([deltas_x, deltas_x_log], dim=-1)
+                deltas_y = torch.cat([deltas_y, deltas_y_log], dim=-1)
+
+        if self.training:
+            assert self.use_act_checkpoint, "activation ckpt not enabled in decoder"
+        deltas_x = activation_ckpt_wrapper(self.boxRPB_embed_x)(
+            x=deltas_x,
+            act_ckpt_enable=self.training and self.use_act_checkpoint,
+        )  # bs, num_queries, W, n_heads
+        deltas_y = activation_ckpt_wrapper(self.boxRPB_embed_y)(
+            x=deltas_y,
+            act_ckpt_enable=self.training and self.use_act_checkpoint,
+        )  # bs, num_queries, H, n_heads
+
+        if not torch.compiler.is_dynamo_compiling():
+            assert deltas_x.shape[:3] == (bs, num_queries, W)
+            assert deltas_y.shape[:3] == (bs, num_queries, H)
+
+        B = deltas_y.unsqueeze(3) + deltas_x.unsqueeze(
+            2
+        )  # bs, num_queries, H, W, n_heads
+        if not torch.compiler.is_dynamo_compiling():
+            assert B.shape[:4] == (bs, num_queries, H, W)
+        B = B.flatten(2, 3)  # bs, num_queries, H*W, n_heads
+        B = B.permute(0, 3, 1, 2)  # bs, n_heads, num_queries, H*W
+        B = B.contiguous()  # memeff attn likes ordered strides
+        if not torch.compiler.is_dynamo_compiling():
+            assert B.shape[2:] == (num_queries, H * W)
+        return B
+
+    def forward(
+        self,
+        tgt,
+        memory,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+        reference_boxes: Optional[Tensor] = None,  # num_queries, bs, 4
+        # for memory
+        level_start_index: Optional[Tensor] = None,  # num_levels
+        spatial_shapes: Optional[Tensor] = None,  # bs, num_levels, 2
+        valid_ratios: Optional[Tensor] = None,
+        # for text
+        memory_text: Optional[Tensor] = None,
+        text_attention_mask: Optional[Tensor] = None,
+        # if `apply_dac` is None, it will default to `self.dac`
+        apply_dac: Optional[bool] = None,
+        is_instance_prompt=False,
+        decoder_extra_kwargs: Optional[Dict] = None,
+        # ROI memory bank
+        obj_roi_memory_feat=None,
+        obj_roi_memory_mask=None,
+        box_head_trk=None,
+    ):
+        """
+        Input:
+            - tgt: nq, bs, d_model
+            - memory: \\sum{hw}, bs, d_model
+            - pos: \\sum{hw}, bs, d_model
+            - reference_boxes: nq, bs, 4 (after sigmoid)
+            - valid_ratios/spatial_shapes: bs, nlevel, 2
+        """
+        if memory_mask is not None:
+            assert (
+                self.boxRPB == "none"
+            ), "inputting a memory_mask in the presence of boxRPB is unexpected/not implemented"
+
+        apply_dac = apply_dac if apply_dac is not None else self.dac
+        if apply_dac:
+            assert (tgt.shape[0] == self.num_queries) or (
+                self.use_instance_query
+                and (tgt.shape[0] == self.instance_query_embed.num_embeddings)
+            )
+
+            tgt = tgt.repeat(2, 1, 1)
+            # note that we don't tile tgt_mask, since DAC doesn't
+            # use self-attention in o2m queries
+            if reference_boxes is not None:
+                assert (reference_boxes.shape[0] == self.num_queries) or (
+                    self.use_instance_query
+                    and (
+                        reference_boxes.shape[0]
+                        == self.instance_query_embed.num_embeddings
+                    )
+                )
+                reference_boxes = reference_boxes.repeat(2, 1, 1)
+
+        bs = tgt.shape[1]
+        intermediate = []
+        intermediate_presence_logits = []
+        presence_feats = None
+
+        if self.box_refine:
+            if reference_boxes is None:
+                # In this case, we're in a one-stage model, so we generate the reference boxes
+                reference_boxes = self.reference_points.weight.unsqueeze(1)
+                reference_boxes = (
+                    reference_boxes.repeat(2, bs, 1)
+                    if apply_dac
+                    else reference_boxes.repeat(1, bs, 1)
+                )
+                reference_boxes = reference_boxes.sigmoid()
+            intermediate_ref_boxes = [reference_boxes]
+        else:
+            reference_boxes = None
+            intermediate_ref_boxes = None
+
+        output = tgt
+        presence_out = None
+        if self.presence_token is not None and is_instance_prompt is False:
+            # expand to batch dim
+            presence_out = self.presence_token.weight[None].expand(1, bs, -1)
+
+        box_head = self.bbox_embed
+        if is_instance_prompt and self.instance_bbox_embed is not None:
+            box_head = self.instance_bbox_embed
+
+        out_norm = self.norm
+        if is_instance_prompt and self.instance_norm is not None:
+            out_norm = self.instance_norm
+
+        for layer_idx, layer in enumerate(self.layers):
+            reference_points_input = (
+                reference_boxes[:, :, None]
+                * torch.cat([valid_ratios, valid_ratios], -1)[None, :]
+            )  # nq, bs, nlevel, 4
+
+            query_sine_embed = gen_sineembed_for_position(
+                reference_points_input[:, :, 0, :], self.d_model
+            )  # nq, bs, d_model*2
+
+            # conditional query
+            query_pos = self.ref_point_head(query_sine_embed)  # nq, bs, d_model
+
+            if self.boxRPB != "none" and reference_boxes is not None:
+                assert (
+                    spatial_shapes.shape[0] == 1
+                ), "only single scale support implemented"
+                memory_mask = self._get_rpb_matrix(
+                    reference_boxes,
+                    (spatial_shapes[0, 0], spatial_shapes[0, 1]),
+                )
+                memory_mask = memory_mask.flatten(0, 1)  # (bs*n_heads, nq, H*W)
+            if self.training:
+                assert (
+                    self.use_act_checkpoint
+                ), "Activation checkpointing not enabled in the decoder"
+            output, presence_out = activation_ckpt_wrapper(layer)(
+                tgt=output,
+                tgt_query_pos=query_pos,
+                tgt_query_sine_embed=query_sine_embed,
+                tgt_key_padding_mask=tgt_key_padding_mask,
+                tgt_reference_points=reference_points_input,
+                memory_text=memory_text,
+                text_attention_mask=text_attention_mask,
+                memory=memory,
+                memory_key_padding_mask=memory_key_padding_mask,
+                memory_level_start_index=level_start_index,
+                memory_spatial_shapes=spatial_shapes,
+                memory_pos=pos,
+                self_attn_mask=tgt_mask,
+                cross_attn_mask=memory_mask,
+                dac=apply_dac,
+                dac_use_selfatt_ln=self.dac_use_selfatt_ln,
+                presence_token=presence_out,
+                **(decoder_extra_kwargs or {}),
+                act_ckpt_enable=self.training and self.use_act_checkpoint,
+                # ROI memory bank
+                obj_roi_memory_feat=obj_roi_memory_feat,
+                obj_roi_memory_mask=obj_roi_memory_mask,
+            )
+
+            # iter update
+            if self.box_refine:
+                reference_before_sigmoid = inverse_sigmoid(reference_boxes)
+                if box_head_trk is None:
+                    # delta_unsig = self.bbox_embed(output)
+                    if not self.use_normed_output_consistently:
+                        delta_unsig = box_head(output)
+                    else:
+                        delta_unsig = box_head(out_norm(output))
+                else:
+                    # box_head_trk use a separate box head for tracking queries
+                    Q_det = decoder_extra_kwargs["Q_det"]
+                    assert output.size(0) >= Q_det
+                    delta_unsig_det = self.bbox_embed(output[:Q_det])
+                    delta_unsig_trk = box_head_trk(output[Q_det:])
+                    delta_unsig = torch.cat([delta_unsig_det, delta_unsig_trk], dim=0)
+                outputs_unsig = delta_unsig + reference_before_sigmoid
+                new_reference_points = outputs_unsig.sigmoid()
+
+                reference_boxes = new_reference_points.detach()
+                if layer_idx != self.num_layers - 1:
+                    intermediate_ref_boxes.append(new_reference_points)
+            else:
+                raise NotImplementedError("not implemented yet")
+
+            intermediate.append(out_norm(output))
+            if self.presence_token is not None and is_instance_prompt is False:
+                # norm, mlp head
+                intermediate_layer_presence_logits = self.presence_token_head(
+                    self.presence_token_out_norm(presence_out)
+                ).squeeze(-1)
+
+                # clamp to mitigate numerical issues
+                if self.clamp_presence_logits:
+                    intermediate_layer_presence_logits.clamp(
+                        min=-self.clamp_presence_logit_max_val,
+                        max=self.clamp_presence_logit_max_val,
+                    )
+
+                intermediate_presence_logits.append(intermediate_layer_presence_logits)
+                presence_feats = presence_out.clone()
+
+        if not self.compiled and self.compile_mode is not None:
+            self.forward = torch.compile(
+                self.forward, mode=self.compile_mode, fullgraph=True
+            )
+            self.compiled = True
+
+        return (
+            torch.stack(intermediate),
+            torch.stack(intermediate_ref_boxes),
+            (
+                torch.stack(intermediate_presence_logits)
+                if self.presence_token is not None and is_instance_prompt is False
+                else None
+            ),
+            presence_feats,
+        )
+
+
+class TransformerEncoderCrossAttention(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        frozen: bool,
+        pos_enc_at_input: bool,
+        layer,
+        num_layers: int,
+        use_act_checkpoint: bool = False,
+        batch_first: bool = False,  # Do layers expect batch first input?
+        # which layers to exclude cross attention? default: None, means all
+        # layers use cross attention
+        remove_cross_attention_layers: Optional[list] = None,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.layers = get_clones(layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = nn.LayerNorm(d_model)
+        self.pos_enc_at_input = pos_enc_at_input
+        self.use_act_checkpoint = use_act_checkpoint
+
+        if frozen:
+            for p in self.parameters():
+                p.requires_grad_(False)
+
+        self.batch_first = batch_first
+
+        # remove cross attention layers if specified
+        self.remove_cross_attention_layers = [False] * self.num_layers
+        if remove_cross_attention_layers is not None:
+            for i in remove_cross_attention_layers:
+                self.remove_cross_attention_layers[i] = True
+        assert len(self.remove_cross_attention_layers) == len(self.layers)
+
+        for i, remove_cross_attention in enumerate(self.remove_cross_attention_layers):
+            if remove_cross_attention:
+                self.layers[i].cross_attn_image = None
+                self.layers[i].norm2 = None
+                self.layers[i].dropout2 = None
+
+    def forward(
+        self,
+        src,  # self-attention inputs
+        prompt,  # cross-attention inputs
+        src_mask: Optional[Tensor] = None,  # att.mask for self-attention inputs
+        prompt_mask: Optional[Tensor] = None,  # att.mask for cross-attention inputs
+        src_key_padding_mask: Optional[Tensor] = None,
+        prompt_key_padding_mask: Optional[Tensor] = None,
+        src_pos: Optional[Tensor] = None,  # pos_enc for self-attention inputs
+        prompt_pos: Optional[Tensor] = None,  # pos_enc for cross-attention inputs
+        feat_sizes: Optional[list] = None,
+        num_obj_ptr_tokens: int = 0,  # number of object pointer *tokens*
+    ):
+        if isinstance(src, list):
+            assert isinstance(src_key_padding_mask, list) and isinstance(src_pos, list)
+            assert len(src) == len(src_key_padding_mask) == len(src_pos) == 1
+            src, src_key_padding_mask, src_pos = (
+                src[0],
+                src_key_padding_mask[0],
+                src_pos[0],
+            )
+
+        assert (
+            src.shape[1] == prompt.shape[1]
+        ), "Batch size must be the same for src and prompt"
+
+        output = src
+
+        if self.pos_enc_at_input and src_pos is not None:
+            output = output + 0.1 * src_pos
+
+        if self.batch_first:
+            # Convert to batch first
+            output = output.transpose(0, 1)
+            src_pos = src_pos.transpose(0, 1)
+            prompt = prompt.transpose(0, 1)
+            prompt_pos = prompt_pos.transpose(0, 1)
+
+        for layer in self.layers:
+            kwds = {}
+            if isinstance(layer.cross_attn_image, RoPEAttention):
+                kwds = {"num_k_exclude_rope": num_obj_ptr_tokens}
+
+            output = activation_ckpt_wrapper(layer)(
+                tgt=output,
+                memory=prompt,
+                tgt_mask=src_mask,
+                memory_mask=prompt_mask,
+                tgt_key_padding_mask=src_key_padding_mask,
+                memory_key_padding_mask=prompt_key_padding_mask,
+                pos=prompt_pos,
+                query_pos=src_pos,
+                dac=False,
+                attn_bias=None,
+                act_ckpt_enable=self.training and self.use_act_checkpoint,
+                **kwds,
+            )
+            normed_output = self.norm(output)
+
+        if self.batch_first:
+            # Convert back to seq first
+            normed_output = normed_output.transpose(0, 1)
+            src_pos = src_pos.transpose(0, 1)
+
+        return {
+            "memory": normed_output,
+            "pos_embed": src_pos,
+            "padding_mask": src_key_padding_mask,
+        }
+
+
+class TransformerDecoderLayerv1(nn.Module):
+    def __init__(
+        self,
+        activation: str,
+        cross_attention: nn.Module,
+        d_model: int,
+        dim_feedforward: int,
+        dropout: float,
+        pos_enc_at_attn: bool,
+        pos_enc_at_cross_attn_keys: bool,
+        pos_enc_at_cross_attn_queries: bool,
+        pre_norm: bool,
+        self_attention: nn.Module,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.dim_feedforward = dim_feedforward
+        self.dropout_value = dropout
+        self.self_attn = self_attention
+        self.cross_attn_image = cross_attention
+
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        self.activation_str = activation
+        self.activation = get_activation_fn(activation)
+        self.pre_norm = pre_norm
+
+        self.pos_enc_at_attn = pos_enc_at_attn
+        self.pos_enc_at_cross_attn_queries = pos_enc_at_cross_attn_queries
+        self.pos_enc_at_cross_attn_keys = pos_enc_at_cross_attn_keys
+
+    def forward_post(
+        self,
+        tgt,
+        memory,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+        query_pos: Optional[Tensor] = None,
+        **kwargs,
+    ):
+        q = k = tgt + query_pos if self.pos_enc_at_attn else tgt
+
+        # Self attention
+        tgt2 = self.self_attn(
+            q,
+            k,
+            value=tgt,
+            attn_mask=tgt_mask,
+            key_padding_mask=tgt_key_padding_mask,
+        )[0]
+        tgt = tgt + self.dropout1(tgt2)
+        tgt = self.norm1(tgt)
+
+        # Cross attention to image
+        tgt2 = self.cross_attn_image(
+            query=tgt + query_pos if self.pos_enc_at_cross_attn_queries else tgt,
+            key=memory + pos if self.pos_enc_at_cross_attn_keys else memory,
+            value=memory,
+            attn_mask=memory_mask,
+            key_padding_mask=memory_key_padding_mask,
+        )[0]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt = self.norm2(tgt)
+
+        # FFN
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+        tgt = tgt + self.dropout3(tgt2)
+        tgt = self.norm3(tgt)
+        return tgt
+
+    def forward_pre(
+        self,
+        tgt,
+        memory,
+        dac: bool = False,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+        query_pos: Optional[Tensor] = None,
+        attn_bias: Optional[Tensor] = None,
+        **kwargs,
+    ):
+        if dac:
+            # we only apply self attention to the first half of the queries
+            assert tgt.shape[0] % 2 == 0
+            other_tgt = tgt[tgt.shape[0] // 2 :]
+            tgt = tgt[: tgt.shape[0] // 2]
+        tgt2 = self.norm1(tgt)
+        q = k = tgt2 + query_pos if self.pos_enc_at_attn else tgt2
+        tgt2 = self.self_attn(
+            q,
+            k,
+            value=tgt2,
+            attn_mask=tgt_mask,
+            key_padding_mask=tgt_key_padding_mask,
+        )[0]
+        tgt = tgt + self.dropout1(tgt2)
+        if dac:
+            # Recombine
+            tgt = torch.cat((tgt, other_tgt), dim=0)
+        tgt2 = self.norm2(tgt)
+        tgt2 = self.cross_attn_image(
+            query=tgt2 + query_pos if self.pos_enc_at_cross_attn_queries else tgt2,
+            key=memory + pos if self.pos_enc_at_cross_attn_keys else memory,
+            value=memory,
+            attn_mask=memory_mask,
+            key_padding_mask=memory_key_padding_mask,
+            attn_bias=attn_bias,
+        )[0]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt2 = self.norm3(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
+        tgt = tgt + self.dropout3(tgt2)
+        return tgt
+
+    def forward(
+        self,
+        tgt,
+        memory,
+        dac: bool = False,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+        query_pos: Optional[Tensor] = None,
+        attn_bias: Optional[Tensor] = None,
+        **kwds: Any,
+    ) -> torch.Tensor:
+        fwd_fn = self.forward_pre if self.pre_norm else self.forward_post
+        return fwd_fn(
+            tgt,
+            memory,
+            dac=dac,
+            tgt_mask=tgt_mask,
+            memory_mask=memory_mask,
+            tgt_key_padding_mask=tgt_key_padding_mask,
+            memory_key_padding_mask=memory_key_padding_mask,
+            pos=pos,
+            query_pos=query_pos,
+            attn_bias=attn_bias,
+            **kwds,
+        )
+
+
+class TransformerDecoderLayerv2(TransformerDecoderLayerv1):
+    def __init__(self, cross_attention_first=False, *args: Any, **kwds: Any):
+        super().__init__(*args, **kwds)
+        self.cross_attention_first = cross_attention_first
+
+    def _forward_sa(self, tgt, query_pos):
+        # Self-Attention
+        tgt2 = self.norm1(tgt)
+        q = k = tgt2 + query_pos if self.pos_enc_at_attn else tgt2
+        tgt2 = self.self_attn(q, k, v=tgt2)
+        tgt = tgt + self.dropout1(tgt2)
+        return tgt
+
+    def _forward_ca(self, tgt, memory, query_pos, pos, num_k_exclude_rope=0):
+        if self.cross_attn_image is None:
+            return tgt
+
+        kwds = {}
+        if num_k_exclude_rope > 0:
+            assert isinstance(self.cross_attn_image, RoPEAttention)
+            kwds = {"num_k_exclude_rope": num_k_exclude_rope}
+
+        # Cross-Attention
+        tgt2 = self.norm2(tgt)
+        tgt2 = self.cross_attn_image(
+            q=tgt2 + query_pos if self.pos_enc_at_cross_attn_queries else tgt2,
+            k=memory + pos if self.pos_enc_at_cross_attn_keys else memory,
+            v=memory,
+            **kwds,
+        )
+        tgt = tgt + self.dropout2(tgt2)
+        return tgt
+
+    def forward_pre(
+        self,
+        tgt,
+        memory,
+        dac: bool,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+        query_pos: Optional[Tensor] = None,
+        attn_bias: Optional[Tensor] = None,
+        num_k_exclude_rope: int = 0,
+    ):
+        assert dac is False
+        assert tgt_mask is None
+        assert memory_mask is None
+        assert tgt_key_padding_mask is None
+        assert memory_key_padding_mask is None
+        assert attn_bias is None
+
+        if self.cross_attention_first:
+            tgt = self._forward_ca(tgt, memory, query_pos, pos, num_k_exclude_rope)
+            tgt = self._forward_sa(tgt, query_pos)
+        else:
+            tgt = self._forward_sa(tgt, query_pos)
+            tgt = self._forward_ca(tgt, memory, query_pos, pos, num_k_exclude_rope)
+
+        # MLP
+        tgt2 = self.norm3(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
+        tgt = tgt + self.dropout3(tgt2)
+        return tgt
+
+    def forward(self, *args: Any, **kwds: Any) -> torch.Tensor:
+        if self.pre_norm:
+            return self.forward_pre(*args, **kwds)
+        raise NotImplementedError

sam3/model/edt.py

@@ -0,0 +1,173 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+"""Triton kernel for euclidean distance transform (EDT)"""
+
+import torch
+import triton
+import triton.language as tl
+
+"""
+Disclaimer: This implementation is not meant to be extremely efficient. A CUDA kernel would likely be more efficient.
+Even in Triton, there may be more suitable algorithms.
+
+The goal of this kernel is to mimic cv2.distanceTransform(input, cv2.DIST_L2, 0).
+Recall that the euclidean distance transform (EDT) calculates the L2 distance to the closest zero pixel for each pixel of the source image.
+
+For images of size NxN, the naive algorithm would be to compute pairwise distances between every pair of points, leading to a O(N^4) algorithm, which is obviously impractical.
+One can do better using the following approach:
+- First, compute the distance to the closest point in the same row. We can write it as Row_EDT[i,j] = min_k (sqrt((k-j)^2) if input[i,k]==0 else +infinity). With a naive implementation, this step has a O(N^3) complexity
+- Then, because of triangular inequality, we notice that the EDT for a given location [i,j] is the min of the row EDTs in the same column. EDT[i,j] = min_k Row_EDT[k, j]. This is also O(N^3)
+
+Overall, this algorithm is quite amenable to parallelization, and has a complexity O(N^3). Can we do better?
+
+It turns out that we can leverage the structure of the L2 distance (nice and convex) to find the minimum in a more efficient way.
+We follow the algorithm from "Distance Transforms of Sampled Functions" (https://cs.brown.edu/people/pfelzens/papers/dt-final.pdf), which is also what's implemented in opencv
+
+For a single dimension EDT, we can compute the EDT of an arbitrary function F, that we discretize over the grid. Note that for the binary EDT that we're interested in, we can set F(i,j) = 0 if input[i,j]==0 else +infinity
+For now, we'll compute the EDT squared, and will take the sqrt only at the very end.
+The basic idea is that each point at location i spawns a parabola around itself, with a bias equal to F(i). So specifically, we're looking at the parabola (x - i)^2 + F(i)
+When we're looking for the row EDT at location j, we're effectively looking for min_i (x-i)^2 + F(i). In other word we want to find the lowest parabola at location j.
+
+To do this efficiently, we need to maintain the lower envelope of the union of parabolas. This can be constructed on the fly using a sort of stack approach:
+ - every time we want to add a new parabola, we check if it may be covering the current right-most parabola. If so, then that parabola was useless, so we can pop it from the stack
+ - repeat until we can't find any more parabola to pop. Then push the new one.
+
+This algorithm runs in O(N) for a single row, so overall O(N^2) when applied to all rows
+Similarly as before, we notice that we can decompose the algorithm for rows and columns, leading to an overall run-time of O(N^2)
+
+This algorithm is less suited for to GPUs, since the one-dimensional EDT computation is quite sequential in nature. However, we can parallelize over batch and row dimensions.
+In Triton, things are particularly bad at the moment, since there is no support for reading/writing to the local memory at a specific index (a local gather is coming soon, see https://github.com/triton-lang/triton/issues/974, but no mention of writing, ie scatter)
+One could emulate these operations with masking, but in initial tests, it proved to be worst than naively reading and writing to the global memory. My guess is that the cache is compensating somewhat for the repeated single-point accesses.
+
+
+The timing obtained on a H100 for a random batch of masks of dimension 256 x 1024 x 1024 are as follows:
+- OpenCV: 1780ms (including round-trip to cpu, but discounting the fact that it introduces a synchronization point)
+- triton, O(N^3) algo: 627ms
+- triton, O(N^2) algo: 322ms
+
+Overall, despite being quite naive, this implementation is roughly 5.5x faster than the openCV cpu implem
+
+"""
+
+
+@triton.jit
+def edt_kernel(inputs_ptr, outputs_ptr, v, z, height, width, horizontal: tl.constexpr):
+    # This is a somewhat verbatim implementation of the efficient 1D EDT algorithm described above
+    # It can be applied horizontally or vertically depending if we're doing the first or second stage.
+    # It's parallelized across batch+row (or batch+col if horizontal=False)
+    # TODO: perhaps the implementation can be revisited if/when local gather/scatter become available in triton
+    batch_id = tl.program_id(axis=0)
+    if horizontal:
+        row_id = tl.program_id(axis=1)
+        block_start = (batch_id * height * width) + row_id * width
+        length = width
+        stride = 1
+    else:
+        col_id = tl.program_id(axis=1)
+        block_start = (batch_id * height * width) + col_id
+        length = height
+        stride = width
+
+    # This will be the index of the right most parabola in the envelope ("the top of the stack")
+    k = 0
+    for q in range(1, length):
+        # Read the function value at the current location. Note that we're doing a singular read, not very efficient
+        cur_input = tl.load(inputs_ptr + block_start + (q * stride))
+        # location of the parabola on top of the stack
+        r = tl.load(v + block_start + (k * stride))
+        # associated boundary
+        z_k = tl.load(z + block_start + (k * stride))
+        # value of the function at the parabola location
+        previous_input = tl.load(inputs_ptr + block_start + (r * stride))
+        # intersection between the two parabolas
+        s = (cur_input - previous_input + q * q - r * r) / (q - r) / 2
+
+        # we'll pop as many parabolas as required
+        while s <= z_k and k - 1 >= 0:
+            k = k - 1
+            r = tl.load(v + block_start + (k * stride))
+            z_k = tl.load(z + block_start + (k * stride))
+            previous_input = tl.load(inputs_ptr + block_start + (r * stride))
+            s = (cur_input - previous_input + q * q - r * r) / (q - r) / 2
+
+        # Store the new one
+        k = k + 1
+        tl.store(v + block_start + (k * stride), q)
+        tl.store(z + block_start + (k * stride), s)
+        if k + 1 < length:
+            tl.store(z + block_start + ((k + 1) * stride), 1e9)
+
+    # Last step, we read the envelope to find the min in every location
+    k = 0
+    for q in range(length):
+        while (
+            k + 1 < length
+            and tl.load(
+                z + block_start + ((k + 1) * stride), mask=(k + 1) < length, other=q
+            )
+            < q
+        ):
+            k += 1
+        r = tl.load(v + block_start + (k * stride))
+        d = q - r
+        old_value = tl.load(inputs_ptr + block_start + (r * stride))
+        tl.store(outputs_ptr + block_start + (q * stride), old_value + d * d)
+
+
+def edt_triton(data: torch.Tensor):
+    """
+    Computes the Euclidean Distance Transform (EDT) of a batch of binary images.
+
+    Args:
+        data: A tensor of shape (B, H, W) representing a batch of binary images.
+
+    Returns:
+        A tensor of the same shape as data containing the EDT.
+        It should be equivalent to a batched version of cv2.distanceTransform(input, cv2.DIST_L2, 0)
+    """
+    assert data.dim() == 3
+    assert data.is_cuda
+    B, H, W = data.shape
+    data = data.contiguous()
+
+    # Allocate the "function" tensor. Implicitly the function is 0 if data[i,j]==0 else +infinity
+    output = torch.where(data, 1e18, 0.0)
+    assert output.is_contiguous()
+
+    # Scratch tensors for the parabola stacks
+    parabola_loc = torch.zeros(B, H, W, dtype=torch.uint32, device=data.device)
+    parabola_inter = torch.empty(B, H, W, dtype=torch.float, device=data.device)
+    parabola_inter[:, :, 0] = -1e18
+    parabola_inter[:, :, 1] = 1e18
+
+    # Grid size (number of blocks)
+    grid = (B, H)
+
+    # Launch initialization kernel
+    edt_kernel[grid](
+        output.clone(),
+        output,
+        parabola_loc,
+        parabola_inter,
+        H,
+        W,
+        horizontal=True,
+    )
+
+    # reset the parabola stacks
+    parabola_loc.zero_()
+    parabola_inter[:, :, 0] = -1e18
+    parabola_inter[:, :, 1] = 1e18
+
+    grid = (B, W)
+    edt_kernel[grid](
+        output.clone(),
+        output,
+        parabola_loc,
+        parabola_inter,
+        H,
+        W,
+        horizontal=False,
+    )
+    # don't forget to take sqrt at the end
+    return output.sqrt()

sam3/model/encoder.py

@@ -0,0 +1,594 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+# Based on https://github.com/IDEA-Research/GroundingDINO
+
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+from torch import nn, Tensor
+
+from .act_ckpt_utils import activation_ckpt_wrapper
+from .model_misc import get_activation_fn, get_clones, get_valid_ratio
+
+
+class TransformerEncoderLayer(nn.Module):
+    """
+    Transformer encoder layer that performs self-attention followed by cross-attention.
+
+    This layer was previously called TransformerDecoderLayer but was renamed to better
+    reflect its role in the architecture. It processes input sequences through self-attention
+    and then cross-attention with another input (typically image features).
+
+    The layer supports both pre-norm and post-norm configurations, as well as
+    positional encoding at different stages of the attention mechanism.
+    """
+
+    def __init__(
+        self,
+        activation: str,
+        cross_attention: nn.Module,
+        d_model: int,
+        dim_feedforward: int,
+        dropout: float,
+        pos_enc_at_attn: bool,
+        pos_enc_at_cross_attn_keys: bool,
+        pos_enc_at_cross_attn_queries: bool,
+        pre_norm: bool,
+        self_attention: nn.Module,
+    ):
+        """
+        Initialize a transformer encoder layer.
+
+        Args:
+            activation: Activation function to use in the feedforward network
+            cross_attention: Cross-attention module for attending to image features
+            d_model: Model dimension/hidden size
+            dim_feedforward: Dimension of the feedforward network
+            dropout: Dropout probability
+            pos_enc_at_attn: Whether to add positional encodings at self-attention
+            pos_enc_at_cross_attn_keys: Whether to add positional encodings to keys in cross-attention
+            pos_enc_at_cross_attn_queries: Whether to add positional encodings to queries in cross-attention
+            pre_norm: Whether to use pre-norm (True) or post-norm (False) architecture
+            self_attention: Self-attention module
+        """
+        super().__init__()
+        self.d_model = d_model
+        self.dim_feedforward = dim_feedforward
+        self.dropout_value = dropout
+        self.self_attn = self_attention
+        self.cross_attn_image = cross_attention
+
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        self.activation_str = activation
+        self.activation = get_activation_fn(activation)
+        self.pre_norm = pre_norm
+
+        self.pos_enc_at_attn = pos_enc_at_attn
+        self.pos_enc_at_cross_attn_queries = pos_enc_at_cross_attn_queries
+        self.pos_enc_at_cross_attn_keys = pos_enc_at_cross_attn_keys
+
+        self.layer_idx = None
+
+    def forward_post(
+        self,
+        tgt: Tensor,
+        memory: Tensor,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+        query_pos: Optional[Tensor] = None,
+        **kwargs,
+    ) -> Tensor:
+        """
+        Forward pass for post-norm architecture.
+
+        In post-norm architecture, normalization is applied after attention and feedforward operations.
+
+        Args:
+            tgt: Input tensor to be processed
+            memory: Memory tensor for cross-attention
+            tgt_mask: Mask for self-attention
+            memory_mask: Mask for cross-attention
+            tgt_key_padding_mask: Key padding mask for self-attention
+            memory_key_padding_mask: Key padding mask for cross-attention
+            pos: Positional encoding for memory
+            query_pos: Positional encoding for query
+            **kwargs: Additional keyword arguments
+
+        Returns:
+            Processed tensor
+        """
+        q = k = tgt + query_pos if self.pos_enc_at_attn else tgt
+
+        # Self attention
+        tgt2 = self.self_attn(
+            q, k, value=tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask
+        )[0]
+        tgt = tgt + self.dropout1(tgt2)
+        tgt = self.norm1(tgt)
+
+        # Cross attention to image
+        tgt2 = self.cross_attn_image(
+            query=tgt + query_pos if self.pos_enc_at_cross_attn_queries else tgt,
+            key=memory + pos if self.pos_enc_at_cross_attn_keys else memory,
+            value=memory,
+            attn_mask=memory_mask,
+            key_padding_mask=memory_key_padding_mask,
+        )[0]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt = self.norm2(tgt)
+
+        # FFN
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+        tgt = tgt + self.dropout3(tgt2)
+        tgt = self.norm3(tgt)
+        return tgt
+
+    def forward_pre(
+        self,
+        tgt: Tensor,
+        memory: Tensor,
+        dac: bool = False,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+        query_pos: Optional[Tensor] = None,
+        # attn_bias: Optional[Tensor] = None,
+        # **kwargs,
+    ) -> Tensor:
+        """
+        Forward pass for pre-norm architecture.
+
+        In pre-norm architecture, normalization is applied before attention and feedforward operations.
+
+        Args:
+            tgt: Input tensor to be processed
+            memory: Memory tensor for cross-attention
+            dac: Whether to use Divide-and-Conquer attention
+            tgt_mask: Mask for self-attention
+            memory_mask: Mask for cross-attention
+            tgt_key_padding_mask: Key padding mask for self-attention
+            memory_key_padding_mask: Key padding mask for cross-attention
+            pos: Positional encoding for memory
+            query_pos: Positional encoding for query
+            attn_bias: Optional attention bias tensor
+            **kwargs: Additional keyword arguments
+
+        Returns:
+            Processed tensor
+        """
+        if dac:
+            # we only apply self attention to the first half of the queries
+            assert tgt.shape[0] % 2 == 0
+            other_tgt = tgt[tgt.shape[0] // 2 :]
+            tgt = tgt[: tgt.shape[0] // 2]
+        tgt2 = self.norm1(tgt)
+        q = k = tgt2 + query_pos if self.pos_enc_at_attn else tgt2
+        tgt2 = self.self_attn(
+            q, k, value=tgt2, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask
+        )[0]
+        tgt = tgt + self.dropout1(tgt2)
+        if dac:
+            # Recombine
+            tgt = torch.cat((tgt, other_tgt), dim=0)
+        tgt2 = self.norm2(tgt)
+        tgt2 = self.cross_attn_image(
+            query=tgt2 + query_pos if self.pos_enc_at_cross_attn_queries else tgt2,
+            key=memory + pos if self.pos_enc_at_cross_attn_keys else memory,
+            value=memory,
+            attn_mask=memory_mask,
+            key_padding_mask=memory_key_padding_mask,
+            # attn_bias=attn_bias,
+        )[0]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt2 = self.norm3(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
+        tgt = tgt + self.dropout3(tgt2)
+        return tgt
+
+    def forward(
+        self,
+        tgt: Tensor,
+        memory: Tensor,
+        dac: bool = False,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+        query_pos: Optional[Tensor] = None,
+        # attn_bias: Optional[Tensor] = None,
+        # **kwds: Any,
+    ) -> torch.Tensor:
+        """
+        Forward pass for the transformer encoder layer.
+
+        Args:
+            tgt: Input tensor to be processed
+            memory: Memory tensor (e.g., image features) for cross-attention
+            dac: Whether to use Divide-and-Conquer attention (only apply self-attention to first half)
+            tgt_mask: Mask for self-attention
+            memory_mask: Mask for cross-attention
+            tgt_key_padding_mask: Key padding mask for self-attention
+            memory_key_padding_mask: Key padding mask for cross-attention
+            pos: Positional encoding for memory
+            query_pos: Positional encoding for query
+            attn_bias: Optional attention bias tensor
+            **kwds: Additional keyword arguments
+
+        Returns:
+            Processed tensor after self-attention, cross-attention, and feedforward network
+        """
+        fwd_fn = self.forward_pre if self.pre_norm else self.forward_post
+        return fwd_fn(
+            tgt,
+            memory,
+            dac=dac,
+            tgt_mask=tgt_mask,
+            memory_mask=memory_mask,
+            tgt_key_padding_mask=tgt_key_padding_mask,
+            memory_key_padding_mask=memory_key_padding_mask,
+            pos=pos,
+            query_pos=query_pos,
+            # attn_bias=attn_bias,
+            # **kwds,
+        )
+
+
+class TransformerEncoder(nn.Module):
+    """
+    Transformer encoder that processes multi-level features.
+
+    This encoder takes multi-level features (e.g., from a backbone network) and processes
+    them through a stack of transformer encoder layers. It supports features from multiple
+    levels (e.g., different resolutions) and can apply activation checkpointing for memory
+    efficiency during training.
+
+    Args:
+        layer: The encoder layer to be stacked multiple times
+        num_layers: Number of encoder layers to stack
+        d_model: Model dimension/hidden size
+        num_feature_levels: Number of feature levels to process
+        frozen: Whether to freeze the parameters of this module
+        use_act_checkpoint: Whether to use activation checkpointing during training
+    """
+
+    def __init__(
+        self,
+        layer: nn.Module,
+        num_layers: int,
+        d_model: int,
+        num_feature_levels: int,
+        frozen: bool = False,
+        use_act_checkpoint: bool = False,
+    ):
+        super().__init__()
+        self.layers = get_clones(layer, num_layers)
+        self.num_layers = num_layers
+
+        self.num_feature_levels = num_feature_levels
+        self.level_embed = None
+        if num_feature_levels > 1:
+            self.level_embed = nn.Parameter(torch.Tensor(num_feature_levels, d_model))
+
+        if frozen:
+            for p in self.parameters():
+                p.requires_grad_(False)
+
+        self.use_act_checkpoint = use_act_checkpoint
+
+        # assign layer index to each layer so that some layers can decide what to do
+        # based on which layer index they are (e.g. cross attention to memory bank only
+        # in selected layers)
+        for layer_idx, layer in enumerate(self.layers):
+            layer.layer_idx = layer_idx
+
+    @staticmethod
+    def get_reference_points(spatial_shapes, valid_ratios, device):
+        with torch.no_grad():
+            reference_points_list = []
+            for lvl, (H_, W_) in enumerate(spatial_shapes):
+                ref_y, ref_x = torch.meshgrid(
+                    torch.linspace(
+                        0.5, H_ - 0.5, H_, dtype=torch.float32, device=device
+                    ),
+                    torch.linspace(
+                        0.5, W_ - 0.5, W_, dtype=torch.float32, device=device
+                    ),
+                )
+                ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, lvl, 1] * H_)
+                ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, lvl, 0] * W_)
+                ref = torch.stack((ref_x, ref_y), -1)
+                reference_points_list.append(ref)
+            reference_points = torch.cat(reference_points_list, 1)
+            reference_points = reference_points[:, :, None] * valid_ratios[:, None]
+
+        return reference_points
+
+    def _prepare_multilevel_features(self, srcs, masks, pos_embeds):
+        assert (
+            len(srcs) == self.num_feature_levels
+        ), "mismatch between expected and received # of feature levels"
+
+        src_flatten = []
+        mask_flatten = []
+        lvl_pos_embed_flatten = []
+        spatial_shapes = []
+        has_mask = masks is not None and masks[0] is not None
+        for lvl, (src, mask, pos_embed) in enumerate(zip(srcs, masks, pos_embeds)):
+            bs, c, h, w = src.shape
+            spatial_shape = (h, w)
+            spatial_shapes.append(spatial_shape)
+
+            src = src.flatten(2).transpose(1, 2)  # bs, hw, c
+            if has_mask:
+                mask = mask.flatten(1)
+            pos_embed = pos_embed.flatten(2).transpose(1, 2)  # bs, hw, c
+            if self.level_embed is not None:
+                lvl_pos_embed = pos_embed + self.level_embed[lvl].view(1, 1, -1)
+            else:
+                lvl_pos_embed = pos_embed
+            lvl_pos_embed_flatten.append(lvl_pos_embed)
+            src_flatten.append(src)
+            if has_mask:
+                mask_flatten.append(mask)
+        src_flatten = torch.cat(src_flatten, 1)  # bs, \sum{hxw}, c
+        mask_flatten = torch.cat(mask_flatten, 1) if has_mask else None  # bs, \sum{hxw}
+        lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1)  # bs, \sum{hxw}, c
+        spatial_shapes = torch.tensor(
+            spatial_shapes, dtype=torch.long, device=src_flatten.device
+        )
+        level_start_index = torch.cat(
+            (
+                spatial_shapes.new_zeros((1,)),
+                spatial_shapes.prod(1).cumsum(0)[:-1],
+            )
+        )
+        if has_mask:
+            valid_ratios = torch.stack([get_valid_ratio(m) for m in masks], 1)
+        else:
+            valid_ratios = torch.ones(
+                (src_flatten.shape[0], self.num_feature_levels, 2),
+                device=src_flatten.device,
+            )
+
+        return (
+            src_flatten,
+            mask_flatten,
+            lvl_pos_embed_flatten,
+            level_start_index,
+            valid_ratios,
+            spatial_shapes,
+        )
+
+    def forward(
+        self,
+        src: List[Tensor],
+        src_key_padding_masks: Optional[List[Tensor]] = None,
+        pos: Optional[List[Tensor]] = None,
+        prompt: Optional[Tensor] = None,
+        prompt_key_padding_mask: Optional[Tensor] = None,
+        encoder_extra_kwargs: Optional[Dict] = None,
+    ) -> Tuple[Tensor, Optional[Tensor], Tensor, Tensor, Tensor, Tensor]:
+        """
+        Process multi-level features through the transformer encoder.
+
+        Args:
+            src: List of multi-level features, each with shape (batch_size, channels, height, width)
+            src_key_padding_masks: List of padding masks for each feature level, each with shape (batch_size, height, width)
+            pos: List of positional embeddings for each feature level, each with shape (batch_size, channels, height, width)
+            prompt: Optional text/prompt features to attend to, with shape (seq_len, batch_size, d_model)
+            prompt_key_padding_mask: Optional padding mask for prompt, with shape (batch_size, seq_len)
+            encoder_extra_kwargs: Optional additional arguments to pass to each encoder layer
+
+        Returns:
+            A tuple containing:
+            - output: Processed features with shape (seq_len, batch_size, d_model)
+            - key_padding_masks_flatten: Flattened padding masks
+            - lvl_pos_embed_flatten: Flattened positional embeddings
+            - level_start_index: Starting indices for each feature level
+            - spatial_shapes: Spatial dimensions of each feature level
+            - valid_ratios: Valid ratios for each feature level
+        """
+        assert (
+            len(src) == self.num_feature_levels
+        ), "must be equal to num_feature_levels"
+        if src_key_padding_masks is not None:
+            assert len(src_key_padding_masks) == self.num_feature_levels
+        if pos is not None:
+            assert len(pos) == self.num_feature_levels
+        # Flatten multilevel feats and add level pos embeds
+        (
+            src_flatten,
+            key_padding_masks_flatten,
+            lvl_pos_embed_flatten,
+            level_start_index,
+            valid_ratios,
+            spatial_shapes,
+        ) = self._prepare_multilevel_features(src, src_key_padding_masks, pos)
+
+        reference_points = self.get_reference_points(
+            spatial_shapes, valid_ratios, device=src_flatten.device
+        )
+
+        output = src_flatten
+        for layer in self.layers:
+            layer_kwargs = {}
+
+            assert isinstance(layer, TransformerEncoderLayer)
+            layer_kwargs["memory"] = prompt
+            layer_kwargs["memory_key_padding_mask"] = prompt_key_padding_mask
+            layer_kwargs["query_pos"] = lvl_pos_embed_flatten
+            layer_kwargs["tgt"] = output
+            layer_kwargs["tgt_key_padding_mask"] = key_padding_masks_flatten
+
+            if self.training:
+                assert self.use_act_checkpoint, "activation ckpt not enabled in encoder"
+            if encoder_extra_kwargs is not None:
+                layer_kwargs.update(encoder_extra_kwargs)
+            output = activation_ckpt_wrapper(layer)(
+                **layer_kwargs,
+                act_ckpt_enable=self.training and self.use_act_checkpoint,
+            )
+        # return as seq first
+        return (
+            output.transpose(0, 1),
+            (
+                key_padding_masks_flatten.transpose(0, 1)
+                if key_padding_masks_flatten is not None
+                else None
+            ),
+            lvl_pos_embed_flatten.transpose(0, 1),
+            level_start_index,
+            spatial_shapes,
+            valid_ratios,
+        )
+
+
+class TransformerEncoderFusion(TransformerEncoder):
+    """
+    Transformer encoder that fuses text and image features.
+
+    This encoder extends TransformerEncoder to handle both text and image features,
+    with the ability to add pooled text features to image features for better
+    cross-modal fusion. It supports torch.compile for performance optimization.
+
+    Args:
+        layer: The encoder layer to be stacked multiple times
+        num_layers: Number of encoder layers to stack
+        d_model: Model dimension/hidden size
+        num_feature_levels: Number of feature levels to process
+        add_pooled_text_to_img_feat: Whether to add pooled text features to image features
+        pool_text_with_mask: Whether to use the mask when pooling text features
+        compile_mode: Mode for torch.compile, or None to disable compilation
+        **kwargs: Additional arguments to pass to the parent class
+    """
+
+    def __init__(
+        self,
+        layer: nn.Module,
+        num_layers: int,
+        d_model: int,
+        num_feature_levels: int,
+        add_pooled_text_to_img_feat: bool = True,
+        pool_text_with_mask: bool = False,
+        compile_mode: Optional[str] = None,
+        **kwargs,
+    ):
+        super().__init__(
+            layer,
+            num_layers,
+            d_model,
+            num_feature_levels,
+            **kwargs,
+        )
+        self.add_pooled_text_to_img_feat = add_pooled_text_to_img_feat
+        if self.add_pooled_text_to_img_feat:
+            self.text_pooling_proj = nn.Linear(d_model, d_model)
+        self.pool_text_with_mask = pool_text_with_mask
+        if compile_mode is not None:
+            self.forward = torch.compile(
+                self.forward, mode=compile_mode, fullgraph=True
+            )
+
+    @staticmethod
+    def get_reference_points(spatial_shapes, valid_ratios, device):
+        # Not needed here
+        return None
+
+    def forward(
+        self,
+        src: List[Tensor],
+        prompt: Tensor,
+        src_key_padding_mask: Optional[List[Tensor]] = None,
+        src_pos: Optional[List[Tensor]] = None,
+        prompt_key_padding_mask: Optional[Tensor] = None,
+        prompt_pos: Optional[Tensor] = None,
+        feat_sizes: Optional[List[int]] = None,
+        encoder_extra_kwargs: Optional[Dict] = None,
+    ):
+        # Restore spatial shapes of vision
+        bs = src[0].shape[1]  # seq first
+        if feat_sizes is not None:
+            assert len(feat_sizes) == len(src)
+            if src_key_padding_mask is None:
+                src_key_padding_mask = [None] * len(src)
+            for i, (h, w) in enumerate(feat_sizes):
+                src[i] = src[i].reshape(h, w, bs, -1).permute(2, 3, 0, 1)
+                src_pos[i] = src_pos[i].reshape(h, w, bs, -1).permute(2, 3, 0, 1)
+                src_key_padding_mask[i] = (
+                    src_key_padding_mask[i].reshape(h, w, bs).permute(2, 0, 1)
+                    if src_key_padding_mask[i] is not None
+                    else None
+                )
+        else:
+            assert all(
+                x.dim == 4 for x in src
+            ), "expected list of (bs, c, h, w) tensors"
+
+        if self.add_pooled_text_to_img_feat:
+            # Fusion: Add mean pooled text to image features
+            pooled_text = pool_text_feat(
+                prompt, prompt_key_padding_mask, self.pool_text_with_mask
+            )
+            pooled_text = self.text_pooling_proj(pooled_text)[
+                ..., None, None
+            ]  # prompt is seq first
+            src = [x.add_(pooled_text) for x in src]
+
+        (
+            out,
+            key_padding_masks_flatten,
+            lvl_pos_embed_flatten,
+            level_start_index,
+            spatial_shapes,
+            valid_ratios,
+        ) = super().forward(
+            src,
+            src_key_padding_masks=src_key_padding_mask,
+            pos=src_pos,
+            prompt=prompt.transpose(0, 1),
+            prompt_key_padding_mask=prompt_key_padding_mask,
+            encoder_extra_kwargs=encoder_extra_kwargs,
+        )
+
+        return {
+            "memory": out,
+            "padding_mask": key_padding_masks_flatten,
+            "pos_embed": lvl_pos_embed_flatten,
+            "memory_text": prompt,
+            "level_start_index": level_start_index,
+            "spatial_shapes": spatial_shapes,
+            "valid_ratios": valid_ratios,
+        }
+
+
+def pool_text_feat(prompt, prompt_mask, pool_with_mask):
+    # prompt has shape (seq, bs, dim)
+    if not pool_with_mask:
+        return prompt.mean(dim=0)
+
+    # prompt_mask has shape (bs, seq), where False is valid and True is padding
+    assert prompt_mask.dim() == 2
+    # is_valid has shape (seq, bs, 1), where 1 is valid and 0 is padding
+    is_valid = (~prompt_mask).float().permute(1, 0)[..., None]
+    # num_valid has shape (bs, 1)
+    num_valid = torch.clamp(torch.sum(is_valid, dim=0), min=1.0)
+
+    # mean pool over all the valid tokens
+    pooled_text = (prompt * is_valid).sum(dim=0) / num_valid
+    return pooled_text

sam3/model/geometry_encoders.py

@@ -0,0 +1,850 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+import torchvision
+from typing_extensions import override
+
+from .act_ckpt_utils import activation_ckpt_wrapper
+from .box_ops import box_cxcywh_to_xyxy
+
+from .model_misc import get_clones
+
+
+def is_right_padded(mask):
+    """Given a padding mask (following pytorch convention, 1s for padded values),
+    returns whether the padding is on the right or not."""
+    return (mask.long() == torch.sort(mask.long(), dim=-1)[0]).all()
+
+
+def concat_padded_sequences(seq1, mask1, seq2, mask2, return_index: bool = False):
+    """
+    Concatenates two right-padded sequences, such that the resulting sequence
+    is contiguous and also right-padded.
+
+    Following pytorch's convention, tensors are sequence first, and the mask are
+    batch first, with 1s for padded values.
+
+    :param seq1: A tensor of shape (seq1_length, batch_size, hidden_size).
+    :param mask1: A tensor of shape (batch_size, seq1_length).
+    :param seq2: A tensor of shape (seq2_length, batch_size,  hidden_size).
+    :param mask2: A tensor of shape (batch_size, seq2_length).
+    :param return_index: If True, also returns the index of the ids of the element of seq2
+        in the concatenated sequence. This can be used to retrieve the elements of seq2
+    :return: A tuple (concatenated_sequence, concatenated_mask) if return_index is False,
+        otherwise (concatenated_sequence, concatenated_mask, index).
+    """
+    seq1_length, batch_size, hidden_size = seq1.shape
+    seq2_length, batch_size, hidden_size = seq2.shape
+
+    assert batch_size == seq1.size(1) == seq2.size(1) == mask1.size(0) == mask2.size(0)
+    assert hidden_size == seq1.size(2) == seq2.size(2)
+    assert seq1_length == mask1.size(1)
+    assert seq2_length == mask2.size(1)
+
+    torch._assert_async(is_right_padded(mask1))
+    torch._assert_async(is_right_padded(mask2))
+
+    actual_seq1_lengths = (~mask1).sum(dim=-1)
+    actual_seq2_lengths = (~mask2).sum(dim=-1)
+
+    final_lengths = actual_seq1_lengths + actual_seq2_lengths
+    max_length = seq1_length + seq2_length
+    concatenated_mask = (
+        torch.arange(max_length, device=seq2.device)[None].repeat(batch_size, 1)
+        >= final_lengths[:, None]
+    )
+
+    # (max_len, batch_size, hidden_size)
+    concatenated_sequence = torch.zeros(
+        (max_length, batch_size, hidden_size), device=seq2.device, dtype=seq2.dtype
+    )
+    concatenated_sequence[:seq1_length, :, :] = seq1
+
+    # At this point, the element of seq1 are in the right place
+    # We just need to shift the elements of seq2
+
+    index = torch.arange(seq2_length, device=seq2.device)[:, None].repeat(1, batch_size)
+    index = index + actual_seq1_lengths[None]
+
+    concatenated_sequence = concatenated_sequence.scatter(
+        0, index[:, :, None].expand(-1, -1, hidden_size), seq2
+    )
+
+    if return_index:
+        return concatenated_sequence, concatenated_mask, index
+
+    return concatenated_sequence, concatenated_mask
+
+
+class Prompt:
+    """Utility class to manipulate geometric prompts.
+
+    We expect the sequences in pytorch convention, that is sequence first, batch second
+    The dimensions are expected as follows:
+    box_embeddings shape: N_boxes x B x C_box
+    box_mask shape: B x N_boxes. Can be None if nothing is masked out
+    point_embeddings shape: N_points x B x C_point
+    point_mask shape: B x N_points. Can be None if nothing is masked out
+    mask_embeddings shape: N_masks x B x 1 x H_mask x W_mask
+    mask_mask shape: B x N_masks. Can be None if nothing is masked out
+
+    We also store positive/negative labels. These tensors are also stored batch-first
+    If they are None, we'll assume positive labels everywhere
+    box_labels: long tensor of shape N_boxes x B
+    point_labels: long tensor of shape N_points x B
+    mask_labels: long tensor of shape N_masks x B
+    """
+
+    def __init__(
+        self,
+        box_embeddings=None,
+        box_mask=None,
+        point_embeddings=None,
+        point_mask=None,
+        box_labels=None,
+        point_labels=None,
+        mask_embeddings=None,
+        mask_mask=None,  # Attention mask for mask prompt
+        mask_labels=None,
+    ):
+        # Check for null prompt
+        if (
+            box_embeddings is None
+            and point_embeddings is None
+            and mask_embeddings is None
+        ):
+            self.box_embeddings = None
+            self.box_labels = None
+            self.box_mask = None
+            self.point_embeddings = None
+            self.point_labels = None
+            self.point_mask = None
+            self.mask_embeddings = None
+            self.mask_mask = None
+            # Masks are assumed positive only for now.
+            self.mask_labels = None
+            return
+        # Get sequence lengths and device
+        box_seq_len, point_seq_len, mask_seq_len, bs, device = (
+            self._init_seq_len_and_device(
+                box_embeddings, point_embeddings, mask_embeddings
+            )
+        )
+
+        # Initialize embeds, labels, attention masks.
+        box_embeddings, box_labels, box_mask = self._init_box(
+            box_embeddings, box_labels, box_mask, box_seq_len, bs, device
+        )
+        point_embeddings, point_labels, point_mask = self._init_point(
+            point_embeddings, point_labels, point_mask, point_seq_len, bs, device
+        )
+        mask_embeddings, mask_labels, mask_mask = self._init_mask(
+            mask_embeddings, mask_labels, mask_mask, mask_seq_len, bs, device
+        )
+
+        # Dimension checks
+        assert (
+            box_embeddings is not None
+            and list(box_embeddings.shape[:2])
+            == [
+                box_seq_len,
+                bs,
+            ]
+        ), f"Wrong dimension for box embeddings. Expected [{box_seq_len}, {bs}, *] got {box_embeddings.shape}"
+        assert (
+            box_mask is not None
+            and list(box_mask.shape)
+            == [
+                bs,
+                box_seq_len,
+            ]
+        ), f"Wrong dimension for box mask. Expected [{bs}, {box_seq_len}] got {box_mask.shape}"
+        assert (
+            point_embeddings is not None
+            and list(point_embeddings.shape[:2])
+            == [
+                point_seq_len,
+                bs,
+            ]
+        ), f"Wrong dimension for point embeddings. Expected [{point_seq_len}, {bs}, *] got {point_embeddings.shape}"
+        assert (
+            point_mask is not None
+            and list(point_mask.shape)
+            == [
+                bs,
+                point_seq_len,
+            ]
+        ), f"Wrong dimension for point mask. Expected [{bs}, {point_seq_len}] got {point_mask.shape}"
+        assert (
+            box_labels is not None
+            and list(box_labels.shape)
+            == [
+                box_seq_len,
+                bs,
+            ]
+        ), f"Wrong dimension for box labels. Expected [{box_seq_len}, {bs}] got {box_labels.shape}"
+        assert (
+            point_labels is not None
+            and list(point_labels.shape)
+            == [
+                point_seq_len,
+                bs,
+            ]
+        ), f"Wrong dimension for point labels. Expected [{point_seq_len}, {bs}] got {point_labels.shape}"
+        assert (
+            # Allowed to be None, we leave it to the encoder to check for validity before encoding.
+            mask_embeddings is None
+            or list(mask_embeddings.shape[:2])
+            == [
+                mask_seq_len,
+                bs,
+            ]
+        ), f"Wrong dimension for mask embeddings. Expected [{mask_seq_len}, {bs}, *] got {mask_embeddings.shape}"
+        assert (
+            mask_mask is None
+            or list(mask_mask.shape)
+            == [
+                bs,
+                mask_seq_len,
+            ]
+        ), f"Wrong dimension for mask attn. mask. Expected [{bs}, {mask_seq_len}] got {mask_mask.shape}"
+
+        # Device checks
+        assert (
+            box_embeddings is not None and box_embeddings.device == device
+        ), f"Expected box embeddings to be on device {device}, got {box_embeddings.device}"
+        assert (
+            box_mask is not None and box_mask.device == device
+        ), f"Expected box mask to be on device {device}, got {box_mask.device}"
+        assert (
+            box_labels is not None and box_labels.device == device
+        ), f"Expected box labels to be on device {device}, got {box_labels.device}"
+        assert (
+            point_embeddings is not None and point_embeddings.device == device
+        ), f"Expected point embeddings to be on device {device}, got {point_embeddings.device}"
+        assert (
+            point_mask is not None and point_mask.device == device
+        ), f"Expected point mask to be on device {device}, got {point_mask.device}"
+        assert (
+            point_labels is not None and point_labels.device == device
+        ), f"Expected point labels to be on device {device}, got {point_labels.device}"
+        assert (
+            mask_embeddings is None or mask_embeddings.device == device
+        ), f"Expected mask embeddings to be on device {device}, got {mask_embeddings.device}"
+        assert (
+            mask_mask is None or mask_mask.device == device
+        ), f"Expected mask attn. mask to be on device {device}, got {mask_mask.device}"
+
+        self.box_embeddings = box_embeddings
+        self.point_embeddings = point_embeddings
+        self.box_mask = box_mask
+        self.point_mask = point_mask
+        self.box_labels = box_labels
+        self.point_labels = point_labels
+        self.mask_embeddings = mask_embeddings
+        self.mask_labels = mask_labels
+        self.mask_mask = mask_mask
+
+    def _init_seq_len_and_device(
+        self, box_embeddings, point_embeddings, mask_embeddings
+    ):
+        box_seq_len = point_seq_len = mask_seq_len = 0
+        bs = None
+        device = None
+        if box_embeddings is not None:
+            bs = box_embeddings.shape[1]
+            box_seq_len = box_embeddings.shape[0]
+            device = box_embeddings.device
+
+        if point_embeddings is not None:
+            point_seq_len = point_embeddings.shape[0]
+            if bs is not None:
+                assert (
+                    bs == point_embeddings.shape[1]
+                ), f"Batch size mismatch between box and point embeddings. Got {bs} and {point_embeddings.shape[1]}."
+            else:
+                bs = point_embeddings.shape[1]
+            if device is not None:
+                assert (
+                    device == point_embeddings.device
+                ), "Device mismatch between box and point embeddings"
+            else:
+                device = point_embeddings.device
+
+        if mask_embeddings is not None:
+            mask_seq_len = mask_embeddings.shape[0]
+            if bs is not None:
+                assert (
+                    bs == mask_embeddings.shape[1]
+                ), f"Batch size mismatch between box/point and mask embedding. Got {bs} and {mask_embeddings.shape[1]}"
+            else:
+                bs = mask_embeddings.shape[1]
+            if device is not None:
+                assert (
+                    device == mask_embeddings.device
+                ), "Device mismatch between box/point and mask embeddings."
+            else:
+                device = mask_embeddings.device
+
+        return box_seq_len, point_seq_len, mask_seq_len, bs, device
+
+    def _init_box(self, box_embeddings, box_labels, box_mask, box_seq_len, bs, device):
+        if box_embeddings is None:
+            box_embeddings = torch.zeros(box_seq_len, bs, 4, device=device)
+        if box_labels is None:
+            box_labels = torch.ones(box_seq_len, bs, device=device, dtype=torch.long)
+        if box_mask is None:
+            box_mask = torch.zeros(bs, box_seq_len, device=device, dtype=torch.bool)
+        return box_embeddings, box_labels, box_mask
+
+    def _init_point(
+        self, point_embeddings, point_labels, point_mask, point_seq_len, bs, device
+    ):
+        """
+        Identical to _init_box. Except that C=2 for points (vs. 4 for boxes).
+        """
+        if point_embeddings is None:
+            point_embeddings = torch.zeros(point_seq_len, bs, 2, device=device)
+        if point_labels is None:
+            point_labels = torch.ones(
+                point_seq_len, bs, device=device, dtype=torch.long
+            )
+        if point_mask is None:
+            point_mask = torch.zeros(bs, point_seq_len, device=device, dtype=torch.bool)
+        return point_embeddings, point_labels, point_mask
+
+    def _init_mask(
+        self, mask_embeddings, mask_labels, mask_mask, mask_seq_len, bs, device
+    ):
+        # NOTE: Mask embeddings can be of arbitrary resolution, so we don't initialize it here.
+        # In case we append new mask, we check that its resolution matches exisiting ones (if any).
+        # In case mask_embeddings is None, we should never encode it.
+        if mask_labels is None:
+            mask_labels = torch.ones(mask_seq_len, bs, device=device, dtype=torch.long)
+        if mask_mask is None:
+            mask_mask = torch.zeros(bs, mask_seq_len, device=device, dtype=torch.bool)
+        return mask_embeddings, mask_labels, mask_mask
+
+    def append_boxes(self, boxes, labels, mask=None):
+        if self.box_embeddings is None:
+            self.box_embeddings = boxes
+            self.box_labels = labels
+            self.box_mask = mask
+            return
+
+        bs = self.box_embeddings.shape[1]
+        assert boxes.shape[1] == labels.shape[1] == bs
+        assert list(boxes.shape[:2]) == list(labels.shape[:2])
+        if mask is None:
+            mask = torch.zeros(
+                bs, boxes.shape[0], dtype=torch.bool, device=boxes.device
+            )
+
+        self.box_labels, _ = concat_padded_sequences(
+            self.box_labels.unsqueeze(-1), self.box_mask, labels.unsqueeze(-1), mask
+        )
+        self.box_labels = self.box_labels.squeeze(-1)
+        self.box_embeddings, self.box_mask = concat_padded_sequences(
+            self.box_embeddings, self.box_mask, boxes, mask
+        )
+
+    def append_points(self, points, labels, mask=None):
+        if self.point_embeddings is None:
+            self.point_embeddings = points
+            self.point_labels = labels
+            self.point_mask = mask
+            return
+
+        bs = self.point_embeddings.shape[1]
+        assert points.shape[1] == labels.shape[1] == bs
+        assert list(points.shape[:2]) == list(labels.shape[:2])
+        if mask is None:
+            mask = torch.zeros(
+                bs, points.shape[0], dtype=torch.bool, device=points.device
+            )
+
+        self.point_labels, _ = concat_padded_sequences(
+            self.point_labels.unsqueeze(-1), self.point_mask, labels.unsqueeze(-1), mask
+        )
+        self.point_labels = self.point_labels.squeeze(-1)
+        self.point_embeddings, self.point_mask = concat_padded_sequences(
+            self.point_embeddings, self.point_mask, points, mask
+        )
+
+    def append_masks(self, masks, labels=None, attn_mask=None):
+        if labels is not None:
+            assert list(masks.shape[:2]) == list(labels.shape[:2])
+        if self.mask_embeddings is None:
+            self.mask_embeddings = masks
+            mask_seq_len, bs = masks.shape[:2]
+            if labels is None:
+                self.mask_labels = torch.ones(
+                    mask_seq_len, bs, device=masks.device, dtype=torch.long
+                )
+            else:
+                self.mask_labels = labels
+            if attn_mask is None:
+                self.mask_mask = torch.zeros(
+                    bs, mask_seq_len, device=masks.device, dtype=torch.bool
+                )
+            else:
+                self.mask_mask = attn_mask
+        else:
+            raise NotImplementedError("Only one mask per prompt is supported.")
+
+    def clone(self):
+        return Prompt(
+            box_embeddings=(
+                None if self.box_embeddings is None else self.box_embeddings.clone()
+            ),
+            box_mask=None if self.box_mask is None else self.box_mask.clone(),
+            point_embeddings=(
+                None if self.point_embeddings is None else self.point_embeddings.clone()
+            ),
+            point_mask=None if self.point_mask is None else self.point_mask.clone(),
+            box_labels=None if self.box_labels is None else self.box_labels.clone(),
+            point_labels=(
+                None if self.point_labels is None else self.point_labels.clone()
+            ),
+        )
+
+
+class MaskEncoder(nn.Module):
+    """
+    Base class for mask encoders.
+    """
+
+    def __init__(
+        self,
+        mask_downsampler: nn.Module,
+        position_encoding: nn.Module,
+    ):
+        super().__init__()
+        self.mask_downsampler = mask_downsampler
+        self.position_encoding = position_encoding
+
+    def forward(self, masks, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
+        masks = self.mask_downsampler(masks)
+        masks_pos = self.position_encoding(masks).to(masks.dtype)
+
+        return masks, masks_pos
+
+
+class FusedMaskEncoder(MaskEncoder):
+    """
+    Identical to memory.SimpleMaskEncoder but follows the interface of geometry_encoders.MaskEncoder.
+    We also remove the `skip_mask_sigmoid` option (to be handled outside the MaskEncoder).
+    Fuses backbone image features with mask features.
+    """
+
+    def __init__(
+        self,
+        mask_downsampler: nn.Module,
+        position_encoding: nn.Module,
+        fuser: nn.Module,
+        in_dim: int = 256,
+        out_dim: int = 256,
+    ):
+        super().__init__(mask_downsampler, position_encoding)
+        self.fuser = fuser
+        self.out_proj = nn.Identity()
+        if out_dim != in_dim:
+            self.out_proj = nn.Conv2d(in_dim, out_dim, kernel_size=1)
+        self.pix_feat_proj = nn.Conv2d(in_dim, in_dim, kernel_size=1)
+
+    @override
+    def forward(
+        self,
+        masks: torch.Tensor,
+        pix_feat: torch.Tensor,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        masks = self.mask_downsampler(masks)
+
+        ## Fuse pix_feats and downsampled masks
+        # in case the visual features are on CPU, cast them to CUDA
+        pix_feat = pix_feat.to(masks.device)
+
+        x = self.pix_feat_proj(pix_feat)
+        x = x + masks
+        x = self.fuser(x)
+        x = self.out_proj(x)
+
+        pos = self.position_encoding(x).to(x.dtype)
+
+        return x, pos
+
+
+class SequenceGeometryEncoder(nn.Module):
+    """
+    This a fully fledged encoder for geometric prompts.
+    It assumes boxes are passed in the "normalized CxCyWH" format, and points in normalized xy
+    This allows flexibility in how to encode the features (eg do pooling)
+
+    Points and boxes can be encoded with any of the three possibilities:
+     - direct projection: we just compute a linear from coordinate space to d_model
+     - pooling: pool features from the backbone in the requested location.
+                For boxes, it's a roi align
+                For points it's a grid sample
+     - pos encoder: Take the position encoding of the point or box center
+
+    These three options are mutually compatible. If several are selected, we'll take a simple addition
+
+    As an alternative, we offer the possibility to encode points only.
+    In that case, the boxes are converted to two points for the top left and bottom right corners (with appropriate labels)
+
+    On top of these encodings, we offer the possibility to further encode the prompt sequence with a transformer.
+    """
+
+    def __init__(
+        self,
+        encode_boxes_as_points: bool,
+        points_direct_project: bool,
+        points_pool: bool,
+        points_pos_enc: bool,
+        boxes_direct_project: bool,
+        boxes_pool: bool,
+        boxes_pos_enc: bool,
+        d_model: int,
+        pos_enc,
+        num_layers: int,
+        layer: nn.Module,
+        roi_size: int = 7,  # for boxes pool
+        add_cls: bool = True,
+        add_post_encode_proj: bool = True,
+        mask_encoder: MaskEncoder = None,
+        add_mask_label: bool = False,
+        use_act_ckpt: bool = False,
+    ):
+        super().__init__()
+
+        self.d_model = d_model
+        self.pos_enc = pos_enc
+        self.encode_boxes_as_points = encode_boxes_as_points
+        self.roi_size = roi_size
+        # There usually are two labels: positive and negatives.
+        # If we encode boxes as points, we have 3 types of points: regular, top left, bottom right
+        # These 3 types can be positives or negatives, hence 2*3 = 6 labels
+        num_labels = 6 if self.encode_boxes_as_points else 2
+        self.label_embed = torch.nn.Embedding(num_labels, self.d_model)
+
+        # This is a cls token, can be used for pooling if need be.
+        # It also ensures that the encoded sequences are always non-empty
+        self.cls_embed = None
+        if add_cls:
+            self.cls_embed = torch.nn.Embedding(1, self.d_model)
+
+        assert (
+            points_direct_project or points_pos_enc or points_pool
+        ), "Error: need at least one way to encode points"
+        assert (
+            encode_boxes_as_points
+            or boxes_direct_project
+            or boxes_pos_enc
+            or boxes_pool
+        ), "Error: need at least one way to encode boxes"
+
+        self.points_direct_project = None
+        if points_direct_project:
+            self.points_direct_project = nn.Linear(2, self.d_model)
+        self.points_pool_project = None
+        if points_pool:
+            self.points_pool_project = nn.Linear(self.d_model, self.d_model)
+        self.points_pos_enc_project = None
+        if points_pos_enc:
+            self.points_pos_enc_project = nn.Linear(self.d_model, self.d_model)
+
+        self.boxes_direct_project = None
+        self.boxes_pool_project = None
+        self.boxes_pos_enc_project = None
+        if not encode_boxes_as_points:
+            if boxes_direct_project:
+                self.boxes_direct_project = nn.Linear(4, self.d_model)
+            if boxes_pool:
+                self.boxes_pool_project = nn.Conv2d(
+                    self.d_model, self.d_model, self.roi_size
+                )
+            if boxes_pos_enc:
+                self.boxes_pos_enc_project = nn.Linear(self.d_model + 2, self.d_model)
+
+        self.final_proj = None
+        if add_post_encode_proj:
+            self.final_proj = nn.Linear(self.d_model, self.d_model)
+            self.norm = nn.LayerNorm(self.d_model)
+
+        self.img_pre_norm = nn.Identity()
+        if self.points_pool_project is not None or self.boxes_pool_project is not None:
+            self.img_pre_norm = nn.LayerNorm(self.d_model)
+
+        self.encode = None
+        if num_layers > 0:
+            assert (
+                add_cls
+            ), "It's currently highly recommended to add a CLS when using a transformer"
+            self.encode = get_clones(layer, num_layers)
+            self.encode_norm = nn.LayerNorm(self.d_model)
+
+        if mask_encoder is not None:
+            assert isinstance(
+                mask_encoder, MaskEncoder
+            ), f"Expected mask_encoder of type MaskEncoder. Got {type(mask_encoder)}."
+            if add_mask_label:
+                self.mask_label_embed = torch.nn.Embedding(2, self.d_model)
+        self.add_mask_label = add_mask_label
+        self.mask_encoder = mask_encoder
+        self.use_act_ckpt = use_act_ckpt
+
+    def _encode_points(self, points, points_mask, points_labels, img_feats):
+        points_embed = None
+        n_points, bs = points.shape[:2]
+
+        if self.points_direct_project is not None:
+            proj = self.points_direct_project(points)
+            assert points_embed is None
+            points_embed = proj
+
+        if self.points_pool_project is not None:
+            # points are [Num_points, bs, 2], normalized in [0, 1]
+            # the grid needs to be [Bs, H_out, W_out, 2] normalized in [-1,1]
+            # Will take H_out = num_points, w_out = 1
+            grid = points.transpose(0, 1).unsqueeze(2)
+            # re normalize to [-1, 1]
+            grid = (grid * 2) - 1
+            sampled = torch.nn.functional.grid_sample(
+                img_feats, grid, align_corners=False
+            )
+            assert list(sampled.shape) == [bs, self.d_model, n_points, 1]
+            sampled = sampled.squeeze(-1).permute(2, 0, 1)
+            proj = self.points_pool_project(sampled)
+            if points_embed is None:
+                points_embed = proj
+            else:
+                points_embed = points_embed + proj
+
+        if self.points_pos_enc_project is not None:
+            x, y = points.unbind(-1)
+            enc_x, enc_y = self.pos_enc._encode_xy(x.flatten(), y.flatten())
+            enc_x = enc_x.view(n_points, bs, enc_x.shape[-1])
+            enc_y = enc_y.view(n_points, bs, enc_y.shape[-1])
+            enc = torch.cat([enc_x, enc_y], -1)
+
+            proj = self.points_pos_enc_project(enc)
+            if points_embed is None:
+                points_embed = proj
+            else:
+                points_embed = points_embed + proj
+
+        type_embed = self.label_embed(points_labels.long())
+        return type_embed + points_embed, points_mask
+
+    def _encode_boxes(self, boxes, boxes_mask, boxes_labels, img_feats):
+        boxes_embed = None
+        n_boxes, bs = boxes.shape[:2]
+
+        if self.boxes_direct_project is not None:
+            proj = self.boxes_direct_project(boxes)
+            assert boxes_embed is None
+            boxes_embed = proj
+
+        if self.boxes_pool_project is not None:
+            H, W = img_feats.shape[-2:]
+
+            # boxes are [Num_boxes, bs, 4], normalized in [0, 1]
+            # We need to denormalize, and convert to [x, y, x, y]
+            boxes_xyxy = box_cxcywh_to_xyxy(boxes)
+            scale = torch.tensor([W, H, W, H], dtype=boxes_xyxy.dtype)
+            scale = scale.pin_memory().to(device=boxes_xyxy.device, non_blocking=True)
+            scale = scale.view(1, 1, 4)
+            boxes_xyxy = boxes_xyxy * scale
+            sampled = torchvision.ops.roi_align(
+                img_feats, boxes_xyxy.float().transpose(0, 1).unbind(0), self.roi_size
+            )
+            assert list(sampled.shape) == [
+                bs * n_boxes,
+                self.d_model,
+                self.roi_size,
+                self.roi_size,
+            ]
+            proj = self.boxes_pool_project(sampled)
+            proj = proj.view(bs, n_boxes, self.d_model).transpose(0, 1)
+            if boxes_embed is None:
+                boxes_embed = proj
+            else:
+                boxes_embed = boxes_embed + proj
+
+        if self.boxes_pos_enc_project is not None:
+            cx, cy, w, h = boxes.unbind(-1)
+            enc = self.pos_enc.encode_boxes(
+                cx.flatten(), cy.flatten(), w.flatten(), h.flatten()
+            )
+            enc = enc.view(boxes.shape[0], boxes.shape[1], enc.shape[-1])
+
+            proj = self.boxes_pos_enc_project(enc)
+            if boxes_embed is None:
+                boxes_embed = proj
+            else:
+                boxes_embed = boxes_embed + proj
+
+        type_embed = self.label_embed(boxes_labels.long())
+        return type_embed + boxes_embed, boxes_mask
+
+    def _encode_masks(
+        self,
+        masks: torch.Tensor,
+        attn_mask: torch.Tensor,
+        mask_labels: torch.Tensor,
+        img_feats: torch.Tensor = None,
+    ):
+        n_masks, bs = masks.shape[:2]
+        assert (
+            n_masks == 1
+        ), "We assume one mask per prompt for now. Code should still be functional if this assertion is removed."
+        assert (
+            list(attn_mask.shape)
+            == [
+                bs,
+                n_masks,
+            ]
+        ), f"Expected attn_mask to be of shape {bs}x{n_masks}. Got {list(attn_mask.shape)}."
+        masks, pos = self.mask_encoder(
+            masks=masks.flatten(0, 1).float(),
+            pix_feat=img_feats,
+        )
+        H, W = masks.shape[-2:]
+        n_tokens_per_mask = H * W
+        # NOTE: We directly add pos enc here as we usually don't keep track of pos encoding for the concatenated prompt (text, other geometric prompts). Might need to do some refactoring for more flexibility.
+        masks = masks + pos
+        masks = masks.view(n_masks, bs, *masks.shape[1:]).flatten(
+            -2
+        )  # n_masks x bs x C x H*W
+        masks = masks.permute(0, 3, 1, 2).flatten(0, 1)  # n_masks * H*W x bs x C
+        attn_mask = attn_mask.repeat_interleave(n_tokens_per_mask, dim=1)
+        if self.add_mask_label:
+            masks = masks + self.mask_label_embed(mask_labels.long())
+        return masks, attn_mask
+
+    def forward(self, geo_prompt: Prompt, img_feats, img_sizes, img_pos_embeds=None):
+        points = geo_prompt.point_embeddings
+        points_mask = geo_prompt.point_mask
+        points_labels = geo_prompt.point_labels
+        boxes = geo_prompt.box_embeddings
+        boxes_mask = geo_prompt.box_mask
+        boxes_labels = geo_prompt.box_labels
+        masks = geo_prompt.mask_embeddings
+        masks_mask = geo_prompt.mask_mask
+        masks_labels = geo_prompt.mask_labels
+        seq_first_img_feats = img_feats[-1]  # [H*W, B, C]
+        seq_first_img_pos_embeds = (
+            img_pos_embeds[-1]
+            if img_pos_embeds is not None
+            else torch.zeros_like(seq_first_img_feats)
+        )
+
+        if self.points_pool_project or self.boxes_pool_project:
+            assert len(img_feats) == len(img_sizes)
+            cur_img_feat = img_feats[-1]
+            cur_img_feat = self.img_pre_norm(cur_img_feat)
+            H, W = img_sizes[-1]
+            assert cur_img_feat.shape[0] == H * W
+            N, C = cur_img_feat.shape[-2:]
+            # Put back in NxCxHxW
+            cur_img_feat = cur_img_feat.permute(1, 2, 0)
+            cur_img_feat = cur_img_feat.view(N, C, H, W)
+            img_feats = cur_img_feat
+
+        if self.encode_boxes_as_points:
+            assert boxes is not None
+            assert geo_prompt.box_mask is not None
+            assert geo_prompt.box_labels is not None
+            assert boxes.shape[-1] == 4
+
+            boxes_xyxy = box_cxcywh_to_xyxy(boxes)
+            top_left, bottom_right = boxes_xyxy.split(split_size=2, dim=-1)
+
+            labels_tl = geo_prompt.box_labels + 2
+            labels_br = geo_prompt.box_labels + 4
+
+            # Append to the existing points
+            points, _ = concat_padded_sequences(
+                points, points_mask, top_left, boxes_mask
+            )
+            points_labels, points_mask = concat_padded_sequences(
+                points_labels.unsqueeze(-1),
+                points_mask,
+                labels_tl.unsqueeze(-1),
+                boxes_mask,
+            )
+            points_labels = points_labels.squeeze(-1)
+
+            points, _ = concat_padded_sequences(
+                points, points_mask, bottom_right, boxes_mask
+            )
+            points_labels, points_mask = concat_padded_sequences(
+                points_labels.unsqueeze(-1),
+                points_mask,
+                labels_br.unsqueeze(-1),
+                boxes_mask,
+            )
+            points_labels = points_labels.squeeze(-1)
+
+        final_embeds, final_mask = self._encode_points(
+            points=points,
+            points_mask=points_mask,
+            points_labels=points_labels,
+            img_feats=img_feats,
+        )
+
+        if not self.encode_boxes_as_points:
+            boxes_embeds, boxes_mask = self._encode_boxes(
+                boxes=boxes,
+                boxes_mask=boxes_mask,
+                boxes_labels=boxes_labels,
+                img_feats=img_feats,
+            )
+
+            final_embeds, final_mask = concat_padded_sequences(
+                final_embeds, final_mask, boxes_embeds, boxes_mask
+            )
+
+        if masks is not None and self.mask_encoder is not None:
+            masks_embed, masks_mask = self._encode_masks(
+                masks=masks,
+                attn_mask=masks_mask,
+                mask_labels=masks_labels,
+                img_feats=img_feats,
+            )
+            if points.size(0) == boxes.size(0) == 0:
+                return masks_embed, masks_mask
+        bs = final_embeds.shape[1]
+        assert final_mask.shape[0] == bs
+        if self.cls_embed is not None:
+            cls = self.cls_embed.weight.view(1, 1, self.d_model).repeat(1, bs, 1)
+            cls_mask = torch.zeros(
+                bs, 1, dtype=final_mask.dtype, device=final_mask.device
+            )
+            final_embeds, final_mask = concat_padded_sequences(
+                final_embeds, final_mask, cls, cls_mask
+            )
+
+        if self.final_proj is not None:
+            final_embeds = self.norm(self.final_proj(final_embeds))
+
+        if self.encode is not None:
+            for lay in self.encode:
+                final_embeds = activation_ckpt_wrapper(lay)(
+                    tgt=final_embeds,
+                    memory=seq_first_img_feats,
+                    tgt_key_padding_mask=final_mask,
+                    pos=seq_first_img_pos_embeds,
+                    act_ckpt_enable=self.training and self.use_act_ckpt,
+                )
+            final_embeds = self.encode_norm(final_embeds)
+        # Finally, concat mask embeddings if any
+        if masks is not None and self.mask_encoder is not None:
+            final_embeds, final_mask = concat_padded_sequences(
+                final_embeds, final_mask, masks_embed, masks_mask
+            )
+        return final_embeds, final_mask

sam3/model/io_utils.py

@@ -0,0 +1,709 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import contextlib
+import os
+import queue
+import re
+import time
+from threading import Condition, get_ident, Lock, Thread
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+
+from PIL import Image
+
+from sam3.logger import get_logger
+from tqdm import tqdm
+
+logger = get_logger(__name__)
+
+IS_MAIN_PROCESS = os.getenv("IS_MAIN_PROCESS", "1") == "1"
+RANK = int(os.getenv("RANK", "0"))
+
+IMAGE_EXTS = [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".webp"]
+VIDEO_EXTS = [".mp4", ".mov", ".avi", ".mkv", ".webm"]
+
+
+def load_resource_as_video_frames(
+    resource_path,
+    image_size,
+    offload_video_to_cpu,
+    img_mean=(0.5, 0.5, 0.5),
+    img_std=(0.5, 0.5, 0.5),
+    async_loading_frames=False,
+    video_loader_type="cv2",
+):
+    """
+    Load video frames from either a video or an image (as a single-frame video).
+    Alternatively, if input is a list of PIL images, convert its format
+    """
+    if isinstance(resource_path, list):
+        img_mean = torch.tensor(img_mean, dtype=torch.float16)[:, None, None]
+        img_std = torch.tensor(img_std, dtype=torch.float16)[:, None, None]
+        assert all(isinstance(img_pil, Image.Image) for img_pil in resource_path)
+        assert len(resource_path) is not None
+        orig_height, orig_width = resource_path[0].size
+        orig_height, orig_width = (
+            orig_width,
+            orig_height,
+        )  # For some reason, this method returns these swapped
+        images = []
+        for img_pil in resource_path:
+            img_np = np.array(img_pil.convert("RGB").resize((image_size, image_size)))
+            assert img_np.dtype == np.uint8, "np.uint8 is expected for JPEG images"
+            img_np = img_np / 255.0
+            img = torch.from_numpy(img_np).permute(2, 0, 1)
+            # float16 precision should be sufficient for image tensor storage
+            img = img.to(dtype=torch.float16)
+            # normalize by mean and std
+            img -= img_mean
+            img /= img_std
+            images.append(img)
+        images = torch.stack(images)
+        if not offload_video_to_cpu:
+            images = images.cuda()
+        return images, orig_height, orig_width
+
+    is_image = (
+        isinstance(resource_path, str)
+        and os.path.splitext(resource_path)[-1].lower() in IMAGE_EXTS
+    )
+    if is_image:
+        return load_image_as_single_frame_video(
+            image_path=resource_path,
+            image_size=image_size,
+            offload_video_to_cpu=offload_video_to_cpu,
+            img_mean=img_mean,
+            img_std=img_std,
+        )
+    else:
+        return load_video_frames(
+            video_path=resource_path,
+            image_size=image_size,
+            offload_video_to_cpu=offload_video_to_cpu,
+            img_mean=img_mean,
+            img_std=img_std,
+            async_loading_frames=async_loading_frames,
+            video_loader_type=video_loader_type,
+        )
+
+
+def load_image_as_single_frame_video(
+    image_path,
+    image_size,
+    offload_video_to_cpu,
+    img_mean=(0.5, 0.5, 0.5),
+    img_std=(0.5, 0.5, 0.5),
+):
+    """Load an image as a single-frame video."""
+    images, image_height, image_width = _load_img_as_tensor(image_path, image_size)
+    images = images.unsqueeze(0).half()
+
+    img_mean = torch.tensor(img_mean, dtype=torch.float16)[:, None, None]
+    img_std = torch.tensor(img_std, dtype=torch.float16)[:, None, None]
+    if not offload_video_to_cpu:
+        images = images.cuda()
+        img_mean = img_mean.cuda()
+        img_std = img_std.cuda()
+    # normalize by mean and std
+    images -= img_mean
+    images /= img_std
+    return images, image_height, image_width
+
+
+def load_video_frames(
+    video_path,
+    image_size,
+    offload_video_to_cpu,
+    img_mean=(0.5, 0.5, 0.5),
+    img_std=(0.5, 0.5, 0.5),
+    async_loading_frames=False,
+    video_loader_type="cv2",
+):
+    """
+    Load the video frames from video_path. The frames are resized to image_size as in
+    the model and are loaded to GPU if offload_video_to_cpu=False. This is used by the demo.
+    """
+    assert isinstance(video_path, str)
+    if video_path.startswith("<load-dummy-video"):
+        # Check for pattern <load-dummy-video-N> where N is an integer
+        match = re.match(r"<load-dummy-video-(\d+)>", video_path)
+        num_frames = int(match.group(1)) if match else 60
+        return load_dummy_video(image_size, offload_video_to_cpu, num_frames=num_frames)
+    elif os.path.isdir(video_path):
+        return load_video_frames_from_image_folder(
+            image_folder=video_path,
+            image_size=image_size,
+            offload_video_to_cpu=offload_video_to_cpu,
+            img_mean=img_mean,
+            img_std=img_std,
+            async_loading_frames=async_loading_frames,
+        )
+    elif os.path.splitext(video_path)[-1].lower() in VIDEO_EXTS:
+        return load_video_frames_from_video_file(
+            video_path=video_path,
+            image_size=image_size,
+            offload_video_to_cpu=offload_video_to_cpu,
+            img_mean=img_mean,
+            img_std=img_std,
+            async_loading_frames=async_loading_frames,
+            video_loader_type=video_loader_type,
+        )
+    else:
+        raise NotImplementedError("Only video files and image folders are supported")
+
+
+def load_video_frames_from_image_folder(
+    image_folder,
+    image_size,
+    offload_video_to_cpu,
+    img_mean,
+    img_std,
+    async_loading_frames,
+):
+    """
+    Load the video frames from a directory of image files ("<frame_index>.<img_ext>" format)
+    """
+    frame_names = [
+        p
+        for p in os.listdir(image_folder)
+        if os.path.splitext(p)[-1].lower() in IMAGE_EXTS
+    ]
+    try:
+        frame_names.sort(key=lambda p: int(os.path.splitext(p)[0]))
+    except ValueError:
+        # fallback to lexicographic sort if the format is not "<frame_index>.<img_ext>"
+        logger.warning(
+            f'frame names are not in "<frame_index>.<img_ext>" format: {frame_names[:5]=}, '
+            f"falling back to lexicographic sort."
+        )
+        frame_names.sort()
+    num_frames = len(frame_names)
+    if num_frames == 0:
+        raise RuntimeError(f"no images found in {image_folder}")
+    img_paths = [os.path.join(image_folder, frame_name) for frame_name in frame_names]
+    img_mean = torch.tensor(img_mean, dtype=torch.float16)[:, None, None]
+    img_std = torch.tensor(img_std, dtype=torch.float16)[:, None, None]
+
+    if async_loading_frames:
+        lazy_images = AsyncImageFrameLoader(
+            img_paths, image_size, offload_video_to_cpu, img_mean, img_std
+        )
+        return lazy_images, lazy_images.video_height, lazy_images.video_width
+
+    # float16 precision should be sufficient for image tensor storage
+    images = torch.zeros(num_frames, 3, image_size, image_size, dtype=torch.float16)
+    video_height, video_width = None, None
+    for n, img_path in enumerate(
+        tqdm(img_paths, desc=f"frame loading (image folder) [rank={RANK}]")
+    ):
+        images[n], video_height, video_width = _load_img_as_tensor(img_path, image_size)
+    if not offload_video_to_cpu:
+        images = images.cuda()
+        img_mean = img_mean.cuda()
+        img_std = img_std.cuda()
+    # normalize by mean and std
+    images -= img_mean
+    images /= img_std
+    return images, video_height, video_width
+
+
+def load_video_frames_from_video_file(
+    video_path,
+    image_size,
+    offload_video_to_cpu,
+    img_mean,
+    img_std,
+    async_loading_frames,
+    gpu_acceleration=False,
+    gpu_device=None,
+    video_loader_type="cv2",
+):
+    """Load the video frames from a video file."""
+    if video_loader_type == "cv2":
+        return load_video_frames_from_video_file_using_cv2(
+            video_path=video_path,
+            image_size=image_size,
+            img_mean=img_mean,
+            img_std=img_std,
+            offload_video_to_cpu=offload_video_to_cpu,
+        )
+    elif video_loader_type == "torchcodec":
+        logger.info("Using torchcodec to load video file")
+        lazy_images = AsyncVideoFileLoaderWithTorchCodec(
+            video_path=video_path,
+            image_size=image_size,
+            offload_video_to_cpu=offload_video_to_cpu,
+            img_mean=img_mean,
+            img_std=img_std,
+            gpu_acceleration=gpu_acceleration,
+            gpu_device=gpu_device,
+        )
+        # The `AsyncVideoFileLoaderWithTorchCodec` class always loads the videos asynchronously,
+        # so we just wait for its loading thread to finish if async_loading_frames=False.
+        if not async_loading_frames:
+            async_thread = lazy_images.thread
+            if async_thread is not None:
+                async_thread.join()
+        return lazy_images, lazy_images.video_height, lazy_images.video_width
+    else:
+        raise RuntimeError("video_loader_type must be either 'cv2' or 'torchcodec'")
+
+
+def load_video_frames_from_video_file_using_cv2(
+    video_path: str,
+    image_size: int,
+    img_mean: tuple = (0.5, 0.5, 0.5),
+    img_std: tuple = (0.5, 0.5, 0.5),
+    offload_video_to_cpu: bool = False,
+) -> torch.Tensor:
+    """
+    Load video from path, convert to normalized tensor with specified preprocessing
+
+    Args:
+        video_path: Path to video file
+        image_size: Target size for square frames (height and width)
+        img_mean: Normalization mean (RGB)
+        img_std: Normalization standard deviation (RGB)
+
+    Returns:
+        torch.Tensor: Preprocessed video tensor in shape (T, C, H, W) with float16 dtype
+    """
+    import cv2  # delay OpenCV import to avoid unnecessary dependency
+
+    # Initialize video capture
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        raise ValueError(f"Could not open video: {video_path}")
+
+    original_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    original_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    num_frames = num_frames if num_frames > 0 else None
+
+    frames = []
+    pbar = tqdm(desc=f"frame loading (OpenCV) [rank={RANK}]", total=num_frames)
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        # Convert BGR to RGB and resize
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        frame_resized = cv2.resize(
+            frame_rgb, (image_size, image_size), interpolation=cv2.INTER_CUBIC
+        )
+        frames.append(frame_resized)
+        pbar.update(1)
+    cap.release()
+    pbar.close()
+
+    # Convert to tensor
+    frames_np = np.stack(frames, axis=0).astype(np.float32)  # (T, H, W, C)
+    video_tensor = torch.from_numpy(frames_np).permute(0, 3, 1, 2)  # (T, C, H, W)
+
+    img_mean = torch.tensor(img_mean, dtype=torch.float16).view(1, 3, 1, 1)
+    img_std = torch.tensor(img_std, dtype=torch.float16).view(1, 3, 1, 1)
+    if not offload_video_to_cpu:
+        video_tensor = video_tensor.cuda()
+        img_mean = img_mean.cuda()
+        img_std = img_std.cuda()
+    # normalize by mean and std
+    video_tensor -= img_mean
+    video_tensor /= img_std
+    return video_tensor, original_height, original_width
+
+
+def load_dummy_video(image_size, offload_video_to_cpu, num_frames=60):
+    """
+    Load a dummy video with random frames for testing and compilation warmup purposes.
+    """
+    video_height, video_width = 480, 640  # dummy original video sizes
+    images = torch.randn(num_frames, 3, image_size, image_size, dtype=torch.float16)
+    if not offload_video_to_cpu:
+        images = images.cuda()
+    return images, video_height, video_width
+
+
+def _load_img_as_tensor(img_path, image_size):
+    """Load and resize an image and convert it into a PyTorch tensor."""
+    img = Image.open(img_path).convert("RGB")
+    orig_width, orig_height = img.width, img.height
+    img = TF.resize(img, size=(image_size, image_size))
+    img = TF.to_tensor(img)
+    return img, orig_height, orig_width
+
+
+class AsyncImageFrameLoader:
+    """
+    A list of video frames to be load asynchronously without blocking session start.
+    """
+
+    def __init__(self, img_paths, image_size, offload_video_to_cpu, img_mean, img_std):
+        self.img_paths = img_paths
+        self.image_size = image_size
+        self.offload_video_to_cpu = offload_video_to_cpu
+        self.img_mean = img_mean
+        self.img_std = img_std
+        # items in `self._images` will be loaded asynchronously
+        self.images = [None] * len(img_paths)
+        # catch and raise any exceptions in the async loading thread
+        self.exception = None
+        # video_height and video_width be filled when loading the first image
+        self.video_height = None
+        self.video_width = None
+
+        # load the first frame to fill video_height and video_width and also
+        # to cache it (since it's most likely where the user will click)
+        self.__getitem__(0)
+
+        # load the rest of frames asynchronously without blocking the session start
+        def _load_frames():
+            try:
+                for n in tqdm(
+                    range(len(self.images)),
+                    desc=f"frame loading (image folder) [rank={RANK}]",
+                ):
+                    self.__getitem__(n)
+            except Exception as e:
+                self.exception = e
+
+        self.thread = Thread(target=_load_frames, daemon=True)
+        self.thread.start()
+
+    def __getitem__(self, index):
+        if self.exception is not None:
+            raise RuntimeError("Failure in frame loading thread") from self.exception
+
+        img = self.images[index]
+        if img is not None:
+            return img
+
+        img, video_height, video_width = _load_img_as_tensor(
+            self.img_paths[index], self.image_size
+        )
+        self.video_height = video_height
+        self.video_width = video_width
+        # float16 precision should be sufficient for image tensor storage
+        img = img.to(dtype=torch.float16)
+        # normalize by mean and std
+        img -= self.img_mean
+        img /= self.img_std
+        if not self.offload_video_to_cpu:
+            img = img.cuda()
+        self.images[index] = img
+        return img
+
+    def __len__(self):
+        return len(self.images)
+
+
+class TorchCodecDecoder:
+    """
+    A wrapper to support GPU device and num_threads in TorchCodec decoder,
+    which are not supported by `torchcodec.decoders.SimpleVideoDecoder` yet.
+    """
+
+    def __init__(self, source, dimension_order="NCHW", device="cpu", num_threads=1):
+        from torchcodec import _core as core
+
+        self._source = source  # hold a reference to the source to prevent it from GC
+        if isinstance(source, str):
+            self._decoder = core.create_from_file(source, "exact")
+        elif isinstance(source, bytes):
+            self._decoder = core.create_from_bytes(source, "exact")
+        else:
+            raise TypeError(f"Unknown source type: {type(source)}.")
+        assert dimension_order in ("NCHW", "NHWC")
+
+        device_string = str(device)
+        core.scan_all_streams_to_update_metadata(self._decoder)
+        core.add_video_stream(
+            self._decoder,
+            dimension_order=dimension_order,
+            device=device_string,
+            num_threads=(1 if "cuda" in device_string else num_threads),
+        )
+        video_metadata = core.get_container_metadata(self._decoder)
+        best_stream_index = video_metadata.best_video_stream_index
+        assert best_stream_index is not None
+        self.metadata = video_metadata.streams[best_stream_index]
+        assert self.metadata.num_frames_from_content is not None
+        self._num_frames = self.metadata.num_frames_from_content
+
+    def __len__(self) -> int:
+        return self._num_frames
+
+    def __getitem__(self, key: int):
+        from torchcodec import _core as core
+
+        if key < 0:
+            key += self._num_frames
+        if key >= self._num_frames or key < 0:
+            raise IndexError(
+                f"Index {key} is out of bounds; length is {self._num_frames}"
+            )
+        frame_data, *_ = core.get_frame_at_index(
+            self._decoder,
+            frame_index=key,
+        )
+        return frame_data
+
+
+class FIFOLock:
+    """A lock that ensures FIFO ordering of lock acquisitions."""
+
+    def __init__(self):
+        self._lock = Lock()
+        self._waiters = queue.Queue()
+        self._condition = Condition()
+
+    def acquire(self):
+        ident = get_ident()
+        with self._condition:
+            self._waiters.put(ident)
+            while self._waiters.queue[0] != ident or not self._lock.acquire(
+                blocking=False
+            ):
+                self._condition.wait()
+                # got the lock and it's our turn
+
+    def release(self):
+        with self._condition:
+            self._lock.release()
+            self._waiters.get()
+            self._condition.notify_all()
+
+    def __enter__(self):
+        self.acquire()
+
+    def __exit__(self, t, v, tb):
+        self.release()
+
+
+class AsyncVideoFileLoaderWithTorchCodec:
+    """
+    Loading frames from video files asynchronously without blocking session start.
+
+    Unlike `AsyncVideoFileLoader`, this class uses PyTorch's offical TorchCodec library
+    for video decoding, which is more efficient and supports more video formats.
+    """
+
+    def __init__(
+        self,
+        video_path,
+        image_size,
+        offload_video_to_cpu,
+        img_mean,
+        img_std,
+        gpu_acceleration=True,
+        gpu_device=None,
+        use_rand_seek_in_loading=False,
+    ):
+        # Check and possibly infer the output device (and also get its GPU id when applicable)
+        assert gpu_device is None or gpu_device.type == "cuda"
+        gpu_id = (
+            gpu_device.index
+            if gpu_device is not None and gpu_device.index is not None
+            else torch.cuda.current_device()
+        )
+        if offload_video_to_cpu:
+            out_device = torch.device("cpu")
+        else:
+            out_device = torch.device("cuda") if gpu_device is None else gpu_device
+        self.out_device = out_device
+        self.gpu_acceleration = gpu_acceleration
+        self.gpu_id = gpu_id
+        self.image_size = image_size
+        self.offload_video_to_cpu = offload_video_to_cpu
+        if not isinstance(img_mean, torch.Tensor):
+            img_mean = torch.tensor(img_mean, dtype=torch.float16)[:, None, None]
+        self.img_mean = img_mean
+        if not isinstance(img_std, torch.Tensor):
+            img_std = torch.tensor(img_std, dtype=torch.float16)[:, None, None]
+        self.img_std = img_std
+
+        if gpu_acceleration:
+            self.img_mean = self.img_mean.to(f"cuda:{self.gpu_id}")
+            self.img_std = self.img_std.to(f"cuda:{self.gpu_id}")
+            decoder_option = {"device": f"cuda:{self.gpu_id}"}
+        else:
+            self.img_mean = self.img_mean.cpu()
+            self.img_std = self.img_std.cpu()
+            decoder_option = {"num_threads": 1}  # use a single thread to save memory
+
+        self.rank = int(os.environ.get("RANK", "0"))
+        self.world_size = int(os.environ.get("WORLD_SIZE", "1"))
+        self.async_reader = TorchCodecDecoder(video_path, **decoder_option)
+
+        # `num_frames_from_content` is the true number of frames in the video content
+        # from the scan operation (rather than from the metadata, which could be wrong)
+        self.num_frames = self.async_reader.metadata.num_frames_from_content
+        self.video_height = self.async_reader.metadata.height
+        self.video_width = self.async_reader.metadata.width
+
+        # items in `self._images` will be loaded asynchronously
+        self.images_loaded = [False] * self.num_frames
+        self.images = torch.zeros(
+            self.num_frames,
+            3,
+            self.image_size,
+            self.image_size,
+            dtype=torch.float16,
+            device=self.out_device,
+        )
+        # catch and raise any exceptions in the async loading thread
+        self.exception = None
+        self.use_rand_seek_in_loading = use_rand_seek_in_loading
+        self.rand_seek_idx_queue = queue.Queue()
+        # use a lock to avoid race condition between concurrent access to torchcodec
+        # libs (which are not thread-safe); the lock is replaced with a nullcontext
+        # when the video is fully loaded
+        self.torchcodec_access_lock = FIFOLock()
+        self._start_video_loading()
+
+    def _load_one_frame(self, idx):
+        frame_resized = self._transform_frame(self.async_reader[idx])
+        return frame_resized
+
+    @torch.inference_mode()
+    def _start_video_loading(self):
+        desc = f"frame loading (TorchCodec w/ {'GPU' if self.gpu_acceleration else 'CPU'}) [rank={RANK}]"
+        pbar = tqdm(desc=desc, total=self.num_frames)
+        self.num_loaded_frames = 0
+        # load the first frame synchronously to cache it before the session is opened
+        idx = self.num_loaded_frames
+        self.images[idx] = self._load_one_frame(idx)
+        self.images_loaded[idx] = True
+        self.num_loaded_frames += 1
+        pbar.update(n=1)
+        self.all_frames_loaded = self.num_loaded_frames == self.num_frames
+
+        # load the frames asynchronously without blocking the session start
+        def _load_frames():
+            finished = self.all_frames_loaded
+            chunk_size = 16
+            while not finished:
+                # asynchronously load `chunk_size` frames each time we acquire the lock
+                with self.torchcodec_access_lock, torch.inference_mode():
+                    for _ in range(chunk_size):
+                        try:
+                            idx = self.num_loaded_frames
+                            self.images[idx] = self._load_one_frame(idx)
+                            self.images_loaded[idx] = True
+                            self.num_loaded_frames += 1
+                            pbar.update(n=1)
+                            if self.num_loaded_frames >= self.num_frames:
+                                finished = True
+                                break
+                        except Exception as e:
+                            self.exception = e
+                            raise
+
+                    # also read the frame that is being randomly seeked to
+                    while True:
+                        try:
+                            idx = self.rand_seek_idx_queue.get_nowait()
+                            if not self.images_loaded[idx]:
+                                self.images[idx] = self._load_one_frame(idx)
+                                self.images_loaded[idx] = True
+                        except queue.Empty:
+                            break
+                        except Exception as e:
+                            self.exception = e
+                            raise
+
+            # finished -- check whether we have loaded the total number of frames
+            if self.num_loaded_frames != self.num_frames:
+                raise RuntimeError(
+                    f"There are {self.num_frames} frames in the video, but only "
+                    f"{self.num_loaded_frames} frames can be loaded successfully."
+                )
+            else:
+                self.all_frames_loaded = True
+                pbar.close()
+                with self.torchcodec_access_lock:
+                    import gc
+
+                    # all frames have been loaded, so we can release the readers and free their memory
+                    # also remove pbar and thread (which shouldn't be a part of session saving)
+                    reader = self.async_reader
+                    if reader is not None:
+                        reader._source = None
+                    self.async_reader = None
+                    self.pbar = None
+                    self.thread = None
+                    self.rand_seek_idx_queue = None
+                    gc.collect()
+                # remove the lock (replace it with nullcontext) when the video is fully loaded
+                self.torchcodec_access_lock = contextlib.nullcontext()
+
+        self.thread = Thread(target=_load_frames, daemon=True)
+        self.thread.start()
+
+    def _transform_frame(self, frame):
+        frame = frame.clone()  # make a copy to avoid modifying the original frame bytes
+        frame = frame.float()  # convert to float32 before interpolation
+        frame_resized = F.interpolate(
+            frame[None, :],
+            size=(self.image_size, self.image_size),
+            mode="bicubic",
+            align_corners=False,
+        )[0]
+        # float16 precision should be sufficient for image tensor storage
+        frame_resized = frame_resized.half()  # uint8 -> float16
+        frame_resized /= 255
+        frame_resized -= self.img_mean
+        frame_resized /= self.img_std
+        if self.offload_video_to_cpu:
+            frame_resized = frame_resized.cpu()
+        elif frame_resized.device != self.out_device:
+            frame_resized = frame_resized.to(device=self.out_device, non_blocking=True)
+        return frame_resized
+
+    def __getitem__(self, index):
+        if self.exception is not None:
+            raise RuntimeError("Failure in frame loading thread") from self.exception
+
+        max_tries = 1200
+        for _ in range(max_tries):
+            # use a lock to avoid race condition between concurrent access to torchcodec
+            # libs (which are not thread-safe); the lock is replaced with a nullcontext
+            # when the video is fully loaded
+            with self.torchcodec_access_lock:
+                if self.images_loaded[index]:
+                    return self.images[index]
+
+                if self.use_rand_seek_in_loading:
+                    # async loading hasn't reached this frame yet, so we load this frame individually
+                    # (it will be loaded by in _load_frames thread and added to self.images[index])
+                    self.rand_seek_idx_queue.put(index)
+
+            time.sleep(0.1)
+
+        raise RuntimeError(f"Failed to load frame {index} after {max_tries} tries")
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getstate__(self):
+        """
+        Remove a few attributes during pickling, so that this async video loader can be
+        saved and loaded as a part of the model session.
+        """
+        # wait for async video loading to finish before pickling
+        async_thread = self.thread
+        if async_thread is not None:
+            async_thread.join()
+        # release a few objects that cannot be pickled
+        reader = self.async_reader
+        if reader is not None:
+            reader._source = None
+        self.async_reader = None
+        self.pbar = None
+        self.thread = None
+        self.rand_seek_idx_queue = None
+        self.torchcodec_access_lock = contextlib.nullcontext()
+        return self.__dict__.copy()

sam3/model/maskformer_segmentation.py

@@ -0,0 +1,323 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import math
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+
+from .model_misc import MLP
+
+
+class LinearPresenceHead(nn.Sequential):
+    def __init__(self, d_model):
+        # a hack to make `LinearPresenceHead` compatible with old checkpoints
+        super().__init__(nn.Identity(), nn.Identity(), nn.Linear(d_model, 1))
+
+    def forward(self, hs, prompt, prompt_mask):
+        return super().forward(hs)
+
+
+class MaskPredictor(nn.Module):
+    def __init__(self, hidden_dim, mask_dim):
+        super().__init__()
+        self.mask_embed = MLP(hidden_dim, hidden_dim, mask_dim, 3)
+
+    def forward(self, obj_queries, pixel_embed):
+        if len(obj_queries.shape) == 3:
+            if pixel_embed.ndim == 3:
+                # batch size was omitted
+                mask_preds = torch.einsum(
+                    "bqc,chw->bqhw", self.mask_embed(obj_queries), pixel_embed
+                )
+            else:
+                mask_preds = torch.einsum(
+                    "bqc,bchw->bqhw", self.mask_embed(obj_queries), pixel_embed
+                )
+        else:
+            # Assumed to have aux masks
+            if pixel_embed.ndim == 3:
+                # batch size was omitted
+                mask_preds = torch.einsum(
+                    "lbqc,chw->lbqhw", self.mask_embed(obj_queries), pixel_embed
+                )
+            else:
+                mask_preds = torch.einsum(
+                    "lbqc,bchw->lbqhw", self.mask_embed(obj_queries), pixel_embed
+                )
+
+        return mask_preds
+
+
+class SegmentationHead(nn.Module):
+    def __init__(
+        self,
+        hidden_dim,
+        upsampling_stages,
+        use_encoder_inputs=False,
+        aux_masks=False,
+        no_dec=False,
+        pixel_decoder=None,
+        act_ckpt=False,
+        shared_conv=False,
+        compile_mode_pixel_decoder=None,
+    ):
+        super().__init__()
+        self.use_encoder_inputs = use_encoder_inputs
+        self.aux_masks = aux_masks
+        if pixel_decoder is not None:
+            self.pixel_decoder = pixel_decoder
+        else:
+            self.pixel_decoder = PixelDecoder(
+                hidden_dim,
+                upsampling_stages,
+                shared_conv=shared_conv,
+                compile_mode=compile_mode_pixel_decoder,
+            )
+        self.no_dec = no_dec
+        if no_dec:
+            self.mask_predictor = nn.Conv2d(
+                hidden_dim, 1, kernel_size=3, stride=1, padding=1
+            )
+        else:
+            self.mask_predictor = MaskPredictor(hidden_dim, mask_dim=hidden_dim)
+
+        self.act_ckpt = act_ckpt
+
+        # used to update the output dictionary
+        self.instance_keys = ["pred_masks"]
+
+    @property
+    def device(self):
+        self._device = getattr(self, "_device", None) or next(self.parameters()).device
+        return self._device
+
+    def to(self, *args, **kwargs):
+        # clear cached _device in case the model is moved to a different device
+        self._device = None
+        return super().to(*args, **kwargs)
+
+    def _embed_pixels(
+        self,
+        backbone_feats: List[torch.Tensor],
+        image_ids,
+        encoder_hidden_states,
+    ) -> torch.Tensor:
+        feature_device = backbone_feats[0].device  # features could be on CPU
+        model_device = self.device
+        image_ids_ = image_ids.to(feature_device)
+        if self.use_encoder_inputs:
+            if backbone_feats[0].shape[0] > 1:
+                # For bs > 1, we construct the per query backbone features
+                backbone_visual_feats = []
+                for feat in backbone_feats:
+                    # Copy the img features per query (pixel decoder won't share img feats)
+                    backbone_visual_feats.append(feat[image_ids_, ...].to(model_device))
+            else:
+                # Bs=1, we rely on broadcasting for query-based processing
+                backbone_visual_feats = [bb_feat.clone() for bb_feat in backbone_feats]
+            # Extract visual embeddings
+            encoder_hidden_states = encoder_hidden_states.permute(1, 2, 0)
+            spatial_dim = math.prod(backbone_feats[-1].shape[-2:])
+            encoder_visual_embed = encoder_hidden_states[..., :spatial_dim].reshape(
+                -1, *backbone_feats[-1].shape[1:]
+            )
+
+            backbone_visual_feats[-1] = encoder_visual_embed
+            if self.act_ckpt:
+                pixel_embed = checkpoint.checkpoint(
+                    self.pixel_decoder, backbone_visual_feats, use_reentrant=False
+                )
+            else:
+                pixel_embed = self.pixel_decoder(backbone_visual_feats)
+        else:
+            backbone_feats = [x.to(model_device) for x in backbone_feats]
+            pixel_embed = self.pixel_decoder(backbone_feats)
+            if pixel_embed.shape[0] == 1:
+                # For batch_size=1 training, we can avoid the indexing to save memory
+                pixel_embed = pixel_embed.squeeze(0)
+            else:
+                pixel_embed = pixel_embed[image_ids, ...]
+        return pixel_embed
+
+    def forward(
+        self,
+        backbone_feats: List[torch.Tensor],
+        obj_queries: torch.Tensor,
+        image_ids,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Dict[str, torch.Tensor]:
+        if self.use_encoder_inputs:
+            assert encoder_hidden_states is not None
+
+        pixel_embed = self._embed_pixels(
+            backbone_feats=backbone_feats,
+            image_ids=image_ids,
+            encoder_hidden_states=encoder_hidden_states,
+        )
+
+        if self.no_dec:
+            mask_pred = self.mask_predictor(pixel_embed)
+        elif self.aux_masks:
+            mask_pred = self.mask_predictor(obj_queries, pixel_embed)
+        else:
+            mask_pred = self.mask_predictor(obj_queries[-1], pixel_embed)
+
+        return {"pred_masks": mask_pred}
+
+
+class PixelDecoder(nn.Module):
+    def __init__(
+        self,
+        hidden_dim,
+        num_upsampling_stages,
+        interpolation_mode="nearest",
+        shared_conv=False,
+        compile_mode=None,
+    ):
+        super().__init__()
+        self.hidden_dim = hidden_dim
+        self.num_upsampling_stages = num_upsampling_stages
+        self.interpolation_mode = interpolation_mode
+        conv_layers = []
+        norms = []
+        num_convs = 1 if shared_conv else num_upsampling_stages
+        for _ in range(num_convs):
+            conv_layers.append(nn.Conv2d(self.hidden_dim, self.hidden_dim, 3, 1, 1))
+            norms.append(nn.GroupNorm(8, self.hidden_dim))
+
+        self.conv_layers = nn.ModuleList(conv_layers)
+        self.norms = nn.ModuleList(norms)
+        self.shared_conv = shared_conv
+        self.out_dim = self.conv_layers[-1].out_channels
+        if compile_mode is not None:
+            self.forward = torch.compile(
+                self.forward, mode=compile_mode, dynamic=True, fullgraph=True
+            )
+            # Needed to make checkpointing happy. But we don't know if the module is checkpointed, so we disable it by default.
+            torch._dynamo.config.optimize_ddp = False
+
+    def forward(self, backbone_feats: List[torch.Tensor]):
+        # Assumes backbone features are already projected (C == hidden dim)
+
+        prev_fpn = backbone_feats[-1]
+        fpn_feats = backbone_feats[:-1]
+        for layer_idx, bb_feat in enumerate(fpn_feats[::-1]):
+            curr_fpn = bb_feat
+            prev_fpn = curr_fpn + F.interpolate(
+                prev_fpn, size=curr_fpn.shape[-2:], mode=self.interpolation_mode
+            )
+            if self.shared_conv:
+                # only one conv layer
+                layer_idx = 0
+            prev_fpn = self.conv_layers[layer_idx](prev_fpn)
+            prev_fpn = F.relu(self.norms[layer_idx](prev_fpn))
+
+        return prev_fpn
+
+
+class UniversalSegmentationHead(SegmentationHead):
+    """This module handles semantic+instance segmentation"""
+
+    def __init__(
+        self,
+        hidden_dim,
+        upsampling_stages,
+        pixel_decoder,
+        aux_masks=False,
+        no_dec=False,
+        act_ckpt=False,
+        presence_head: bool = False,
+        dot_product_scorer=None,
+        cross_attend_prompt=None,
+    ):
+        super().__init__(
+            hidden_dim=hidden_dim,
+            upsampling_stages=upsampling_stages,
+            use_encoder_inputs=True,
+            aux_masks=aux_masks,
+            no_dec=no_dec,
+            pixel_decoder=pixel_decoder,
+            act_ckpt=act_ckpt,
+        )
+        self.d_model = hidden_dim
+
+        if dot_product_scorer is not None:
+            assert presence_head, "Specifying a dot product scorer without a presence head is likely a mistake"
+
+        self.presence_head = None
+        if presence_head:
+            self.presence_head = (
+                dot_product_scorer
+                if dot_product_scorer is not None
+                else LinearPresenceHead(self.d_model)
+            )
+
+        self.cross_attend_prompt = cross_attend_prompt
+        if self.cross_attend_prompt is not None:
+            self.cross_attn_norm = nn.LayerNorm(self.d_model)
+
+        self.semantic_seg_head = nn.Conv2d(self.pixel_decoder.out_dim, 1, kernel_size=1)
+        self.instance_seg_head = nn.Conv2d(
+            self.pixel_decoder.out_dim, self.d_model, kernel_size=1
+        )
+
+    def forward(
+        self,
+        backbone_feats: List[torch.Tensor],
+        obj_queries: torch.Tensor,
+        image_ids,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        prompt: Optional[torch.Tensor] = None,
+        prompt_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Dict[str, Optional[torch.Tensor]]:
+        assert encoder_hidden_states is not None
+        bs = encoder_hidden_states.shape[1]
+
+        if self.cross_attend_prompt is not None:
+            tgt2 = self.cross_attn_norm(encoder_hidden_states)
+            tgt2 = self.cross_attend_prompt(
+                query=tgt2,
+                key=prompt,
+                value=prompt,
+                key_padding_mask=prompt_mask,
+            )[0]
+            encoder_hidden_states = tgt2 + encoder_hidden_states
+
+        presence_logit = None
+        if self.presence_head is not None:
+            pooled_enc = encoder_hidden_states.mean(0)
+            presence_logit = (
+                self.presence_head(
+                    pooled_enc.view(1, bs, 1, self.d_model),
+                    prompt=prompt,
+                    prompt_mask=prompt_mask,
+                )
+                .squeeze(0)
+                .squeeze(1)
+            )
+
+        pixel_embed = self._embed_pixels(
+            backbone_feats=backbone_feats,
+            image_ids=image_ids,
+            encoder_hidden_states=encoder_hidden_states,
+        )
+
+        instance_embeds = self.instance_seg_head(pixel_embed)
+
+        if self.no_dec:
+            mask_pred = self.mask_predictor(instance_embeds)
+        elif self.aux_masks:
+            mask_pred = self.mask_predictor(obj_queries, instance_embeds)
+        else:
+            mask_pred = self.mask_predictor(obj_queries[-1], instance_embeds)
+
+        return {
+            "pred_masks": mask_pred,
+            "semantic_seg": self.semantic_seg_head(pixel_embed),
+            "presence_logit": presence_logit,
+        }

sam3/model/memory.py

@@ -0,0 +1,201 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import math
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from timm.layers import DropPath
+except ModuleNotFoundError:
+    # compatibility for older timm versions
+    from timm.models.layers import DropPath
+
+from .model_misc import get_clones, LayerNorm2d
+
+
+class SimpleMaskDownSampler(nn.Module):
+    """
+    Progressively downsample a mask by total_stride, each time by stride.
+    Note that LayerNorm is applied per *token*, like in ViT.
+
+    With each downsample (by a factor stride**2), channel capacity increases by the same factor.
+    In the end, we linearly project to embed_dim channels.
+    """
+
+    def __init__(
+        self,
+        embed_dim=256,
+        kernel_size=4,
+        stride=4,
+        padding=0,
+        total_stride=16,
+        activation=nn.GELU,
+        # Option to interpolate the input mask first before downsampling using convs. In that case, the total_stride is assumed to be after interpolation.
+        # If set to input resolution or None, we don't interpolate. We default to None to be safe (for older configs or if not explicitly set)
+        interpol_size=None,
+    ):
+        super().__init__()
+        num_layers = int(math.log2(total_stride) // math.log2(stride))
+        assert stride**num_layers == total_stride
+        self.encoder = nn.Sequential()
+        mask_in_chans, mask_out_chans = 1, 1
+        for _ in range(num_layers):
+            mask_out_chans = mask_in_chans * (stride**2)
+            self.encoder.append(
+                nn.Conv2d(
+                    mask_in_chans,
+                    mask_out_chans,
+                    kernel_size=kernel_size,
+                    stride=stride,
+                    padding=padding,
+                )
+            )
+            self.encoder.append(LayerNorm2d(mask_out_chans))
+            self.encoder.append(activation())
+            mask_in_chans = mask_out_chans
+
+        self.encoder.append(nn.Conv2d(mask_out_chans, embed_dim, kernel_size=1))
+        self.interpol_size = interpol_size
+        if self.interpol_size is not None:
+            assert isinstance(
+                self.interpol_size, (list, tuple)
+            ), f"Unsupported type {type(self.interpol_size)}. Should be a list or tuple."
+            self.interpol_size = list(interpol_size)
+            assert len(self.interpol_size) == 2
+
+    def forward(self, x: torch.Tensor):
+        if self.interpol_size is not None and self.interpol_size != list(x.shape[-2:]):
+            x = F.interpolate(
+                x.float(),
+                size=self.interpol_size,
+                align_corners=False,
+                mode="bilinear",
+                antialias=True,
+            )
+        return self.encoder(x)
+
+
+# Lightly adapted from ConvNext (https://github.com/facebookresearch/ConvNeXt)
+class CXBlock(nn.Module):
+    r"""ConvNeXt Block. There are two equivalent implementations:
+    (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
+    (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
+    We use (2) as we find it slightly faster in PyTorch
+
+    Args:
+        dim (int): Number of input channels.
+        drop_path (float): Stochastic depth rate. Default: 0.0
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+    """
+
+    def __init__(
+        self,
+        dim,
+        kernel_size=7,
+        padding=3,
+        drop_path=0.0,
+        layer_scale_init_value=1e-6,
+        use_dwconv=True,
+    ):
+        super().__init__()
+        self.dwconv = nn.Conv2d(
+            dim,
+            dim,
+            kernel_size=kernel_size,
+            padding=padding,
+            groups=dim if use_dwconv else 1,
+        )  # depthwise conv
+        self.norm = LayerNorm2d(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(
+            dim, 4 * dim
+        )  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.gamma = (
+            nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
+            if layer_scale_init_value > 0
+            else None
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x):
+        input = x
+        x = self.dwconv(x)
+        x = self.norm(x)
+        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
+
+        x = input + self.drop_path(x)
+        return x
+
+
+class SimpleFuser(nn.Module):
+    def __init__(self, layer, num_layers, dim=None, input_projection=False):
+        super().__init__()
+        self.proj = nn.Identity()
+        self.layers = get_clones(layer, num_layers)
+
+        if input_projection:
+            assert dim is not None
+            self.proj = nn.Conv2d(dim, dim, kernel_size=1)
+
+    def forward(self, x):
+        # normally x: (N, C, H, W)
+        x = self.proj(x)
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+class SimpleMaskEncoder(nn.Module):
+    def __init__(
+        self,
+        out_dim,
+        mask_downsampler,
+        fuser,
+        position_encoding,
+        in_dim=256,  # in_dim of pix_feats
+    ):
+        super().__init__()
+
+        self.mask_downsampler = mask_downsampler
+
+        self.pix_feat_proj = nn.Conv2d(in_dim, in_dim, kernel_size=1)
+        self.fuser = fuser
+        self.position_encoding = position_encoding
+        self.out_proj = nn.Identity()
+        if out_dim != in_dim:
+            self.out_proj = nn.Conv2d(in_dim, out_dim, kernel_size=1)
+
+    def forward(
+        self,
+        pix_feat: torch.Tensor,
+        masks: torch.Tensor,
+        skip_mask_sigmoid: bool = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        ## Process masks
+        # sigmoid, so that less domain shift from gt masks which are bool
+        if not skip_mask_sigmoid:
+            masks = F.sigmoid(masks)
+        masks = self.mask_downsampler(masks)
+
+        ## Fuse pix_feats and downsampled masks
+        # in case the visual features are on CPU, cast them to CUDA
+        pix_feat = pix_feat.to(masks.device)
+
+        x = self.pix_feat_proj(pix_feat)
+        x = x + masks
+        x = self.fuser(x)
+        x = self.out_proj(x)
+
+        pos = self.position_encoding(x).to(x.dtype)
+
+        return {"vision_features": x, "vision_pos_enc": [pos]}

sam3/model/model_misc.py

@@ -0,0 +1,428 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+"""Various utility models"""
+
+import copy
+import math
+import weakref
+from collections.abc import Iterator
+from contextlib import AbstractContextManager
+from enum import auto, Enum
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+from typing_extensions import override
+
+
+def inverse_sigmoid(x, eps=1e-3):
+    """
+    The inverse function for sigmoid activation function.
+    Note: It might face numberical issues with fp16 small eps.
+    """
+    x = x.clamp(min=0, max=1)
+    x1 = x.clamp(min=eps)
+    x2 = (1 - x).clamp(min=eps)
+    return torch.log(x1 / x2)
+
+
+class MultiheadAttentionWrapper(nn.MultiheadAttention):
+    def forward(self, *args, **kwargs):
+        kwargs["need_weights"] = False
+        return super().forward(*args, **kwargs)
+
+
+class DotProductScoring(torch.nn.Module):
+    def __init__(
+        self,
+        d_model,
+        d_proj,
+        prompt_mlp=None,
+        clamp_logits=True,
+        clamp_max_val=12.0,
+    ):
+        super().__init__()
+        self.d_proj = d_proj
+        assert isinstance(prompt_mlp, torch.nn.Module) or prompt_mlp is None
+        self.prompt_mlp = prompt_mlp  # an optional MLP projection for prompt
+        self.prompt_proj = torch.nn.Linear(d_model, d_proj)
+        self.hs_proj = torch.nn.Linear(d_model, d_proj)
+        self.scale = float(1.0 / np.sqrt(d_proj))
+        self.clamp_logits = clamp_logits
+        if self.clamp_logits:
+            self.clamp_max_val = clamp_max_val
+
+    def mean_pool_text(self, prompt, prompt_mask):
+        # is_valid has shape (seq, bs, 1), where 1 is valid and 0 is padding
+        is_valid = (~prompt_mask).float().permute(1, 0)[..., None]
+        # num_valid has shape (bs, 1)
+        num_valid = torch.clamp(torch.sum(is_valid, dim=0), min=1.0)
+        # mean pool over all the valid tokens -- pooled_prompt has shape (bs, proj_dim)
+        pooled_prompt = (prompt * is_valid).sum(dim=0) / num_valid
+        return pooled_prompt
+
+    def forward(self, hs, prompt, prompt_mask):
+        # hs has shape (num_layer, bs, num_query, d_model)
+        # prompt has shape (seq, bs, d_model)
+        # prompt_mask has shape (bs, seq), where 1 is valid and 0 is padding
+        assert hs.dim() == 4 and prompt.dim() == 3 and prompt_mask.dim() == 2
+
+        # apply MLP on prompt if specified
+        if self.prompt_mlp is not None:
+            prompt = self.prompt_mlp(prompt)
+
+        # first, get the mean-pooled version of the prompt
+        pooled_prompt = self.mean_pool_text(prompt, prompt_mask)
+
+        # then, project pooled_prompt and hs to d_proj dimensions
+        proj_pooled_prompt = self.prompt_proj(pooled_prompt)  # (bs, d_proj)
+        proj_hs = self.hs_proj(hs)  # (num_layer, bs, num_query, d_proj)
+
+        # finally, get dot-product scores of shape (num_layer, bs, num_query, 1)
+        scores = torch.matmul(proj_hs, proj_pooled_prompt.unsqueeze(-1))
+        scores *= self.scale
+
+        # clamp scores to a max value to avoid numerical issues in loss or matcher
+        if self.clamp_logits:
+            scores.clamp_(min=-self.clamp_max_val, max=self.clamp_max_val)
+
+        return scores
+
+
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        init_values: Union[float, Tensor] = 1e-5,
+        inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+class LayerNorm2d(nn.Module):
+    def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(num_channels))
+        self.bias = nn.Parameter(torch.zeros(num_channels))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        u = x.mean(1, keepdim=True)
+        s = (x - u).pow(2).mean(1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.eps)
+        x = self.weight[:, None, None] * x + self.bias[:, None, None]
+        return x
+
+
+class TransformerWrapper(nn.Module):
+    def __init__(
+        self,
+        encoder,
+        decoder,
+        d_model: int,
+        two_stage_type="none",  # ["none"] only for now
+        pos_enc_at_input_dec=True,
+    ):
+        super().__init__()
+
+        self.encoder = encoder
+        self.decoder = decoder
+        self.num_queries = decoder.num_queries if decoder is not None else None
+        self.pos_enc_at_input_dec = pos_enc_at_input_dec
+
+        # for two stage
+        assert two_stage_type in ["none"], "unknown param {} of two_stage_type".format(
+            two_stage_type
+        )
+        self.two_stage_type = two_stage_type
+
+        self._reset_parameters()
+        self.d_model = d_model
+
+    def _reset_parameters(self):
+        for n, p in self.named_parameters():
+            if p.dim() > 1:
+                if (
+                    "box_embed" not in n
+                    and "query_embed" not in n
+                    and "reference_points" not in n
+                ):
+                    nn.init.xavier_uniform_(p)
+
+
+class MLP(nn.Module):
+    """Very simple multi-layer perceptron (also called FFN)"""
+
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int,
+        output_dim: int,
+        num_layers: int,
+        dropout: float = 0.0,
+        residual: bool = False,
+        out_norm: Optional[nn.Module] = None,
+    ):
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(
+            nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])
+        )
+        self.drop = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+        # whether to add the output as a residual connection to the input
+        if residual and input_dim != output_dim:
+            raise ValueError("residual is only supported if input_dim == output_dim")
+        self.residual = residual
+        # whether to apply a normalization layer to the output
+        assert isinstance(out_norm, nn.Module) or out_norm is None
+        self.out_norm = out_norm or nn.Identity()
+
+    def forward(self, x):
+        orig_x = x
+        for i, layer in enumerate(self.layers):
+            x = self.drop(F.relu(layer(x))) if i < self.num_layers - 1 else layer(x)
+        if self.residual:
+            x = x + orig_x
+        x = self.out_norm(x)
+        return x
+
+
+def get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def get_clones_seq(module, N):
+    return nn.Sequential(*[copy.deepcopy(module) for i in range(N)])
+
+
+def get_activation_fn(activation):
+    """Return an activation function given a string"""
+    if activation == "relu":
+        return F.relu
+    if activation == "gelu":
+        return F.gelu
+    if activation == "glu":
+        return F.glu
+    raise RuntimeError(f"activation should be relu/gelu, not {activation}.")
+
+
+def get_activation_module(activation):
+    """Return an activation function given a string"""
+    if activation == "relu":
+        return nn.ReLU
+    if activation == "gelu":
+        return nn.GELU
+    if activation == "glu":
+        return nn.GLU
+    raise RuntimeError(f"activation should be relu/gelu, not {activation}.")
+
+
+def get_valid_ratio(mask):
+    _, H, W = mask.shape
+    valid_H = torch.sum(~mask[:, :, 0], 1)
+    valid_W = torch.sum(~mask[:, 0, :], 1)
+    valid_ratio_h = valid_H.float() / H
+    valid_ratio_w = valid_W.float() / W
+    valid_ratio = torch.stack([valid_ratio_w, valid_ratio_h], -1)
+    return valid_ratio
+
+
+def gen_sineembed_for_position(pos_tensor, num_feats=256):
+    assert num_feats % 2 == 0
+    num_feats = num_feats // 2
+    # n_query, bs, _ = pos_tensor.size()
+    # sineembed_tensor = torch.zeros(n_query, bs, 256)
+    scale = 2 * math.pi
+    dim_t = torch.arange(num_feats, dtype=torch.float32, device=pos_tensor.device)
+    dim_t = 10000 ** (2 * (torch.div(dim_t, 2, rounding_mode="floor")) / num_feats)
+    x_embed = pos_tensor[:, :, 0] * scale
+    y_embed = pos_tensor[:, :, 1] * scale
+    pos_x = x_embed[:, :, None] / dim_t
+    pos_y = y_embed[:, :, None] / dim_t
+    pos_x = torch.stack(
+        (pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3
+    ).flatten(2)
+    pos_y = torch.stack(
+        (pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim=3
+    ).flatten(2)
+    if pos_tensor.size(-1) == 2:
+        pos = torch.cat((pos_y, pos_x), dim=2)
+    elif pos_tensor.size(-1) == 4:
+        w_embed = pos_tensor[:, :, 2] * scale
+        pos_w = w_embed[:, :, None] / dim_t
+        pos_w = torch.stack(
+            (pos_w[:, :, 0::2].sin(), pos_w[:, :, 1::2].cos()), dim=3
+        ).flatten(2)
+
+        h_embed = pos_tensor[:, :, 3] * scale
+        pos_h = h_embed[:, :, None] / dim_t
+        pos_h = torch.stack(
+            (pos_h[:, :, 0::2].sin(), pos_h[:, :, 1::2].cos()), dim=3
+        ).flatten(2)
+
+        pos = torch.cat((pos_y, pos_x, pos_w, pos_h), dim=2)
+    else:
+        raise ValueError("Unknown pos_tensor shape(-1):{}".format(pos_tensor.size(-1)))
+    return pos
+
+
+class SAM3Output(list):
+    """
+    A class representing the output of a SAM3 model.
+    It provides an iterable interface that supports different iteration modes, including iterating over all steps per stage,
+    last step per stage, and flattened output.
+    Attributes:
+        output: The output of the SAM3 model, represented as a list of lists.
+        iter_mode: The current iteration mode.
+    Example:
+        >>> output = [[1, 2], [3, 4], [5, 6]]
+        >>> sam3_output = SAM3Output(output)
+        >>> for step in sam3_output:
+        ...     print(step)
+        [1, 2]
+        [3, 4]
+        [5, 6]
+        >>> with SAM3Output.iteration_mode(SAM3Output.IterMode.LAST_STEP_PER_STAGE) as sam3_last_step_out:
+        ...     for step in sam3_last_step_out:
+        ...         print(step)
+        [2]
+        [4]
+        [6]
+        >>> with SAM3Output.iteration_mode(SAM3Output.IterMode.FLATTENED) as sam3_flattened_out:
+        ...     for step in sam3_flattened_out:
+        ...         print(step)
+        1
+        2
+        3
+        4
+        5
+        6
+    """
+
+    class IterMode(Enum):
+        # Defines the type of iterator over ouptuts.
+        ALL_STEPS_PER_STAGE = auto()
+        LAST_STEP_PER_STAGE = auto()
+        FLATTENED = auto()  # Returns each interactivity step as if it is a separate stage (this is used in SAM3Image model)
+
+    def __init__(
+        self,
+        output: List[List[Dict]] = None,
+        iter_mode: IterMode = IterMode.ALL_STEPS_PER_STAGE,
+        loss_stages: Optional[List[int]] = None,
+    ):
+        if output is not None:
+            assert (
+                isinstance(output, list)
+                and len(output) > 0
+                and isinstance(output[0], list)
+            ), "Expected output to be a list of lists"
+            self.output = output
+        else:
+            self.output = []
+        assert isinstance(
+            iter_mode, SAM3Output.IterMode
+        ), f"iter_mode shoulf be of enum type 'SAM3Output.IterMode'. Got {type(iter_mode)}"
+
+        self.iter_mode = iter_mode
+        # We create a weak reference to self to be used in the lambda functions.
+        # This is to avoid cyclic references and let SAM3Output be garabge collected.
+        self_ref = weakref.ref(self)
+        self._mode2iter = {
+            SAM3Output.IterMode.ALL_STEPS_PER_STAGE: lambda: iter(self_ref().output),
+            SAM3Output.IterMode.LAST_STEP_PER_STAGE: lambda: (
+                inner_list[-1] for inner_list in self_ref().output
+            ),
+            SAM3Output.IterMode.FLATTENED: lambda: (
+                element for inner_list in self_ref().output for element in inner_list
+            ),
+        }
+        self.loss_stages = loss_stages
+
+    @override
+    def __iter__(self) -> Iterator:
+        return self._mode2iter[self.iter_mode]()
+
+    def __getitem__(self, index):
+        """
+        Returns the item at the specified index.
+        Args:
+            index (int): The index of the item to return.
+        Returns:
+            list or element: The item at the specified index.
+        """
+        assert isinstance(index, int), f"index should be an integer. Got {type(index)}"
+        if self.iter_mode == SAM3Output.IterMode.ALL_STEPS_PER_STAGE:
+            return self.output[index]
+        elif self.iter_mode == SAM3Output.IterMode.LAST_STEP_PER_STAGE:
+            return self.output[index][-1]
+        elif self.iter_mode == SAM3Output.IterMode.FLATTENED:
+            if index == -1:
+                return self.self.output[-1][-1]
+            else:
+                flattened_output = sum(self.output, [])
+                return flattened_output[index]
+
+    class _IterationMode(AbstractContextManager):
+        """
+        A context manager that temporarily changes the iteration mode of a SAM3Output object.
+        This class is used internally by the SAM3Output.iteration_mode method.
+        """
+
+        def __init__(
+            self, model_output: "SAM3Output", iter_mode: "SAM3Output.IterMode"
+        ):
+            self._model_output = model_output
+            self._orig_iter_mode = model_output.iter_mode
+            self._new_iter_mode = iter_mode
+
+        @override
+        def __enter__(self) -> "SAM3Output":
+            self._model_output.iter_mode = self._new_iter_mode
+            return self._model_output
+
+        @override
+        def __exit__(self, exc_type, exc_value, traceback):
+            self._model_output.iter_mode = self._orig_iter_mode
+            return super().__exit__(exc_type, exc_value, traceback)
+
+    @staticmethod
+    def iteration_mode(
+        model_output: "SAM3Output", iter_mode: IterMode
+    ) -> _IterationMode:
+        """
+        Returns a context manager that allows you to temporarily change the iteration mode of the SAM3Output object.
+        Args:
+            model_output: The SAM3Output object.
+            iter_mode: The new iteration mode.
+        Returns:
+            SAM3Output._IterationMode: A context manager that changes the iteration mode of the SAM3Output object.
+        """
+        return SAM3Output._IterationMode(model_output=model_output, iter_mode=iter_mode)
+
+    def append(self, item: list):
+        assert isinstance(
+            item, list
+        ), f"Only list items are supported. Got {type(item)}"
+        self.output.append(item)
+
+    def __repr__(self):
+        return self.output.__repr__()
+
+    def __len__(self):
+        if self.iter_mode in [
+            SAM3Output.IterMode.ALL_STEPS_PER_STAGE,
+            SAM3Output.IterMode.LAST_STEP_PER_STAGE,
+        ]:
+            return len(self.output)
+        elif self.iter_mode == SAM3Output.IterMode.FLATTENED:
+            flattened_output = sum(self.output, [])
+            return len(flattened_output)

sam3/model/necks.py

@@ -0,0 +1,125 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+"""Necks are the interface between a vision backbone and the rest of the detection model"""
+
+from copy import deepcopy
+from typing import List, Optional, Tuple
+
+import torch
+
+import torch.nn as nn
+
+
+class Sam3DualViTDetNeck(nn.Module):
+    def __init__(
+        self,
+        trunk: nn.Module,
+        position_encoding: nn.Module,
+        d_model: int,
+        scale_factors=(4.0, 2.0, 1.0, 0.5),
+        add_sam2_neck: bool = False,
+    ):
+        """
+        SimpleFPN neck a la ViTDet
+        (From detectron2, very lightly adapted)
+        It supports a "dual neck" setting, where we have two identical necks (for SAM3 and SAM2), with different weights
+
+        :param trunk: the backbone
+        :param position_encoding: the positional encoding to use
+        :param d_model: the dimension of the model
+        """
+        super().__init__()
+        self.trunk = trunk
+        self.position_encoding = position_encoding
+        self.convs = nn.ModuleList()
+
+        self.scale_factors = scale_factors
+        use_bias = True
+        dim: int = self.trunk.channel_list[-1]
+
+        for _, scale in enumerate(scale_factors):
+            current = nn.Sequential()
+
+            if scale == 4.0:
+                current.add_module(
+                    "dconv_2x2_0",
+                    nn.ConvTranspose2d(dim, dim // 2, kernel_size=2, stride=2),
+                )
+                current.add_module(
+                    "gelu",
+                    nn.GELU(),
+                )
+                current.add_module(
+                    "dconv_2x2_1",
+                    nn.ConvTranspose2d(dim // 2, dim // 4, kernel_size=2, stride=2),
+                )
+                out_dim = dim // 4
+            elif scale == 2.0:
+                current.add_module(
+                    "dconv_2x2",
+                    nn.ConvTranspose2d(dim, dim // 2, kernel_size=2, stride=2),
+                )
+                out_dim = dim // 2
+            elif scale == 1.0:
+                out_dim = dim
+            elif scale == 0.5:
+                current.add_module(
+                    "maxpool_2x2",
+                    nn.MaxPool2d(kernel_size=2, stride=2),
+                )
+                out_dim = dim
+            else:
+                raise NotImplementedError(f"scale_factor={scale} is not supported yet.")
+
+            current.add_module(
+                "conv_1x1",
+                nn.Conv2d(
+                    in_channels=out_dim,
+                    out_channels=d_model,
+                    kernel_size=1,
+                    bias=use_bias,
+                ),
+            )
+            current.add_module(
+                "conv_3x3",
+                nn.Conv2d(
+                    in_channels=d_model,
+                    out_channels=d_model,
+                    kernel_size=3,
+                    padding=1,
+                    bias=use_bias,
+                ),
+            )
+            self.convs.append(current)
+
+        self.sam2_convs = None
+        if add_sam2_neck:
+            # Assumes sam2 neck is just a clone of the original neck
+            self.sam2_convs = deepcopy(self.convs)
+
+    def forward(
+        self, tensor_list: List[torch.Tensor]
+    ) -> Tuple[
+        List[torch.Tensor],
+        List[torch.Tensor],
+        Optional[List[torch.Tensor]],
+        Optional[List[torch.Tensor]],
+    ]:
+        xs = self.trunk(tensor_list)
+        sam3_out, sam3_pos = [], []
+        sam2_out, sam2_pos = None, None
+        if self.sam2_convs is not None:
+            sam2_out, sam2_pos = [], []
+        x = xs[-1]  # simpleFPN
+        for i in range(len(self.convs)):
+            sam3_x_out = self.convs[i](x)
+            sam3_pos_out = self.position_encoding(sam3_x_out).to(sam3_x_out.dtype)
+            sam3_out.append(sam3_x_out)
+            sam3_pos.append(sam3_pos_out)
+
+            if self.sam2_convs is not None:
+                sam2_x_out = self.sam2_convs[i](x)
+                sam2_pos_out = self.position_encoding(sam2_x_out).to(sam2_x_out.dtype)
+                sam2_out.append(sam2_x_out)
+                sam2_pos.append(sam2_pos_out)
+        return sam3_out, sam3_pos, sam2_out, sam2_pos

sam3/model/position_encoding.py

@@ -0,0 +1,124 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import math
+from typing import Optional
+
+import torch
+from torch import nn
+
+
+class PositionEmbeddingSine(nn.Module):
+    """
+    This is a more standard version of the position embedding, very similar to the one
+    used by the Attention is all you need paper, generalized to work on images.
+    """
+
+    def __init__(
+        self,
+        num_pos_feats,
+        temperature: int = 10000,
+        normalize: bool = True,
+        scale: Optional[float] = None,
+        precompute_resolution: Optional[int] = None,
+    ):
+        super().__init__()
+        assert num_pos_feats % 2 == 0, "Expecting even model width"
+        self.num_pos_feats = num_pos_feats // 2
+        self.temperature = temperature
+        self.normalize = normalize
+        if scale is not None and normalize is False:
+            raise ValueError("normalize should be True if scale is passed")
+        if scale is None:
+            scale = 2 * math.pi
+        self.scale = scale
+
+        self.cache = {}
+        # Precompute positional encodings under `precompute_resolution` to fill the cache
+        # and avoid symbolic shape tracing errors in torch.compile in PyTorch 2.4 nightly.
+        if precompute_resolution is not None:
+            # We precompute pos enc for stride 4, 8, 16 and 32 to fill `self.cache`.
+            precompute_sizes = [
+                (precompute_resolution // 4, precompute_resolution // 4),
+                (precompute_resolution // 8, precompute_resolution // 8),
+                (precompute_resolution // 16, precompute_resolution // 16),
+                (precompute_resolution // 32, precompute_resolution // 32),
+            ]
+            for size in precompute_sizes:
+                tensors = torch.zeros((1, 1) + size, device="cuda")
+                self.forward(tensors)
+                # further clone and detach it in the cache (just to be safe)
+                self.cache[size] = self.cache[size].clone().detach()
+
+    def _encode_xy(self, x, y):
+        # The positions are expected to be normalized
+        assert len(x) == len(y) and x.ndim == y.ndim == 1
+        x_embed = x * self.scale
+        y_embed = y * self.scale
+
+        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
+        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
+
+        pos_x = x_embed[:, None] / dim_t
+        pos_y = y_embed[:, None] / dim_t
+        pos_x = torch.stack(
+            (pos_x[:, 0::2].sin(), pos_x[:, 1::2].cos()), dim=2
+        ).flatten(1)
+        pos_y = torch.stack(
+            (pos_y[:, 0::2].sin(), pos_y[:, 1::2].cos()), dim=2
+        ).flatten(1)
+        return pos_x, pos_y
+
+    @torch.no_grad()
+    def encode_boxes(self, x, y, w, h):
+        pos_x, pos_y = self._encode_xy(x, y)
+        pos = torch.cat((pos_y, pos_x, h[:, None], w[:, None]), dim=1)
+        return pos
+
+    encode = encode_boxes  # Backwards compatibility
+
+    @torch.no_grad()
+    def encode_points(self, x, y, labels):
+        (bx, nx), (by, ny), (bl, nl) = x.shape, y.shape, labels.shape
+        assert bx == by and nx == ny and bx == bl and nx == nl
+        pos_x, pos_y = self._encode_xy(x.flatten(), y.flatten())
+        pos_x, pos_y = pos_x.reshape(bx, nx, -1), pos_y.reshape(by, ny, -1)
+        pos = torch.cat((pos_y, pos_x, labels[:, :, None]), dim=2)
+        return pos
+
+    @torch.no_grad()
+    def forward(self, x):
+        cache_key = None
+        cache_key = (x.shape[-2], x.shape[-1])
+        if cache_key in self.cache:
+            return self.cache[cache_key][None].repeat(x.shape[0], 1, 1, 1)
+        y_embed = (
+            torch.arange(1, x.shape[-2] + 1, dtype=torch.float32, device=x.device)
+            .view(1, -1, 1)
+            .repeat(x.shape[0], 1, x.shape[-1])
+        )
+        x_embed = (
+            torch.arange(1, x.shape[-1] + 1, dtype=torch.float32, device=x.device)
+            .view(1, 1, -1)
+            .repeat(x.shape[0], x.shape[-2], 1)
+        )
+
+        if self.normalize:
+            eps = 1e-6
+            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
+            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
+
+        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
+        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
+
+        pos_x = x_embed[:, :, :, None] / dim_t
+        pos_y = y_embed[:, :, :, None] / dim_t
+        pos_x = torch.stack(
+            (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4
+        ).flatten(3)
+        pos_y = torch.stack(
+            (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4
+        ).flatten(3)
+        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
+        if cache_key is not None:
+            self.cache[cache_key] = pos[0]
+        return pos

sam3/model/sam1_task_predictor.py

@@ -0,0 +1,458 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+import torch.nn as nn
+from PIL.Image import Image
+
+from sam3.model.sam3_tracker_base import Sam3TrackerBase
+from sam3.model.utils.sam1_utils import SAM2Transforms
+
+
+# Adapted from https://github.com/facebookresearch/sam2/blob/main/sam2/sam2_image_predictor.py
+class SAM3InteractiveImagePredictor(nn.Module):
+    def __init__(
+        self,
+        sam_model: Sam3TrackerBase,
+        mask_threshold=0.0,
+        max_hole_area=256.0,
+        max_sprinkle_area=0.0,
+        **kwargs,
+    ) -> None:
+        """
+        Uses SAM-3 to calculate the image embedding for an image, and then
+        allow repeated, efficient mask prediction given prompts.
+
+        Arguments:
+          sam_model : The model to use for mask prediction.
+          mask_threshold (float): The threshold to use when converting mask logits
+            to binary masks. Masks are thresholded at 0 by default.
+          max_hole_area (int): If max_hole_area > 0, we fill small holes in up to
+            the maximum area of max_hole_area in low_res_masks.
+          max_sprinkle_area (int): If max_sprinkle_area > 0, we remove small sprinkles up to
+            the maximum area of max_sprinkle_area in low_res_masks.
+        """
+        super().__init__()
+        self.model = sam_model
+        self._transforms = SAM2Transforms(
+            resolution=self.model.image_size,
+            mask_threshold=mask_threshold,
+            max_hole_area=max_hole_area,
+            max_sprinkle_area=max_sprinkle_area,
+        )
+
+        # Predictor state
+        self._is_image_set = False
+        self._features = None
+        self._orig_hw = None
+        # Whether the predictor is set for single image or a batch of images
+        self._is_batch = False
+
+        # Predictor config
+        self.mask_threshold = mask_threshold
+
+        # Spatial dim for backbone feature maps
+        self._bb_feat_sizes = [
+            (288, 288),
+            (144, 144),
+            (72, 72),
+        ]
+
+    @torch.no_grad()
+    def set_image(
+        self,
+        image: Union[np.ndarray, Image],
+    ) -> None:
+        """
+        Calculates the image embeddings for the provided image, allowing
+        masks to be predicted with the 'predict' method.
+
+        Arguments:
+          image (np.ndarray or PIL Image): The input image to embed in RGB format. The image should be in HWC format if np.ndarray, or WHC format if PIL Image
+          with pixel values in [0, 255].
+          image_format (str): The color format of the image, in ['RGB', 'BGR'].
+        """
+        self.reset_predictor()
+        # Transform the image to the form expected by the model
+        if isinstance(image, np.ndarray):
+            logging.info("For numpy array image, we assume (HxWxC) format")
+            self._orig_hw = [image.shape[:2]]
+        elif isinstance(image, Image):
+            w, h = image.size
+            self._orig_hw = [(h, w)]
+        else:
+            raise NotImplementedError("Image format not supported")
+
+        input_image = self._transforms(image)
+        input_image = input_image[None, ...].to(self.device)
+
+        assert (
+            len(input_image.shape) == 4 and input_image.shape[1] == 3
+        ), f"input_image must be of size 1x3xHxW, got {input_image.shape}"
+        logging.info("Computing image embeddings for the provided image...")
+        backbone_out = self.model.forward_image(input_image)
+        (
+            _,
+            vision_feats,
+            _,
+            _,
+        ) = self.model._prepare_backbone_features(backbone_out)
+        # Add no_mem_embed, which is added to the lowest rest feat. map during training on videos
+        vision_feats[-1] = vision_feats[-1] + self.model.no_mem_embed
+
+        feats = [
+            feat.permute(1, 2, 0).view(1, -1, *feat_size)
+            for feat, feat_size in zip(vision_feats[::-1], self._bb_feat_sizes[::-1])
+        ][::-1]
+        self._features = {"image_embed": feats[-1], "high_res_feats": feats[:-1]}
+        self._is_image_set = True
+        logging.info("Image embeddings computed.")
+
+    @torch.no_grad()
+    def set_image_batch(
+        self,
+        image_list: List[Union[np.ndarray]],
+    ) -> None:
+        """
+        Calculates the image embeddings for the provided image batch, allowing
+        masks to be predicted with the 'predict_batch' method.
+
+        Arguments:
+          image_list (List[np.ndarray]): The input images to embed in RGB format. The image should be in HWC format if np.ndarray
+          with pixel values in [0, 255].
+        """
+        self.reset_predictor()
+        assert isinstance(image_list, list)
+        self._orig_hw = []
+        for image in image_list:
+            assert isinstance(
+                image, np.ndarray
+            ), "Images are expected to be an np.ndarray in RGB format, and of shape  HWC"
+            self._orig_hw.append(image.shape[:2])
+        # Transform the image to the form expected by the model
+        img_batch = self._transforms.forward_batch(image_list)
+        img_batch = img_batch.to(self.device)
+        batch_size = img_batch.shape[0]
+        assert (
+            len(img_batch.shape) == 4 and img_batch.shape[1] == 3
+        ), f"img_batch must be of size Bx3xHxW, got {img_batch.shape}"
+        logging.info("Computing image embeddings for the provided images...")
+        backbone_out = self.model.forward_image(img_batch)
+        (
+            _,
+            vision_feats,
+            _,
+            _,
+        ) = self.model._prepare_backbone_features(backbone_out)
+        # Add no_mem_embed, which is added to the lowest rest feat. map during training on videos
+        vision_feats[-1] = vision_feats[-1] + self.model.no_mem_embed
+
+        feats = [
+            feat.permute(1, 2, 0).view(batch_size, -1, *feat_size)
+            for feat, feat_size in zip(vision_feats[::-1], self._bb_feat_sizes[::-1])
+        ][::-1]
+        self._features = {"image_embed": feats[-1], "high_res_feats": feats[:-1]}
+        self._is_image_set = True
+        self._is_batch = True
+        logging.info("Image embeddings computed.")
+
+    def predict_batch(
+        self,
+        point_coords_batch: List[np.ndarray] = None,
+        point_labels_batch: List[np.ndarray] = None,
+        box_batch: List[np.ndarray] = None,
+        mask_input_batch: List[np.ndarray] = None,
+        multimask_output: bool = True,
+        return_logits: bool = False,
+        normalize_coords=True,
+    ) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray]]:
+        """This function is very similar to predict(...), however it is used for batched mode, when the model is expected to generate predictions on multiple images.
+        It returns a tuple of lists of masks, ious, and low_res_masks_logits.
+        """
+        assert self._is_batch, "This function should only be used when in batched mode"
+        if not self._is_image_set:
+            raise RuntimeError(
+                "An image must be set with .set_image_batch(...) before mask prediction."
+            )
+        num_images = len(self._features["image_embed"])
+        all_masks = []
+        all_ious = []
+        all_low_res_masks = []
+        for img_idx in range(num_images):
+            # Transform input prompts
+            point_coords = (
+                point_coords_batch[img_idx] if point_coords_batch is not None else None
+            )
+            point_labels = (
+                point_labels_batch[img_idx] if point_labels_batch is not None else None
+            )
+            box = box_batch[img_idx] if box_batch is not None else None
+            mask_input = (
+                mask_input_batch[img_idx] if mask_input_batch is not None else None
+            )
+            mask_input, unnorm_coords, labels, unnorm_box = self._prep_prompts(
+                point_coords,
+                point_labels,
+                box,
+                mask_input,
+                normalize_coords,
+                img_idx=img_idx,
+            )
+            masks, iou_predictions, low_res_masks = self._predict(
+                unnorm_coords,
+                labels,
+                unnorm_box,
+                mask_input,
+                multimask_output,
+                return_logits=return_logits,
+                img_idx=img_idx,
+            )
+            masks_np = masks.squeeze(0).float().detach().cpu().numpy()
+            iou_predictions_np = (
+                iou_predictions.squeeze(0).float().detach().cpu().numpy()
+            )
+            low_res_masks_np = low_res_masks.squeeze(0).float().detach().cpu().numpy()
+            all_masks.append(masks_np)
+            all_ious.append(iou_predictions_np)
+            all_low_res_masks.append(low_res_masks_np)
+
+        return all_masks, all_ious, all_low_res_masks
+
+    def predict(
+        self,
+        point_coords: Optional[np.ndarray] = None,
+        point_labels: Optional[np.ndarray] = None,
+        box: Optional[np.ndarray] = None,
+        mask_input: Optional[np.ndarray] = None,
+        multimask_output: bool = True,
+        return_logits: bool = False,
+        normalize_coords=True,
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """
+        Predict masks for the given input prompts, using the currently set image.
+
+        Arguments:
+          point_coords (np.ndarray or None): A Nx2 array of point prompts to the
+            model. Each point is in (X,Y) in pixels.
+          point_labels (np.ndarray or None): A length N array of labels for the
+            point prompts. 1 indicates a foreground point and 0 indicates a
+            background point.
+          box (np.ndarray or None): A length 4 array given a box prompt to the
+            model, in XYXY format.
+          mask_input (np.ndarray): A low resolution mask input to the model, typically
+            coming from a previous prediction iteration. Has form 1xHxW, where
+            for SAM, H=W=256.
+          multimask_output (bool): If true, the model will return three masks.
+            For ambiguous input prompts (such as a single click), this will often
+            produce better masks than a single prediction. If only a single
+            mask is needed, the model's predicted quality score can be used
+            to select the best mask. For non-ambiguous prompts, such as multiple
+            input prompts, multimask_output=False can give better results.
+          return_logits (bool): If true, returns un-thresholded masks logits
+            instead of a binary mask.
+          normalize_coords (bool): If true, the point coordinates will be normalized to the range [0,1] and point_coords is expected to be wrt. image dimensions.
+
+        Returns:
+          (np.ndarray): The output masks in CxHxW format, where C is the
+            number of masks, and (H, W) is the original image size.
+          (np.ndarray): An array of length C containing the model's
+            predictions for the quality of each mask.
+          (np.ndarray): An array of shape CxHxW, where C is the number
+            of masks and H=W=256. These low resolution logits can be passed to
+            a subsequent iteration as mask input.
+        """
+        if not self._is_image_set:
+            raise RuntimeError(
+                "An image must be set with .set_image(...) before mask prediction."
+            )
+
+        # Transform input prompts
+
+        mask_input, unnorm_coords, labels, unnorm_box = self._prep_prompts(
+            point_coords, point_labels, box, mask_input, normalize_coords
+        )
+
+        masks, iou_predictions, low_res_masks = self._predict(
+            unnorm_coords,
+            labels,
+            unnorm_box,
+            mask_input,
+            multimask_output,
+            return_logits=return_logits,
+        )
+
+        masks_np = masks.squeeze(0).float().detach().cpu().numpy()
+        iou_predictions_np = iou_predictions.squeeze(0).float().detach().cpu().numpy()
+        low_res_masks_np = low_res_masks.squeeze(0).float().detach().cpu().numpy()
+        return masks_np, iou_predictions_np, low_res_masks_np
+
+    def _prep_prompts(
+        self, point_coords, point_labels, box, mask_logits, normalize_coords, img_idx=-1
+    ):
+        unnorm_coords, labels, unnorm_box, mask_input = None, None, None, None
+        if point_coords is not None:
+            assert (
+                point_labels is not None
+            ), "point_labels must be supplied if point_coords is supplied."
+            point_coords = torch.as_tensor(
+                point_coords, dtype=torch.float, device=self.device
+            )
+            unnorm_coords = self._transforms.transform_coords(
+                point_coords, normalize=normalize_coords, orig_hw=self._orig_hw[img_idx]
+            )
+            labels = torch.as_tensor(point_labels, dtype=torch.int, device=self.device)
+            if len(unnorm_coords.shape) == 2:
+                unnorm_coords, labels = unnorm_coords[None, ...], labels[None, ...]
+        if box is not None:
+            box = torch.as_tensor(box, dtype=torch.float, device=self.device)
+            unnorm_box = self._transforms.transform_boxes(
+                box, normalize=normalize_coords, orig_hw=self._orig_hw[img_idx]
+            )  # Bx2x2
+        if mask_logits is not None:
+            mask_input = torch.as_tensor(
+                mask_logits, dtype=torch.float, device=self.device
+            )
+            if len(mask_input.shape) == 3:
+                mask_input = mask_input[None, :, :, :]
+        return mask_input, unnorm_coords, labels, unnorm_box
+
+    @torch.no_grad()
+    def _predict(
+        self,
+        point_coords: Optional[torch.Tensor],
+        point_labels: Optional[torch.Tensor],
+        boxes: Optional[torch.Tensor] = None,
+        mask_input: Optional[torch.Tensor] = None,
+        multimask_output: bool = True,
+        return_logits: bool = False,
+        img_idx: int = -1,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Predict masks for the given input prompts, using the currently set image.
+        Input prompts are batched torch tensors and are expected to already be
+        transformed to the input frame using SAM2Transforms.
+
+        Arguments:
+          point_coords (torch.Tensor or None): A BxNx2 array of point prompts to the
+            model. Each point is in (X,Y) in pixels.
+          point_labels (torch.Tensor or None): A BxN array of labels for the
+            point prompts. 1 indicates a foreground point and 0 indicates a
+            background point.
+          boxes (np.ndarray or None): A Bx4 array given a box prompt to the
+            model, in XYXY format.
+          mask_input (np.ndarray): A low resolution mask input to the model, typically
+            coming from a previous prediction iteration. Has form Bx1xHxW, where
+            for SAM, H=W=256. Masks returned by a previous iteration of the
+            predict method do not need further transformation.
+          multimask_output (bool): If true, the model will return three masks.
+            For ambiguous input prompts (such as a single click), this will often
+            produce better masks than a single prediction. If only a single
+            mask is needed, the model's predicted quality score can be used
+            to select the best mask. For non-ambiguous prompts, such as multiple
+            input prompts, multimask_output=False can give better results.
+          return_logits (bool): If true, returns un-thresholded masks logits
+            instead of a binary mask.
+
+        Returns:
+          (torch.Tensor): The output masks in BxCxHxW format, where C is the
+            number of masks, and (H, W) is the original image size.
+          (torch.Tensor): An array of shape BxC containing the model's
+            predictions for the quality of each mask.
+          (torch.Tensor): An array of shape BxCxHxW, where C is the number
+            of masks and H=W=256. These low res logits can be passed to
+            a subsequent iteration as mask input.
+        """
+        if not self._is_image_set:
+            raise RuntimeError(
+                "An image must be set with .set_image(...) before mask prediction."
+            )
+
+        if point_coords is not None:
+            concat_points = (point_coords, point_labels)
+        else:
+            concat_points = None
+
+        # Embed prompts
+        if boxes is not None:
+            box_coords = boxes.reshape(-1, 2, 2)
+            box_labels = torch.tensor([[2, 3]], dtype=torch.int, device=boxes.device)
+            box_labels = box_labels.repeat(boxes.size(0), 1)
+            # we merge "boxes" and "points" into a single "concat_points" input (where
+            # boxes are added at the beginning) to sam_prompt_encoder
+            if concat_points is not None:
+                concat_coords = torch.cat([box_coords, concat_points[0]], dim=1)
+                concat_labels = torch.cat([box_labels, concat_points[1]], dim=1)
+                concat_points = (concat_coords, concat_labels)
+            else:
+                concat_points = (box_coords, box_labels)
+
+        sparse_embeddings, dense_embeddings = self.model.sam_prompt_encoder(
+            points=concat_points,
+            boxes=None,
+            masks=mask_input,
+        )
+
+        # Predict masks
+        batched_mode = (
+            concat_points is not None and concat_points[0].shape[0] > 1
+        )  # multi object prediction
+        high_res_features = [
+            feat_level[img_idx].unsqueeze(0)
+            for feat_level in self._features["high_res_feats"]
+        ]
+        low_res_masks, iou_predictions, _, _ = self.model.sam_mask_decoder(
+            image_embeddings=self._features["image_embed"][img_idx].unsqueeze(0),
+            image_pe=self.model.sam_prompt_encoder.get_dense_pe(),
+            sparse_prompt_embeddings=sparse_embeddings,
+            dense_prompt_embeddings=dense_embeddings,
+            multimask_output=multimask_output,
+            repeat_image=batched_mode,
+            high_res_features=high_res_features,
+        )
+
+        # Upscale the masks to the original image resolution
+        masks = self._transforms.postprocess_masks(
+            low_res_masks, self._orig_hw[img_idx]
+        )
+        low_res_masks = torch.clamp(low_res_masks, -32.0, 32.0)
+        if not return_logits:
+            masks = masks > self.mask_threshold
+
+        return masks, iou_predictions, low_res_masks
+
+    def get_image_embedding(self) -> torch.Tensor:
+        """
+        Returns the image embeddings for the currently set image, with
+        shape 1xCxHxW, where C is the embedding dimension and (H,W) are
+        the embedding spatial dimension of SAM (typically C=256, H=W=64).
+        """
+        if not self._is_image_set:
+            raise RuntimeError(
+                "An image must be set with .set_image(...) to generate an embedding."
+            )
+        assert (
+            self._features is not None
+        ), "Features must exist if an image has been set."
+        return self._features["image_embed"]
+
+    @property
+    def device(self) -> torch.device:
+        return self.model.device
+
+    def reset_predictor(self) -> None:
+        """
+        Resets the image embeddings and other state variables.
+        """
+        self._is_image_set = False
+        self._features = None
+        self._orig_hw = None
+        self._is_batch = False

sam3/model/sam3_image.py

@@ -0,0 +1,883 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import os
+from copy import deepcopy
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+
+from sam3.model.model_misc import SAM3Output
+
+from sam3.model.sam1_task_predictor import SAM3InteractiveImagePredictor
+from sam3.model.vl_combiner import SAM3VLBackbone
+from sam3.perflib.nms import nms_masks
+
+from sam3.train.data.collator import BatchedDatapoint
+
+from .act_ckpt_utils import activation_ckpt_wrapper
+
+from .box_ops import box_cxcywh_to_xyxy
+
+from .geometry_encoders import Prompt
+from .model_misc import inverse_sigmoid
+
+
+def _update_out(out, out_name, out_value, auxiliary=True, update_aux=True):
+    out[out_name] = out_value[-1] if auxiliary else out_value
+    if auxiliary and update_aux:
+        if "aux_outputs" not in out:
+            out["aux_outputs"] = [{} for _ in range(len(out_value) - 1)]
+        assert len(out["aux_outputs"]) == len(out_value) - 1
+        for aux_output, aux_value in zip(out["aux_outputs"], out_value[:-1]):
+            aux_output[out_name] = aux_value
+
+
+class Sam3Image(torch.nn.Module):
+    TEXT_ID_FOR_TEXT = 0
+    TEXT_ID_FOR_VISUAL = 1
+    TEXT_ID_FOR_GEOMETRIC = 2
+
+    def __init__(
+        self,
+        backbone: SAM3VLBackbone,
+        transformer,
+        input_geometry_encoder,
+        segmentation_head=None,
+        num_feature_levels=1,
+        o2m_mask_predict=True,
+        dot_prod_scoring=None,
+        use_instance_query: bool = True,
+        multimask_output: bool = True,
+        use_act_checkpoint_seg_head: bool = True,
+        interactivity_in_encoder: bool = True,
+        matcher=None,
+        use_dot_prod_scoring=True,
+        supervise_joint_box_scores: bool = False,  # only relevant if using presence token/score
+        detach_presence_in_joint_score: bool = False,  # only relevant if using presence token/score
+        separate_scorer_for_instance: bool = False,
+        num_interactive_steps_val: int = 0,
+        inst_interactive_predictor: SAM3InteractiveImagePredictor = None,
+        **kwargs,
+    ):
+        super().__init__()
+        self.backbone = backbone
+        self.geometry_encoder = input_geometry_encoder
+        self.transformer = transformer
+        self.hidden_dim = transformer.d_model
+        self.num_feature_levels = num_feature_levels
+        self.segmentation_head = segmentation_head
+
+        self.o2m_mask_predict = o2m_mask_predict
+
+        self.dot_prod_scoring = dot_prod_scoring
+        self.use_act_checkpoint_seg_head = use_act_checkpoint_seg_head
+        self.interactivity_in_encoder = interactivity_in_encoder
+        self.matcher = matcher
+
+        self.num_interactive_steps_val = num_interactive_steps_val
+        self.use_dot_prod_scoring = use_dot_prod_scoring
+
+        if self.use_dot_prod_scoring:
+            assert dot_prod_scoring is not None
+            self.dot_prod_scoring = dot_prod_scoring
+            self.instance_dot_prod_scoring = None
+            if separate_scorer_for_instance:
+                self.instance_dot_prod_scoring = deepcopy(dot_prod_scoring)
+        else:
+            self.class_embed = torch.nn.Linear(self.hidden_dim, 1)
+            self.instance_class_embed = None
+            if separate_scorer_for_instance:
+                self.instance_class_embed = deepcopy(self.class_embed)
+
+        self.supervise_joint_box_scores = supervise_joint_box_scores
+        self.detach_presence_in_joint_score = detach_presence_in_joint_score
+
+        # verify the number of queries for O2O and O2M
+        num_o2o_static = self.transformer.decoder.num_queries
+        num_o2m_static = self.transformer.decoder.num_o2m_queries
+        assert num_o2m_static == (num_o2o_static if self.transformer.decoder.dac else 0)
+        self.dac = self.transformer.decoder.dac
+
+        self.use_instance_query = use_instance_query
+        self.multimask_output = multimask_output
+
+        self.inst_interactive_predictor = inst_interactive_predictor
+
+    @property
+    def device(self):
+        self._device = getattr(self, "_device", None) or next(self.parameters()).device
+        return self._device
+
+    def to(self, *args, **kwargs):
+        # clear cached _device in case the model is moved to a different device
+        self._device = None
+        return super().to(*args, **kwargs)
+
+    def _get_img_feats(self, backbone_out, img_ids):
+        """Retrieve correct image features from backbone output."""
+        if "backbone_fpn" in backbone_out:
+            if "id_mapping" in backbone_out and backbone_out["id_mapping"] is not None:
+                img_ids = backbone_out["id_mapping"][img_ids]
+                # If this assert fails, it likely means we're requesting different img_ids (perhaps a different frame?)
+                # We currently don't expect this to happen. We could technically trigger a recompute here,
+                # but likely at the cost of a cpu<->gpu sync point, which would deteriorate perf
+                torch._assert_async((img_ids >= 0).all())
+
+            vis_feats = backbone_out["backbone_fpn"][-self.num_feature_levels :]
+            vis_pos_enc = backbone_out["vision_pos_enc"][-self.num_feature_levels :]
+            vis_feat_sizes = [x.shape[-2:] for x in vis_pos_enc]  # (H, W) shapes
+            # index and flatten visual features NxCxHxW => HWxNxC (batch-first => seq-first)
+            img_feats = [x[img_ids].flatten(2).permute(2, 0, 1) for x in vis_feats]
+            img_pos_embeds = [
+                x[img_ids].flatten(2).permute(2, 0, 1) for x in vis_pos_enc
+            ]
+            return backbone_out, img_feats, img_pos_embeds, vis_feat_sizes
+
+        # Image features not available in backbone output, so we compute them on the fly
+        # This case likely occurs for video. In that case, we want to forward only the current frame
+        img_batch = backbone_out["img_batch_all_stages"]
+        if img_ids.numel() > 1:
+            # Only forward backbone on unique image ids to avoid repetitive computation
+            unique_ids, _ = torch.unique(img_ids, return_inverse=True)
+        else:
+            unique_ids, _ = img_ids, slice(None)
+        # Compute the image features on those unique image ids
+        # note: we allow using a list (or other indexable types) of tensors as img_batch
+        # (e.g. for async frame loading in demo). In this case we index img_batch.tensors directly
+        if isinstance(img_batch, torch.Tensor):
+            image = img_batch[unique_ids]
+        elif unique_ids.numel() == 1:
+            image = img_batch[unique_ids.item()].unsqueeze(0)
+        else:
+            image = torch.stack([img_batch[i] for i in unique_ids.tolist()])
+        # `img_batch` might be fp16 and offloaded to CPU
+        image = image.to(dtype=torch.float32, device=self.device)
+        # Next time we call this function, we want to remember which indices we computed
+        id_mapping = torch.full(
+            (len(img_batch),), -1, dtype=torch.long, device=self.device
+        )
+        id_mapping[unique_ids] = torch.arange(len(unique_ids), device=self.device)
+        backbone_out = {
+            **backbone_out,
+            **self.backbone.forward_image(image),
+            "id_mapping": id_mapping,
+        }
+        assert "backbone_fpn" in backbone_out
+        return self._get_img_feats(backbone_out, img_ids=img_ids)
+
+    def _encode_prompt(
+        self,
+        backbone_out,
+        find_input,
+        geometric_prompt,
+        visual_prompt_embed=None,
+        visual_prompt_mask=None,
+        encode_text=True,
+        prev_mask_pred=None,
+    ):
+        # index text features (note that regardless of early or late fusion, the batch size of
+        # `txt_feats` is always the number of *prompts* in the encoder)
+        txt_ids = find_input.text_ids
+        txt_feats = backbone_out["language_features"][:, txt_ids]
+        txt_masks = backbone_out["language_mask"][txt_ids]
+
+        feat_tuple = self._get_img_feats(backbone_out, find_input.img_ids)
+        backbone_out, img_feats, img_pos_embeds, vis_feat_sizes = feat_tuple
+
+        if prev_mask_pred is not None:
+            img_feats = [img_feats[-1] + prev_mask_pred]
+        # Encode geometry
+        geo_feats, geo_masks = self.geometry_encoder(
+            geo_prompt=geometric_prompt,
+            img_feats=img_feats,
+            img_sizes=vis_feat_sizes,
+            img_pos_embeds=img_pos_embeds,
+        )
+        if visual_prompt_embed is None:
+            visual_prompt_embed = torch.zeros(
+                (0, *geo_feats.shape[1:]), device=geo_feats.device
+            )
+            visual_prompt_mask = torch.zeros(
+                (*geo_masks.shape[:-1], 0),
+                device=geo_masks.device,
+                dtype=geo_masks.dtype,
+            )
+        if encode_text:
+            prompt = torch.cat([txt_feats, geo_feats, visual_prompt_embed], dim=0)
+            prompt_mask = torch.cat([txt_masks, geo_masks, visual_prompt_mask], dim=1)
+        else:
+            prompt = torch.cat([geo_feats, visual_prompt_embed], dim=0)
+            prompt_mask = torch.cat([geo_masks, visual_prompt_mask], dim=1)
+        return prompt, prompt_mask, backbone_out
+
+    def _run_encoder(
+        self,
+        backbone_out,
+        find_input,
+        prompt,
+        prompt_mask,
+        encoder_extra_kwargs: Optional[Dict] = None,
+    ):
+        feat_tuple = self._get_img_feats(backbone_out, find_input.img_ids)
+        backbone_out, img_feats, img_pos_embeds, vis_feat_sizes = feat_tuple
+
+        # Run the encoder
+        prompt_pos_embed = torch.zeros_like(prompt)
+        # make a copy of the image feature lists since the encoder may modify these lists in-place
+        memory = self.transformer.encoder(
+            src=img_feats.copy(),
+            src_key_padding_mask=None,
+            src_pos=img_pos_embeds.copy(),
+            prompt=prompt,
+            prompt_pos=prompt_pos_embed,
+            prompt_key_padding_mask=prompt_mask,
+            feat_sizes=vis_feat_sizes,
+            encoder_extra_kwargs=encoder_extra_kwargs,
+        )
+        encoder_out = {
+            # encoded image features
+            "encoder_hidden_states": memory["memory"],
+            "pos_embed": memory["pos_embed"],
+            "padding_mask": memory["padding_mask"],
+            "level_start_index": memory["level_start_index"],
+            "spatial_shapes": memory["spatial_shapes"],
+            "valid_ratios": memory["valid_ratios"],
+            "vis_feat_sizes": vis_feat_sizes,
+            # encoded text features (or other prompts)
+            "prompt_before_enc": prompt,
+            "prompt_after_enc": memory.get("memory_text", prompt),
+            "prompt_mask": prompt_mask,
+        }
+        return backbone_out, encoder_out, feat_tuple
+
+    def _run_decoder(
+        self,
+        pos_embed,
+        memory,
+        src_mask,
+        out,
+        prompt,
+        prompt_mask,
+        encoder_out,
+    ):
+        bs = memory.shape[1]
+        query_embed = self.transformer.decoder.query_embed.weight
+        tgt = query_embed.unsqueeze(1).repeat(1, bs, 1)
+
+        apply_dac = self.transformer.decoder.dac and self.training
+        hs, reference_boxes, dec_presence_out, dec_presence_feats = (
+            self.transformer.decoder(
+                tgt=tgt,
+                memory=memory,
+                memory_key_padding_mask=src_mask,
+                pos=pos_embed,
+                reference_boxes=None,
+                level_start_index=encoder_out["level_start_index"],
+                spatial_shapes=encoder_out["spatial_shapes"],
+                valid_ratios=encoder_out["valid_ratios"],
+                tgt_mask=None,
+                memory_text=prompt,
+                text_attention_mask=prompt_mask,
+                apply_dac=apply_dac,
+            )
+        )
+        hs = hs.transpose(1, 2)  # seq-first to batch-first
+        reference_boxes = reference_boxes.transpose(1, 2)  # seq-first to batch-first
+        if dec_presence_out is not None:
+            # seq-first to batch-first
+            dec_presence_out = dec_presence_out.transpose(1, 2)
+
+        out["presence_feats"] = dec_presence_feats
+        self._update_scores_and_boxes(
+            out,
+            hs,
+            reference_boxes,
+            prompt,
+            prompt_mask,
+            dec_presence_out=dec_presence_out,
+        )
+        return out, hs
+
+    def _update_scores_and_boxes(
+        self,
+        out,
+        hs,
+        reference_boxes,
+        prompt,
+        prompt_mask,
+        dec_presence_out=None,
+        is_instance_prompt=False,
+    ):
+        apply_dac = self.transformer.decoder.dac and self.training
+        num_o2o = (hs.size(2) // 2) if apply_dac else hs.size(2)
+        num_o2m = hs.size(2) - num_o2o
+        assert num_o2m == (num_o2o if apply_dac else 0)
+        out["queries"] = hs[-1][:, :num_o2o]  # remove o2m queries if there are any
+        # score prediction
+        if self.use_dot_prod_scoring:
+            dot_prod_scoring_head = self.dot_prod_scoring
+            if is_instance_prompt and self.instance_dot_prod_scoring is not None:
+                dot_prod_scoring_head = self.instance_dot_prod_scoring
+            outputs_class = dot_prod_scoring_head(hs, prompt, prompt_mask)
+        else:
+            class_embed_head = self.class_embed
+            if is_instance_prompt and self.instance_class_embed is not None:
+                class_embed_head = self.instance_class_embed
+            outputs_class = class_embed_head(hs)
+
+        # box prediction
+        box_head = self.transformer.decoder.bbox_embed
+        if (
+            is_instance_prompt
+            and self.transformer.decoder.instance_bbox_embed is not None
+        ):
+            box_head = self.transformer.decoder.instance_bbox_embed
+        anchor_box_offsets = box_head(hs)
+        reference_boxes_inv_sig = inverse_sigmoid(reference_boxes)
+        outputs_coord = (reference_boxes_inv_sig + anchor_box_offsets).sigmoid()
+        outputs_boxes_xyxy = box_cxcywh_to_xyxy(outputs_coord)
+
+        if dec_presence_out is not None:
+            _update_out(
+                out, "presence_logit_dec", dec_presence_out, update_aux=self.training
+            )
+
+        if self.supervise_joint_box_scores:
+            assert dec_presence_out is not None
+            prob_dec_presence_out = dec_presence_out.clone().sigmoid()
+            if self.detach_presence_in_joint_score:
+                prob_dec_presence_out = prob_dec_presence_out.detach()
+
+            outputs_class = inverse_sigmoid(
+                outputs_class.sigmoid() * prob_dec_presence_out.unsqueeze(2)
+            ).clamp(min=-10.0, max=10.0)
+
+        _update_out(
+            out, "pred_logits", outputs_class[:, :, :num_o2o], update_aux=self.training
+        )
+        _update_out(
+            out, "pred_boxes", outputs_coord[:, :, :num_o2o], update_aux=self.training
+        )
+        _update_out(
+            out,
+            "pred_boxes_xyxy",
+            outputs_boxes_xyxy[:, :, :num_o2o],
+            update_aux=self.training,
+        )
+        if num_o2m > 0 and self.training:
+            _update_out(
+                out,
+                "pred_logits_o2m",
+                outputs_class[:, :, num_o2o:],
+                update_aux=self.training,
+            )
+            _update_out(
+                out,
+                "pred_boxes_o2m",
+                outputs_coord[:, :, num_o2o:],
+                update_aux=self.training,
+            )
+            _update_out(
+                out,
+                "pred_boxes_xyxy_o2m",
+                outputs_boxes_xyxy[:, :, num_o2o:],
+                update_aux=self.training,
+            )
+
+    def _run_segmentation_heads(
+        self,
+        out,
+        backbone_out,
+        img_ids,
+        vis_feat_sizes,
+        encoder_hidden_states,
+        prompt,
+        prompt_mask,
+        hs,
+    ):
+        apply_dac = self.transformer.decoder.dac and self.training
+        if self.segmentation_head is not None:
+            num_o2o = (hs.size(2) // 2) if apply_dac else hs.size(2)
+            num_o2m = hs.size(2) - num_o2o
+            obj_queries = hs if self.o2m_mask_predict else hs[:, :, :num_o2o]
+            seg_head_outputs = activation_ckpt_wrapper(self.segmentation_head)(
+                backbone_feats=backbone_out["backbone_fpn"],
+                obj_queries=obj_queries,
+                image_ids=img_ids,
+                encoder_hidden_states=encoder_hidden_states,
+                act_ckpt_enable=self.training and self.use_act_checkpoint_seg_head,
+                prompt=prompt,
+                prompt_mask=prompt_mask,
+            )
+            aux_masks = False  # self.aux_loss and self.segmentation_head.aux_masks
+            for k, v in seg_head_outputs.items():
+                if k in self.segmentation_head.instance_keys:
+                    _update_out(out, k, v[:, :num_o2o], auxiliary=aux_masks)
+                    if (
+                        self.o2m_mask_predict and num_o2m > 0
+                    ):  # handle o2m mask prediction
+                        _update_out(
+                            out, f"{k}_o2m", v[:, num_o2o:], auxiliary=aux_masks
+                        )
+                else:
+                    out[k] = v
+        else:
+            backbone_out.pop("backbone_fpn", None)
+
+    def _get_best_mask(self, out):
+        prev_mask_idx = out["pred_logits"].argmax(dim=1).squeeze(1)
+        batch_idx = torch.arange(
+            out["pred_logits"].shape[0], device=prev_mask_idx.device
+        )
+        prev_mask_pred = out["pred_masks"][batch_idx, prev_mask_idx][:, None]
+        # Downsample mask to match image resolution.
+        prev_mask_pred = self.geometry_encoder.mask_encoder.mask_downsampler(
+            prev_mask_pred
+        )
+        prev_mask_pred = prev_mask_pred.flatten(-2).permute(2, 0, 1)
+
+        return prev_mask_pred
+
+    def forward_grounding(
+        self,
+        backbone_out,
+        find_input,
+        find_target,
+        geometric_prompt: Prompt,
+    ):
+        with torch.profiler.record_function("SAM3Image._encode_prompt"):
+            prompt, prompt_mask, backbone_out = self._encode_prompt(
+                backbone_out, find_input, geometric_prompt
+            )
+        # Run the encoder
+        with torch.profiler.record_function("SAM3Image._run_encoder"):
+            backbone_out, encoder_out, _ = self._run_encoder(
+                backbone_out, find_input, prompt, prompt_mask
+            )
+        out = {
+            "encoder_hidden_states": encoder_out["encoder_hidden_states"],
+            "prev_encoder_out": {
+                "encoder_out": encoder_out,
+                "backbone_out": backbone_out,
+            },
+        }
+
+        # Run the decoder
+        with torch.profiler.record_function("SAM3Image._run_decoder"):
+            out, hs = self._run_decoder(
+                memory=out["encoder_hidden_states"],
+                pos_embed=encoder_out["pos_embed"],
+                src_mask=encoder_out["padding_mask"],
+                out=out,
+                prompt=prompt,
+                prompt_mask=prompt_mask,
+                encoder_out=encoder_out,
+            )
+
+        # Run segmentation heads
+        with torch.profiler.record_function("SAM3Image._run_segmentation_heads"):
+            self._run_segmentation_heads(
+                out=out,
+                backbone_out=backbone_out,
+                img_ids=find_input.img_ids,
+                vis_feat_sizes=encoder_out["vis_feat_sizes"],
+                encoder_hidden_states=out["encoder_hidden_states"],
+                prompt=prompt,
+                prompt_mask=prompt_mask,
+                hs=hs,
+            )
+
+        if self.training or self.num_interactive_steps_val > 0:
+            self._compute_matching(out, self.back_convert(find_target))
+        return out
+
+    def _postprocess_out(self, out: Dict, multimask_output: bool = False):
+        # For multimask output, during eval we return the single best mask with the dict keys expected by the evaluators, but also return the multimasks output with new keys.
+        num_mask_boxes = out["pred_boxes"].size(1)
+        if not self.training and multimask_output and num_mask_boxes > 1:
+            out["multi_pred_logits"] = out["pred_logits"]
+            if "pred_masks" in out:
+                out["multi_pred_masks"] = out["pred_masks"]
+            out["multi_pred_boxes"] = out["pred_boxes"]
+            out["multi_pred_boxes_xyxy"] = out["pred_boxes_xyxy"]
+
+            best_mask_idx = out["pred_logits"].argmax(1).squeeze(1)
+            batch_idx = torch.arange(len(best_mask_idx), device=best_mask_idx.device)
+
+            out["pred_logits"] = out["pred_logits"][batch_idx, best_mask_idx].unsqueeze(
+                1
+            )
+            if "pred_masks" in out:
+                out["pred_masks"] = out["pred_masks"][
+                    batch_idx, best_mask_idx
+                ].unsqueeze(1)
+            out["pred_boxes"] = out["pred_boxes"][batch_idx, best_mask_idx].unsqueeze(1)
+            out["pred_boxes_xyxy"] = out["pred_boxes_xyxy"][
+                batch_idx, best_mask_idx
+            ].unsqueeze(1)
+
+        return out
+
+    def _get_dummy_prompt(self, num_prompts=1):
+        device = self.device
+        geometric_prompt = Prompt(
+            box_embeddings=torch.zeros(0, num_prompts, 4, device=device),
+            box_mask=torch.zeros(num_prompts, 0, device=device, dtype=torch.bool),
+        )
+        return geometric_prompt
+
+    def forward(self, input: BatchedDatapoint):
+        device = self.device
+        backbone_out = {"img_batch_all_stages": input.img_batch}
+        backbone_out.update(self.backbone.forward_image(input.img_batch))
+        num_frames = len(input.find_inputs)
+        assert num_frames == 1
+
+        text_outputs = self.backbone.forward_text(input.find_text_batch, device=device)
+        backbone_out.update(text_outputs)
+
+        previous_stages_out = SAM3Output(
+            iter_mode=SAM3Output.IterMode.LAST_STEP_PER_STAGE
+        )
+
+        find_input = input.find_inputs[0]
+        find_target = input.find_targets[0]
+
+        if find_input.input_points is not None and find_input.input_points.numel() > 0:
+            print("Warning: Point prompts are ignored in PCS.")
+
+        num_interactive_steps = 0 if self.training else self.num_interactive_steps_val
+        geometric_prompt = Prompt(
+            box_embeddings=find_input.input_boxes,
+            box_mask=find_input.input_boxes_mask,
+            box_labels=find_input.input_boxes_label,
+        )
+
+        # Init vars that are shared across the loop.
+        stage_outs = []
+        for cur_step in range(num_interactive_steps + 1):
+            if cur_step > 0:
+                # We sample interactive geometric prompts (boxes, points)
+                geometric_prompt, _ = self.interactive_prompt_sampler.sample(
+                    geo_prompt=geometric_prompt,
+                    find_target=find_target,
+                    previous_out=stage_outs[-1],
+                )
+            out = self.forward_grounding(
+                backbone_out=backbone_out,
+                find_input=find_input,
+                find_target=find_target,
+                geometric_prompt=geometric_prompt.clone(),
+            )
+            stage_outs.append(out)
+
+        previous_stages_out.append(stage_outs)
+        return previous_stages_out
+
+    def _compute_matching(self, out, targets):
+        out["indices"] = self.matcher(out, targets)
+        for aux_out in out.get("aux_outputs", []):
+            aux_out["indices"] = self.matcher(aux_out, targets)
+
+    def back_convert(self, targets):
+        batched_targets = {
+            "boxes": targets.boxes.view(-1, 4),
+            "boxes_xyxy": box_cxcywh_to_xyxy(targets.boxes.view(-1, 4)),
+            "boxes_padded": targets.boxes_padded,
+            "positive_map": targets.boxes.new_ones(len(targets.boxes), 1),
+            "num_boxes": targets.num_boxes,
+            "masks": targets.segments,
+            "semantic_masks": targets.semantic_segments,
+            "is_valid_mask": targets.is_valid_segment,
+            "is_exhaustive": targets.is_exhaustive,
+            "object_ids_packed": targets.object_ids,
+            "object_ids_padded": targets.object_ids_padded,
+        }
+        return batched_targets
+
+    def predict_inst(
+        self,
+        inference_state,
+        **kwargs,
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        orig_h, orig_w = (
+            inference_state["original_height"],
+            inference_state["original_width"],
+        )
+        backbone_out = inference_state["backbone_out"]["sam2_backbone_out"]
+        (
+            _,
+            vision_feats,
+            _,
+            _,
+        ) = self.inst_interactive_predictor.model._prepare_backbone_features(
+            backbone_out
+        )
+        # Add no_mem_embed, which is added to the lowest rest feat. map during training on videos
+        vision_feats[-1] = (
+            vision_feats[-1] + self.inst_interactive_predictor.model.no_mem_embed
+        )
+        feats = [
+            feat.permute(1, 2, 0).view(1, -1, *feat_size)
+            for feat, feat_size in zip(
+                vision_feats[::-1], self.inst_interactive_predictor._bb_feat_sizes[::-1]
+            )
+        ][::-1]
+        self.inst_interactive_predictor._features = {
+            "image_embed": feats[-1],
+            "high_res_feats": feats[:-1],
+        }
+        self.inst_interactive_predictor._is_image_set = True
+        self.inst_interactive_predictor._orig_hw = [(orig_h, orig_w)]
+        res = self.inst_interactive_predictor.predict(**kwargs)
+        self.inst_interactive_predictor._features = None
+        self.inst_interactive_predictor._is_image_set = False
+        return res
+
+    def predict_inst_batch(
+        self,
+        inference_state,
+        *args,
+        **kwargs,
+    ) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray]]:
+        backbone_out = inference_state["backbone_out"]["sam2_backbone_out"]
+        (
+            _,
+            vision_feats,
+            _,
+            _,
+        ) = self.inst_interactive_predictor.model._prepare_backbone_features(
+            backbone_out
+        )
+        # Add no_mem_embed, which is added to the lowest res feat. map during training on videos
+        vision_feats[-1] = (
+            vision_feats[-1] + self.inst_interactive_predictor.model.no_mem_embed
+        )
+        batch_size = vision_feats[-1].shape[1]
+        orig_heights, orig_widths = (
+            inference_state["original_heights"],
+            inference_state["original_widths"],
+        )
+        assert (
+            batch_size == len(orig_heights) == len(orig_widths)
+        ), f"Batch size mismatch in predict_inst_batch. Got {batch_size}, {len(orig_heights)}, {len(orig_widths)}"
+        feats = [
+            feat.permute(1, 2, 0).view(batch_size, -1, *feat_size)
+            for feat, feat_size in zip(
+                vision_feats[::-1], self.inst_interactive_predictor._bb_feat_sizes[::-1]
+            )
+        ][::-1]
+        self.inst_interactive_predictor._features = {
+            "image_embed": feats[-1],
+            "high_res_feats": feats[:-1],
+        }
+        self.inst_interactive_predictor._is_image_set = True
+        self.inst_interactive_predictor._is_batch = True
+        self.inst_interactive_predictor._orig_hw = [
+            (orig_h, orig_w) for orig_h, orig_w in zip(orig_heights, orig_widths)
+        ]
+        res = self.inst_interactive_predictor.predict_batch(*args, **kwargs)
+        self.inst_interactive_predictor._features = None
+        self.inst_interactive_predictor._is_image_set = False
+        self.inst_interactive_predictor._is_batch = False
+        return res
+
+
+class Sam3ImageOnVideoMultiGPU(Sam3Image):
+    def __init__(
+        self, *args, async_all_gather=True, gather_backbone_out=None, **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.rank = int(os.getenv("RANK", "0"))
+        self.world_size = int(os.getenv("WORLD_SIZE", "1"))
+        self.async_all_gather = async_all_gather
+
+        # if gather_backbone is not set, default to gathering only for `SAM3VLBackbone`
+        if gather_backbone_out is None:
+            gather_backbone_out = isinstance(self.backbone, SAM3VLBackbone)
+        self.gather_backbone_out = gather_backbone_out
+
+    def forward_video_grounding_multigpu(
+        self,
+        backbone_out,
+        find_inputs,
+        geometric_prompt: Prompt,
+        frame_idx,
+        num_frames,
+        # `multigpu_buffer` is a dict to cache detector's outputs in a chunk between different calls
+        multigpu_buffer,
+        track_in_reverse=False,
+        # whether to also return the SAM2 backbone features
+        return_sam2_backbone_feats=False,
+        # whether to perform NMS and suppress the scores of those detections removed by NMS
+        run_nms=False,
+        nms_prob_thresh=None,
+        nms_iou_thresh=None,
+        **kwargs,
+    ):
+        """
+        Compute the detector's detection outputs in a distributed manner, where all GPUs process
+        a chunk of frames (equal to the number of GPUs) at once and store them in cache.
+        """
+        # Step 1: fetch the detector outputs in the current chunk from buffer
+        frame_idx_curr_b = frame_idx - frame_idx % self.world_size
+        frame_idx_curr_e = min(frame_idx_curr_b + self.world_size, num_frames)
+        # in case the current frame's detection results are not in the buffer yet, build the current chunk
+        # (this should only happen on the first chunk, since we are also building the next chunk below)
+        if frame_idx not in multigpu_buffer:
+            with torch.profiler.record_function("build_multigpu_buffer_next_chunk1"):
+                self._build_multigpu_buffer_next_chunk(
+                    backbone_out=backbone_out,
+                    find_inputs=find_inputs,
+                    geometric_prompt=geometric_prompt,
+                    frame_idx_begin=frame_idx_curr_b,
+                    frame_idx_end=frame_idx_curr_e,
+                    num_frames=num_frames,
+                    multigpu_buffer=multigpu_buffer,
+                    run_nms=run_nms,
+                    nms_prob_thresh=nms_prob_thresh,
+                    nms_iou_thresh=nms_iou_thresh,
+                )
+
+        # read out the current frame's results from `multigpu_buffer`
+        out = {}
+        for k, (v, handle) in multigpu_buffer[frame_idx].items():
+            if k.startswith("sam2_backbone_") and not return_sam2_backbone_feats:
+                continue
+            if handle is not None:
+                handle.wait()  # wait for async all-gather to finish
+            out[k] = v
+
+        # Step 2: remove detection outputs of the previous chunk from cache to save GPU memory
+        if not track_in_reverse and frame_idx_curr_b - self.world_size >= 0:
+            frame_idx_prev_e = frame_idx_curr_b
+            frame_idx_prev_b = frame_idx_curr_b - self.world_size
+        elif track_in_reverse and frame_idx_curr_e < num_frames:
+            frame_idx_prev_b = frame_idx_curr_e
+            frame_idx_prev_e = min(frame_idx_prev_b + self.world_size, num_frames)
+        else:
+            frame_idx_prev_b = frame_idx_prev_e = None
+        if frame_idx_prev_b is not None:
+            for frame_idx_rm in range(frame_idx_prev_b, frame_idx_prev_e):
+                multigpu_buffer.pop(frame_idx_rm, None)
+
+        # Step 3: compute and cache detection outputs of the next chunk ahead of time
+        # (so that we can overlap computation with all-gather transfer)
+        if not track_in_reverse and frame_idx_curr_e < num_frames:
+            frame_idx_next_b = frame_idx_curr_e
+            frame_idx_next_e = min(frame_idx_next_b + self.world_size, num_frames)
+        elif track_in_reverse and frame_idx_curr_b - self.world_size >= 0:
+            frame_idx_next_e = frame_idx_curr_b
+            frame_idx_next_b = frame_idx_curr_b - self.world_size
+        else:
+            frame_idx_next_b = frame_idx_next_e = None
+        if frame_idx_next_b is not None and frame_idx_next_b not in multigpu_buffer:
+            with torch.profiler.record_function("build_multigpu_buffer_next_chunk2"):
+                self._build_multigpu_buffer_next_chunk(
+                    backbone_out=backbone_out,
+                    find_inputs=find_inputs,
+                    geometric_prompt=geometric_prompt,
+                    frame_idx_begin=frame_idx_next_b,
+                    frame_idx_end=frame_idx_next_e,
+                    num_frames=num_frames,
+                    multigpu_buffer=multigpu_buffer,
+                    run_nms=run_nms,
+                    nms_prob_thresh=nms_prob_thresh,
+                    nms_iou_thresh=nms_iou_thresh,
+                )
+
+        return out, backbone_out
+
+    def _build_multigpu_buffer_next_chunk(
+        self,
+        backbone_out,
+        find_inputs,
+        geometric_prompt: Prompt,
+        frame_idx_begin,
+        frame_idx_end,
+        num_frames,
+        multigpu_buffer,
+        run_nms=False,
+        nms_prob_thresh=None,
+        nms_iou_thresh=None,
+    ):
+        """Compute detection outputs on a chunk of frames and store their results in multigpu_buffer."""
+        # each GPU computes detections on one frame in the chunk (in a round-robin manner)
+        frame_idx_local_gpu = min(frame_idx_begin + self.rank, frame_idx_end - 1)
+        # `forward_grounding` (from base class `Sam3ImageOnVideo`) runs the detector on a single frame
+        with torch.profiler.record_function("forward_grounding"):
+            out_local = self.forward_grounding(
+                backbone_out=backbone_out,
+                find_input=find_inputs[frame_idx_local_gpu],
+                find_target=None,
+                geometric_prompt=geometric_prompt,
+            )
+        if run_nms:
+            with torch.profiler.record_function("nms_masks"):
+                # run NMS as a post-processing step on top of the detection outputs
+                assert nms_prob_thresh is not None and nms_iou_thresh is not None
+                pred_probs = out_local["pred_logits"].squeeze(-1).sigmoid()
+                pred_masks = out_local["pred_masks"]
+                # loop over text prompts (not an overhead for demo where there's only 1 prompt)
+                for prompt_idx in range(pred_probs.size(0)):
+                    keep = nms_masks(
+                        pred_probs=pred_probs[prompt_idx],
+                        pred_masks=pred_masks[prompt_idx],
+                        prob_threshold=nms_prob_thresh,
+                        iou_threshold=nms_iou_thresh,
+                    )
+                    # set a very low threshold for those detections removed by NMS
+                    out_local["pred_logits"][prompt_idx, :, 0] -= 1e4 * (~keep).float()
+
+        if self.gather_backbone_out:
+            # gather the SAM 2 backbone features across GPUs
+            feats = out_local["prev_encoder_out"]["backbone_out"]["sam2_backbone_out"]
+            assert len(feats["backbone_fpn"]) == 3  # SAM2 backbone always have 3 levels
+            # cast the SAM2 backbone features to bfloat16 for all-gather (this is usually
+            # a no-op, SAM2 backbone features are likely already in bfloat16 due to AMP)
+            backbone_fpn_bf16 = [x.to(torch.bfloat16) for x in feats["backbone_fpn"]]
+            fpn0, fpn_handle0 = self._gather_tensor(backbone_fpn_bf16[0])
+            fpn1, fpn_handle1 = self._gather_tensor(backbone_fpn_bf16[1])
+            fpn2, fpn_handle2 = self._gather_tensor(backbone_fpn_bf16[2])
+            # vision_pos_enc is the same on all frames, so no need to all-gather them
+            vision_pos_enc = feats["vision_pos_enc"]
+
+        # trim the detector output to only include the necessary keys
+        out_local = {
+            "pred_logits": out_local["pred_logits"],
+            "pred_boxes": out_local["pred_boxes"],
+            "pred_boxes_xyxy": out_local["pred_boxes_xyxy"],
+            "pred_masks": out_local["pred_masks"],
+        }
+
+        # gather the results: after this step, each GPU will receive detector outputs on
+        # all frames in the chunk and store them in `multigpu_buffer`
+        out_gathered = {k: self._gather_tensor(v) for k, v in out_local.items()}
+        for rank in range(self.world_size):
+            frame_idx_to_save = frame_idx_begin + rank
+            if frame_idx_to_save >= num_frames:
+                continue
+            frame_buffer = {
+                k: (v[rank], handle) for k, (v, handle) in out_gathered.items()
+            }
+            if self.gather_backbone_out:
+                # also add gathered SAM 2 backbone features to frame_buffer
+                frame_buffer["tracker_backbone_fpn_0"] = (fpn0[rank], fpn_handle0)
+                frame_buffer["tracker_backbone_fpn_1"] = (fpn1[rank], fpn_handle1)
+                frame_buffer["tracker_backbone_fpn_2"] = (fpn2[rank], fpn_handle2)
+                frame_buffer["tracker_backbone_pos_enc"] = (vision_pos_enc, None)
+
+            multigpu_buffer[frame_idx_to_save] = frame_buffer
+
+    def _gather_tensor(self, x):
+        if self.world_size == 1:
+            return [x], None
+
+        async_op = self.async_all_gather
+        # here `.contiguous()` is required -- otherwise NCCL all_gather
+        # sometimes gives wrong results
+        x = x.contiguous()  # ensure contiguous memory for NCCL
+        output_list = [torch.empty_like(x) for _ in range(self.world_size)]
+        handle = torch.distributed.all_gather(output_list, x, async_op=async_op)
+        return output_list, handle

sam3/model/sam3_image_processor.py

@@ -0,0 +1,222 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+from typing import Dict, List
+
+import numpy as np
+import PIL
+import torch
+
+from sam3.model import box_ops
+
+from sam3.model.data_misc import FindStage, interpolate
+from torchvision.transforms import v2
+
+
+class Sam3Processor:
+    """ """
+
+    def __init__(self, model, resolution=1008, device="cuda", confidence_threshold=0.5):
+        self.model = model
+        self.resolution = resolution
+        self.device = device
+        self.transform = v2.Compose(
+            [
+                v2.ToDtype(torch.uint8, scale=True),
+                v2.Resize(size=(resolution, resolution)),
+                v2.ToDtype(torch.float32, scale=True),
+                v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+            ]
+        )
+        self.confidence_threshold = confidence_threshold
+
+        self.find_stage = FindStage(
+            img_ids=torch.tensor([0], device=device, dtype=torch.long),
+            text_ids=torch.tensor([0], device=device, dtype=torch.long),
+            input_boxes=None,
+            input_boxes_mask=None,
+            input_boxes_label=None,
+            input_points=None,
+            input_points_mask=None,
+        )
+
+    @torch.inference_mode()
+    def set_image(self, image, state=None):
+        """Sets the image on which we want to do predictions."""
+        if state is None:
+            state = {}
+
+        if isinstance(image, PIL.Image.Image):
+            width, height = image.size
+        elif isinstance(image, (torch.Tensor, np.ndarray)):
+            height, width = image.shape[-2:]
+        else:
+            raise ValueError("Image must be a PIL image or a tensor")
+
+        image = v2.functional.to_image(image).to(self.device)
+        image = self.transform(image).unsqueeze(0)
+
+        state["original_height"] = height
+        state["original_width"] = width
+        state["backbone_out"] = self.model.backbone.forward_image(image)
+        inst_interactivity_en = self.model.inst_interactive_predictor is not None
+        if inst_interactivity_en and "sam2_backbone_out" in state["backbone_out"]:
+            sam2_backbone_out = state["backbone_out"]["sam2_backbone_out"]
+            sam2_backbone_out["backbone_fpn"][0] = (
+                self.model.inst_interactive_predictor.model.sam_mask_decoder.conv_s0(
+                    sam2_backbone_out["backbone_fpn"][0]
+                )
+            )
+            sam2_backbone_out["backbone_fpn"][1] = (
+                self.model.inst_interactive_predictor.model.sam_mask_decoder.conv_s1(
+                    sam2_backbone_out["backbone_fpn"][1]
+                )
+            )
+        return state
+
+    @torch.inference_mode()
+    def set_image_batch(self, images: List[np.ndarray], state=None):
+        """Sets the image batch on which we want to do predictions."""
+        if state is None:
+            state = {}
+
+        if not isinstance(images, list):
+            raise ValueError("Images must be a list of PIL images or tensors")
+        assert len(images) > 0, "Images list must not be empty"
+        assert isinstance(
+            images[0], PIL.Image.Image
+        ), "Images must be a list of PIL images"
+
+        state["original_heights"] = [image.height for image in images]
+        state["original_widths"] = [image.width for image in images]
+
+        images = [
+            self.transform(v2.functional.to_image(image).to(self.device))
+            for image in images
+        ]
+        images = torch.stack(images, dim=0)
+        state["backbone_out"] = self.model.backbone.forward_image(images)
+        inst_interactivity_en = self.model.inst_interactive_predictor is not None
+        if inst_interactivity_en and "sam2_backbone_out" in state["backbone_out"]:
+            sam2_backbone_out = state["backbone_out"]["sam2_backbone_out"]
+            sam2_backbone_out["backbone_fpn"][0] = (
+                self.model.inst_interactive_predictor.model.sam_mask_decoder.conv_s0(
+                    sam2_backbone_out["backbone_fpn"][0]
+                )
+            )
+            sam2_backbone_out["backbone_fpn"][1] = (
+                self.model.inst_interactive_predictor.model.sam_mask_decoder.conv_s1(
+                    sam2_backbone_out["backbone_fpn"][1]
+                )
+            )
+        return state
+
+    @torch.inference_mode()
+    def set_text_prompt(self, prompt: str, state: Dict):
+        """Sets the text prompt and run the inference"""
+
+        if "backbone_out" not in state:
+            raise ValueError("You must call set_image before set_text_prompt")
+
+        text_outputs = self.model.backbone.forward_text([prompt], device=self.device)
+        # will erase the previous text prompt if any
+        state["backbone_out"].update(text_outputs)
+        if "geometric_prompt" not in state:
+            state["geometric_prompt"] = self.model._get_dummy_prompt()
+
+        return self._forward_grounding(state)
+
+    @torch.inference_mode()
+    def add_geometric_prompt(self, box: List, label: bool, state: Dict):
+        """Adds a box prompt and run the inference.
+        The image needs to be set, but not necessarily the text prompt.
+        The box is assumed to be in [center_x, center_y, width, height] format and normalized in [0, 1] range.
+        The label is True for a positive box, False for a negative box.
+        """
+        if "backbone_out" not in state:
+            raise ValueError("You must call set_image before set_text_prompt")
+
+        if "language_features" not in state["backbone_out"]:
+            # Looks like we don't have a text prompt yet. This is allowed, but we need to set the text prompt to "visual" for the model to rely only on the geometric prompt
+            dummy_text_outputs = self.model.backbone.forward_text(
+                ["visual"], device=self.device
+            )
+            state["backbone_out"].update(dummy_text_outputs)
+
+        if "geometric_prompt" not in state:
+            state["geometric_prompt"] = self.model._get_dummy_prompt()
+
+        # adding a batch and sequence dimension
+        boxes = torch.tensor(box, device=self.device, dtype=torch.float32).view(1, 1, 4)
+        labels = torch.tensor([label], device=self.device, dtype=torch.bool).view(1, 1)
+        state["geometric_prompt"].append_boxes(boxes, labels)
+
+        return self._forward_grounding(state)
+
+    def reset_all_prompts(self, state: Dict):
+        """Removes all the prompts and results"""
+        if "backbone_out" in state:
+            backbone_keys_to_del = [
+                "language_features",
+                "language_mask",
+                "language_embeds",
+            ]
+            for key in backbone_keys_to_del:
+                if key in state["backbone_out"]:
+                    del state["backbone_out"][key]
+
+        keys_to_del = ["geometric_prompt", "boxes", "masks", "masks_logits", "scores"]
+        for key in keys_to_del:
+            if key in state:
+                del state[key]
+
+    @torch.inference_mode()
+    def set_confidence_threshold(self, threshold: float, state=None):
+        """Sets the confidence threshold for the masks"""
+        self.confidence_threshold = threshold
+        if state is not None and "boxes" in state:
+            # we need to filter the boxes again
+            # In principle we could do this more efficiently since we would only need
+            # to rerun the heads. But this is simpler and not too inefficient
+            return self._forward_grounding(state)
+        return state
+
+    @torch.inference_mode()
+    def _forward_grounding(self, state: Dict):
+        outputs = self.model.forward_grounding(
+            backbone_out=state["backbone_out"],
+            find_input=self.find_stage,
+            geometric_prompt=state["geometric_prompt"],
+            find_target=None,
+        )
+
+        out_bbox = outputs["pred_boxes"]
+        out_logits = outputs["pred_logits"]
+        out_masks = outputs["pred_masks"]
+        out_probs = out_logits.sigmoid()
+        presence_score = outputs["presence_logit_dec"].sigmoid().unsqueeze(1)
+        out_probs = (out_probs * presence_score).squeeze(-1)
+
+        keep = out_probs > self.confidence_threshold
+        out_probs = out_probs[keep]
+        out_masks = out_masks[keep]
+        out_bbox = out_bbox[keep]
+
+        # convert to [x0, y0, x1, y1] format
+        boxes = box_ops.box_cxcywh_to_xyxy(out_bbox)
+
+        img_h = state["original_height"]
+        img_w = state["original_width"]
+        scale_fct = torch.tensor([img_w, img_h, img_w, img_h]).to(self.device)
+        boxes = boxes * scale_fct[None, :]
+
+        out_masks = interpolate(
+            out_masks.unsqueeze(1),
+            (img_h, img_w),
+            mode="bilinear",
+            align_corners=False,
+        ).sigmoid()
+
+        state["masks_logits"] = out_masks
+        state["masks"] = out_masks > 0.5
+        state["boxes"] = boxes
+        state["scores"] = out_probs
+        return state

sam3/model/sam3_tracker_utils.py

@@ -0,0 +1,427 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from numpy.typing import NDArray
+
+from sam3.model.edt import edt_triton
+
+
+def sample_box_points(
+    masks: torch.Tensor,
+    noise: float = 0.1,  # SAM default
+    noise_bound: int = 20,  # SAM default
+    top_left_label: int = 2,
+    bottom_right_label: int = 3,
+) -> tuple[NDArray, NDArray]:
+    """
+    Sample a noised version of the top left and bottom right corners of a given `bbox`
+
+    Inputs:
+    - masks: [B, 1, H, W] tensor
+    - noise: noise as a fraction of box width and height, dtype=float
+    - noise_bound: maximum amount of noise (in pure pixels), dtype=int
+
+    Returns:
+    - box_coords: [B, num_pt, 2], contains (x, y) coordinates of top left and bottom right box corners, dtype=torch.float
+    - box_labels: [B, num_pt], label 2 is reserverd for top left and 3 for bottom right corners, dtype=torch.int32
+    """
+    device = masks.device
+    box_coords = mask_to_box(masks)
+    B, _, H, W = masks.shape
+    box_labels = torch.tensor(
+        [top_left_label, bottom_right_label], dtype=torch.int, device=device
+    ).repeat(B)
+    if noise > 0.0:
+        if not isinstance(noise_bound, torch.Tensor):
+            noise_bound = torch.tensor(noise_bound, device=device)
+        bbox_w = box_coords[..., 2] - box_coords[..., 0]
+        bbox_h = box_coords[..., 3] - box_coords[..., 1]
+        max_dx = torch.min(bbox_w * noise, noise_bound)
+        max_dy = torch.min(bbox_h * noise, noise_bound)
+        box_noise = 2 * torch.rand(B, 1, 4, device=device) - 1
+        box_noise = box_noise * torch.stack((max_dx, max_dy, max_dx, max_dy), dim=-1)
+
+        box_coords = box_coords + box_noise
+        img_bounds = (
+            torch.tensor([W, H, W, H], device=device) - 1
+        )  # uncentered pixel coords
+        box_coords.clamp_(torch.zeros_like(img_bounds), img_bounds)  # In place clamping
+
+    box_coords = box_coords.reshape(-1, 2, 2)  # always 2 points
+    box_labels = box_labels.reshape(-1, 2)
+    return box_coords, box_labels
+
+
+def mask_to_box(masks: torch.Tensor):
+    """
+    compute bounding box given an input mask
+
+    Inputs:
+    - masks: [B, 1, H, W] tensor
+
+    Returns:
+    - box_coords: [B, 1, 4], contains (x, y) coordinates of top left and bottom right box corners, dtype=torch.Tensor
+    """
+    B, _, h, w = masks.shape
+    device = masks.device
+    mask_area = masks.sum(dim=(-1, -2))
+    xs = torch.arange(w, device=device, dtype=torch.int32)
+    ys = torch.arange(h, device=device, dtype=torch.int32)
+    grid_xs, grid_ys = torch.meshgrid(xs, ys, indexing="xy")
+    grid_xs = grid_xs[None, None, ...].expand(B, 1, h, w)
+    grid_ys = grid_ys[None, None, ...].expand(B, 1, h, w)
+    min_xs, _ = torch.min(torch.where(masks, grid_xs, w).flatten(-2), dim=-1)
+    max_xs, _ = torch.max(torch.where(masks, grid_xs, -1).flatten(-2), dim=-1)
+    min_ys, _ = torch.min(torch.where(masks, grid_ys, h).flatten(-2), dim=-1)
+    max_ys, _ = torch.max(torch.where(masks, grid_ys, -1).flatten(-2), dim=-1)
+    bbox_coords = torch.stack((min_xs, min_ys, max_xs, max_ys), dim=-1)
+    bbox_coords = torch.where(
+        mask_area[..., None] > 0, bbox_coords, torch.zeros_like(bbox_coords)
+    )
+    return bbox_coords
+
+
+def sample_random_points_from_errors(gt_masks, pred_masks, num_pt=1):
+    """
+    Sample `num_pt` random points (along with their labels) independently from the error regions.
+
+    Inputs:
+    - gt_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool
+    - pred_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool or None
+    - num_pt: int, number of points to sample independently for each of the B error maps
+
+    Outputs:
+    - points: [B, num_pt, 2], dtype=torch.float, contains (x, y) coordinates of each sampled point
+    - labels: [B, num_pt], dtype=torch.int32, where 1 means positive clicks and 0 means
+      negative clicks
+    """
+    if pred_masks is None:  # if pred_masks is not provided, treat it as empty
+        pred_masks = torch.zeros_like(gt_masks)
+    assert gt_masks.dtype == torch.bool and gt_masks.size(1) == 1
+    assert pred_masks.dtype == torch.bool and pred_masks.shape == gt_masks.shape
+    assert num_pt >= 0
+
+    B, _, H_im, W_im = gt_masks.shape
+    device = gt_masks.device
+
+    # false positive region, a new point sampled in this region should have
+    # negative label to correct the FP error
+    fp_masks = ~gt_masks & pred_masks
+    # false negative region, a new point sampled in this region should have
+    # positive label to correct the FN error
+    fn_masks = gt_masks & ~pred_masks
+    # whether the prediction completely match the ground-truth on each mask
+    all_correct = torch.all((gt_masks == pred_masks).flatten(2), dim=2)
+    all_correct = all_correct[..., None, None]
+
+    # channel 0 is FP map, while channel 1 is FN map
+    pts_noise = torch.rand(B, num_pt, H_im, W_im, 2, device=device)
+    # sample a negative new click from FP region or a positive new click
+    # from FN region, depend on where the maximum falls,
+    # and in case the predictions are all correct (no FP or FN), we just
+    # sample a negative click from the background region
+    pts_noise[..., 0] *= fp_masks | (all_correct & ~gt_masks)
+    pts_noise[..., 1] *= fn_masks
+    pts_idx = pts_noise.flatten(2).argmax(dim=2)
+    labels = (pts_idx % 2).to(torch.int32)
+    pts_idx = pts_idx // 2
+    pts_x = pts_idx % W_im
+    pts_y = pts_idx // W_im
+    points = torch.stack([pts_x, pts_y], dim=2).to(torch.float)
+    return points, labels
+
+
+def sample_one_point_from_error_center(gt_masks, pred_masks, padding=True):
+    """
+    Sample 1 random point (along with its label) from the center of each error region,
+    that is, the point with the largest distance to the boundary of each error region.
+    This is the RITM sampling method from https://github.com/saic-vul/ritm_interactive_segmentation/blob/master/isegm/inference/clicker.py
+
+    Inputs:
+    - gt_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool
+    - pred_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool or None
+    - padding: if True, pad with boundary of 1 px for distance transform
+
+    Outputs:
+    - points: [B, 1, 2], dtype=torch.float, contains (x, y) coordinates of each sampled point
+    - labels: [B, 1], dtype=torch.int32, where 1 means positive clicks and 0 means negative clicks
+    """
+    if pred_masks is None:
+        pred_masks = torch.zeros_like(gt_masks)
+    assert gt_masks.dtype == torch.bool and gt_masks.size(1) == 1
+    assert pred_masks.dtype == torch.bool and pred_masks.shape == gt_masks.shape
+
+    B, _, H, W = gt_masks.shape
+
+    # false positive region, a new point sampled in this region should have
+    # negative label to correct the FP error
+    fp_masks = (~gt_masks & pred_masks).squeeze(1)
+    # false negative region, a new point sampled in this region should have
+    # positive label to correct the FN error
+    fn_masks = (gt_masks & ~pred_masks).squeeze(1)
+
+    if padding:
+        padded_fp_masks = torch.zeros(
+            B, H + 2, W + 2, dtype=fp_masks.dtype, device=fp_masks.device
+        )
+        padded_fp_masks[:, 1 : H + 1, 1 : W + 1] = fp_masks
+        padded_fn_masks = torch.zeros(
+            B, H + 2, W + 2, dtype=fp_masks.dtype, device=fp_masks.device
+        )
+        padded_fn_masks[:, 1 : H + 1, 1 : W + 1] = fn_masks
+    else:
+        padded_fp_masks = fp_masks
+        padded_fn_masks = fn_masks
+
+    fn_mask_dt = edt_triton(padded_fn_masks)
+    fp_mask_dt = edt_triton(padded_fp_masks)
+    if padding:
+        fn_mask_dt = fn_mask_dt[:, 1:-1, 1:-1]
+        fp_mask_dt = fp_mask_dt[:, 1:-1, 1:-1]
+
+    fn_max, fn_argmax = fn_mask_dt.reshape(B, -1).max(dim=-1)
+    fp_max, fp_argmax = fp_mask_dt.reshape(B, -1).max(dim=-1)
+    is_positive = fn_max > fp_max
+    chosen = torch.where(is_positive, fn_argmax, fp_argmax)
+    points_x = chosen % W
+    points_y = chosen // W
+
+    labels = is_positive.long()
+    points = torch.stack([points_x, points_y], -1)
+    return points.unsqueeze(1), labels.unsqueeze(1)
+
+
+def sample_one_point_from_error_center_slow(gt_masks, pred_masks, padding=True):
+    """
+    Sample 1 random point (along with its label) from the center of each error region,
+    that is, the point with the largest distance to the boundary of each error region.
+    This is the RITM sampling method from https://github.com/saic-vul/ritm_interactive_segmentation/blob/master/isegm/inference/clicker.py
+
+    Inputs:
+    - gt_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool
+    - pred_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool or None
+    - padding: if True, pad with boundary of 1 px for distance transform
+
+    Outputs:
+    - points: [B, 1, 2], dtype=torch.float, contains (x, y) coordinates of each sampled point
+    - labels: [B, 1], dtype=torch.int32, where 1 means positive clicks and 0 means negative clicks
+    """
+    import cv2  # delay OpenCV import to avoid unnecessary dependency
+
+    if pred_masks is None:
+        pred_masks = torch.zeros_like(gt_masks)
+    assert gt_masks.dtype == torch.bool and gt_masks.size(1) == 1
+    assert pred_masks.dtype == torch.bool and pred_masks.shape == gt_masks.shape
+
+    B, _, _, W_im = gt_masks.shape
+    device = gt_masks.device
+
+    # false positive region, a new point sampled in this region should have
+    # negative label to correct the FP error
+    fp_masks = ~gt_masks & pred_masks
+    # false negative region, a new point sampled in this region should have
+    # positive label to correct the FN error
+    fn_masks = gt_masks & ~pred_masks
+
+    fp_masks = fp_masks.cpu().numpy()
+    fn_masks = fn_masks.cpu().numpy()
+    points = torch.zeros(B, 1, 2, dtype=torch.float)
+    labels = torch.ones(B, 1, dtype=torch.int32)
+    for b in range(B):
+        fn_mask = fn_masks[b, 0]
+        fp_mask = fp_masks[b, 0]
+        if padding:
+            fn_mask = np.pad(fn_mask, ((1, 1), (1, 1)), "constant")
+            fp_mask = np.pad(fp_mask, ((1, 1), (1, 1)), "constant")
+        # compute the distance of each point in FN/FP region to its boundary
+        fn_mask_dt = cv2.distanceTransform(fn_mask.astype(np.uint8), cv2.DIST_L2, 0)
+        fp_mask_dt = cv2.distanceTransform(fp_mask.astype(np.uint8), cv2.DIST_L2, 0)
+        if padding:
+            fn_mask_dt = fn_mask_dt[1:-1, 1:-1]
+            fp_mask_dt = fp_mask_dt[1:-1, 1:-1]
+
+        # take the point in FN/FP region with the largest distance to its boundary
+        fn_mask_dt_flat = fn_mask_dt.reshape(-1)
+        fp_mask_dt_flat = fp_mask_dt.reshape(-1)
+        fn_argmax = np.argmax(fn_mask_dt_flat)
+        fp_argmax = np.argmax(fp_mask_dt_flat)
+        is_positive = fn_mask_dt_flat[fn_argmax] > fp_mask_dt_flat[fp_argmax]
+        pt_idx = fn_argmax if is_positive else fp_argmax
+        points[b, 0, 0] = pt_idx % W_im  # x
+        points[b, 0, 1] = pt_idx // W_im  # y
+        labels[b, 0] = int(is_positive)
+
+    points = points.to(device)
+    labels = labels.to(device)
+    return points, labels
+
+
+def get_next_point(gt_masks, pred_masks, method):
+    if method == "uniform":
+        return sample_random_points_from_errors(gt_masks, pred_masks)
+    elif method == "center":
+        return sample_one_point_from_error_center(gt_masks, pred_masks)
+    else:
+        raise ValueError(f"unknown sampling method {method}")
+
+
+def select_closest_cond_frames(
+    frame_idx, cond_frame_outputs, max_cond_frame_num, keep_first_cond_frame=False
+):
+    """
+    Select up to `max_cond_frame_num` conditioning frames from `cond_frame_outputs`
+    that are temporally closest to the current frame at `frame_idx`. Here, we take
+    - a) the closest conditioning frame before `frame_idx` (if any);
+    - b) the closest conditioning frame after `frame_idx` (if any);
+    - c) any other temporally closest conditioning frames until reaching a total
+         of `max_cond_frame_num` conditioning frames.
+
+    Outputs:
+    - selected_outputs: selected items (keys & values) from `cond_frame_outputs`.
+    - unselected_outputs: items (keys & values) not selected in `cond_frame_outputs`.
+    """
+    if max_cond_frame_num == -1 or len(cond_frame_outputs) <= max_cond_frame_num:
+        selected_outputs = cond_frame_outputs
+        unselected_outputs = {}
+    else:
+        assert max_cond_frame_num >= 2, "we should allow using 2+ conditioning frames"
+        selected_outputs = {}
+        if keep_first_cond_frame:
+            idx_first = min(
+                (t for t in cond_frame_outputs if t < frame_idx), default=None
+            )
+            if idx_first is None:
+                # Maybe we are tracking in reverse
+                idx_first = max(
+                    (t for t in cond_frame_outputs if t > frame_idx), default=None
+                )
+            if idx_first is not None:
+                selected_outputs[idx_first] = cond_frame_outputs[idx_first]
+        # the closest conditioning frame before `frame_idx` (if any)
+        idx_before = max((t for t in cond_frame_outputs if t < frame_idx), default=None)
+        if idx_before is not None:
+            selected_outputs[idx_before] = cond_frame_outputs[idx_before]
+
+        # the closest conditioning frame after `frame_idx` (if any)
+        idx_after = min((t for t in cond_frame_outputs if t >= frame_idx), default=None)
+        if idx_after is not None:
+            selected_outputs[idx_after] = cond_frame_outputs[idx_after]
+
+        # add other temporally closest conditioning frames until reaching a total
+        # of `max_cond_frame_num` conditioning frames.
+        num_remain = max_cond_frame_num - len(selected_outputs)
+        inds_remain = sorted(
+            (t for t in cond_frame_outputs if t not in selected_outputs),
+            key=lambda x: abs(x - frame_idx),
+        )[:num_remain]
+        selected_outputs.update((t, cond_frame_outputs[t]) for t in inds_remain)
+        unselected_outputs = {
+            t: v for t, v in cond_frame_outputs.items() if t not in selected_outputs
+        }
+
+    return selected_outputs, unselected_outputs
+
+
+def get_1d_sine_pe(pos_inds, dim, temperature=10000):
+    """
+    Get 1D sine positional embedding as in the original Transformer paper.
+    """
+    pe_dim = dim // 2
+    dim_t = torch.arange(pe_dim, dtype=torch.float32, device=pos_inds.device)
+    dim_t = temperature ** (2 * (dim_t // 2) / pe_dim)
+
+    pos_embed = pos_inds.unsqueeze(-1) / dim_t
+    pos_embed = torch.cat([pos_embed.sin(), pos_embed.cos()], dim=-1)
+    return pos_embed
+
+
+def get_best_gt_match_from_multimasks(pred_multimasks, gt_masks, pred_scores=None):
+    """
+    Get the mask with the best match to GT masks (based on IoU) from pred_multimasks.
+    Optionally, use `pred_scores` to break ties in case all IoUs are zeros.
+    """
+    assert pred_multimasks.ndim == 4 and gt_masks.ndim == 4
+    if pred_multimasks.size(1) == 1:
+        return pred_multimasks  # only a single mask channel, nothing to select
+
+    pred_multimasks_binary = pred_multimasks > 0
+    area_i = torch.sum(pred_multimasks_binary & gt_masks, dim=(2, 3)).float()
+    area_u = torch.sum(pred_multimasks_binary | gt_masks, dim=(2, 3)).float()
+    ious = area_i / torch.clamp(area_u, min=1.0)
+
+    # In case all IoUs are zeros (e.g. because the GT mask is empty), use pred_scores
+    # to break ties and select the best mask
+    if pred_scores is not None:
+        has_nonzero_ious = torch.any(ious > 0).expand_as(ious)
+        scores = torch.where(has_nonzero_ious, ious, pred_scores)
+    else:
+        scores = ious
+
+    # Finally, take the best mask prediction (with the highest score)
+    best_scores_inds = torch.argmax(scores, dim=-1)
+    batch_inds = torch.arange(scores.size(0), device=scores.device)
+    best_pred_mask = pred_multimasks[batch_inds, best_scores_inds].unsqueeze(1)
+    return best_pred_mask
+
+
+def fill_holes_in_mask_scores(mask, max_area, fill_holes=True, remove_sprinkles=True):
+    """
+    A post processor to fill small holes in mask scores with area under `max_area`.
+    Holes are those small connected components in either background or foreground.
+
+    Note that it relies on the "cc_torch" package to find connected components fast. You can
+    install it via the following command (`TORCH_CUDA_ARCH_LIST=8.0` is for A100 GPUs):
+    ```
+    pip uninstall -y cc_torch; TORCH_CUDA_ARCH_LIST=8.0 9.0 pip install git+https://github.com/ronghanghu/cc_torch
+    ```
+    Otherwise, it will fallback to a slightly slower triton implementation, or skimage if the tensor is on cpu
+    """
+
+    if max_area <= 0:
+        return mask  # nothing to fill in this case
+
+    if fill_holes:
+        # We remove small connected components in background by changing them to foreground
+        # with a small positive mask score (0.1).
+        mask_bg = mask <= 0
+        bg_area_thresh = max_area
+        _, areas_bg = _get_connected_components_with_padding(mask_bg)
+        small_components_bg = mask_bg & (areas_bg <= bg_area_thresh)
+        mask = torch.where(small_components_bg, 0.1, mask)
+
+    if remove_sprinkles:
+        # We remove small connected components in foreground by changing them to background
+        # with a small negative mask score (-0.1). Here we only remove connected components
+        # whose areas are under both `max_area` and half of the entire mask's area. This
+        # removes sprinkles while avoids filtering out tiny objects that we want to track.
+        mask_fg = mask > 0
+        fg_area_thresh = torch.sum(mask_fg, dim=(2, 3), keepdim=True, dtype=torch.int32)
+        fg_area_thresh.floor_divide_(2).clamp_(max=max_area)
+        _, areas_fg = _get_connected_components_with_padding(mask_fg)
+        small_components_fg = mask_fg & (areas_fg <= fg_area_thresh)
+        mask = torch.where(small_components_fg, -0.1, mask)
+    return mask
+
+
+def _get_connected_components_with_padding(mask):
+    """Get connected components from masks (possibly padding them to an even size)."""
+    from sam3.perflib.connected_components import connected_components
+
+    mask = mask.to(torch.uint8)
+    _, _, H, W = mask.shape
+    # make sure both height and width are even (to be compatible with cc_torch)
+    pad_h = H % 2
+    pad_w = W % 2
+    if pad_h == 0 and pad_w == 0:
+        labels, counts = connected_components(mask)
+    else:
+        # pad the mask to make its height and width even
+        # padding format is (padding_left,padding_right,padding_top,padding_bottom)
+        mask_pad = F.pad(mask, (0, pad_w, 0, pad_h), mode="constant", value=0)
+        labels, counts = connected_components(mask_pad)
+        labels = labels[:, :, :H, :W]
+        counts = counts[:, :, :H, :W]
+
+    return labels, counts

sam3/model/sam3_video_predictor.py

@@ -0,0 +1,521 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import datetime
+import gc
+import multiprocessing as mp
+import os
+import queue
+import socket
+import sys
+import time
+import uuid
+from contextlib import closing
+from typing import List, Optional
+
+import psutil
+import torch
+
+from sam3.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class Sam3VideoPredictor:
+    # a global dictionary that holds all inference states for this model (key is session_id)
+    _ALL_INFERENCE_STATES = {}
+
+    def __init__(
+        self,
+        checkpoint_path=None,
+        bpe_path=None,
+        has_presence_token=True,
+        geo_encoder_use_img_cross_attn=True,
+        strict_state_dict_loading=True,
+        async_loading_frames=False,
+        video_loader_type="cv2",
+        apply_temporal_disambiguation: bool = True,
+    ):
+        self.async_loading_frames = async_loading_frames
+        self.video_loader_type = video_loader_type
+        from sam3.model_builder import build_sam3_video_model
+
+        self.model = (
+            build_sam3_video_model(
+                checkpoint_path=checkpoint_path,
+                bpe_path=bpe_path,
+                has_presence_token=has_presence_token,
+                geo_encoder_use_img_cross_attn=geo_encoder_use_img_cross_attn,
+                strict_state_dict_loading=strict_state_dict_loading,
+                apply_temporal_disambiguation=apply_temporal_disambiguation,
+            )
+            .cuda()
+            .eval()
+        )
+
+    @torch.inference_mode()
+    def handle_request(self, request):
+        """Dispatch a request based on its type."""
+        request_type = request["type"]
+        if request_type == "start_session":
+            return self.start_session(
+                resource_path=request["resource_path"],
+                session_id=request.get("session_id", None),
+            )
+        elif request_type == "add_prompt":
+            return self.add_prompt(
+                session_id=request["session_id"],
+                frame_idx=request["frame_index"],
+                text=request.get("text", None),
+                points=request.get("points", None),
+                point_labels=request.get("point_labels", None),
+                bounding_boxes=request.get("bounding_boxes", None),
+                bounding_box_labels=request.get("bounding_box_labels", None),
+                obj_id=request.get("obj_id", None),
+            )
+        elif request_type == "remove_object":
+            return self.remove_object(
+                session_id=request["session_id"],
+                obj_id=request["obj_id"],
+                is_user_action=request.get("is_user_action", True),
+            )
+        elif request_type == "reset_session":
+            return self.reset_session(session_id=request["session_id"])
+        elif request_type == "close_session":
+            return self.close_session(session_id=request["session_id"])
+        else:
+            raise RuntimeError(f"invalid request type: {request_type}")
+
+    @torch.inference_mode()
+    def handle_stream_request(self, request):
+        """Dispatch a stream request based on its type."""
+        request_type = request["type"]
+        if request_type == "propagate_in_video":
+            yield from self.propagate_in_video(
+                session_id=request["session_id"],
+                propagation_direction=request.get("propagation_direction", "both"),
+                start_frame_idx=request.get("start_frame_index", None),
+                max_frame_num_to_track=request.get("max_frame_num_to_track", None),
+            )
+        else:
+            raise RuntimeError(f"invalid request type: {request_type}")
+
+    def start_session(self, resource_path, session_id=None):
+        """
+        Start a new inference session on an image or a video. Here `resource_path`
+        can be either a path to an image file (for image inference) or an MP4 file
+        or directory with JPEG video frames (for video inference).
+
+        If `session_id` is defined, it will be used as identifier for the
+        session. If it is not defined, the start_session function will create
+        a session id and return it.
+        """
+        # get an initial inference_state from the model
+        inference_state = self.model.init_state(
+            resource_path=resource_path,
+            async_loading_frames=self.async_loading_frames,
+            video_loader_type=self.video_loader_type,
+        )
+        if not session_id:
+            session_id = str(uuid.uuid4())
+        self._ALL_INFERENCE_STATES[session_id] = {
+            "state": inference_state,
+            "session_id": session_id,
+            "start_time": time.time(),
+        }
+        logger.debug(
+            f"started new session {session_id}; {self._get_session_stats()}; "
+            f"{self._get_torch_and_gpu_properties()}"
+        )
+        return {"session_id": session_id}
+
+    def add_prompt(
+        self,
+        session_id: str,
+        frame_idx: int,
+        text: Optional[str] = None,
+        points: Optional[List[List[float]]] = None,
+        point_labels: Optional[List[int]] = None,
+        bounding_boxes: Optional[List[List[float]]] = None,
+        bounding_box_labels: Optional[List[int]] = None,
+        obj_id: Optional[int] = None,
+    ):
+        """Add text, box and/or point prompt on a specific video frame."""
+        logger.debug(
+            f"add prompt on frame {frame_idx} in session {session_id}: "
+            f"{text=}, {points=}, {point_labels=}, "
+            f"{bounding_boxes=}, {bounding_box_labels=}"
+        )
+        session = self._get_session(session_id)
+        inference_state = session["state"]
+
+        frame_idx, outputs = self.model.add_prompt(
+            inference_state=inference_state,
+            frame_idx=frame_idx,
+            text_str=text,
+            points=points,
+            point_labels=point_labels,
+            boxes_xywh=bounding_boxes,
+            box_labels=bounding_box_labels,
+            obj_id=obj_id,
+        )
+        return {"frame_index": frame_idx, "outputs": outputs}
+
+    def remove_object(
+        self,
+        session_id: str,
+        obj_id: int,
+        is_user_action: bool = True,
+    ):
+        """Remove an object from tracking."""
+        logger.debug(
+            f"remove object {obj_id} in session {session_id}: " f"{is_user_action=}"
+        )
+        session = self._get_session(session_id)
+        inference_state = session["state"]
+
+        self.model.remove_object(
+            inference_state=inference_state,
+            obj_id=obj_id,
+            is_user_action=is_user_action,
+        )
+        return {"is_success": True}
+
+    def propagate_in_video(
+        self,
+        session_id,
+        propagation_direction,
+        start_frame_idx,
+        max_frame_num_to_track,
+    ):
+        """Propagate the added prompts to get grounding results on all video frames."""
+        logger.debug(
+            f"propagate in video in session {session_id}: "
+            f"{propagation_direction=}, {start_frame_idx=}, {max_frame_num_to_track=}"
+        )
+        try:
+            session = self._get_session(session_id)
+            inference_state = session["state"]
+            if propagation_direction not in ["both", "forward", "backward"]:
+                raise ValueError(
+                    f"invalid propagation direction: {propagation_direction}"
+                )
+
+            # First doing the forward propagation
+            if propagation_direction in ["both", "forward"]:
+                for frame_idx, outputs in self.model.propagate_in_video(
+                    inference_state=inference_state,
+                    start_frame_idx=start_frame_idx,
+                    max_frame_num_to_track=max_frame_num_to_track,
+                    reverse=False,
+                ):
+                    yield {"frame_index": frame_idx, "outputs": outputs}
+            # Then doing the backward propagation (reverse in time)
+            if propagation_direction in ["both", "backward"]:
+                for frame_idx, outputs in self.model.propagate_in_video(
+                    inference_state=inference_state,
+                    start_frame_idx=start_frame_idx,
+                    max_frame_num_to_track=max_frame_num_to_track,
+                    reverse=True,
+                ):
+                    yield {"frame_index": frame_idx, "outputs": outputs}
+        finally:
+            # Log upon completion (so that e.g. we can see if two propagations happen in parallel).
+            # Using `finally` here to log even when the tracking is aborted with GeneratorExit.
+            logger.debug(
+                f"propagation ended in session {session_id}; {self._get_session_stats()}"
+            )
+
+    def reset_session(self, session_id):
+        """Reset the session to its initial state (as when it's initial opened)."""
+        logger.debug(f"reset session {session_id}")
+        session = self._get_session(session_id)
+        inference_state = session["state"]
+        self.model.reset_state(inference_state)
+        return {"is_success": True}
+
+    def close_session(self, session_id):
+        """
+        Close a session. This method is idempotent and can be called multiple
+        times on the same "session_id".
+        """
+        session = self._ALL_INFERENCE_STATES.pop(session_id, None)
+        if session is None:
+            logger.warning(
+                f"cannot close session {session_id} as it does not exist (it might have expired); "
+                f"{self._get_session_stats()}"
+            )
+        else:
+            del session
+            gc.collect()
+            logger.info(f"removed session {session_id}; {self._get_session_stats()}")
+        return {"is_success": True}
+
+    def _get_session(self, session_id):
+        session = self._ALL_INFERENCE_STATES.get(session_id, None)
+        if session is None:
+            raise RuntimeError(
+                f"Cannot find session {session_id}; it might have expired"
+            )
+        return session
+
+    def _get_session_stats(self):
+        """Get a statistics string for live sessions and their GPU usage."""
+        # print both the session ids and their video frame numbers
+        live_session_strs = [
+            f"'{session_id}' ({session['state']['num_frames']} frames)"
+            for session_id, session in self._ALL_INFERENCE_STATES.items()
+        ]
+        session_stats_str = (
+            f"live sessions: [{', '.join(live_session_strs)}], GPU memory: "
+            f"{torch.cuda.memory_allocated() // 1024**2} MiB used and "
+            f"{torch.cuda.memory_reserved() // 1024**2} MiB reserved"
+            f" (max over time: {torch.cuda.max_memory_allocated() // 1024**2} MiB used "
+            f"and {torch.cuda.max_memory_reserved() // 1024**2} MiB reserved)"
+        )
+        return session_stats_str
+
+    def _get_torch_and_gpu_properties(self):
+        """Get a string for PyTorch and GPU properties (for logging and debugging)."""
+        torch_and_gpu_str = (
+            f"torch: {torch.__version__} with CUDA arch {torch.cuda.get_arch_list()}, "
+            f"GPU device: {torch.cuda.get_device_properties(torch.cuda.current_device())}"
+        )
+        return torch_and_gpu_str
+
+    def shutdown(self):
+        """Shutdown the predictor and clear all sessions."""
+        self._ALL_INFERENCE_STATES.clear()
+
+
+class Sam3VideoPredictorMultiGPU(Sam3VideoPredictor):
+    def __init__(self, *model_args, gpus_to_use=None, **model_kwargs):
+        if gpus_to_use is None:
+            # if not specified, use only the current GPU by default
+            gpus_to_use = [torch.cuda.current_device()]
+
+        IS_MAIN_PROCESS = os.getenv("IS_MAIN_PROCESS", "1") == "1"
+        if IS_MAIN_PROCESS:
+            gpus_to_use = sorted(set(gpus_to_use))
+            logger.info(f"using the following GPU IDs: {gpus_to_use}")
+            assert len(gpus_to_use) > 0 and all(isinstance(i, int) for i in gpus_to_use)
+            assert all(0 <= i < torch.cuda.device_count() for i in gpus_to_use)
+            os.environ["MASTER_ADDR"] = "localhost"
+            os.environ["MASTER_PORT"] = f"{self._find_free_port()}"
+            os.environ["RANK"] = "0"
+            os.environ["WORLD_SIZE"] = f"{len(gpus_to_use)}"
+
+        self.gpus_to_use = gpus_to_use
+        self.rank = int(os.environ["RANK"])
+        self.world_size = int(os.environ["WORLD_SIZE"])
+        self.rank_str = f"rank={self.rank} with world_size={self.world_size}"
+        self.device = torch.device(f"cuda:{self.gpus_to_use[self.rank]}")
+        torch.cuda.set_device(self.device)
+        self.has_shutdown = False
+        if self.rank == 0:
+            logger.info("\n\n\n\t*** START loading model on all ranks ***\n\n")
+
+        logger.info(f"loading model on {self.rank_str} -- this could take a while ...")
+        super().__init__(*model_args, **model_kwargs)
+        logger.info(f"loading model on {self.rank_str} -- DONE locally")
+
+        if self.world_size > 1 and self.rank == 0:
+            # start the worker processes *after* the model is loaded in the main process
+            # so that the main process can run torch.compile and fill the cache first
+            self._start_worker_processes(*model_args, **model_kwargs)
+            for rank in range(1, self.world_size):
+                self.command_queues[rank].put(("start_nccl_process_group", None))
+            self._start_nccl_process_group()
+
+        if self.rank == 0:
+            logger.info("\n\n\n\t*** DONE loading model on all ranks ***\n\n")
+
+    @torch.inference_mode()
+    def handle_request(self, request):
+        """Dispatch a request based on its type."""
+        if self.has_shutdown:
+            raise RuntimeError(
+                "cannot handle request after the predictor has shutdown; please create a new predictor"
+            )
+
+        # when starting a session, we need to create a session id before dispatching
+        # the request to the workers
+        if request["type"] == "start_session" and request.get("session_id") is None:
+            request["session_id"] = str(uuid.uuid4())
+        # dispatch the request to all worker processes
+        if self.world_size > 1 and self.rank == 0:
+            for rank in range(1, self.world_size):
+                self.command_queues[rank].put((request, False))
+
+        response = super().handle_request(request)
+
+        if self.world_size > 1:
+            torch.distributed.barrier()  # wait for all ranks to finish
+        return response
+
+    @torch.inference_mode()
+    def handle_stream_request(self, request):
+        """Dispatch a stream request based on its type."""
+        if self.has_shutdown:
+            raise RuntimeError(
+                "cannot handle request after the predictor has shutdown; please create a new predictor"
+            )
+
+        # dispatch the request to all worker processes
+        if self.world_size > 1 and self.rank == 0:
+            for rank in range(1, self.world_size):
+                self.command_queues[rank].put((request, True))
+
+        yield from super().handle_stream_request(request)
+
+        if self.world_size > 1:
+            torch.distributed.barrier()  # wait for all ranks to finish
+
+    def _start_worker_processes(self, *model_args, **model_kwargs):
+        """Start worker processes for handling model inference."""
+        world_size = self.world_size
+        logger.info(f"spawning {world_size - 1} worker processes")
+        # Use "spawn" (instead of "fork") for different PyTorch or CUDA context
+        mp_ctx = mp.get_context("spawn")
+        self.command_queues = {rank: mp_ctx.Queue() for rank in range(1, world_size)}
+        self.result_queues = {rank: mp_ctx.Queue() for rank in range(1, world_size)}
+        parent_pid = os.getpid()
+        for rank in range(1, world_size):
+            # set the environment variables for each worker process
+            os.environ["IS_MAIN_PROCESS"] = "0"  # mark this as a worker process
+            os.environ["RANK"] = f"{rank}"
+            worker_process = mp_ctx.Process(
+                target=Sam3VideoPredictorMultiGPU._worker_process_command_loop,
+                args=(
+                    rank,
+                    world_size,
+                    self.command_queues[rank],
+                    self.result_queues[rank],
+                    model_args,
+                    model_kwargs,
+                    self.gpus_to_use,
+                    parent_pid,
+                ),
+                daemon=True,
+            )
+            worker_process.start()
+        # revert the environment variables for the main process
+        os.environ["IS_MAIN_PROCESS"] = "1"
+        os.environ["RANK"] = "0"
+        # wait for all the worker processes to load the model and collect their PIDs
+        self.worker_pids = {}
+        for rank in range(1, self.world_size):
+            # a large timeout to cover potentially long model loading time due to compilation
+            _, worker_pid = self.result_queues[rank].get(timeout=7200)
+            self.worker_pids[rank] = worker_pid
+        logger.info(f"spawned {world_size - 1} worker processes")
+
+    def _start_nccl_process_group(self):
+        rank = int(os.environ["RANK"])
+        world_size = int(os.environ["WORLD_SIZE"])
+        if world_size == 1:
+            return
+
+        logger.debug(f"starting NCCL process group on {rank=} with {world_size=}")
+        assert not torch.distributed.is_initialized()
+        # use the "env://" init method with environment variables set in start_worker_processes
+        # a short 3-min timeout to quickly detect any synchronization failures
+        timeout_sec = int(os.getenv("SAM3_COLLECTIVE_OP_TIMEOUT_SEC", "180"))
+        timeout = datetime.timedelta(seconds=timeout_sec)
+        torch.distributed.init_process_group(
+            backend="nccl",
+            init_method="env://",
+            timeout=timeout,
+            device_id=self.device,
+        )
+        # warm-up the NCCL process group by running a dummy all-reduce
+        tensor = torch.ones(1024, 1024).cuda()
+        torch.distributed.all_reduce(tensor)
+        logger.debug(f"started NCCL process group on {rank=} with {world_size=}")
+
+    def _find_free_port(self) -> int:
+        """
+        Find a free port (a random free port from 1024 to 65535 will be selected)
+        https://stackoverflow.com/questions/1365265/on-localhost-how-do-i-pick-a-free-port-number)
+        """
+        with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
+            s.bind(("", 0))
+            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            return s.getsockname()[1]
+
+    @staticmethod
+    def _worker_process_command_loop(
+        rank,
+        world_size,
+        command_queue,
+        result_queue,
+        model_args,
+        model_kwargs,
+        gpus_to_use,
+        parent_pid,
+    ):
+        """
+        The command loop for each worker process. It listens to commands from the main process
+        and executes them using the model.
+        """
+        logger.info(f"starting worker process {rank=} with {world_size=}")
+        # verify that the environment variables are set correctly
+        assert int(os.environ["IS_MAIN_PROCESS"]) == 0
+        assert int(os.environ["RANK"]) == rank
+        assert int(os.environ["WORLD_SIZE"]) == world_size
+        # load the model in this worker process
+        predictor = Sam3VideoPredictorMultiGPU(
+            *model_args, gpus_to_use=gpus_to_use, **model_kwargs
+        )
+        logger.info(f"started worker {rank=} with {world_size=}")
+        # return the worker process id to the main process for bookkeeping
+        worker_pid = os.getpid()
+        result_queue.put(("load_model", worker_pid))
+
+        # wait for the command to start the NCCL process group
+        request_type, _ = command_queue.get(timeout=7200)
+        assert request_type == "start_nccl_process_group"
+        predictor._start_nccl_process_group()
+
+        # keep listening to commands from the main process
+        while True:
+            try:
+                request, is_stream_request = command_queue.get(timeout=5.0)
+                if request == "shutdown":
+                    logger.info(f"worker {rank=} shutting down")
+                    torch.distributed.destroy_process_group()
+                    result_queue.put(("shutdown", True))  # acknowledge the shutdown
+                    sys.exit(0)
+
+                logger.debug(f"worker {rank=} received request {request['type']=}")
+                if is_stream_request:
+                    for _ in predictor.handle_stream_request(request):
+                        pass  # handle stream requests in a generator fashion
+                else:
+                    predictor.handle_request(request)
+            except queue.Empty:
+                # Usually Python's multiprocessing module will shutdown all the daemon worker
+                # processes when the main process exits gracefully. However, the user may kill
+                # the main process using SIGKILL and thereby leaving no chance for the main process
+                # to clean up its daemon child processes. So here we manually check whether the
+                # parent process still exists (every 5 sec as in `command_queue.get` timeout).
+                if not psutil.pid_exists(parent_pid):
+                    logger.info(
+                        f"stopping worker {rank=} as its parent process has exited"
+                    )
+                    sys.exit(1)
+            except Exception as e:
+                logger.error(f"worker {rank=} exception: {e}", exc_info=True)
+
+    def shutdown(self):
+        """Shutdown all worker processes."""
+        if self.rank == 0 and self.world_size > 1:
+            logger.info(f"shutting down {self.world_size - 1} worker processes")
+            for rank in range(1, self.world_size):
+                self.command_queues[rank].put(("shutdown", False))
+            torch.distributed.destroy_process_group()
+            for rank in range(1, self.world_size):
+                self.result_queues[rank].get()  # wait for the worker to acknowledge
+            logger.info(f"shut down {self.world_size - 1} worker processes")
+        self.has_shutdown = True
+
+        super().shutdown()

sam3/model/text_encoder_ve.py

@@ -0,0 +1,328 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+from collections import OrderedDict
+from typing import Callable, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from torch.utils.checkpoint import checkpoint
+
+from .model_misc import LayerScale
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        n_head: int,
+        mlp_ratio: float = 4.0,
+        ls_init_value: Optional[float] = None,
+        act_layer: Callable[[], nn.Module] = nn.GELU,
+        norm_layer: Callable[[int], nn.Module] = nn.LayerNorm,
+    ):
+        super().__init__()
+        # Attention
+        self.attn = nn.MultiheadAttention(d_model, n_head, batch_first=True)
+
+        # LayerNorm, LayerScale
+        self.ln_1 = norm_layer(d_model)
+        self.ln_2 = norm_layer(d_model)
+
+        self.ls_1 = (
+            LayerScale(d_model, ls_init_value)
+            if ls_init_value is not None
+            else nn.Identity()
+        )
+        self.ls_2 = (
+            LayerScale(d_model, ls_init_value)
+            if ls_init_value is not None
+            else nn.Identity()
+        )
+
+        # MLP
+        mlp_width = int(d_model * mlp_ratio)
+        self.mlp = nn.Sequential(
+            OrderedDict(
+                [
+                    ("c_fc", nn.Linear(d_model, mlp_width)),
+                    ("gelu", act_layer()),
+                    ("c_proj", nn.Linear(mlp_width, d_model)),
+                ]
+            )
+        )
+
+    def attention(
+        self,
+        q_x: torch.Tensor,
+        k_x: Optional[torch.Tensor] = None,
+        v_x: Optional[torch.Tensor] = None,
+        attn_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        k_x = k_x if k_x is not None else q_x
+        v_x = v_x if v_x is not None else q_x
+        if attn_mask is not None:
+            # Leave boolean masks as is
+            if not attn_mask.dtype == torch.bool:
+                attn_mask = attn_mask.to(q_x.dtype)
+
+        return self.attn(q_x, k_x, v_x, need_weights=False, attn_mask=attn_mask)[0]
+
+    def forward(
+        self,
+        q_x: torch.Tensor,
+        k_x: Optional[torch.Tensor] = None,
+        v_x: Optional[torch.Tensor] = None,
+        attn_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        k_x = (
+            self.ln_1_kv(k_x) if hasattr(self, "ln_1_kv") and k_x is not None else None
+        )
+        v_x = (
+            self.ln_1_kv(v_x) if hasattr(self, "ln_1_kv") and v_x is not None else None
+        )
+        x = q_x + self.ls_1(
+            self.attention(q_x=self.ln_1(q_x), k_x=k_x, v_x=v_x, attn_mask=attn_mask)
+        )
+        x = x + self.ls_2(self.mlp(self.ln_2(x)))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(
+        self,
+        width: int,
+        layers: int,
+        heads: int,
+        mlp_ratio: float = 4.0,
+        ls_init_value: Optional[float] = None,
+        act_layer: Callable[[], nn.Module] = nn.GELU,
+        norm_layer: Callable[[int], nn.Module] = nn.LayerNorm,
+        compile_mode: Optional[str] = None,
+        use_act_checkpoint: bool = False,
+    ):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.grad_checkpointing = use_act_checkpoint
+        self.resblocks = nn.ModuleList(
+            [
+                ResidualAttentionBlock(
+                    width,
+                    heads,
+                    mlp_ratio,
+                    ls_init_value=ls_init_value,
+                    act_layer=act_layer,
+                    norm_layer=norm_layer,
+                )
+                for _ in range(layers)
+            ]
+        )
+
+        if compile_mode is not None:
+            self.forward = torch.compile(
+                self.forward, mode=compile_mode, fullgraph=True
+            )
+            if self.grad_checkpointing:
+                torch._dynamo.config.optimize_ddp = False
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        for _, r in enumerate(self.resblocks):
+            if (
+                self.grad_checkpointing
+                and not torch.jit.is_scripting()
+                and self.training
+            ):
+                x = checkpoint(r, x, None, None, attn_mask, use_reentrant=False)
+            else:
+                x = r(
+                    x,
+                    attn_mask=attn_mask,
+                )
+        return x
+
+
+def text_global_pool(
+    x: torch.Tensor, text: Optional[torch.Tensor] = None, pool_type: str = "argmax"
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    if pool_type == "first":
+        pooled, tokens = x[:, 0], x[:, 1:]
+    elif pool_type == "last":
+        pooled, tokens = x[:, -1], x[:, :-1]
+    elif pool_type == "argmax":
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        assert text is not None
+        pooled, tokens = x[torch.arange(x.shape[0]), text.argmax(dim=-1)], x
+    else:
+        pooled = tokens = x
+    return pooled, tokens
+
+
+class TextTransformer(nn.Module):
+    def __init__(
+        self,
+        context_length: int = 77,
+        vocab_size: int = 49408,
+        width: int = 512,
+        heads: int = 8,
+        layers: int = 12,
+        mlp_ratio: float = 4.0,
+        ls_init_value: Optional[float] = None,
+        output_dim: int = 512,
+        no_causal_mask: bool = False,
+        pool_type: str = "none",  # no pooling
+        proj_bias: bool = False,
+        act_layer: Callable = nn.GELU,
+        norm_layer: Callable = nn.LayerNorm,
+        output_tokens: bool = False,
+        use_ln_post: bool = True,
+        compile_mode: Optional[str] = None,
+        use_act_checkpoint: bool = False,
+    ):
+        super().__init__()
+        assert pool_type in ("first", "last", "argmax", "none")
+        self.output_tokens = output_tokens
+        self.num_pos = self.context_length = context_length
+        self.vocab_size = vocab_size
+        self.width = width
+        self.output_dim = output_dim
+        self.heads = heads
+        self.pool_type = pool_type
+
+        self.token_embedding = nn.Embedding(self.vocab_size, width)
+        self.positional_embedding = nn.Parameter(torch.empty(self.num_pos, width))
+        self.transformer = Transformer(
+            width=width,
+            layers=layers,
+            heads=heads,
+            mlp_ratio=mlp_ratio,
+            ls_init_value=ls_init_value,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+            compile_mode=compile_mode,
+            use_act_checkpoint=use_act_checkpoint,
+        )
+        self.ln_final = norm_layer(width) if use_ln_post else nn.Identity()
+        if no_causal_mask:
+            self.attn_mask = None
+        else:
+            self.register_buffer(
+                "attn_mask", self.build_causal_mask(), persistent=False
+            )
+        if proj_bias:
+            self.text_projection = nn.Linear(width, output_dim)
+        else:
+            self.text_projection = nn.Parameter(torch.empty(width, output_dim))
+
+    def build_causal_mask(self) -> torch.Tensor:
+        # lazily create causal attention mask, with full attention between the tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.num_pos, self.num_pos)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+
+    def forward(
+        self, text: torch.Tensor
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        seq_len = text.shape[1]
+        x = self.token_embedding(text)  # [batch_size, n_ctx, d_model]
+
+        attn_mask = self.attn_mask
+        if attn_mask is not None:
+            attn_mask = attn_mask[:seq_len, :seq_len]
+
+        x = x + self.positional_embedding[:seq_len]
+        x = self.transformer(x, attn_mask=attn_mask)
+
+        x = self.ln_final(x)
+        pooled, tokens = text_global_pool(x, text, pool_type=self.pool_type)
+        if self.text_projection is not None:
+            if isinstance(self.text_projection, nn.Linear):
+                pooled = self.text_projection(pooled)
+            else:
+                pooled = pooled @ self.text_projection
+        if self.output_tokens:
+            return pooled, tokens
+        return pooled
+
+
+class VETextEncoder(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        tokenizer: Callable,
+        width: int = 1024,
+        heads: int = 16,
+        layers: int = 24,
+        context_length: int = 32,
+        vocab_size: int = 49408,
+        use_ln_post: bool = True,
+        compile_mode: Optional[str] = None,
+        use_act_checkpoint: bool = True,
+    ):
+        super().__init__()
+        self.context_length = context_length
+        self.use_ln_post = use_ln_post
+        self.tokenizer = tokenizer
+
+        self.encoder = TextTransformer(
+            context_length=self.context_length,
+            vocab_size=vocab_size,
+            width=width,
+            heads=heads,
+            layers=layers,
+            # we want the tokens, not just the pooled output
+            output_tokens=True,
+            use_ln_post=use_ln_post,
+            compile_mode=compile_mode,
+            use_act_checkpoint=use_act_checkpoint,
+        )
+        self.resizer = nn.Linear(self.encoder.width, d_model)
+
+    def forward(
+        self,
+        text: Union[List[str], Tuple[torch.Tensor, torch.Tensor, dict]],
+        input_boxes: Optional[List] = None,
+        device: torch.device = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if isinstance(text[0], str):
+            # no use case for this
+            assert input_boxes is None or len(input_boxes) == 0, "not supported"
+
+            # Encode the text
+            tokenized = self.tokenizer(text, context_length=self.context_length).to(
+                device
+            )  # [b, seq_len]
+            text_attention_mask = (tokenized != 0).bool()
+
+            # manually embed the tokens
+            inputs_embeds = self.encoder.token_embedding(
+                tokenized
+            )  # [b, seq_len, d=1024]
+            _, text_memory = self.encoder(tokenized)  # [b, seq_len, d=1024]
+
+            assert text_memory.shape[1] == inputs_embeds.shape[1]
+            # Invert attention mask because its the opposite in pytorch transformer
+            text_attention_mask = text_attention_mask.ne(1)
+            # Transpose memory because pytorch's attention expects sequence first
+            text_memory = text_memory.transpose(0, 1)
+            # Resize the encoder hidden states to be of the same d_model as the decoder
+            text_memory_resized = self.resizer(text_memory)
+        else:
+            # The text is already encoded, use as is.
+            text_attention_mask, text_memory_resized, tokenized = text
+            inputs_embeds = tokenized["inputs_embeds"]
+            assert (
+                input_boxes is None or len(input_boxes) == 0
+            ), "Can't replace boxes in text if it's already encoded"
+
+        # Note that the input_embeds are returned in pytorch's convention (sequence first)
+        return (
+            text_attention_mask,
+            text_memory_resized,
+            inputs_embeds.transpose(0, 1),
+        )

sam3/model/tokenizer_ve.py

@@ -0,0 +1,253 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+"""
+Text Tokenizer.
+
+Copied and lightly adapted from VE repo, which in turn copied
+from open_clip and openAI CLIP.
+"""
+
+import gzip
+import html
+import io
+import os
+import string
+from functools import lru_cache
+from typing import List, Optional, Union
+
+import ftfy
+import regex as re
+import torch
+from iopath.common.file_io import g_pathmgr
+
+
+# https://stackoverflow.com/q/62691279
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+DEFAULT_CONTEXT_LENGTH = 77
+
+
+@lru_cache()
+def bytes_to_unicode():
+    """
+    Returns list of utf-8 byte and a corresponding list of unicode strings.
+    The reversible bpe codes work on unicode strings.
+    This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
+    When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
+    This is a significant percentage of your normal, say, 32K bpe vocab.
+    To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
+    And avoids mapping to whitespace/control characters the bpe code barfs on.
+    """
+    bs = (
+        list(range(ord("!"), ord("~") + 1))
+        + list(range(ord("¡"), ord("¬") + 1))
+        + list(range(ord("®"), ord("ÿ") + 1))
+    )
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8 + n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+def get_pairs(word):
+    """Return set of symbol pairs in a word.
+    Word is represented as tuple of symbols (symbols being variable-length strings).
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r"\s+", " ", text)
+    text = text.strip()
+    return text
+
+
+def _clean_canonicalize(x):
+    # basic, remove whitespace, remove punctuation, lower case
+    return canonicalize_text(basic_clean(x))
+
+
+def _clean_lower(x):
+    # basic, remove whitespace, lower case
+    return whitespace_clean(basic_clean(x)).lower()
+
+
+def _clean_whitespace(x):
+    # basic, remove whitespace
+    return whitespace_clean(basic_clean(x))
+
+
+def get_clean_fn(type: str):
+    if type == "canonicalize":
+        return _clean_canonicalize
+    elif type == "lower":
+        return _clean_lower
+    elif type == "whitespace":
+        return _clean_whitespace
+    else:
+        assert False, f"Invalid clean function ({type})."
+
+
+def canonicalize_text(text, *, keep_punctuation_exact_string=None):
+    """Returns canonicalized `text` (lowercase and punctuation removed).
+    From: https://github.com/google-research/big_vision/blob/53f18caf27a9419231bbf08d3388b07671616d3d/big_vision/evaluators/proj/image_text/prompt_engineering.py#L94
+    Args:
+      text: string to be canonicalized.
+      keep_punctuation_exact_string: If provided, then this exact string kept.
+        For example providing '{}' will keep any occurrences of '{}' (but will
+        still remove '{' and '}' that appear separately).
+    """
+    text = text.replace("_", " ")
+    if keep_punctuation_exact_string:
+        text = keep_punctuation_exact_string.join(
+            part.translate(str.maketrans("", "", string.punctuation))
+            for part in text.split(keep_punctuation_exact_string)
+        )
+    else:
+        text = text.translate(str.maketrans("", "", string.punctuation))
+    text = text.lower()
+    text = re.sub(r"\s+", " ", text)
+    return text.strip()
+
+
+class SimpleTokenizer(object):
+    def __init__(
+        self,
+        bpe_path: Union[str, os.PathLike],
+        additional_special_tokens: Optional[List[str]] = None,
+        context_length: Optional[int] = DEFAULT_CONTEXT_LENGTH,
+        clean: str = "lower",
+    ):
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        with g_pathmgr.open(bpe_path, "rb") as fh:
+            bpe_bytes = io.BytesIO(fh.read())
+            merges = gzip.open(bpe_bytes).read().decode("utf-8").split("\n")
+        # merges = gzip.open(bpe_path).read().decode("utf-8").split("\n")
+        merges = merges[1 : 49152 - 256 - 2 + 1]
+        merges = [tuple(merge.split()) for merge in merges]
+        vocab = list(bytes_to_unicode().values())
+        vocab = vocab + [v + "</w>" for v in vocab]
+        for merge in merges:
+            vocab.append("".join(merge))
+        special_tokens = ["<start_of_text>", "<end_of_text>"]
+        if additional_special_tokens:
+            special_tokens += additional_special_tokens
+        vocab.extend(special_tokens)
+        self.encoder = dict(zip(vocab, range(len(vocab))))
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.bpe_ranks = dict(zip(merges, range(len(merges))))
+        self.cache = {t: t for t in special_tokens}
+        special = "|".join(special_tokens)
+        self.pat = re.compile(
+            special + r"""|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""",
+            re.IGNORECASE,
+        )
+        self.vocab_size = len(self.encoder)
+        self.all_special_ids = [self.encoder[t] for t in special_tokens]
+        self.sot_token_id = self.all_special_ids[0]
+        self.eot_token_id = self.all_special_ids[1]
+        self.context_length = context_length
+        self.clean_fn = get_clean_fn(clean)
+
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token[:-1]) + (token[-1] + "</w>",)
+        pairs = get_pairs(word)
+        if not pairs:
+            return token + "</w>"
+        while True:
+            bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except:
+                    new_word.extend(word[i:])
+                    break
+                if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
+                    new_word.append(first + second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = " ".join(word)
+        self.cache[token] = word
+        return word
+
+    def encode(self, text):
+        bpe_tokens = []
+        text = self.clean_fn(text)
+        for token in re.findall(self.pat, text):
+            token = "".join(self.byte_encoder[b] for b in token.encode("utf-8"))
+            bpe_tokens.extend(
+                self.encoder[bpe_token] for bpe_token in self.bpe(token).split(" ")
+            )
+        return bpe_tokens
+
+    def decode(self, tokens):
+        text = "".join([self.decoder[token] for token in tokens])
+        text = (
+            bytearray([self.byte_decoder[c] for c in text])
+            .decode("utf-8", errors="replace")
+            .replace("</w>", " ")
+        )
+        return text
+
+    def __call__(
+        self, texts: Union[str, List[str]], context_length: Optional[int] = None
+    ) -> torch.LongTensor:
+        """Returns the tokenized representation of given input string(s)
+        Parameters
+        ----------
+        texts : Union[str, List[str]]
+            An input string or a list of input strings to tokenize
+        context_length : int
+            The context length to use; all CLIP models use 77 as the context length
+        Returns
+        -------
+        A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length]
+        """
+        if isinstance(texts, str):
+            texts = [texts]
+        context_length = context_length or self.context_length
+        assert context_length, "Please set a valid context length"
+        all_tokens = [
+            [self.sot_token_id] + self.encode(text) + [self.eot_token_id]
+            for text in texts
+        ]
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+        for i, tokens in enumerate(all_tokens):
+            if len(tokens) > context_length:
+                tokens = tokens[:context_length]  # Truncate
+                tokens[-1] = self.eot_token_id
+            result[i, : len(tokens)] = torch.tensor(tokens)
+        return result

sam3/model/utils/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.

sam3/model/utils/misc.py

@@ -0,0 +1,77 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+from collections import defaultdict
+from dataclasses import fields, is_dataclass
+from typing import Any, Mapping, Protocol, runtime_checkable
+
+import torch
+
+
+def _is_named_tuple(x) -> bool:
+    return isinstance(x, tuple) and hasattr(x, "_asdict") and hasattr(x, "_fields")
+
+
+@runtime_checkable
+class _CopyableData(Protocol):
+    def to(self, device: torch.device, *args: Any, **kwargs: Any):
+        """Copy data to the specified device"""
+        ...
+
+
+def copy_data_to_device(data, device: torch.device, *args: Any, **kwargs: Any):
+    """Function that recursively copies data to a torch.device.
+
+    Args:
+        data: The data to copy to device
+        device: The device to which the data should be copied
+        args: positional arguments that will be passed to the `to` call
+        kwargs: keyword arguments that will be passed to the `to` call
+
+    Returns:
+        The data on the correct device
+    """
+
+    if _is_named_tuple(data):
+        return type(data)(
+            **copy_data_to_device(data._asdict(), device, *args, **kwargs)
+        )
+    elif isinstance(data, (list, tuple)):
+        return type(data)(copy_data_to_device(e, device, *args, **kwargs) for e in data)
+    elif isinstance(data, defaultdict):
+        return type(data)(
+            data.default_factory,
+            {
+                k: copy_data_to_device(v, device, *args, **kwargs)
+                for k, v in data.items()
+            },
+        )
+    elif isinstance(data, Mapping):
+        return type(data)(
+            {
+                k: copy_data_to_device(v, device, *args, **kwargs)
+                for k, v in data.items()
+            }
+        )
+    elif is_dataclass(data) and not isinstance(data, type):
+        new_data_class = type(data)(
+            **{
+                field.name: copy_data_to_device(
+                    getattr(data, field.name), device, *args, **kwargs
+                )
+                for field in fields(data)
+                if field.init
+            }
+        )
+        for field in fields(data):
+            if not field.init:
+                setattr(
+                    new_data_class,
+                    field.name,
+                    copy_data_to_device(
+                        getattr(data, field.name), device, *args, **kwargs
+                    ),
+                )
+        return new_data_class
+    elif isinstance(data, _CopyableData):
+        return data.to(device, *args, **kwargs)
+    return data

sam3/model/utils/sam1_utils.py

@@ -0,0 +1,119 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import warnings
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.transforms import Normalize, Resize, ToTensor
+
+
+# Adapted from https://github.com/facebookresearch/sam2/blob/main/sam2/utils/transforms.py
+class SAM2Transforms(nn.Module):
+    def __init__(
+        self, resolution, mask_threshold, max_hole_area=0.0, max_sprinkle_area=0.0
+    ):
+        """
+        Transforms for SAM2.
+        """
+        super().__init__()
+        self.resolution = resolution
+        self.mask_threshold = mask_threshold
+        self.max_hole_area = max_hole_area
+        self.max_sprinkle_area = max_sprinkle_area
+        self.mean = [0.5, 0.5, 0.5]
+        self.std = [0.5, 0.5, 0.5]
+        self.to_tensor = ToTensor()
+        self.transforms = torch.jit.script(
+            nn.Sequential(
+                Resize((self.resolution, self.resolution)),
+                Normalize(self.mean, self.std),
+            )
+        )
+
+    def __call__(self, x):
+        x = self.to_tensor(x)
+        return self.transforms(x)
+
+    def forward_batch(self, img_list):
+        img_batch = [self.transforms(self.to_tensor(img)) for img in img_list]
+        img_batch = torch.stack(img_batch, dim=0)
+        return img_batch
+
+    def transform_coords(
+        self, coords: torch.Tensor, normalize=False, orig_hw=None
+    ) -> torch.Tensor:
+        """
+        Expects a torch tensor with length 2 in the last dimension. The coordinates can be in absolute image or normalized coordinates,
+        If the coords are in absolute image coordinates, normalize should be set to True and original image size is required.
+
+        Returns
+            Un-normalized coordinates in the range of [0, 1] which is expected by the SAM2 model.
+        """
+        if normalize:
+            assert orig_hw is not None
+            h, w = orig_hw
+            coords = coords.clone()
+            coords[..., 0] = coords[..., 0] / w
+            coords[..., 1] = coords[..., 1] / h
+
+        coords = coords * self.resolution  # unnormalize coords
+        return coords
+
+    def transform_boxes(
+        self, boxes: torch.Tensor, normalize=False, orig_hw=None
+    ) -> torch.Tensor:
+        """
+        Expects a tensor of shape Bx4. The coordinates can be in absolute image or normalized coordinates,
+        if the coords are in absolute image coordinates, normalize should be set to True and original image size is required.
+        """
+        boxes = self.transform_coords(boxes.reshape(-1, 2, 2), normalize, orig_hw)
+        return boxes
+
+    def postprocess_masks(self, masks: torch.Tensor, orig_hw) -> torch.Tensor:
+        """
+        Perform PostProcessing on output masks.
+        """
+        masks = masks.float()
+        input_masks = masks
+        mask_flat = masks.flatten(0, 1).unsqueeze(1)  # flatten as 1-channel image
+        try:
+            from sam3.perflib.connected_components import connected_components
+
+            if self.max_hole_area > 0:
+                # Holes are those connected components in background with area <= self.fill_hole_area
+                # (background regions are those with mask scores <= self.mask_threshold)
+                labels, areas = connected_components(
+                    (mask_flat <= self.mask_threshold).to(torch.uint8)
+                )
+                is_hole = (labels > 0) & (areas <= self.max_hole_area)
+                is_hole = is_hole.reshape_as(masks)
+                # We fill holes with a small positive mask score (10.0) to change them to foreground.
+                masks = torch.where(is_hole, self.mask_threshold + 10.0, masks)
+
+            if self.max_sprinkle_area > 0:
+                labels, areas = connected_components(
+                    (mask_flat > self.mask_threshold).to(torch.uint8)
+                )
+                is_hole = (labels > 0) & (areas <= self.max_sprinkle_area)
+                is_hole = is_hole.reshape_as(masks)
+                # We fill holes with negative mask score (-10.0) to change them to background.
+                masks = torch.where(is_hole, self.mask_threshold - 10.0, masks)
+        except Exception as e:
+            # Skip the post-processing step if the CUDA kernel fails
+            warnings.warn(
+                f"{e}\n\nSkipping the post-processing step due to the error above. You can "
+                "still use SAM 3 and it's OK to ignore the error above, although some post-processing "
+                "functionality may be limited (which doesn't affect the results in most cases; see "
+                "https://github.com/facebookresearch/sam3/blob/main/INSTALL.md).",
+                category=UserWarning,
+                stacklevel=2,
+            )
+            masks = input_masks
+
+        masks = F.interpolate(masks, orig_hw, mode="bilinear", align_corners=False)
+        return masks

sam3/model/utils/sam2_utils.py

@@ -0,0 +1,233 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+from threading import Thread
+
+import numpy as np
+import torch
+from PIL import Image
+from tqdm import tqdm
+
+
+def _load_img_as_tensor(img_path, image_size):
+    img_pil = Image.open(img_path)
+    img_np = np.array(img_pil.convert("RGB").resize((image_size, image_size)))
+    if img_np.dtype == np.uint8:  # np.uint8 is expected for JPEG images
+        img_np = img_np / 255.0
+    else:
+        raise RuntimeError(f"Unknown image dtype: {img_np.dtype} on {img_path}")
+    img = torch.from_numpy(img_np).permute(2, 0, 1)
+    video_width, video_height = img_pil.size  # the original video size
+    return img, video_height, video_width
+
+
+class AsyncVideoFrameLoader:
+    """
+    A list of video frames to be load asynchronously without blocking session start.
+    """
+
+    def __init__(
+        self,
+        img_paths,
+        image_size,
+        offload_video_to_cpu,
+        img_mean,
+        img_std,
+        compute_device,
+    ):
+        self.img_paths = img_paths
+        self.image_size = image_size
+        self.offload_video_to_cpu = offload_video_to_cpu
+        self.img_mean = img_mean
+        self.img_std = img_std
+        # items in `self.images` will be loaded asynchronously
+        self.images = [None] * len(img_paths)
+        # catch and raise any exceptions in the async loading thread
+        self.exception = None
+        # video_height and video_width be filled when loading the first image
+        self.video_height = None
+        self.video_width = None
+        self.compute_device = compute_device
+
+        # load the first frame to fill video_height and video_width and also
+        # to cache it (since it's most likely where the user will click)
+        self.__getitem__(0)
+
+        # load the rest of frames asynchronously without blocking the session start
+        def _load_frames():
+            try:
+                for n in tqdm(range(len(self.images)), desc="frame loading (JPEG)"):
+                    self.__getitem__(n)
+            except Exception as e:
+                self.exception = e
+
+        self.thread = Thread(target=_load_frames, daemon=True)
+        self.thread.start()
+
+    def __getitem__(self, index):
+        if self.exception is not None:
+            raise RuntimeError("Failure in frame loading thread") from self.exception
+
+        img = self.images[index]
+        if img is not None:
+            return img
+
+        img, video_height, video_width = _load_img_as_tensor(
+            self.img_paths[index], self.image_size
+        )
+        self.video_height = video_height
+        self.video_width = video_width
+        # normalize by mean and std
+        img -= self.img_mean
+        img /= self.img_std
+        if not self.offload_video_to_cpu:
+            img = img.to(self.compute_device, non_blocking=True)
+        self.images[index] = img
+        return img
+
+    def __len__(self):
+        return len(self.images)
+
+
+def load_video_frames(
+    video_path,
+    image_size,
+    offload_video_to_cpu,
+    img_mean=(0.485, 0.456, 0.406),
+    img_std=(0.229, 0.224, 0.225),
+    async_loading_frames=False,
+    compute_device=torch.device("cuda"),
+):
+    """
+    Load the video frames from video_path. The frames are resized to image_size as in
+    the model and are loaded to GPU if offload_video_to_cpu=False. This is used by the demo.
+    """
+    is_bytes = isinstance(video_path, bytes)
+    is_str = isinstance(video_path, str)
+    is_mp4_path = is_str and os.path.splitext(video_path)[-1] in [".mp4", ".MP4"]
+    if is_bytes or is_mp4_path:
+        return load_video_frames_from_video_file(
+            video_path=video_path,
+            image_size=image_size,
+            offload_video_to_cpu=offload_video_to_cpu,
+            img_mean=img_mean,
+            img_std=img_std,
+            compute_device=compute_device,
+        )
+    elif is_str and os.path.isdir(video_path):
+        return load_video_frames_from_jpg_images(
+            video_path=video_path,
+            image_size=image_size,
+            offload_video_to_cpu=offload_video_to_cpu,
+            img_mean=img_mean,
+            img_std=img_std,
+            async_loading_frames=async_loading_frames,
+            compute_device=compute_device,
+        )
+    else:
+        raise NotImplementedError(
+            "Only MP4 video and JPEG folder are supported at this moment"
+        )
+
+
+def load_video_frames_from_jpg_images(
+    video_path,
+    image_size,
+    offload_video_to_cpu,
+    img_mean=(0.485, 0.456, 0.406),
+    img_std=(0.229, 0.224, 0.225),
+    async_loading_frames=False,
+    compute_device=torch.device("cuda"),
+):
+    """
+    Load the video frames from a directory of JPEG files ("<frame_index>.jpg" format).
+
+    The frames are resized to image_size x image_size and are loaded to GPU if
+    `offload_video_to_cpu` is `False` and to CPU if `offload_video_to_cpu` is `True`.
+
+    You can load a frame asynchronously by setting `async_loading_frames` to `True`.
+    """
+    if isinstance(video_path, str) and os.path.isdir(video_path):
+        jpg_folder = video_path
+    else:
+        raise NotImplementedError(
+            "Only JPEG frames are supported at this moment. For video files, you may use "
+            "ffmpeg (https://ffmpeg.org/) to extract frames into a folder of JPEG files, such as \n"
+            "```\n"
+            "ffmpeg -i <your_video>.mp4 -q:v 2 -start_number 0 <output_dir>/'%05d.jpg'\n"
+            "```\n"
+            "where `-q:v` generates high-quality JPEG frames and `-start_number 0` asks "
+            "ffmpeg to start the JPEG file from 00000.jpg."
+        )
+
+    frame_names = [
+        p
+        for p in os.listdir(jpg_folder)
+        if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"]
+    ]
+    frame_names.sort(key=lambda p: int(os.path.splitext(p)[0]))
+    num_frames = len(frame_names)
+    if num_frames == 0:
+        raise RuntimeError(f"no images found in {jpg_folder}")
+    img_paths = [os.path.join(jpg_folder, frame_name) for frame_name in frame_names]
+    img_mean = torch.tensor(img_mean, dtype=torch.float32)[:, None, None]
+    img_std = torch.tensor(img_std, dtype=torch.float32)[:, None, None]
+
+    if async_loading_frames:
+        lazy_images = AsyncVideoFrameLoader(
+            img_paths,
+            image_size,
+            offload_video_to_cpu,
+            img_mean,
+            img_std,
+            compute_device,
+        )
+        return lazy_images, lazy_images.video_height, lazy_images.video_width
+
+    images = torch.zeros(num_frames, 3, image_size, image_size, dtype=torch.float32)
+    for n, img_path in enumerate(tqdm(img_paths, desc="frame loading (JPEG)")):
+        images[n], video_height, video_width = _load_img_as_tensor(img_path, image_size)
+    if not offload_video_to_cpu:
+        images = images.to(compute_device)
+        img_mean = img_mean.to(compute_device)
+        img_std = img_std.to(compute_device)
+    # normalize by mean and std
+    images -= img_mean
+    images /= img_std
+    return images, video_height, video_width
+
+
+def load_video_frames_from_video_file(
+    video_path,
+    image_size,
+    offload_video_to_cpu,
+    img_mean=(0.485, 0.456, 0.406),
+    img_std=(0.229, 0.224, 0.225),
+    compute_device=torch.device("cuda"),
+):
+    """Load the video frames from a video file."""
+    import decord
+
+    img_mean = torch.tensor(img_mean, dtype=torch.float32)[:, None, None]
+    img_std = torch.tensor(img_std, dtype=torch.float32)[:, None, None]
+    # Get the original video height and width
+    decord.bridge.set_bridge("torch")
+    video_height, video_width, _ = decord.VideoReader(video_path).next().shape
+    # Iterate over all frames in the video
+    images = []
+    for frame in decord.VideoReader(video_path, width=image_size, height=image_size):
+        images.append(frame.permute(2, 0, 1))
+
+    images = torch.stack(images, dim=0).float() / 255.0
+    if not offload_video_to_cpu:
+        images = images.to(compute_device)
+        img_mean = img_mean.to(compute_device)
+        img_std = img_std.to(compute_device)
+    # normalize by mean and std
+    images -= img_mean
+    images /= img_std
+    return images, video_height, video_width

sam3/model/vitdet.py

@@ -0,0 +1,879 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+"""
+ViTDet backbone adapted from Detectron2.
+This module implements Vision Transformer (ViT) backbone for object detection.
+
+Rope embedding code adopted from:
+1. https://github.com/meta-llama/codellama/blob/main/llama/model.py
+2. https://github.com/naver-ai/rope-vit
+3. https://github.com/lucidrains/rotary-embedding-torch
+"""
+
+import math
+from functools import partial
+from typing import Callable, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+
+try:
+    from timm.layers import DropPath, Mlp, trunc_normal_
+except ModuleNotFoundError:
+    # compatibility for older timm versions
+    from timm.models.layers import DropPath, Mlp, trunc_normal_
+from torch import Tensor
+
+from .model_misc import LayerScale
+
+
+def init_t_xy(
+    end_x: int, end_y: int, scale: float = 1.0, offset: int = 0
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    t = torch.arange(end_x * end_y, dtype=torch.float32)
+    t_x = (t % end_x).float()
+    t_y = torch.div(t, end_x, rounding_mode="floor").float()
+    return t_x * scale + offset, t_y * scale + offset
+
+
+def compute_axial_cis(
+    dim: int,
+    end_x: int,
+    end_y: int,
+    theta: float = 10000.0,
+    scale_pos: float = 1.0,
+    offset: int = 0,
+) -> torch.Tensor:
+    freqs_x = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))
+    freqs_y = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))
+
+    t_x, t_y = init_t_xy(end_x, end_y, scale_pos, offset)
+    freqs_x = torch.outer(t_x, freqs_x)
+    freqs_y = torch.outer(t_y, freqs_y)
+    freqs_cis_x = torch.polar(torch.ones_like(freqs_x), freqs_x)
+    freqs_cis_y = torch.polar(torch.ones_like(freqs_y), freqs_y)
+    return torch.cat([freqs_cis_x, freqs_cis_y], dim=-1)
+
+
+def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+    ndim = x.ndim
+    assert 0 <= 1 < ndim
+    assert freqs_cis.shape == (x.shape[-2], x.shape[-1])
+    shape = [d if i >= ndim - 2 else 1 for i, d in enumerate(x.shape)]
+    return freqs_cis.view(*shape)
+
+
+def apply_rotary_enc(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cis: torch.Tensor,
+    repeat_freqs_k: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
+    xk_ = (
+        torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+        if xk.shape[-2] != 0
+        else None
+    )
+    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
+    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
+    if xk_ is None:
+        # no keys to rotate, due to dropout
+        return xq_out.type_as(xq).to(xq.device), xk
+    # repeat freqs along seq_len dim to match k seq_len
+    if repeat_freqs_k:
+        r = xk_.shape[-2] // xq_.shape[-2]
+        freqs_cis = freqs_cis.repeat(*([1] * (freqs_cis.ndim - 2)), r, 1)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
+    return xq_out.type_as(xq).to(xq.device), xk_out.type_as(xk).to(xk.device)
+
+
+def window_partition(x: Tensor, window_size: int) -> Tuple[Tensor, Tuple[int, int]]:
+    """
+    Partition into non-overlapping windows with padding if needed.
+    Args:
+        x (tensor): input tokens with [B, H, W, C].
+        window_size (int): window size.
+    Returns:
+        windows: windows after partition with [B * num_windows, window_size, window_size, C].
+        (Hp, Wp): padded height and width before partition
+    """
+    B, H, W, C = x.shape
+
+    pad_h = (window_size - H % window_size) % window_size
+    pad_w = (window_size - W % window_size) % window_size
+    if pad_h > 0 or pad_w > 0:
+        x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))
+    Hp, Wp = H + pad_h, W + pad_w
+
+    x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C)
+    windows = x.permute(0, 1, 3, 2, 4, 5).reshape(-1, window_size, window_size, C)
+    return windows, (Hp, Wp)
+
+
+def window_unpartition(
+    windows: Tensor, window_size: int, pad_hw: Tuple[int, int], hw: Tuple[int, int]
+) -> Tensor:
+    """
+    Window unpartition into original sequences and removing padding.
+    Args:
+        x (tensor): input tokens with [B * num_windows, window_size, window_size, C].
+        window_size (int): window size.
+        pad_hw (Tuple): padded height and width (Hp, Wp).
+        hw (Tuple): original height and width (H, W) before padding.
+    Returns:
+        x: unpartitioned sequences with [B, H, W, C].
+    """
+    Hp, Wp = pad_hw
+    H, W = hw
+    B = windows.shape[0] // (Hp * Wp // window_size // window_size)
+    x = windows.reshape(
+        B, Hp // window_size, Wp // window_size, window_size, window_size, -1
+    )
+    x = x.permute(0, 1, 3, 2, 4, 5).reshape(B, Hp, Wp, -1)
+
+    if Hp > H or Wp > W:
+        x = x[:, :H, :W, :]
+    return x
+
+
+def get_rel_pos(q_size: int, k_size: int, rel_pos: Tensor) -> Tensor:
+    """
+    Get relative positional embeddings according to the relative positions of
+        query and key sizes.
+    Args:
+        q_size (int): size of query q.
+        k_size (int): size of key k.
+        rel_pos (Tensor): relative position embeddings (L, C).
+    Returns:
+        Extracted positional embeddings according to relative positions.
+    """
+    max_rel_dist = int(2 * max(q_size, k_size) - 1)
+    # Interpolate rel pos if needed.
+    if rel_pos.shape[0] != max_rel_dist:
+        # Interpolate rel pos.
+        rel_pos_resized = F.interpolate(
+            rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),
+            size=max_rel_dist,
+            mode="linear",
+            align_corners=False,
+        )
+        rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)
+    else:
+        rel_pos_resized = rel_pos
+
+    # Scale the coords with short length if shapes for q and k are different.
+    q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0)
+    k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0)
+    relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)
+
+    return rel_pos_resized[relative_coords.long()]
+
+
+def get_abs_pos(
+    abs_pos: Tensor,
+    has_cls_token: bool,
+    hw: Tuple[int, int],
+    retain_cls_token: bool = False,
+    tiling: bool = False,
+) -> Tensor:
+    """
+    Calculate absolute positional embeddings. If needed, resize embeddings and remove cls_token
+        dimension for the original embeddings.
+    Args:
+        abs_pos (Tensor): absolute positional embeddings with (1, num_position, C).
+        has_cls_token (bool): If true, has 1 embedding in abs_pos for cls token.
+        hw (Tuple): size of input image tokens.
+        retain_cls_token: whether to retain the cls_token
+        tiling: whether to tile the embeddings, *instead* of interpolation (a la abs_win)
+    Returns:
+        Absolute positional embeddings after processing with shape (1, H, W, C),
+        if retain_cls_token is False, otherwise (1, 1+H*W, C)
+    """
+    if retain_cls_token:
+        assert has_cls_token
+
+    h, w = hw
+    if has_cls_token:
+        cls_pos = abs_pos[:, :1]
+        abs_pos = abs_pos[:, 1:]
+
+    xy_num = abs_pos.shape[1]
+    size = int(math.sqrt(xy_num))
+    assert size * size == xy_num
+
+    if size != h or size != w:
+        new_abs_pos = abs_pos.reshape(1, size, size, -1).permute(0, 3, 1, 2)
+        if tiling:
+            new_abs_pos = new_abs_pos.tile(
+                [1, 1] + [x // y + 1 for x, y in zip((h, w), new_abs_pos.shape[2:])]
+            )[:, :, :h, :w]
+        else:
+            new_abs_pos = F.interpolate(
+                new_abs_pos,
+                size=(h, w),
+                mode="bicubic",
+                align_corners=False,
+            )
+
+        if not retain_cls_token:
+            return new_abs_pos.permute(0, 2, 3, 1)
+        else:
+            # add cls_token back, flatten spatial dims
+            assert has_cls_token
+            return torch.cat(
+                [cls_pos, new_abs_pos.permute(0, 2, 3, 1).reshape(1, h * w, -1)],
+                dim=1,
+            )
+
+    else:
+        if not retain_cls_token:
+            return abs_pos.reshape(1, h, w, -1)
+        else:
+            assert has_cls_token
+            return torch.cat([cls_pos, abs_pos], dim=1)
+
+
+def concat_rel_pos(
+    q: Tensor,
+    k: Tensor,
+    q_hw: Tuple[int, int],
+    k_hw: Tuple[int, int],
+    rel_pos_h: Tensor,
+    rel_pos_w: Tensor,
+    rescale: bool = False,
+    relative_coords: Optional[Tensor] = None,
+) -> Tuple[Tensor, Tensor]:
+    """
+    Concatenate rel pos coeffs to the q & k tensors, so that qk^T is now
+    effectively including rel pos biases.
+    Args:
+        q (Tensor): q tensor with shape (B, L_q, C).
+        k (Tensor): k tensor with shape (B, L_k, C).
+        q_hw, k_hw: These are spatial size of q & k tensors.
+        rel_pos_h, rel_pos_w: These are relative pos embeddings/params of height, width.
+        rescale (bool): whether to rescale. e.g. for use when using sdpa, pytorch will
+            scale by the wrong factor due to the concat.
+    Returns:
+        q, k: But, padded so that qk^T accounts for rel pos biases
+    """
+    q_h, q_w = q_hw
+    k_h, k_w = k_hw
+
+    assert (q_h == q_w) and (k_h == k_w), "only square inputs supported"
+
+    if relative_coords is not None:
+        Rh = rel_pos_h[relative_coords]
+        Rw = rel_pos_w[relative_coords]
+    else:
+        Rh = get_rel_pos(q_h, k_h, rel_pos_h)
+        Rw = get_rel_pos(q_w, k_w, rel_pos_w)
+
+    B, _, dim = q.shape
+    r_q = q.reshape(B, q_h, q_w, dim)
+
+    old_scale = dim**0.5
+    new_scale = (dim + k_h + k_w) ** 0.5 if rescale else old_scale  # for sdpa
+    # attn will be divided by new_scale, but we want to divide q by old_scale
+    scale_ratio = new_scale / old_scale
+
+    rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh) * new_scale  # (B, q_h, q_w, k_h)
+    rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw) * new_scale  # (B, q_h, q_w, k_w)
+
+    eye_h = torch.eye(k_h, dtype=q.dtype, device=q.device)
+    eye_w = torch.eye(k_w, dtype=q.dtype, device=q.device)
+
+    eye_h = eye_h.view(1, k_h, 1, k_h).expand([B, k_h, k_w, k_h])
+    eye_w = eye_w.view(1, 1, k_w, k_w).expand([B, k_h, k_w, k_w])
+
+    q = torch.cat([r_q * scale_ratio, rel_h, rel_w], dim=-1).view(B, q_h * q_w, -1)
+    k = torch.cat([k.view(B, k_h, k_w, -1), eye_h, eye_w], dim=-1).view(
+        B, k_h * k_w, -1
+    )
+
+    return q, k
+
+
+class PatchEmbed(nn.Module):
+    """
+    Image to Patch Embedding.
+    """
+
+    def __init__(
+        self,
+        kernel_size: Tuple[int, int] = (16, 16),
+        stride: Tuple[int, int] = (16, 16),
+        padding: Tuple[int, int] = (0, 0),
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        bias: bool = True,
+    ):
+        """
+        Args:
+            kernel_size (Tuple): kernel size of the projection layer.
+            stride (Tuple): stride of the projection layer.
+            padding (Tuple): padding size of the projection layer.
+            in_chans (int): Number of input image channels.
+            embed_dim (int):  embed_dim (int): Patch embedding dimension.
+        """
+        super().__init__()
+
+        self.proj = nn.Conv2d(
+            in_chans,
+            embed_dim,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            bias=bias,
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.proj(x)
+        # B C H W -> B H W C
+        x = x.permute(0, 2, 3, 1)
+        return x
+
+
+class Attention(nn.Module):
+    """Multi-head Attention block with relative position embeddings and 2d-rope."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = True,
+        use_rel_pos: bool = False,
+        rel_pos_zero_init: bool = True,
+        input_size: Optional[Tuple[int, int]] = None,
+        cls_token: bool = False,
+        use_rope: bool = False,
+        rope_theta: float = 10000.0,
+        rope_pt_size: Optional[Tuple[int, int]] = None,
+        rope_interp: bool = False,
+    ):
+        """
+        Args:
+            dim (int): Number of input channels.
+            num_heads (int): Number of attention heads.
+            qkv_bias (bool:  If True, add a learnable bias to query, key, value.
+            rel_pos (bool): If True, add relative positional embeddings to the attention map.
+            rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+            input_size (int or None): Input resolution for calculating the relative positional
+                parameter size or rope size.
+            attn_type: Type of attention operation, e.g. "vanilla", "vanilla-xformer".
+            cls_token: whether a cls_token is present.
+            use_rope: whether to use rope 2d (indep of use_rel_pos, as it can be used together)
+            rope_theta: control frequencies of rope
+            rope_pt_size: size of rope in previous stage of training, needed for interpolation or tiling
+            rope_interp: whether to interpolate (or extrapolate) rope to match input size
+        """
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim**-0.5
+        self.cls_token = cls_token
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.proj = nn.Linear(dim, dim)
+
+        # rel_pos embeddings and rope
+        self.use_rel_pos = use_rel_pos
+        self.input_size = input_size
+
+        self.use_rope = use_rope
+        self.rope_theta = rope_theta
+        self.rope_pt_size = rope_pt_size
+        self.rope_interp = rope_interp
+
+        # init rel_pos embeddings and rope
+        self._setup_rel_pos(rel_pos_zero_init)
+        self._setup_rope_freqs()
+
+    def _setup_rel_pos(self, rel_pos_zero_init: bool = True) -> None:
+        if not self.use_rel_pos:
+            self.rel_pos_h = None
+            self.rel_pos_w = None
+            return
+
+        assert self.input_size is not None
+        assert self.cls_token is False, "not supported"
+        # initialize relative positional embeddings
+        self.rel_pos_h = nn.Parameter(
+            torch.zeros(2 * self.input_size[0] - 1, self.head_dim)
+        )
+        self.rel_pos_w = nn.Parameter(
+            torch.zeros(2 * self.input_size[1] - 1, self.head_dim)
+        )
+
+        if not rel_pos_zero_init:
+            trunc_normal_(self.rel_pos_h, std=0.02)
+            trunc_normal_(self.rel_pos_w, std=0.02)
+
+        # Precompute the relative coords
+        H, W = self.input_size
+        q_coords = torch.arange(H)[:, None]
+        k_coords = torch.arange(W)[None, :]
+        relative_coords = (q_coords - k_coords) + (H - 1)
+        self.register_buffer("relative_coords", relative_coords.long())
+
+    def _setup_rope_freqs(self) -> None:
+        if not self.use_rope:
+            self.freqs_cis = None
+            return
+
+        assert self.input_size is not None
+        # determine rope input size
+        if self.rope_pt_size is None:
+            self.rope_pt_size = self.input_size
+
+        # initialize 2d rope freqs
+        self.compute_cis = partial(
+            compute_axial_cis,
+            dim=self.head_dim,
+            theta=self.rope_theta,
+        )
+
+        # interpolate rope
+        scale_pos = 1.0
+        if self.rope_interp:
+            scale_pos = self.rope_pt_size[0] / self.input_size[0]
+        # get scaled freqs_cis
+        freqs_cis = self.compute_cis(
+            end_x=self.input_size[0],
+            end_y=self.input_size[1],
+            scale_pos=scale_pos,
+        )
+        if self.cls_token:
+            t = torch.zeros(
+                self.head_dim // 2,
+                dtype=torch.float32,
+                device=freqs_cis.device,
+            )
+            cls_freqs_cis = torch.polar(torch.ones_like(t), t)[None, :]
+            freqs_cis = torch.cat([cls_freqs_cis, freqs_cis], dim=0)
+
+        self.register_buffer("freqs_cis", freqs_cis)
+
+    def _apply_rope(self, q, k) -> Tuple[Tensor, Tensor]:
+        if not self.use_rope:
+            return q, k
+
+        assert self.freqs_cis is not None
+        return apply_rotary_enc(q, k, freqs_cis=self.freqs_cis)
+
+    def forward(self, x: Tensor) -> Tensor:
+        s = 1 if self.cls_token else 0  # used to exclude cls_token
+        if x.ndim == 4:
+            B, H, W, _ = x.shape
+            assert s == 0  # no cls_token
+            L = H * W
+            ndim = 4
+        else:
+            assert x.ndim == 3
+            B, L, _ = x.shape
+            ndim = 3
+            H = W = math.sqrt(L - s)
+
+        # qkv with shape (3, B, nHead, L, C)
+        qkv = self.qkv(x).reshape(B, L, 3, self.num_heads, -1)
+        # q, k, v with shape (B, nHead, L, C)
+        q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)
+
+        # handle rope and rel pos embeddings
+        q, k = self._apply_rope(q, k)
+        if self.use_rel_pos:
+            q, k = concat_rel_pos(
+                q.flatten(0, 1),
+                k.flatten(0, 1),
+                (H, W),
+                x.shape[1:3],
+                self.rel_pos_h,
+                self.rel_pos_w,
+                rescale=True,
+                relative_coords=self.relative_coords,
+            )
+
+            # sdpa expects [B, nheads, H*W, C] so we transpose back
+            q = q.reshape(B, self.num_heads, H * W, -1)
+            k = k.reshape(B, self.num_heads, H * W, -1)
+
+        x = F.scaled_dot_product_attention(q, k, v)
+
+        if ndim == 4:
+            x = (
+                x.view(B, self.num_heads, H, W, -1)
+                .permute(0, 2, 3, 1, 4)
+                .reshape(B, H, W, -1)
+            )
+        else:
+            x = x.view(B, self.num_heads, L, -1).permute(0, 2, 1, 3).reshape(B, L, -1)
+
+        x = self.proj(x)
+
+        return x
+
+
+class Block(nn.Module):
+    """Transformer blocks with support of window attention"""
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = True,
+        drop_path: float = 0.0,
+        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        use_rel_pos: bool = False,
+        rel_pos_zero_init: bool = True,
+        window_size: int = 0,
+        input_size: Optional[Tuple[int, int]] = None,
+        use_rope: bool = False,
+        rope_pt_size: Optional[Tuple[int, int]] = None,
+        rope_tiled: bool = False,
+        rope_interp: bool = False,
+        use_ve_rope: bool = False,
+        cls_token: bool = False,
+        dropout: float = 0.0,
+        init_values: Optional[float] = None,
+    ):
+        """
+        Args:
+            dim (int): Number of input channels.
+            num_heads (int): Number of attention heads in each ViT block.
+            mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+            qkv_bias (bool): If True, add a learnable bias to query, key, value.
+            drop_path (float): Stochastic depth rate.
+            norm_layer (nn.Module): Normalization layer.
+            act_layer (nn.Module): Activation layer.
+            use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
+            rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+            window_size (int): Window size for window attention blocks. If it equals 0, then not
+                use window attention.
+            input_size (int or None): Input resolution for calculating the relative positional
+                parameter size.
+            dropout (float): Dropout rate.
+            cls_token: whether a cls_token is present.
+            use_rope: whether to use rope 2d (indep of use_rel_pos, as it can be used together)
+            rope_pt_size: size of rope in previous stage of training, needed for interpolation or tiling
+            rope_interp: whether to interpolate (or extrapolate) rope to match target input size,
+                expected to specify source size as rope_pt_size.
+        """
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            use_rel_pos=use_rel_pos,
+            rel_pos_zero_init=rel_pos_zero_init,
+            input_size=input_size if window_size == 0 else (window_size, window_size),
+            use_rope=use_rope,
+            rope_pt_size=rope_pt_size,
+            rope_interp=rope_interp,
+            cls_token=cls_token,
+        )
+        self.ls1 = (
+            LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=int(dim * mlp_ratio),
+            act_layer=act_layer,
+            drop=(dropout, 0.0),
+        )
+        self.ls2 = (
+            LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        )
+        self.dropout = nn.Dropout(dropout)
+        self.window_size = window_size
+
+    def forward(self, x: Tensor) -> Tensor:
+        shortcut = x
+        x = self.norm1(x)
+        # Window partition
+        if self.window_size > 0:
+            H, W = x.shape[1], x.shape[2]
+            x, pad_hw = window_partition(x, self.window_size)
+
+        x = self.ls1(self.attn(x))
+        # Reverse window partition
+        if self.window_size > 0:
+            x = window_unpartition(x, self.window_size, pad_hw, (H, W))
+
+        x = shortcut + self.dropout(self.drop_path(x))
+        x = x + self.dropout(self.drop_path(self.ls2(self.mlp(self.norm2(x)))))
+
+        return x
+
+
+class ViT(nn.Module):
+    """
+    This module implements Vision Transformer (ViT) backbone in :paper:`vitdet`.
+    "Exploring Plain Vision Transformer Backbones for Object Detection",
+    https://arxiv.org/abs/2203.16527
+    """
+
+    def __init__(
+        self,
+        img_size: int = 1024,
+        patch_size: int = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        depth: int = 12,
+        num_heads: int = 12,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = True,
+        drop_path_rate: float = 0.0,
+        norm_layer: Union[Callable[..., nn.Module], str] = "LayerNorm",
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        use_abs_pos: bool = True,
+        tile_abs_pos: bool = True,
+        rel_pos_blocks: Union[Tuple[int, ...], bool] = (2, 5, 8, 11),
+        rel_pos_zero_init: bool = True,
+        window_size: int = 14,
+        global_att_blocks: Tuple[int, ...] = (2, 5, 8, 11),
+        use_rope: bool = False,
+        rope_pt_size: Optional[int] = None,
+        use_interp_rope: bool = False,
+        pretrain_img_size: int = 224,
+        pretrain_use_cls_token: bool = True,
+        retain_cls_token: bool = True,
+        dropout: float = 0.0,
+        return_interm_layers: bool = False,
+        init_values: Optional[float] = None,  # for layerscale
+        ln_pre: bool = False,
+        ln_post: bool = False,
+        bias_patch_embed: bool = True,
+        compile_mode: Optional[str] = None,
+        use_act_checkpoint: bool = True,
+    ):
+        """
+        Args:
+            img_size (int): Input image size. Only relevant for rel pos or rope.
+            patch_size (int): Patch size.
+            in_chans (int): Number of input image channels.
+            embed_dim (int): Patch embedding dimension.
+            depth (int): Depth of ViT.
+            num_heads (int): Number of attention heads in each ViT block.
+            mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+            qkv_bias (bool): If True, add a learnable bias to query, key, value.
+            drop_path_rate (float): Stochastic depth rate.
+            norm_layer (nn.Module): Normalization layer.
+            act_layer (nn.Module): Activation layer.
+            use_abs_pos (bool): If True, use absolute positional embeddings.
+            tile_abs_pos (bool): If True, tile absolute positional embeddings instead of interpolation.
+            rel_pos_blocks (list): Blocks which have rel pos embeddings.
+            rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+            window_size (int): Window size for window attention blocks.
+            global_att_blocks (list): Indexes for blocks using global attention (other blocks use window attention).
+            use_rope (bool): whether to use rope 2d (indep of rel_pos_blocks, as it can be used together).
+            rope_pt_size (int): size of rope in previous stage of training, needed for interpolation or tiling.
+            use_interp_rope: whether to interpolate (or extrapolate) rope to match target input size,
+                expected to specify source size as rope_pt_size.
+            use_act_checkpoint (bool): If True, use activation checkpointing.
+            pretrain_img_size (int): input image size for pretraining models.
+            pretrain_use_cls_token (bool): If True, pretraining models use class token.
+            retain_cls_token: whether cls_token should be retained.
+            dropout (float): Dropout rate. Applied in residual blocks of attn, mlp and inside the mlp.
+
+            return_interm_layers (bool): Whether to return intermediate layers (all global attention blocks).
+            init_values: layer scale init, None for no layer scale.
+
+            ln_pre (bool): If True, apply layer norm before transformer blocks.
+            ln_post (bool): If True, apply layer norm after transformer blocks.
+            bias_patch_embed (bool): bias in conv for patch embed?
+            compile_mode (str): mode to compile the forward
+        """
+        super().__init__()
+        self.pretrain_use_cls_token = pretrain_use_cls_token
+
+        window_block_indexes = [i for i in range(depth) if i not in global_att_blocks]
+        self.full_attn_ids = list(global_att_blocks)
+        self.rel_pos_blocks = [False] * depth
+        if isinstance(rel_pos_blocks, bool) and rel_pos_blocks:
+            self.rel_pos_blocks = [True] * depth
+        else:
+            for i in rel_pos_blocks:
+                self.rel_pos_blocks[i] = True
+
+        self.retain_cls_token = retain_cls_token
+        if self.retain_cls_token:
+            assert pretrain_use_cls_token
+            assert (
+                len(window_block_indexes) == 0
+            ), "windowing not supported with cls token"
+
+            assert sum(self.rel_pos_blocks) == 0, "rel pos not supported with cls token"
+
+            scale = embed_dim**-0.5
+            self.class_embedding = nn.Parameter(scale * torch.randn(1, 1, embed_dim))
+
+        if isinstance(norm_layer, str):
+            norm_layer = partial(getattr(nn, norm_layer), eps=1e-5)
+
+        self.patch_embed = PatchEmbed(
+            kernel_size=(patch_size, patch_size),
+            stride=(patch_size, patch_size),
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+            bias=bias_patch_embed,
+        )
+
+        # Handle absolute positional embedding
+        self.tile_abs_pos = tile_abs_pos
+        self.use_abs_pos = use_abs_pos
+        if self.tile_abs_pos:
+            assert self.use_abs_pos
+
+        if self.use_abs_pos:
+            # Initialize absolute positional embedding with pretrain image size.
+            num_patches = (pretrain_img_size // patch_size) * (
+                pretrain_img_size // patch_size
+            )
+            num_positions = (num_patches + 1) if pretrain_use_cls_token else num_patches
+            self.pos_embed = nn.Parameter(torch.zeros(1, num_positions, embed_dim))
+        else:
+            self.pos_embed = None
+
+        # stochastic depth decay rule
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
+
+        self.blocks = nn.ModuleList()
+        cur_stage = 1
+        for i in range(depth):
+            block = Block(
+                dim=embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                drop_path=dpr[i],
+                norm_layer=norm_layer,
+                act_layer=act_layer,
+                use_rel_pos=self.rel_pos_blocks[i],
+                rel_pos_zero_init=rel_pos_zero_init,
+                window_size=window_size if i in window_block_indexes else 0,
+                input_size=(img_size // patch_size, img_size // patch_size),
+                use_rope=use_rope,
+                rope_pt_size=(
+                    (window_size, window_size)
+                    if rope_pt_size is None
+                    else (rope_pt_size, rope_pt_size)
+                ),
+                rope_interp=use_interp_rope,
+                cls_token=self.retain_cls_token,
+                dropout=dropout,
+                init_values=init_values,
+            )
+
+            if i not in window_block_indexes:
+                cur_stage += 1
+
+            self.use_act_checkpoint = use_act_checkpoint
+
+            self.blocks.append(block)
+
+        self.return_interm_layers = return_interm_layers
+        self.channel_list = (
+            [embed_dim] * len(self.full_attn_ids)
+            if return_interm_layers
+            else [embed_dim]
+        )
+
+        if self.pos_embed is not None:
+            trunc_normal_(self.pos_embed, std=0.02)
+
+        self.ln_pre = norm_layer(embed_dim) if ln_pre else nn.Identity()
+        self.ln_post = norm_layer(embed_dim) if ln_post else nn.Identity()
+
+        self.apply(self._init_weights)
+
+        if compile_mode is not None:
+            self.forward = torch.compile(
+                self.forward, mode=compile_mode, fullgraph=True
+            )
+            if self.use_act_checkpoint and self.training:
+                torch._dynamo.config.optimize_ddp = False
+
+    def _init_weights(self, m: nn.Module) -> None:
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+        x = self.patch_embed(x)
+        h, w = x.shape[1], x.shape[2]
+
+        s = 0
+        if self.retain_cls_token:
+            # If cls_token is retained, we don't
+            # maintain spatial shape
+            x = torch.cat([self.class_embedding, x.flatten(1, 2)], dim=1)
+            s = 1
+
+        if self.pos_embed is not None:
+            x = x + get_abs_pos(
+                self.pos_embed,
+                self.pretrain_use_cls_token,
+                (h, w),
+                self.retain_cls_token,
+                tiling=self.tile_abs_pos,
+            )
+
+        x = self.ln_pre(x)
+
+        outputs = []
+        for i, blk in enumerate(self.blocks):
+            if self.use_act_checkpoint and self.training:
+                x = checkpoint.checkpoint(blk, x, use_reentrant=False)
+            else:
+                x = blk(x)
+            if (i == self.full_attn_ids[-1]) or (
+                self.return_interm_layers and i in self.full_attn_ids
+            ):
+                if i == self.full_attn_ids[-1]:
+                    x = self.ln_post(x)
+
+                feats = x[:, s:]
+                if feats.ndim == 4:
+                    feats = feats.permute(0, 3, 1, 2)
+                else:
+                    assert feats.ndim == 3
+                    h = w = math.sqrt(feats.shape[1])
+                    feats = feats.reshape(
+                        feats.shape[0], h, w, feats.shape[-1]
+                    ).permute(0, 3, 1, 2)
+
+                outputs.append(feats)
+
+        return outputs
+
+    def get_layer_id(self, layer_name: str) -> int:
+        # https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L33
+        num_layers = self.get_num_layers()
+
+        if layer_name.find("rel_pos") != -1:
+            return num_layers + 1
+        elif layer_name.find("ln_pre") != -1:
+            return 0
+        elif layer_name.find("pos_embed") != -1 or layer_name.find("cls_token") != -1:
+            return 0
+        elif layer_name.find("patch_embed") != -1:
+            return 0
+        elif layer_name.find("blocks") != -1:
+            return int(layer_name.split("blocks")[1].split(".")[1]) + 1
+        else:
+            return num_layers + 1
+
+    def get_num_layers(self) -> int:
+        return len(self.blocks)

sam3/model/vl_combiner.py

@@ -0,0 +1,176 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+"""Provides utility to combine a vision backbone with a language backbone."""
+
+from copy import copy
+from typing import List, Optional
+
+import torch
+import torch.nn as nn
+
+from torch.nn.attention import sdpa_kernel, SDPBackend
+
+from .act_ckpt_utils import activation_ckpt_wrapper
+from .necks import Sam3DualViTDetNeck
+
+
+class SAM3VLBackbone(nn.Module):
+    """This backbone combines a vision backbone and a language backbone without fusion.
+    As such it is more of a convenience wrapper to handle the two backbones together.
+
+    It adds support for activation checkpointing and compilation.
+    """
+
+    def __init__(
+        self,
+        visual: Sam3DualViTDetNeck,
+        text,
+        compile_visual: bool = False,
+        act_ckpt_whole_vision_backbone: bool = False,
+        act_ckpt_whole_language_backbone: bool = False,
+        scalp=0,
+    ):
+        """Initialize the backbone combiner.
+
+        :param visual: The vision backbone to use
+        :param text: The text encoder to use
+        """
+        super().__init__()
+        self.vision_backbone: Sam3DualViTDetNeck = (
+            torch.compile(visual) if compile_visual else visual
+        )
+        self.language_backbone = text
+        self.scalp = scalp
+        # allow running activation checkpointing on the entire vision and language backbones
+        self.act_ckpt_whole_vision_backbone = act_ckpt_whole_vision_backbone
+        self.act_ckpt_whole_language_backbone = act_ckpt_whole_language_backbone
+
+    def forward(
+        self,
+        samples: torch.Tensor,
+        captions: List[str],
+        input_boxes: Optional[torch.Tensor] = None,
+        additional_text: Optional[List[str]] = None,
+    ):
+        """Forward pass of the backbone combiner.
+
+        :param samples: The input images
+        :param captions: The input captions
+        :param input_boxes: If the text contains place-holders for boxes, this
+            parameter contains the tensor containing their spatial features
+        :param additional_text: This can be used to encode some additional text
+            (different from the captions) in the same forward of the backbone
+        :return: Output dictionary with the following keys:
+            - vision_features: The output of the vision backbone
+            - language_features: The output of the language backbone
+            - language_mask: The attention mask of the language backbone
+            - vision_pos_enc: The positional encoding of the vision backbone
+            - (optional) additional_text_features: The output of the language
+                backbone for the additional text
+            - (optional) additional_text_mask: The attention mask of the
+                language backbone for the additional text
+        """
+        output = self.forward_image(samples)
+        device = output["vision_features"].device
+        output.update(self.forward_text(captions, input_boxes, additional_text, device))
+        return output
+
+    def forward_image(self, samples: torch.Tensor):
+        return activation_ckpt_wrapper(self._forward_image_no_act_ckpt)(
+            samples=samples,
+            act_ckpt_enable=self.act_ckpt_whole_vision_backbone and self.training,
+        )
+
+    def _forward_image_no_act_ckpt(self, samples):
+        # Forward through backbone
+        sam3_features, sam3_pos, sam2_features, sam2_pos = self.vision_backbone.forward(
+            samples
+        )
+        if self.scalp > 0:
+            # Discard the lowest resolution features
+            sam3_features, sam3_pos = (
+                sam3_features[: -self.scalp],
+                sam3_pos[: -self.scalp],
+            )
+            if sam2_features is not None and sam2_pos is not None:
+                sam2_features, sam2_pos = (
+                    sam2_features[: -self.scalp],
+                    sam2_pos[: -self.scalp],
+                )
+
+        sam2_output = None
+
+        if sam2_features is not None and sam2_pos is not None:
+            sam2_src = sam2_features[-1]
+            sam2_output = {
+                "vision_features": sam2_src,
+                "vision_pos_enc": sam2_pos,
+                "backbone_fpn": sam2_features,
+            }
+
+        sam3_src = sam3_features[-1]
+        output = {
+            "vision_features": sam3_src,
+            "vision_pos_enc": sam3_pos,
+            "backbone_fpn": sam3_features,
+            "sam2_backbone_out": sam2_output,
+        }
+
+        return output
+
+    def forward_text(
+        self, captions, input_boxes=None, additional_text=None, device="cuda"
+    ):
+        return activation_ckpt_wrapper(self._forward_text_no_ack_ckpt)(
+            captions=captions,
+            input_boxes=input_boxes,
+            additional_text=additional_text,
+            device=device,
+            act_ckpt_enable=self.act_ckpt_whole_language_backbone and self.training,
+        )
+
+    def _forward_text_no_ack_ckpt(
+        self,
+        captions,
+        input_boxes=None,
+        additional_text=None,
+        device="cuda",
+    ):
+        output = {}
+
+        # Forward through text_encoder
+        text_to_encode = copy(captions)
+        if additional_text is not None:
+            # if there are additional_text, we piggy-back them into this forward.
+            # They'll be used later for output alignment
+            text_to_encode += additional_text
+
+        sdpa_context = sdpa_kernel(
+            [
+                SDPBackend.MATH,
+                SDPBackend.EFFICIENT_ATTENTION,
+                SDPBackend.FLASH_ATTENTION,
+            ]
+        )
+
+        with sdpa_context:
+            text_attention_mask, text_memory, text_embeds = self.language_backbone(
+                text_to_encode, input_boxes, device=device
+            )
+
+        if additional_text is not None:
+            output["additional_text_features"] = text_memory[:, -len(additional_text) :]
+            output["additional_text_mask"] = text_attention_mask[
+                -len(additional_text) :
+            ]
+
+        text_memory = text_memory[:, : len(captions)]
+        text_attention_mask = text_attention_mask[: len(captions)]
+        text_embeds = text_embeds[:, : len(captions)]
+        output["language_features"] = text_memory
+        output["language_mask"] = text_attention_mask
+        output["language_embeds"] = (
+            text_embeds  # Text embeddings before forward to the encoder
+        )
+
+        return output

sam3/model_builder.py

@@ -0,0 +1,795 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import os
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from huggingface_hub import hf_hub_download
+from iopath.common.file_io import g_pathmgr
+
+from sam3.model.decoder import (
+    TransformerDecoder,
+    TransformerDecoderLayer,
+    TransformerDecoderLayerv2,
+    TransformerEncoderCrossAttention,
+)
+from sam3.model.encoder import TransformerEncoderFusion, TransformerEncoderLayer
+from sam3.model.geometry_encoders import SequenceGeometryEncoder
+from sam3.model.maskformer_segmentation import PixelDecoder, UniversalSegmentationHead
+from sam3.model.memory import (
+    CXBlock,
+    SimpleFuser,
+    SimpleMaskDownSampler,
+    SimpleMaskEncoder,
+)
+from sam3.model.model_misc import (
+    DotProductScoring,
+    MLP,
+    MultiheadAttentionWrapper as MultiheadAttention,
+    TransformerWrapper,
+)
+from sam3.model.necks import Sam3DualViTDetNeck
+from sam3.model.position_encoding import PositionEmbeddingSine
+
+from sam3.model.sam1_task_predictor import SAM3InteractiveImagePredictor
+
+from sam3.model.sam3_image import Sam3Image, Sam3ImageOnVideoMultiGPU
+from sam3.model.sam3_tracking_predictor import Sam3TrackerPredictor
+from sam3.model.sam3_video_inference import Sam3VideoInferenceWithInstanceInteractivity
+from sam3.model.sam3_video_predictor import Sam3VideoPredictorMultiGPU
+from sam3.model.text_encoder_ve import VETextEncoder
+from sam3.model.tokenizer_ve import SimpleTokenizer
+from sam3.model.vitdet import ViT
+from sam3.model.vl_combiner import SAM3VLBackbone
+from sam3.sam.transformer import RoPEAttention
+
+SAM3_MODEL_ID = "facebook/sam3"
+SAM3_CKPT_NAME = "sam3.pt"
+
+
+# Setup TensorFloat-32 for Ampere GPUs if available
+def _setup_tf32() -> None:
+    """Enable TensorFloat-32 for Ampere GPUs if available."""
+    if torch.cuda.is_available():
+        device_props = torch.cuda.get_device_properties(0)
+        if device_props.major >= 8:
+            torch.backends.cuda.matmul.allow_tf32 = True
+            torch.backends.cudnn.allow_tf32 = True
+
+
+_setup_tf32()
+
+
+def _create_position_encoding(precompute_resolution=None):
+    """Create position encoding for visual backbone."""
+    return PositionEmbeddingSine(
+        num_pos_feats=256,
+        normalize=True,
+        scale=None,
+        temperature=10000,
+        precompute_resolution=precompute_resolution,
+    )
+
+
+def _create_vit_backbone(compile_mode=None):
+    """Create ViT backbone for visual feature extraction."""
+    return ViT(
+        img_size=1008,
+        pretrain_img_size=336,
+        patch_size=14,
+        embed_dim=1024,
+        depth=32,
+        num_heads=16,
+        mlp_ratio=4.625,
+        norm_layer="LayerNorm",
+        drop_path_rate=0.1,
+        qkv_bias=True,
+        use_abs_pos=True,
+        tile_abs_pos=True,
+        global_att_blocks=(7, 15, 23, 31),
+        rel_pos_blocks=(),
+        use_rope=True,
+        use_interp_rope=True,
+        window_size=24,
+        pretrain_use_cls_token=True,
+        retain_cls_token=False,
+        ln_pre=True,
+        ln_post=False,
+        return_interm_layers=False,
+        bias_patch_embed=False,
+        compile_mode=compile_mode,
+    )
+
+
+def _create_vit_neck(position_encoding, vit_backbone, enable_inst_interactivity=False):
+    """Create ViT neck for feature pyramid."""
+    return Sam3DualViTDetNeck(
+        position_encoding=position_encoding,
+        d_model=256,
+        scale_factors=[4.0, 2.0, 1.0, 0.5],
+        trunk=vit_backbone,
+        add_sam2_neck=enable_inst_interactivity,
+    )
+
+
+def _create_vl_backbone(vit_neck, text_encoder):
+    """Create visual-language backbone."""
+    return SAM3VLBackbone(visual=vit_neck, text=text_encoder, scalp=1)
+
+
+def _create_transformer_encoder() -> TransformerEncoderFusion:
+    """Create transformer encoder with its layer."""
+    encoder_layer = TransformerEncoderLayer(
+        activation="relu",
+        d_model=256,
+        dim_feedforward=2048,
+        dropout=0.1,
+        pos_enc_at_attn=True,
+        pos_enc_at_cross_attn_keys=False,
+        pos_enc_at_cross_attn_queries=False,
+        pre_norm=True,
+        self_attention=MultiheadAttention(
+            num_heads=8,
+            dropout=0.1,
+            embed_dim=256,
+            batch_first=True,
+        ),
+        cross_attention=MultiheadAttention(
+            num_heads=8,
+            dropout=0.1,
+            embed_dim=256,
+            batch_first=True,
+        ),
+    )
+
+    encoder = TransformerEncoderFusion(
+        layer=encoder_layer,
+        num_layers=6,
+        d_model=256,
+        num_feature_levels=1,
+        frozen=False,
+        use_act_checkpoint=True,
+        add_pooled_text_to_img_feat=False,
+        pool_text_with_mask=True,
+    )
+    return encoder
+
+
+def _create_transformer_decoder() -> TransformerDecoder:
+    """Create transformer decoder with its layer."""
+    decoder_layer = TransformerDecoderLayer(
+        activation="relu",
+        d_model=256,
+        dim_feedforward=2048,
+        dropout=0.1,
+        cross_attention=MultiheadAttention(
+            num_heads=8,
+            dropout=0.1,
+            embed_dim=256,
+        ),
+        n_heads=8,
+        use_text_cross_attention=True,
+    )
+
+    decoder = TransformerDecoder(
+        layer=decoder_layer,
+        num_layers=6,
+        num_queries=200,
+        return_intermediate=True,
+        box_refine=True,
+        num_o2m_queries=0,
+        dac=True,
+        boxRPB="log",
+        d_model=256,
+        frozen=False,
+        interaction_layer=None,
+        dac_use_selfatt_ln=True,
+        resolution=1008,
+        stride=14,
+        use_act_checkpoint=True,
+        presence_token=True,
+    )
+    return decoder
+
+
+def _create_dot_product_scoring():
+    """Create dot product scoring module."""
+    prompt_mlp = MLP(
+        input_dim=256,
+        hidden_dim=2048,
+        output_dim=256,
+        num_layers=2,
+        dropout=0.1,
+        residual=True,
+        out_norm=nn.LayerNorm(256),
+    )
+    return DotProductScoring(d_model=256, d_proj=256, prompt_mlp=prompt_mlp)
+
+
+def _create_segmentation_head(compile_mode=None):
+    """Create segmentation head with pixel decoder."""
+    pixel_decoder = PixelDecoder(
+        num_upsampling_stages=3,
+        interpolation_mode="nearest",
+        hidden_dim=256,
+        compile_mode=compile_mode,
+    )
+
+    cross_attend_prompt = MultiheadAttention(
+        num_heads=8,
+        dropout=0,
+        embed_dim=256,
+    )
+
+    segmentation_head = UniversalSegmentationHead(
+        hidden_dim=256,
+        upsampling_stages=3,
+        aux_masks=False,
+        presence_head=False,
+        dot_product_scorer=None,
+        act_ckpt=True,
+        cross_attend_prompt=cross_attend_prompt,
+        pixel_decoder=pixel_decoder,
+    )
+    return segmentation_head
+
+
+def _create_geometry_encoder():
+    """Create geometry encoder with all its components."""
+    # Create position encoding for geometry encoder
+    geo_pos_enc = _create_position_encoding()
+    # Create CX block for fuser
+    cx_block = CXBlock(
+        dim=256,
+        kernel_size=7,
+        padding=3,
+        layer_scale_init_value=1.0e-06,
+        use_dwconv=True,
+    )
+    # Create geometry encoder layer
+    geo_layer = TransformerEncoderLayer(
+        activation="relu",
+        d_model=256,
+        dim_feedforward=2048,
+        dropout=0.1,
+        pos_enc_at_attn=False,
+        pre_norm=True,
+        self_attention=MultiheadAttention(
+            num_heads=8,
+            dropout=0.1,
+            embed_dim=256,
+            batch_first=False,
+        ),
+        pos_enc_at_cross_attn_queries=False,
+        pos_enc_at_cross_attn_keys=True,
+        cross_attention=MultiheadAttention(
+            num_heads=8,
+            dropout=0.1,
+            embed_dim=256,
+            batch_first=False,
+        ),
+    )
+
+    # Create geometry encoder
+    input_geometry_encoder = SequenceGeometryEncoder(
+        pos_enc=geo_pos_enc,
+        encode_boxes_as_points=False,
+        points_direct_project=True,
+        points_pool=True,
+        points_pos_enc=True,
+        boxes_direct_project=True,
+        boxes_pool=True,
+        boxes_pos_enc=True,
+        d_model=256,
+        num_layers=3,
+        layer=geo_layer,
+        use_act_ckpt=True,
+        add_cls=True,
+        add_post_encode_proj=True,
+    )
+    return input_geometry_encoder
+
+
+def _create_sam3_model(
+    backbone,
+    transformer,
+    input_geometry_encoder,
+    segmentation_head,
+    dot_prod_scoring,
+    inst_interactive_predictor,
+    eval_mode,
+):
+    """Create the SAM3 image model."""
+    common_params = {
+        "backbone": backbone,
+        "transformer": transformer,
+        "input_geometry_encoder": input_geometry_encoder,
+        "segmentation_head": segmentation_head,
+        "num_feature_levels": 1,
+        "o2m_mask_predict": True,
+        "dot_prod_scoring": dot_prod_scoring,
+        "use_instance_query": False,
+        "multimask_output": True,
+        "inst_interactive_predictor": inst_interactive_predictor,
+    }
+
+    matcher = None
+    if not eval_mode:
+        from sam3.train.matcher import BinaryHungarianMatcherV2
+
+        matcher = BinaryHungarianMatcherV2(
+            focal=True,
+            cost_class=2.0,
+            cost_bbox=5.0,
+            cost_giou=2.0,
+            alpha=0.25,
+            gamma=2,
+            stable=False,
+        )
+    common_params["matcher"] = matcher
+    model = Sam3Image(**common_params)
+
+    return model
+
+
+def _create_tracker_maskmem_backbone():
+    """Create the SAM3 Tracker memory encoder."""
+    # Position encoding for mask memory backbone
+    position_encoding = PositionEmbeddingSine(
+        num_pos_feats=64,
+        normalize=True,
+        scale=None,
+        temperature=10000,
+        precompute_resolution=1008,
+    )
+
+    # Mask processing components
+    mask_downsampler = SimpleMaskDownSampler(
+        kernel_size=3, stride=2, padding=1, interpol_size=[1152, 1152]
+    )
+
+    cx_block_layer = CXBlock(
+        dim=256,
+        kernel_size=7,
+        padding=3,
+        layer_scale_init_value=1.0e-06,
+        use_dwconv=True,
+    )
+
+    fuser = SimpleFuser(layer=cx_block_layer, num_layers=2)
+
+    maskmem_backbone = SimpleMaskEncoder(
+        out_dim=64,
+        position_encoding=position_encoding,
+        mask_downsampler=mask_downsampler,
+        fuser=fuser,
+    )
+
+    return maskmem_backbone
+
+
+def _create_tracker_transformer():
+    """Create the SAM3 Tracker transformer components."""
+    # Self attention
+    self_attention = RoPEAttention(
+        embedding_dim=256,
+        num_heads=1,
+        downsample_rate=1,
+        dropout=0.1,
+        rope_theta=10000.0,
+        feat_sizes=[72, 72],
+        use_fa3=False,
+        use_rope_real=False,
+    )
+
+    # Cross attention
+    cross_attention = RoPEAttention(
+        embedding_dim=256,
+        num_heads=1,
+        downsample_rate=1,
+        dropout=0.1,
+        kv_in_dim=64,
+        rope_theta=10000.0,
+        feat_sizes=[72, 72],
+        rope_k_repeat=True,
+        use_fa3=False,
+        use_rope_real=False,
+    )
+
+    # Encoder layer
+    encoder_layer = TransformerDecoderLayerv2(
+        cross_attention_first=False,
+        activation="relu",
+        dim_feedforward=2048,
+        dropout=0.1,
+        pos_enc_at_attn=False,
+        pre_norm=True,
+        self_attention=self_attention,
+        d_model=256,
+        pos_enc_at_cross_attn_keys=True,
+        pos_enc_at_cross_attn_queries=False,
+        cross_attention=cross_attention,
+    )
+
+    # Encoder
+    encoder = TransformerEncoderCrossAttention(
+        remove_cross_attention_layers=[],
+        batch_first=True,
+        d_model=256,
+        frozen=False,
+        pos_enc_at_input=True,
+        layer=encoder_layer,
+        num_layers=4,
+        use_act_checkpoint=False,
+    )
+
+    # Transformer wrapper
+    transformer = TransformerWrapper(
+        encoder=encoder,
+        decoder=None,
+        d_model=256,
+    )
+
+    return transformer
+
+
+def build_tracker(
+    apply_temporal_disambiguation: bool, with_backbone: bool = False, compile_mode=None
+) -> Sam3TrackerPredictor:
+    """
+    Build the SAM3 Tracker module for video tracking.
+
+    Returns:
+        Sam3TrackerPredictor: Wrapped SAM3 Tracker module
+    """
+
+    # Create model components
+    maskmem_backbone = _create_tracker_maskmem_backbone()
+    transformer = _create_tracker_transformer()
+    backbone = None
+    if with_backbone:
+        vision_backbone = _create_vision_backbone(compile_mode=compile_mode)
+        backbone = SAM3VLBackbone(scalp=1, visual=vision_backbone, text=None)
+    # Create the Tracker module
+    model = Sam3TrackerPredictor(
+        image_size=1008,
+        num_maskmem=7,
+        backbone=backbone,
+        backbone_stride=14,
+        transformer=transformer,
+        maskmem_backbone=maskmem_backbone,
+        # SAM parameters
+        multimask_output_in_sam=True,
+        # Evaluation
+        forward_backbone_per_frame_for_eval=True,
+        trim_past_non_cond_mem_for_eval=False,
+        # Multimask
+        multimask_output_for_tracking=True,
+        multimask_min_pt_num=0,
+        multimask_max_pt_num=1,
+        # Additional settings
+        always_start_from_first_ann_frame=False,
+        # Mask overlap
+        non_overlap_masks_for_mem_enc=False,
+        non_overlap_masks_for_output=False,
+        max_cond_frames_in_attn=4,
+        offload_output_to_cpu_for_eval=False,
+        # SAM decoder settings
+        sam_mask_decoder_extra_args={
+            "dynamic_multimask_via_stability": True,
+            "dynamic_multimask_stability_delta": 0.05,
+            "dynamic_multimask_stability_thresh": 0.98,
+        },
+        clear_non_cond_mem_around_input=True,
+        fill_hole_area=0,
+        use_memory_selection=apply_temporal_disambiguation,
+    )
+
+    return model
+
+
+def _create_text_encoder(bpe_path: str) -> VETextEncoder:
+    """Create SAM3 text encoder."""
+    tokenizer = SimpleTokenizer(bpe_path=bpe_path)
+    return VETextEncoder(
+        tokenizer=tokenizer,
+        d_model=256,
+        width=1024,
+        heads=16,
+        layers=24,
+    )
+
+
+def _create_vision_backbone(
+    compile_mode=None, enable_inst_interactivity=True
+) -> Sam3DualViTDetNeck:
+    """Create SAM3 visual backbone with ViT and neck."""
+    # Position encoding
+    position_encoding = _create_position_encoding(precompute_resolution=1008)
+    # ViT backbone
+    vit_backbone: ViT = _create_vit_backbone(compile_mode=compile_mode)
+    vit_neck: Sam3DualViTDetNeck = _create_vit_neck(
+        position_encoding,
+        vit_backbone,
+        enable_inst_interactivity=enable_inst_interactivity,
+    )
+    # Visual neck
+    return vit_neck
+
+
+def _create_sam3_transformer(has_presence_token: bool = True) -> TransformerWrapper:
+    """Create SAM3 transformer encoder and decoder."""
+    encoder: TransformerEncoderFusion = _create_transformer_encoder()
+    decoder: TransformerDecoder = _create_transformer_decoder()
+
+    return TransformerWrapper(encoder=encoder, decoder=decoder, d_model=256)
+
+
+def _load_checkpoint(model, checkpoint_path):
+    """Load model checkpoint from file."""
+    with g_pathmgr.open(checkpoint_path, "rb") as f:
+        ckpt = torch.load(f, map_location="cpu", weights_only=True)
+    if "model" in ckpt and isinstance(ckpt["model"], dict):
+        ckpt = ckpt["model"]
+    sam3_image_ckpt = {
+        k.replace("detector.", ""): v for k, v in ckpt.items() if "detector" in k
+    }
+    if model.inst_interactive_predictor is not None:
+        sam3_image_ckpt.update(
+            {
+                k.replace("tracker.", "inst_interactive_predictor.model."): v
+                for k, v in ckpt.items()
+                if "tracker" in k
+            }
+        )
+    missing_keys, _ = model.load_state_dict(sam3_image_ckpt, strict=False)
+    if len(missing_keys) > 0:
+        print(
+            f"loaded {checkpoint_path} and found "
+            f"missing and/or unexpected keys:\n{missing_keys=}"
+        )
+
+
+def _setup_device_and_mode(model, device, eval_mode):
+    """Setup model device and evaluation mode."""
+    if device == "cuda":
+        model = model.cuda()
+    if eval_mode:
+        model.eval()
+    return model
+
+
+def build_sam3_image_model(
+    bpe_path=None,
+    device="cuda" if torch.cuda.is_available() else "cpu",
+    eval_mode=True,
+    checkpoint_path=None,
+    load_from_HF=True,
+    enable_segmentation=True,
+    enable_inst_interactivity=False,
+    compile=False,
+):
+    """
+    Build SAM3 image model
+
+    Args:
+        bpe_path: Path to the BPE tokenizer vocabulary
+        device: Device to load the model on ('cuda' or 'cpu')
+        eval_mode: Whether to set the model to evaluation mode
+        checkpoint_path: Optional path to model checkpoint
+        enable_segmentation: Whether to enable segmentation head
+        enable_inst_interactivity: Whether to enable instance interactivity (SAM 1 task)
+        compile_mode: To enable compilation, set to "default"
+
+    Returns:
+        A SAM3 image model
+    """
+    if bpe_path is None:
+        bpe_path = os.path.join(
+            os.path.dirname(__file__), "..", "assets", "bpe_simple_vocab_16e6.txt.gz"
+        )
+    # Create visual components
+    compile_mode = "default" if compile else None
+    vision_encoder = _create_vision_backbone(
+        compile_mode=compile_mode, enable_inst_interactivity=enable_inst_interactivity
+    )
+
+    # Create text components
+    text_encoder = _create_text_encoder(bpe_path)
+
+    # Create visual-language backbone
+    backbone = _create_vl_backbone(vision_encoder, text_encoder)
+
+    # Create transformer components
+    transformer = _create_sam3_transformer()
+
+    # Create dot product scoring
+    dot_prod_scoring = _create_dot_product_scoring()
+
+    # Create segmentation head if enabled
+    segmentation_head = (
+        _create_segmentation_head(compile_mode=compile_mode)
+        if enable_segmentation
+        else None
+    )
+
+    # Create geometry encoder
+    input_geometry_encoder = _create_geometry_encoder()
+    if enable_inst_interactivity:
+        sam3_pvs_base = build_tracker(apply_temporal_disambiguation=False)
+        inst_predictor = SAM3InteractiveImagePredictor(sam3_pvs_base)
+    else:
+        inst_predictor = None
+    # Create the SAM3 model
+    model = _create_sam3_model(
+        backbone,
+        transformer,
+        input_geometry_encoder,
+        segmentation_head,
+        dot_prod_scoring,
+        inst_predictor,
+        eval_mode,
+    )
+    if load_from_HF and checkpoint_path is None:
+        checkpoint_path = hf_hub_download(
+            repo_id=SAM3_MODEL_ID, filename=SAM3_CKPT_NAME
+        )
+    # Load checkpoint if provided
+    if checkpoint_path is not None:
+        _load_checkpoint(model, checkpoint_path)
+
+    # Setup device and mode
+    model = _setup_device_and_mode(model, device, eval_mode)
+
+    return model
+
+
+def build_sam3_video_model(
+    checkpoint_path: Optional[str] = None,
+    load_from_HF=True,
+    bpe_path: Optional[str] = None,
+    has_presence_token: bool = True,
+    geo_encoder_use_img_cross_attn: bool = True,
+    strict_state_dict_loading: bool = True,
+    apply_temporal_disambiguation: bool = True,
+    device="cuda" if torch.cuda.is_available() else "cpu",
+    compile=False,
+) -> Sam3VideoInferenceWithInstanceInteractivity:
+    """
+    Build SAM3 dense tracking model.
+
+    Args:
+        checkpoint_path: Optional path to checkpoint file
+        bpe_path: Path to the BPE tokenizer file
+
+    Returns:
+        Sam3VideoInferenceWithInstanceInteractivity: The instantiated dense tracking model
+    """
+    if bpe_path is None:
+        bpe_path = os.path.join(
+            os.path.dirname(__file__), "..", "assets", "bpe_simple_vocab_16e6.txt.gz"
+        )
+
+    # Build Tracker module
+    tracker = build_tracker(apply_temporal_disambiguation=apply_temporal_disambiguation)
+
+    # Build Detector components
+    visual_neck = _create_vision_backbone()
+    text_encoder = _create_text_encoder(bpe_path)
+    backbone = SAM3VLBackbone(scalp=1, visual=visual_neck, text=text_encoder)
+    transformer = _create_sam3_transformer(has_presence_token=has_presence_token)
+    segmentation_head: UniversalSegmentationHead = _create_segmentation_head()
+    input_geometry_encoder = _create_geometry_encoder()
+
+    # Create main dot product scoring
+    main_dot_prod_mlp = MLP(
+        input_dim=256,
+        hidden_dim=2048,
+        output_dim=256,
+        num_layers=2,
+        dropout=0.1,
+        residual=True,
+        out_norm=nn.LayerNorm(256),
+    )
+    main_dot_prod_scoring = DotProductScoring(
+        d_model=256, d_proj=256, prompt_mlp=main_dot_prod_mlp
+    )
+
+    # Build Detector module
+    detector = Sam3ImageOnVideoMultiGPU(
+        num_feature_levels=1,
+        backbone=backbone,
+        transformer=transformer,
+        segmentation_head=segmentation_head,
+        semantic_segmentation_head=None,
+        input_geometry_encoder=input_geometry_encoder,
+        use_early_fusion=True,
+        use_dot_prod_scoring=True,
+        dot_prod_scoring=main_dot_prod_scoring,
+        supervise_joint_box_scores=has_presence_token,
+    )
+
+    # Build the main SAM3 video model
+    if apply_temporal_disambiguation:
+        model = Sam3VideoInferenceWithInstanceInteractivity(
+            detector=detector,
+            tracker=tracker,
+            score_threshold_detection=0.5,
+            assoc_iou_thresh=0.1,
+            det_nms_thresh=0.1,
+            new_det_thresh=0.7,
+            hotstart_delay=15,
+            hotstart_unmatch_thresh=8,
+            hotstart_dup_thresh=8,
+            suppress_unmatched_only_within_hotstart=True,
+            min_trk_keep_alive=-1,
+            max_trk_keep_alive=30,
+            init_trk_keep_alive=30,
+            suppress_overlapping_based_on_recent_occlusion_threshold=0.7,
+            suppress_det_close_to_boundary=False,
+            fill_hole_area=16,
+            recondition_every_nth_frame=16,
+            masklet_confirmation_enable=False,
+            decrease_trk_keep_alive_for_empty_masklets=False,
+            image_size=1008,
+            image_mean=(0.5, 0.5, 0.5),
+            image_std=(0.5, 0.5, 0.5),
+            compile_model=compile,
+        )
+    else:
+        # a version without any heuristics for ablation studies
+        model = Sam3VideoInferenceWithInstanceInteractivity(
+            detector=detector,
+            tracker=tracker,
+            score_threshold_detection=0.5,
+            assoc_iou_thresh=0.1,
+            det_nms_thresh=0.1,
+            new_det_thresh=0.7,
+            hotstart_delay=0,
+            hotstart_unmatch_thresh=0,
+            hotstart_dup_thresh=0,
+            suppress_unmatched_only_within_hotstart=True,
+            min_trk_keep_alive=-1,
+            max_trk_keep_alive=30,
+            init_trk_keep_alive=30,
+            suppress_overlapping_based_on_recent_occlusion_threshold=0.7,
+            suppress_det_close_to_boundary=False,
+            fill_hole_area=16,
+            recondition_every_nth_frame=0,
+            masklet_confirmation_enable=False,
+            decrease_trk_keep_alive_for_empty_masklets=False,
+            image_size=1008,
+            image_mean=(0.5, 0.5, 0.5),
+            image_std=(0.5, 0.5, 0.5),
+            compile_model=compile,
+        )
+
+    # Load checkpoint if provided
+    if load_from_HF and checkpoint_path is None:
+        checkpoint_path = hf_hub_download(
+            repo_id=SAM3_MODEL_ID, filename=SAM3_CKPT_NAME
+        )
+    if checkpoint_path is not None:
+        with g_pathmgr.open(checkpoint_path, "rb") as f:
+            ckpt = torch.load(f, map_location="cpu", weights_only=True)
+        if "model" in ckpt and isinstance(ckpt["model"], dict):
+            ckpt = ckpt["model"]
+
+        missing_keys, unexpected_keys = model.load_state_dict(
+            ckpt, strict=strict_state_dict_loading
+        )
+        if missing_keys:
+            print(f"Missing keys: {missing_keys}")
+        if unexpected_keys:
+            print(f"Unexpected keys: {unexpected_keys}")
+
+    model.to(device=device)
+    return model
+
+
+def build_sam3_video_predictor(*model_args, gpus_to_use=None, **model_kwargs):
+    return Sam3VideoPredictorMultiGPU(
+        *model_args, gpus_to_use=gpus_to_use, **model_kwargs
+    )

sam3/perflib/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import os
+
+is_enabled = False
+if os.getenv("USE_PERFLIB", "1") == "1":
+    # print("Enabled the use of perflib.\n", end="")
+    is_enabled = True

sam3/perflib/associate_det_trk.py

@@ -0,0 +1,137 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+from collections import defaultdict
+
+import torch
+import torch.nn.functional as F
+from sam3.perflib.masks_ops import mask_iou
+from scipy.optimize import linear_sum_assignment
+
+
+def associate_det_trk(
+    det_masks,
+    track_masks,
+    iou_threshold=0.5,
+    iou_threshold_trk=0.5,
+    det_scores=None,
+    new_det_thresh=0.0,
+):
+    """
+    Optimized implementation of detection <-> track association that minimizes DtoH syncs.
+
+    Args:
+        det_masks: (N, H, W) tensor of predicted masks
+        track_masks: (M, H, W) tensor of track masks
+
+    Returns:
+        new_det_indices: list of indices in det_masks considered 'new'
+        unmatched_trk_indices: list of indices in track_masks considered 'unmatched'
+    """
+    with torch.autograd.profiler.record_function("perflib: associate_det_trk"):
+        assert isinstance(det_masks, torch.Tensor), "det_masks should be a tensor"
+        assert isinstance(track_masks, torch.Tensor), "track_masks should be a tensor"
+        if det_masks.size(0) == 0 or track_masks.size(0) == 0:
+            return list(range(det_masks.size(0))), [], {}, {}  # all detections are new
+
+        if list(det_masks.shape[-2:]) != list(track_masks.shape[-2:]):
+            # resize to the smaller size to save GPU memory
+            if torch.numel(det_masks[-2:]) < torch.numel(track_masks[-2:]):
+                track_masks = (
+                    F.interpolate(
+                        track_masks.unsqueeze(1).float(),
+                        size=det_masks.shape[-2:],
+                        mode="bilinear",
+                        align_corners=False,
+                    ).squeeze(1)
+                    > 0
+                )
+            else:
+                # resize detections to track size
+                det_masks = (
+                    F.interpolate(
+                        det_masks.unsqueeze(1).float(),
+                        size=track_masks.shape[-2:],
+                        mode="bilinear",
+                        align_corners=False,
+                    ).squeeze(1)
+                    > 0
+                )
+
+        det_masks = det_masks > 0
+        track_masks = track_masks > 0
+
+        iou = mask_iou(det_masks, track_masks)  # (N, M)
+        igeit = iou >= iou_threshold
+        igeit_any_dim_1 = igeit.any(dim=1)
+        igeit_trk = iou >= iou_threshold_trk
+
+        iou_list = iou.cpu().numpy().tolist()
+        igeit_list = igeit.cpu().numpy().tolist()
+        igeit_any_dim_1_list = igeit_any_dim_1.cpu().numpy().tolist()
+        igeit_trk_list = igeit_trk.cpu().numpy().tolist()
+
+        det_scores_list = (
+            det_scores
+            if det_scores is None
+            else det_scores.cpu().float().numpy().tolist()
+        )
+
+        # Hungarian matching for tracks (one-to-one: each track matches at most one detection)
+        # For detections: allow many tracks to match to the same detection (many-to-one)
+
+        # If either is empty, return all detections as new
+        if det_masks.size(0) == 0 or track_masks.size(0) == 0:
+            return list(range(det_masks.size(0))), [], {}
+
+        # Hungarian matching: maximize IoU for tracks
+        cost_matrix = 1 - iou.cpu().numpy()  # Hungarian solves for minimum cost
+        row_ind, col_ind = linear_sum_assignment(cost_matrix)
+
+        def branchy_hungarian_better_uses_the_cpu(
+            cost_matrix, row_ind, col_ind, iou_list, det_masks, track_masks
+        ):
+            matched_trk = set()
+            matched_det = set()
+            matched_det_scores = {}  # track index -> [det_score, det_score * iou] det score of matched detection mask
+            for d, t in zip(row_ind, col_ind):
+                matched_det_scores[t] = [
+                    det_scores_list[d],
+                    det_scores_list[d] * iou_list[d][t],
+                ]
+                if igeit_trk_list[d][t]:
+                    matched_trk.add(t)
+                    matched_det.add(d)
+
+            # Tracks not matched by Hungarian assignment above threshold are unmatched
+            unmatched_trk_indices = [
+                t for t in range(track_masks.size(0)) if t not in matched_trk
+            ]
+
+            # For detections: allow many tracks to match to the same detection (many-to-one)
+            # So, a detection is 'new' if it does not match any track above threshold
+            assert track_masks.size(0) == igeit.size(
+                1
+            )  # Needed for loop optimizaiton below
+            new_det_indices = []
+            for d in range(det_masks.size(0)):
+                if not igeit_any_dim_1_list[d]:
+                    if det_scores is not None and det_scores[d] >= new_det_thresh:
+                        new_det_indices.append(d)
+
+            # for each detection, which tracks it matched to (above threshold)
+            det_to_matched_trk = defaultdict(list)
+            for d in range(det_masks.size(0)):
+                for t in range(track_masks.size(0)):
+                    if igeit_list[d][t]:
+                        det_to_matched_trk[d].append(t)
+
+            return (
+                new_det_indices,
+                unmatched_trk_indices,
+                det_to_matched_trk,
+                matched_det_scores,
+            )
+
+        return (branchy_hungarian_better_uses_the_cpu)(
+            cost_matrix, row_ind, col_ind, iou_list, det_masks, track_masks
+        )

sam3/perflib/compile.py

@@ -0,0 +1,99 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import torch
+
+
+def recursive_fn_factory(fn):
+    def recursive_fn(b):
+        if isinstance(b, dict):
+            return {k: recursive_fn(b[k]) for k in b}
+        if isinstance(b, list):
+            return [recursive_fn(t) for t in b]
+        if isinstance(b, tuple):
+            return tuple(recursive_fn(t) for t in b)
+        if isinstance(b, torch.Tensor):
+            return fn(b)
+        # Yes, writing out an explicit white list of
+        # trivial types is tedious, but so are bugs that
+        # come from not applying fn, when expected to have
+        # applied it.
+        if b is None:
+            return b
+        trivial_types = [bool, int]
+        for t in trivial_types:
+            if isinstance(b, t):
+                return b
+        raise TypeError(f"Unexpected type {type(b)}")
+
+    return recursive_fn
+
+
+recursive_contiguous = recursive_fn_factory(lambda x: x.contiguous())
+recursive_clone = recursive_fn_factory(torch.clone)
+
+
+def compile_wrapper(
+    fn, *, mode="max-autotune", fullgraph=True, dynamic=False, name=None
+):
+    compiled_fn = torch.compile(fn, mode=mode, fullgraph=fullgraph, dynamic=dynamic)
+
+    def compiled_fn_wrapper(*args, **kwargs):
+        with torch.autograd.profiler.record_function(
+            f"compiled {fn}" if name is None else name
+        ):
+            cont_args = recursive_contiguous(args)
+            cont_kwargs = recursive_contiguous(kwargs)
+            result = compiled_fn(*cont_args, **cont_kwargs)
+            cloned_result = recursive_clone(result)
+            return cloned_result
+
+    return compiled_fn_wrapper
+
+
+def shape_logging_wrapper(fn, keep_kwargs, enable_logging=False):
+    """
+    Wraps a function and prints the shapes of all tensor inputs.
+    Only prints when a new combination of shapes is seen.
+    Thread-safe.
+
+    Args:
+        fn: Function to wrap
+        enable_logging: Boolean flag to enable/disable logging
+    """
+    seen_shapes = set()
+
+    def get_shape(obj):
+        if isinstance(obj, torch.Tensor):
+            return obj.shape
+        elif isinstance(obj, (list, tuple)):
+            if len(obj) > 1:
+                return tuple(get_shape(x) for x in obj)
+            return get_shape(obj[0])
+        elif isinstance(obj, dict):
+            return tuple(sorted((k, get_shape(v)) for k, v in obj.items()))
+        else:
+            return type(obj).__name__
+
+    def wrapper(*args, **kwargs):
+        shapes = tuple(get_shape(arg) for arg in args) + tuple(
+            (k, get_shape(v))
+            for k, v in kwargs.items()
+            if isinstance(v, (torch.Tensor, list))
+            and (len(keep_kwargs) > 0 and k in keep_kwargs)
+        )
+        if shapes not in seen_shapes:
+            seen_shapes.add(shapes)
+            if enable_logging:
+                print(f"[ShapeLogger] New input shapes for {fn.__qualname__}: {shapes}")
+        return fn(*args, **kwargs)
+
+    # Allow toggling the flag at runtime
+    wrapper.enable_logging = enable_logging
+
+    def set_logging(enabled=False):
+        nonlocal enable_logging
+        enable_logging = enabled
+        wrapper.enable_logging = enable_logging
+
+    wrapper.set_logging = set_logging
+    return wrapper

sam3/perflib/connected_components.py

@@ -0,0 +1,84 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+import logging
+
+import torch
+
+try:
+    from cc_torch import get_connected_components
+
+    HAS_CC_TORCH = True
+except ImportError:
+    logging.debug(
+        "cc_torch not found. Consider installing for better performance. Command line:"
+        " pip install git+https://github.com/ronghanghu/cc_torch.git"
+    )
+    HAS_CC_TORCH = False
+
+
+def connected_components_cpu_single(values: torch.Tensor):
+    assert values.dim() == 2
+    from skimage.measure import label
+
+    labels, num = label(values.cpu().numpy(), return_num=True)
+    labels = torch.from_numpy(labels)
+    counts = torch.zeros_like(labels)
+    for i in range(1, num + 1):
+        cur_mask = labels == i
+        cur_count = cur_mask.sum()
+        counts[cur_mask] = cur_count
+    return labels, counts
+
+
+def connected_components_cpu(input_tensor: torch.Tensor):
+    out_shape = input_tensor.shape
+    if input_tensor.dim() == 4 and input_tensor.shape[1] == 1:
+        input_tensor = input_tensor.squeeze(1)
+    else:
+        assert (
+            input_tensor.dim() == 3
+        ), "Input tensor must be (B, H, W) or (B, 1, H, W)."
+
+    batch_size = input_tensor.shape[0]
+    labels_list = []
+    counts_list = []
+    for b in range(batch_size):
+        labels, counts = connected_components_cpu_single(input_tensor[b])
+        labels_list.append(labels)
+        counts_list.append(counts)
+    labels_tensor = torch.stack(labels_list, dim=0).to(input_tensor.device)
+    counts_tensor = torch.stack(counts_list, dim=0).to(input_tensor.device)
+    return labels_tensor.view(out_shape), counts_tensor.view(out_shape)
+
+
+def connected_components(input_tensor: torch.Tensor):
+    """
+    Computes connected components labeling on a batch of 2D tensors, using the best available backend.
+
+    Args:
+        input_tensor (torch.Tensor): A BxHxW integer tensor or Bx1xHxW. Non-zero values are considered foreground. Bool tensor also accepted
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Both tensors have the same shape as input_tensor.
+            - A tensor with dense labels. Background is 0.
+            - A tensor with the size of the connected component for each pixel.
+    """
+    if input_tensor.dim() == 3:
+        input_tensor = input_tensor.unsqueeze(1)
+
+    assert (
+        input_tensor.dim() == 4 and input_tensor.shape[1] == 1
+    ), "Input tensor must be (B, H, W) or (B, 1, H, W)."
+
+    if input_tensor.is_cuda:
+        if HAS_CC_TORCH:
+            return get_connected_components(input_tensor.to(torch.uint8))
+        else:
+            # triton fallback
+            from sam3.perflib.triton.connected_components import (
+                connected_components_triton,
+            )
+
+            return connected_components_triton(input_tensor)
+
+    # CPU fallback
+    return connected_components_cpu(input_tensor)

sam3/perflib/fa3.py

@@ -0,0 +1,27 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import torch
+
+
+@torch.library.custom_op("flash::flash_attn_func", mutates_args=())
+def flash_attn_func_op(
+    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor
+) -> torch.Tensor:
+    from flash_attn_interface import flash_attn_func as fa3
+
+    return fa3(q, k, v)
+
+
+def flash_attn_func(q, k, v):
+    dtype = torch.float8_e4m3fn
+    return flash_attn_func_op(q.to(dtype), k.to(dtype), v.to(dtype)).to(q.dtype)
+
+
+@flash_attn_func_op.register_fake
+def _(q, k, v, **kwargs):
+    # two outputs:
+    # 1. output: (batch, seq_len, num_heads, head_dim)
+    # 2. softmax_lse: (batch, num_heads, seq_len) with dtype=torch.float32
+    # output needs to be bfloat16, not float8!
+    meta_q = torch.empty_like(q, dtype=torch.bfloat16).contiguous()
+    return meta_q

sam3/perflib/masks_ops.py

@@ -0,0 +1,69 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import torch
+
+
+def masks_to_boxes(masks: torch.Tensor, obj_ids: list[int]):
+    with torch.autograd.profiler.record_function("perflib: masks_to_boxes"):
+        # Sanity check based on callsite for replacement
+        assert masks.shape[0] == len(obj_ids)
+        assert masks.dim() == 3
+
+        # Based on torchvision masks_to_boxes
+        if masks.numel() == 0:
+            return torch.zeros((0, 4), device=masks.device, dtype=torch.float)
+
+        N, H, W = masks.shape
+        device = masks.device
+        y = torch.arange(H, device=device).view(1, H)
+        x = torch.arange(W, device=device).view(1, W)
+
+        masks_with_obj = masks != 0  # N, H, W
+        masks_with_obj_x = masks_with_obj.amax(
+            dim=1
+        )  # N, H (which columns have objects)
+        masks_with_obj_y = masks_with_obj.amax(dim=2)  # N, W (which rows have objects)
+        masks_without_obj_x = ~masks_with_obj_x
+        masks_without_obj_y = ~masks_with_obj_y
+
+        bounding_boxes_0 = torch.amin(
+            (masks_without_obj_x * W) + (masks_with_obj_x * x), dim=1
+        )
+        bounding_boxes_1 = torch.amin(
+            (masks_without_obj_y * H) + (masks_with_obj_y * y), dim=1
+        )
+        bounding_boxes_2 = torch.amax(masks_with_obj_x * x, dim=1)
+        bounding_boxes_3 = torch.amax(masks_with_obj_y * y, dim=1)
+
+        bounding_boxes = torch.stack(
+            [bounding_boxes_0, bounding_boxes_1, bounding_boxes_2, bounding_boxes_3],
+            dim=1,
+        ).to(dtype=torch.float)
+        assert bounding_boxes.shape == (N, 4)
+        assert bounding_boxes.device == masks.device
+        assert bounding_boxes.dtype == torch.float
+        return bounding_boxes
+
+
+def mask_iou(pred_masks: torch.Tensor, gt_masks: torch.Tensor) -> torch.Tensor:
+    """
+    Compute the IoU (Intersection over Union) between predicted masks and ground truth masks.
+    Args:
+      - pred_masks: (N, H, W) bool Tensor, containing binary predicted segmentation masks
+      - gt_masks: (M, H, W) bool Tensor, containing binary ground truth segmentation masks
+    Returns:
+      - ious: (N, M) float Tensor, containing IoUs for each pair of predicted and ground truth masks
+    """
+    assert pred_masks.dtype == gt_masks.dtype == torch.bool
+    N, H, W = pred_masks.shape
+    M, _, _ = gt_masks.shape
+
+    # Flatten masks: (N, 1, H*W) and (1, M, H*W)
+    pred_flat = pred_masks.view(N, 1, H * W)
+    gt_flat = gt_masks.view(1, M, H * W)
+
+    # Compute intersection and union: (N, M)
+    intersection = (pred_flat & gt_flat).sum(dim=2).float()
+    union = (pred_flat | gt_flat).sum(dim=2).float()
+    ious = intersection / union.clamp(min=1)
+    return ious  # shape: (N, M)