Add files using upload-large-folder tool

data/audio_embed.tar.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d7cfb23c805a62df237904736403226b60329075a3147d2f93fcc78f2a163843
+size 17678459

data/gt_mask.tar.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c2fd1d8d37ceea45d0b8b7714c06afa0218710061a7da158652617d9107e197a
+size 203344455

data/image_embed.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0b0f5c8ae133bbddbfa558b2052b3aeb757492ffe310650988103d07e24135bb
+size 167486740480

data/media.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6e4d2f70333148331b736f2e65eea4d0d8e38a030fc70b49bee90b63ed5ba7f5
+size 19026513920

models/segment_anything/utils/amg.py

@@ -0,0 +1,346 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# 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 math
+from copy import deepcopy
+from itertools import product
+from typing import Any, Dict, Generator, ItemsView, List, Tuple
+
+import numpy as np
+import torch
+
+
+class MaskData:
+    """
+    A structure for storing masks and their related data in batched format.
+    Implements basic filtering and concatenation.
+    """
+
+    def __init__(self, **kwargs) -> None:
+        for v in kwargs.values():
+            assert isinstance(
+                v, (list, np.ndarray, torch.Tensor)
+            ), "MaskData only supports list, numpy arrays, and torch tensors."
+        self._stats = dict(**kwargs)
+
+    def __setitem__(self, key: str, item: Any) -> None:
+        assert isinstance(
+            item, (list, np.ndarray, torch.Tensor)
+        ), "MaskData only supports list, numpy arrays, and torch tensors."
+        self._stats[key] = item
+
+    def __delitem__(self, key: str) -> None:
+        del self._stats[key]
+
+    def __getitem__(self, key: str) -> Any:
+        return self._stats[key]
+
+    def items(self) -> ItemsView[str, Any]:
+        return self._stats.items()
+
+    def filter(self, keep: torch.Tensor) -> None:
+        for k, v in self._stats.items():
+            if v is None:
+                self._stats[k] = None
+            elif isinstance(v, torch.Tensor):
+                self._stats[k] = v[torch.as_tensor(keep, device=v.device)]
+            elif isinstance(v, np.ndarray):
+                self._stats[k] = v[keep.detach().cpu().numpy()]
+            elif isinstance(v, list) and keep.dtype == torch.bool:
+                self._stats[k] = [a for i, a in enumerate(v) if keep[i]]
+            elif isinstance(v, list):
+                self._stats[k] = [v[i] for i in keep]
+            else:
+                raise TypeError(f"MaskData key {k} has an unsupported type {type(v)}.")
+
+    def cat(self, new_stats: "MaskData") -> None:
+        for k, v in new_stats.items():
+            if k not in self._stats or self._stats[k] is None:
+                self._stats[k] = deepcopy(v)
+            elif isinstance(v, torch.Tensor):
+                self._stats[k] = torch.cat([self._stats[k], v], dim=0)
+            elif isinstance(v, np.ndarray):
+                self._stats[k] = np.concatenate([self._stats[k], v], axis=0)
+            elif isinstance(v, list):
+                self._stats[k] = self._stats[k] + deepcopy(v)
+            else:
+                raise TypeError(f"MaskData key {k} has an unsupported type {type(v)}.")
+
+    def to_numpy(self) -> None:
+        for k, v in self._stats.items():
+            if isinstance(v, torch.Tensor):
+                self._stats[k] = v.detach().cpu().numpy()
+
+
+def is_box_near_crop_edge(
+    boxes: torch.Tensor, crop_box: List[int], orig_box: List[int], atol: float = 20.0
+) -> torch.Tensor:
+    """Filter masks at the edge of a crop, but not at the edge of the original image."""
+    crop_box_torch = torch.as_tensor(crop_box, dtype=torch.float, device=boxes.device)
+    orig_box_torch = torch.as_tensor(orig_box, dtype=torch.float, device=boxes.device)
+    boxes = uncrop_boxes_xyxy(boxes, crop_box).float()
+    near_crop_edge = torch.isclose(boxes, crop_box_torch[None, :], atol=atol, rtol=0)
+    near_image_edge = torch.isclose(boxes, orig_box_torch[None, :], atol=atol, rtol=0)
+    near_crop_edge = torch.logical_and(near_crop_edge, ~near_image_edge)
+    return torch.any(near_crop_edge, dim=1)
+
+
+def box_xyxy_to_xywh(box_xyxy: torch.Tensor) -> torch.Tensor:
+    box_xywh = deepcopy(box_xyxy)
+    box_xywh[2] = box_xywh[2] - box_xywh[0]
+    box_xywh[3] = box_xywh[3] - box_xywh[1]
+    return box_xywh
+
+
+def batch_iterator(batch_size: int, *args) -> Generator[List[Any], None, None]:
+    assert len(args) > 0 and all(
+        len(a) == len(args[0]) for a in args
+    ), "Batched iteration must have inputs of all the same size."
+    n_batches = len(args[0]) // batch_size + int(len(args[0]) % batch_size != 0)
+    for b in range(n_batches):
+        yield [arg[b * batch_size : (b + 1) * batch_size] for arg in args]
+
+
+def mask_to_rle_pytorch(tensor: torch.Tensor) -> List[Dict[str, Any]]:
+    """
+    Encodes masks to an uncompressed RLE, in the format expected by
+    pycoco tools.
+    """
+    # Put in fortran order and flatten h,w
+    b, h, w = tensor.shape
+    tensor = tensor.permute(0, 2, 1).flatten(1)
+
+    # Compute change indices
+    diff = tensor[:, 1:] ^ tensor[:, :-1]
+    change_indices = diff.nonzero()
+
+    # Encode run length
+    out = []
+    for i in range(b):
+        cur_idxs = change_indices[change_indices[:, 0] == i, 1]
+        cur_idxs = torch.cat(
+            [
+                torch.tensor([0], dtype=cur_idxs.dtype, device=cur_idxs.device),
+                cur_idxs + 1,
+                torch.tensor([h * w], dtype=cur_idxs.dtype, device=cur_idxs.device),
+            ]
+        )
+        btw_idxs = cur_idxs[1:] - cur_idxs[:-1]
+        counts = [] if tensor[i, 0] == 0 else [0]
+        counts.extend(btw_idxs.detach().cpu().tolist())
+        out.append({"size": [h, w], "counts": counts})
+    return out
+
+
+def rle_to_mask(rle: Dict[str, Any]) -> np.ndarray:
+    """Compute a binary mask from an uncompressed RLE."""
+    h, w = rle["size"]
+    mask = np.empty(h * w, dtype=bool)
+    idx = 0
+    parity = False
+    for count in rle["counts"]:
+        mask[idx : idx + count] = parity
+        idx += count
+        parity ^= True
+    mask = mask.reshape(w, h)
+    return mask.transpose()  # Put in C order
+
+
+def area_from_rle(rle: Dict[str, Any]) -> int:
+    return sum(rle["counts"][1::2])
+
+
+def calculate_stability_score(
+    masks: torch.Tensor, mask_threshold: float, threshold_offset: float
+) -> torch.Tensor:
+    """
+    Computes the stability score for a batch of masks. The stability
+    score is the IoU between the binary masks obtained by thresholding
+    the predicted mask logits at high and low values.
+    """
+    # One mask is always contained inside the other.
+    # Save memory by preventing unnecessary cast to torch.int64
+    intersections = (
+        (masks > (mask_threshold + threshold_offset))
+        .sum(-1, dtype=torch.int16)
+        .sum(-1, dtype=torch.int32)
+    )
+    unions = (
+        (masks > (mask_threshold - threshold_offset))
+        .sum(-1, dtype=torch.int16)
+        .sum(-1, dtype=torch.int32)
+    )
+    return intersections / unions
+
+
+def build_point_grid(n_per_side: int) -> np.ndarray:
+    """Generates a 2D grid of points evenly spaced in [0,1]x[0,1]."""
+    offset = 1 / (2 * n_per_side)
+    points_one_side = np.linspace(offset, 1 - offset, n_per_side)
+    points_x = np.tile(points_one_side[None, :], (n_per_side, 1))
+    points_y = np.tile(points_one_side[:, None], (1, n_per_side))
+    points = np.stack([points_x, points_y], axis=-1).reshape(-1, 2)
+    return points
+
+
+def build_all_layer_point_grids(
+    n_per_side: int, n_layers: int, scale_per_layer: int
+) -> List[np.ndarray]:
+    """Generates point grids for all crop layers."""
+    points_by_layer = []
+    for i in range(n_layers + 1):
+        n_points = int(n_per_side / (scale_per_layer**i))
+        points_by_layer.append(build_point_grid(n_points))
+    return points_by_layer
+
+
+def generate_crop_boxes(
+    im_size: Tuple[int, ...], n_layers: int, overlap_ratio: float
+) -> Tuple[List[List[int]], List[int]]:
+    """
+    Generates a list of crop boxes of different sizes. Each layer
+    has (2**i)**2 boxes for the ith layer.
+    """
+    crop_boxes, layer_idxs = [], []
+    im_h, im_w = im_size
+    short_side = min(im_h, im_w)
+
+    # Original image
+    crop_boxes.append([0, 0, im_w, im_h])
+    layer_idxs.append(0)
+
+    def crop_len(orig_len, n_crops, overlap):
+        return int(math.ceil((overlap * (n_crops - 1) + orig_len) / n_crops))
+
+    for i_layer in range(n_layers):
+        n_crops_per_side = 2 ** (i_layer + 1)
+        overlap = int(overlap_ratio * short_side * (2 / n_crops_per_side))
+
+        crop_w = crop_len(im_w, n_crops_per_side, overlap)
+        crop_h = crop_len(im_h, n_crops_per_side, overlap)
+
+        crop_box_x0 = [int((crop_w - overlap) * i) for i in range(n_crops_per_side)]
+        crop_box_y0 = [int((crop_h - overlap) * i) for i in range(n_crops_per_side)]
+
+        # Crops in XYWH format
+        for x0, y0 in product(crop_box_x0, crop_box_y0):
+            box = [x0, y0, min(x0 + crop_w, im_w), min(y0 + crop_h, im_h)]
+            crop_boxes.append(box)
+            layer_idxs.append(i_layer + 1)
+
+    return crop_boxes, layer_idxs
+
+
+def uncrop_boxes_xyxy(boxes: torch.Tensor, crop_box: List[int]) -> torch.Tensor:
+    x0, y0, _, _ = crop_box
+    offset = torch.tensor([[x0, y0, x0, y0]], device=boxes.device)
+    # Check if boxes has a channel dimension
+    if len(boxes.shape) == 3:
+        offset = offset.unsqueeze(1)
+    return boxes + offset
+
+
+def uncrop_points(points: torch.Tensor, crop_box: List[int]) -> torch.Tensor:
+    x0, y0, _, _ = crop_box
+    offset = torch.tensor([[x0, y0]], device=points.device)
+    # Check if points has a channel dimension
+    if len(points.shape) == 3:
+        offset = offset.unsqueeze(1)
+    return points + offset
+
+
+def uncrop_masks(
+    masks: torch.Tensor, crop_box: List[int], orig_h: int, orig_w: int
+) -> torch.Tensor:
+    x0, y0, x1, y1 = crop_box
+    if x0 == 0 and y0 == 0 and x1 == orig_w and y1 == orig_h:
+        return masks
+    # Coordinate transform masks
+    pad_x, pad_y = orig_w - (x1 - x0), orig_h - (y1 - y0)
+    pad = (x0, pad_x - x0, y0, pad_y - y0)
+    return torch.nn.functional.pad(masks, pad, value=0)
+
+
+def remove_small_regions(
+    mask: np.ndarray, area_thresh: float, mode: str
+) -> Tuple[np.ndarray, bool]:
+    """
+    Removes small disconnected regions and holes in a mask. Returns the
+    mask and an indicator of if the mask has been modified.
+    """
+    import cv2  # type: ignore
+
+    assert mode in ["holes", "islands"]
+    correct_holes = mode == "holes"
+    working_mask = (correct_holes ^ mask).astype(np.uint8)
+    n_labels, regions, stats, _ = cv2.connectedComponentsWithStats(working_mask, 8)
+    sizes = stats[:, -1][1:]  # Row 0 is background label
+    small_regions = [i + 1 for i, s in enumerate(sizes) if s < area_thresh]
+    if len(small_regions) == 0:
+        return mask, False
+    fill_labels = [0] + small_regions
+    if not correct_holes:
+        fill_labels = [i for i in range(n_labels) if i not in fill_labels]
+        # If every region is below threshold, keep largest
+        if len(fill_labels) == 0:
+            fill_labels = [int(np.argmax(sizes)) + 1]
+    mask = np.isin(regions, fill_labels)
+    return mask, True
+
+
+def coco_encode_rle(uncompressed_rle: Dict[str, Any]) -> Dict[str, Any]:
+    from pycocotools import mask as mask_utils  # type: ignore
+
+    h, w = uncompressed_rle["size"]
+    rle = mask_utils.frPyObjects(uncompressed_rle, h, w)
+    rle["counts"] = rle["counts"].decode("utf-8")  # Necessary to serialize with json
+    return rle
+
+
+def batched_mask_to_box(masks: torch.Tensor) -> torch.Tensor:
+    """
+    Calculates boxes in XYXY format around masks. Return [0,0,0,0] for
+    an empty mask. For input shape C1xC2x...xHxW, the output shape is C1xC2x...x4.
+    """
+    # torch.max below raises an error on empty inputs, just skip in this case
+    if torch.numel(masks) == 0:
+        return torch.zeros(*masks.shape[:-2], 4, device=masks.device)
+
+    # Normalize shape to CxHxW
+    shape = masks.shape
+    h, w = shape[-2:]
+    if len(shape) > 2:
+        masks = masks.flatten(0, -3)
+    else:
+        masks = masks.unsqueeze(0)
+
+    # Get top and bottom edges
+    in_height, _ = torch.max(masks, dim=-1)
+    in_height_coords = in_height * torch.arange(h, device=in_height.device)[None, :]
+    bottom_edges, _ = torch.max(in_height_coords, dim=-1)
+    in_height_coords = in_height_coords + h * (~in_height)
+    top_edges, _ = torch.min(in_height_coords, dim=-1)
+
+    # Get left and right edges
+    in_width, _ = torch.max(masks, dim=-2)
+    in_width_coords = in_width * torch.arange(w, device=in_width.device)[None, :]
+    right_edges, _ = torch.max(in_width_coords, dim=-1)
+    in_width_coords = in_width_coords + w * (~in_width)
+    left_edges, _ = torch.min(in_width_coords, dim=-1)
+
+    # If the mask is empty the right edge will be to the left of the left edge.
+    # Replace these boxes with [0, 0, 0, 0]
+    empty_filter = (right_edges < left_edges) | (bottom_edges < top_edges)
+    out = torch.stack([left_edges, top_edges, right_edges, bottom_edges], dim=-1)
+    out = out * (~empty_filter).unsqueeze(-1)
+
+    # Return to original shape
+    if len(shape) > 2:
+        out = out.reshape(*shape[:-2], 4)
+    else:
+        out = out[0]
+
+    return out

models/segment_anything/utils/onnx.py

@@ -0,0 +1,157 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from ..modeling import Sam
+from .amg import calculate_stability_score
+
+
+class SamOnnxModel(nn.Module):
+    """
+    This model should not be called directly, but is used in ONNX export.
+    It combines the prompt encoder, mask decoder, and mask postprocessing of Sam,
+    with some functions modified to enable model tracing. Also supports extra
+    options controlling what information. See the ONNX export script for details.
+    """
+
+    def __init__(
+        self,
+        model: Sam,
+        return_single_mask: bool,
+        use_stability_score: bool = False,
+        return_extra_metrics: bool = False,
+    ) -> None:
+        super().__init__()
+        self.mask_decoder = model.mask_decoder
+        self.model = model
+        self.img_size = model.image_encoder.img_size
+        self.return_single_mask = return_single_mask
+        self.use_stability_score = use_stability_score
+        self.stability_score_offset = 1.0
+        self.return_extra_metrics = return_extra_metrics
+
+    @staticmethod
+    def resize_longest_image_size(
+        input_image_size: torch.Tensor, longest_side: int
+    ) -> torch.Tensor:
+        input_image_size = input_image_size.to(torch.float32)
+        scale = longest_side / torch.max(input_image_size)
+        transformed_size = scale * input_image_size
+        transformed_size = torch.floor(transformed_size + 0.5).to(torch.int64)
+        return transformed_size
+
+    def _embed_points(
+        self, point_coords: torch.Tensor, point_labels: torch.Tensor
+    ) -> torch.Tensor:
+        point_coords = point_coords + 0.5
+        point_coords = point_coords / self.img_size
+        point_embedding = self.model.prompt_encoder.pe_layer._pe_encoding(point_coords)
+        point_labels = point_labels.unsqueeze(-1).expand_as(point_embedding)
+
+        point_embedding = point_embedding * (point_labels != -1)
+        point_embedding = (
+            point_embedding
+            + self.model.prompt_encoder.not_a_point_embed.weight * (point_labels == -1)
+        )
+
+        for i in range(self.model.prompt_encoder.num_point_embeddings):
+            point_embedding = (
+                point_embedding
+                + self.model.prompt_encoder.point_embeddings[i].weight
+                * (point_labels == i)
+            )
+
+        return point_embedding
+
+    def _embed_masks(
+        self, input_mask: torch.Tensor, has_mask_input: torch.Tensor
+    ) -> torch.Tensor:
+        mask_embedding = has_mask_input * self.model.prompt_encoder.mask_downscaling(
+            input_mask
+        )
+        mask_embedding = mask_embedding + (
+            1 - has_mask_input
+        ) * self.model.prompt_encoder.no_mask_embed.weight.reshape(1, -1, 1, 1)
+        return mask_embedding
+
+    def mask_postprocessing(
+        self, masks: torch.Tensor, orig_im_size: torch.Tensor
+    ) -> torch.Tensor:
+        masks = F.interpolate(
+            masks,
+            size=(self.img_size, self.img_size),
+            mode="bilinear",
+            align_corners=False,
+        )
+
+        prepadded_size = self.resize_longest_image_size(orig_im_size, self.img_size).to(
+            torch.int64
+        )
+        masks = masks[..., : prepadded_size[0], : prepadded_size[1]]  # type: ignore
+
+        orig_im_size = orig_im_size.to(torch.int64)
+        h, w = orig_im_size[0], orig_im_size[1]
+        masks = F.interpolate(masks, size=(h, w), mode="bilinear", align_corners=False)
+        return masks
+
+    def select_masks(
+        self, masks: torch.Tensor, iou_preds: torch.Tensor, num_points: int
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Determine if we should return the multiclick mask or not from the number of points.
+        # The reweighting is used to avoid control flow.
+        score_reweight = torch.tensor(
+            [[1000] + [0] * (self.model.mask_decoder.num_mask_tokens - 1)]
+        ).to(iou_preds.device)
+        score = iou_preds + (num_points - 2.5) * score_reweight
+        best_idx = torch.argmax(score, dim=1)
+        masks = masks[torch.arange(masks.shape[0]), best_idx, :, :].unsqueeze(1)
+        iou_preds = iou_preds[torch.arange(masks.shape[0]), best_idx].unsqueeze(1)
+
+        return masks, iou_preds
+
+    @torch.no_grad()
+    def forward(
+        self,
+        image_embeddings: torch.Tensor,
+        point_coords: torch.Tensor,
+        point_labels: torch.Tensor,
+        mask_input: torch.Tensor,
+        has_mask_input: torch.Tensor,
+        orig_im_size: torch.Tensor,
+    ):
+        sparse_embedding = self._embed_points(point_coords, point_labels)
+        dense_embedding = self._embed_masks(mask_input, has_mask_input)
+
+        masks, scores = self.model.mask_decoder.predict_masks(
+            image_embeddings=image_embeddings,
+            image_pe=self.model.prompt_encoder.get_dense_pe(),
+            sparse_prompt_embeddings=sparse_embedding,
+            dense_prompt_embeddings=dense_embedding,
+        )
+
+        if self.use_stability_score:
+            scores = calculate_stability_score(
+                masks, self.model.mask_threshold, self.stability_score_offset
+            )
+
+        if self.return_single_mask:
+            masks, scores = self.select_masks(masks, scores, point_coords.shape[1])
+
+        upscaled_masks = self.mask_postprocessing(masks, orig_im_size)
+
+        if self.return_extra_metrics:
+            stability_scores = calculate_stability_score(
+                upscaled_masks, self.model.mask_threshold, self.stability_score_offset
+            )
+            areas = (upscaled_masks > self.model.mask_threshold).sum(-1).sum(-1)
+            return upscaled_masks, scores, stability_scores, areas, masks
+
+        return upscaled_masks, scores, masks

models/segment_anything/utils/transforms.py

@@ -0,0 +1,113 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from copy import deepcopy
+from typing import Tuple
+
+import numpy as np
+import torch
+from torch.nn import functional as F
+from torchvision.transforms.functional import resize  # type: ignore
+from torchvision.transforms.functional import to_pil_image
+
+
+class ResizeLongestSide:
+    """
+    Resizes images to the longest side 'target_length', as well as provides
+    methods for resizing coordinates and boxes. Provides methods for
+    transforming both numpy array and batched torch tensors.
+    """
+
+    def __init__(self, target_length: int) -> None:
+        self.target_length = target_length
+
+    def apply_image(self, image: np.ndarray) -> np.ndarray:
+        """
+        Expects a numpy array with shape HxWxC in uint8 format.
+        """
+        target_size = self.get_preprocess_shape(
+            image.shape[0], image.shape[1], self.target_length
+        )
+        return np.array(resize(to_pil_image(image), target_size))
+
+    def apply_coords(
+        self, coords: np.ndarray, original_size: Tuple[int, ...]
+    ) -> np.ndarray:
+        """
+        Expects a numpy array of length 2 in the final dimension. Requires the
+        original image size in (H, W) format.
+        """
+        old_h, old_w = original_size
+        new_h, new_w = self.get_preprocess_shape(
+            original_size[0], original_size[1], self.target_length
+        )
+        coords = deepcopy(coords).astype(float)
+        coords[..., 0] = coords[..., 0] * (new_w / old_w)
+        coords[..., 1] = coords[..., 1] * (new_h / old_h)
+        return coords
+
+    def apply_boxes(
+        self, boxes: np.ndarray, original_size: Tuple[int, ...]
+    ) -> np.ndarray:
+        """
+        Expects a numpy array shape Bx4. Requires the original image size
+        in (H, W) format.
+        """
+        boxes = self.apply_coords(boxes.reshape(-1, 2, 2), original_size)
+        return boxes.reshape(-1, 4)
+
+    def apply_image_torch(self, image: torch.Tensor) -> torch.Tensor:
+        """
+        Expects batched images with shape BxCxHxW and float format. This
+        transformation may not exactly match apply_image. apply_image is
+        the transformation expected by the model.
+        """
+        # Expects an image in BCHW format. May not exactly match apply_image.
+        target_size = self.get_preprocess_shape(
+            image.shape[0], image.shape[1], self.target_length
+        )
+        return F.interpolate(
+            image, target_size, mode="bilinear", align_corners=False, antialias=True
+        )
+
+    def apply_coords_torch(
+        self, coords: torch.Tensor, original_size: Tuple[int, ...]
+    ) -> torch.Tensor:
+        """
+        Expects a torch tensor with length 2 in the last dimension. Requires the
+        original image size in (H, W) format.
+        """
+        old_h, old_w = original_size
+        new_h, new_w = self.get_preprocess_shape(
+            original_size[0], original_size[1], self.target_length
+        )
+        coords = deepcopy(coords).to(torch.float)
+        coords[..., 0] = coords[..., 0] * (new_w / old_w)
+        coords[..., 1] = coords[..., 1] * (new_h / old_h)
+        return coords
+
+    def apply_boxes_torch(
+        self, boxes: torch.Tensor, original_size: Tuple[int, ...]
+    ) -> torch.Tensor:
+        """
+        Expects a torch tensor with shape Bx4. Requires the original image
+        size in (H, W) format.
+        """
+        boxes = self.apply_coords_torch(boxes.reshape(-1, 2, 2), original_size)
+        return boxes.reshape(-1, 4)
+
+    @staticmethod
+    def get_preprocess_shape(
+        oldh: int, oldw: int, long_side_length: int
+    ) -> Tuple[int, int]:
+        """
+        Compute the output size given input size and target long side length.
+        """
+        scale = long_side_length * 1.0 / max(oldh, oldw)
+        newh, neww = oldh * scale, oldw * scale
+        neww = int(neww + 0.5)
+        newh = int(newh + 0.5)
+        return (newh, neww)

models/tf/modeling_outputs.py

@@ -0,0 +1,43 @@
+import torch
+import warnings
+from dataclasses import dataclass
+from typing import Optional, Tuple, Dict, List
+from transformers.utils import ModelOutput
+
+@dataclass
+class CausalLMOutputWithPastAndLabel(ModelOutput):
+    """
+    Base class for causal language model (or autoregressive) outputs.
+
+    Args:
+        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+            Language modeling loss (for next-token prediction).
+        labels (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*, returned when `labels` is provided):
+        logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+        past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+            Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
+            `(batch_size, num_heads, sequence_length, embed_size_per_head)`)
+
+            Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+            `past_key_values` input) to speed up sequential decoding.
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    labels: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+    bs2imgs_token_list: List[List[int]] = None

setup_simtoken.md

@@ -0,0 +1,152 @@
+# SimToken Setup
+
+---
+
+## 1. Create Environment
+
+```bash
+conda create -n simtoken python=3.10 -y
+conda activate simtoken
+
+python -m pip install --upgrade pip wheel "setuptools<81"
+
+pip install \
+  torch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 \
+  --index-url https://download.pytorch.org/whl/cu121
+
+pip install \
+  transformers==4.30.2 \
+  peft==0.2.0 \
+  accelerate==0.21.0 \
+  sentencepiece \
+  protobuf \
+  safetensors \
+  numpy==1.26.4 \
+  pandas \
+  matplotlib \
+  opencv-python \
+  pillow \
+  tqdm \
+  einops \
+  timm \
+  requests \
+  towhee \
+  huggingface_hub
+```
+
+---
+
+## 2. Download from HuggingFace（新机器初始化）
+
+登录 HuggingFace（token 在 https://huggingface.co/settings/tokens 生成）：
+
+```bash
+huggingface-cli login
+```
+
+下载完整 repo（代码 + 权重 + 压缩数据包，共约 190G）：
+
+```bash
+mkdir -p /workspace/SimToken
+cd /workspace/SimToken
+
+huggingface-cli download yfan07/SimToken \
+  --repo-type model \
+  --local-dir . \
+  --local-dir-use-symlinks False
+```
+
+下载完成后解压数据包：
+
+```bash
+cd /workspace/SimToken/data
+
+tar -xf image_embed.tar      # ~5–10 分钟
+tar -xzf gt_mask.tar.gz
+tar -xzf audio_embed.tar.gz
+tar -xf media.tar
+```
+
+
+---
+
+## 3. Pre-download Model Weights（首次使用必做）
+
+`transformers==4.30.2` 与新版 `huggingface_hub` 存在 API 不兼容（`use_auth_token` 已移除）。
+解决方案：先用 CLI 将模型下载到本地缓存，之后运行实验时加 `TRANSFORMERS_OFFLINE=1`，跳过所有网络请求。
+
+```bash
+# Chat-UniVi-7B（~14G）
+huggingface-cli download Chat-UniVi/Chat-UniVi-7B-v1.5
+
+# CLIP ViT-L（~1.6G）
+huggingface-cli download openai/clip-vit-large-patch14
+```
+
+下载完成后即永久缓存，新 session 无需重复下载。
+
+---
+
+## 4. Example Evaluation
+
+所有评测命令统一加 `TRANSFORMERS_OFFLINE=1`：
+
+```bash
+cd /workspace/SimToken
+
+# Unseen split（全量 1656 样本）
+TRANSFORMERS_OFFLINE=1 python -W ignore load_model.py --eval_split test_u
+
+# Seen split
+TRANSFORMERS_OFFLINE=1 python -W ignore load_model.py --eval_split test_s
+
+# Null split（S metric，越低越好）
+TRANSFORMERS_OFFLINE=1 python -W ignore load_model.py --eval_split test_n
+
+# 限制样本数（快速验证）
+TRANSFORMERS_OFFLINE=1 python -W ignore load_model.py --eval_split test_u --max_eval_rows 50
+
+# Stage 0 梯度连通性 + bypass 等价性检查（仅诊断）
+TRANSFORMERS_OFFLINE=1 python -W ignore load_model.py --eval_split test_u --max_eval_rows 0
+```
+
+每次评估依次输出：Baseline + q-LTPO Stage 1 两组结果及诊断统计。
+
+---
+
+## 5. Upload to HuggingFace（实验结束后）
+
+数据目录以压缩包形式存储，可大幅减少文件数量，避免 HuggingFace commit 频率限制。
+
+**第一步：将数据目录压缩为归档文件（如尚未压缩）**
+
+```bash
+cd /workspace/SimToken/data
+
+tar -cf image_embed.tar image_embed/     # 不压缩（.pt 已是二进制）
+tar -czf gt_mask.tar.gz gt_mask/
+tar -czf audio_embed.tar.gz audio_embed/
+tar -cf media.tar media/
+
+# 确认压缩包存在后删除原始目录
+ls -lh *.tar*
+rm -rf image_embed/ gt_mask/ audio_embed/ media/
+```
+
+**第二步：清理缓存并上传**
+
+```bash
+find /workspace/SimToken -name "__pycache__" -exec rm -rf {} + 2>/dev/null
+find /workspace/SimToken -name "*.pyc" -delete
+
+huggingface-cli login   # token 在 https://huggingface.co/settings/tokens 生成（需 Write 权限）
+
+cd /workspace/SimToken
+python upload_hf.py --repo yfan07/SimToken
+```
+
+**注意事项：**
+- 建议在 `tmux` 里运行，防止 SSH 断开：`tmux new -s upload`，完成后 `Ctrl+B D` detach
+- 支持断点续传：中断后重新执行同一命令会自动跳过已上传文件
+- 遇到 rate limit（HTTP 429）时脚本会自动等待约 1 小时后重试
+- 监控进度：`tail -f /workspace/SimToken/upload.log`

utils/__init__.py

@@ -0,0 +1,2 @@
+from .metric import pyutils
+from .metric import utility

utils/metric/pyutils.py

@@ -0,0 +1,160 @@
+
+import numpy as np
+import time
+import sys
+
+class Logger(object):
+    def __init__(self, outfile):
+        self.terminal = sys.stdout
+        self.log = open(outfile, "w")
+        sys.stdout = self
+
+    def write(self, message):
+        self.terminal.write(message)
+        self.log.write(message)
+
+    def flush(self):
+        self.terminal.flush()
+
+
+class AverageMeter:
+    def __init__(self, *keys):
+        self.__data = dict()
+        for k in keys:
+            self.__data[k] = [0.0, 0]
+
+    def add(self, dict):
+        for k, v in dict.items():
+            self.__data[k][0] += v
+            self.__data[k][1] += 1
+
+    def get(self, *keys):
+        if len(keys) == 1:
+            return self.__data[keys[0]][0] / self.__data[keys[0]][1]
+        else:
+            v_list = [self.__data[k][0] / self.__data[k][1] for k in keys]
+            return tuple(v_list)
+
+    def pop(self, key=None):
+        if key is None:
+            for k in self.__data.keys():
+                self.__data[k] = [0.0, 0]
+        else:
+            v = self.get(key)
+            self.__data[key] = [0.0, 0]
+            return v
+
+
+class Timer:
+    def __init__(self, starting_msg = None):
+        self.start = time.time()
+        self.stage_start = self.start
+
+        if starting_msg is not None:
+            print(starting_msg, time.ctime(time.time()))
+
+
+    def update_progress(self, progress):
+        self.elapsed = time.time() - self.start
+        self.est_total = self.elapsed / progress
+        self.est_remaining = self.est_total - self.elapsed
+        self.est_finish = int(self.start + self.est_total)
+
+
+    def str_est_finish(self):
+        return str(time.ctime(self.est_finish))
+
+    def get_stage_elapsed(self):
+        return time.time() - self.stage_start
+
+    def reset_stage(self):
+        self.stage_start = time.time()
+
+
+from multiprocessing.pool import ThreadPool
+
+class BatchThreader:
+
+    def __init__(self, func, args_list, batch_size, prefetch_size=4, processes=12):
+        self.batch_size = batch_size
+        self.prefetch_size = prefetch_size
+
+        self.pool = ThreadPool(processes=processes)
+        self.async_result = []
+
+        self.func = func
+        self.left_args_list = args_list
+        self.n_tasks = len(args_list)
+
+        # initial work
+        self.__start_works(self.__get_n_pending_works())
+
+
+    def __start_works(self, times):
+        for _ in range(times):
+            args = self.left_args_list.pop(0)
+            self.async_result.append(
+                self.pool.apply_async(self.func, args))
+
+
+    def __get_n_pending_works(self):
+        return min((self.prefetch_size + 1) * self.batch_size - len(self.async_result)
+                   , len(self.left_args_list))
+
+
+
+    def pop_results(self):
+
+        n_inwork = len(self.async_result)
+
+        n_fetch = min(n_inwork, self.batch_size)
+        rtn = [self.async_result.pop(0).get()
+                for _ in range(n_fetch)]
+
+        to_fill = self.__get_n_pending_works()
+        if to_fill == 0:
+            self.pool.close()
+        else:
+            self.__start_works(to_fill)
+
+        return rtn
+
+
+
+
+def get_indices_of_pairs(radius, size):
+
+    search_dist = []
+
+    for x in range(1, radius):
+        search_dist.append((0, x))
+
+    for y in range(1, radius):
+        for x in range(-radius + 1, radius):
+            if x * x + y * y < radius * radius:
+                search_dist.append((y, x))
+
+    radius_floor = radius - 1
+
+    full_indices = np.reshape(np.arange(0, size[0]*size[1], dtype=np.int64),
+                                   (size[0], size[1]))
+
+    cropped_height = size[0] - radius_floor
+    cropped_width = size[1] - 2 * radius_floor
+
+    indices_from = np.reshape(full_indices[:-radius_floor, radius_floor:-radius_floor],
+                              [-1])
+
+    indices_to_list = []
+
+    for dy, dx in search_dist:
+        indices_to = full_indices[dy:dy + cropped_height,
+                     radius_floor + dx:radius_floor + dx + cropped_width]
+        indices_to = np.reshape(indices_to, [-1])
+
+        indices_to_list.append(indices_to)
+
+    concat_indices_to = np.concatenate(indices_to_list, axis=0)
+
+    return indices_from, concat_indices_to
+

utils/metric/utility.py

@@ -0,0 +1,125 @@
+import matplotlib.pyplot as plt
+import torch
+from torch.nn import functional as F
+
+import os
+import shutil
+# import logging
+import cv2
+import numpy as np
+from PIL import Image
+
+import sys
+import time
+import pandas as pd
+import pdb
+from torchvision import transforms
+
+def metric_s_for_null(pred):
+    num_seg, T, H, W = pred.shape
+    pred = pred.view(num_seg*T, H, W)
+    assert len(pred.shape) == 3
+
+    N = pred.size(0)
+    num_pixels = pred.view(-1).shape[0] 
+
+    temp_pred = torch.sigmoid(pred)
+    pred = (temp_pred > 0.5).int()
+
+    x = torch.sum(pred.view(-1)) 
+    s = torch.sqrt(x / num_pixels)
+
+    return s
+
+def mask_iou(pred, target, eps=1e-7, size_average=True):
+    r"""
+        param:
+            pred: size [N x H x W]
+            target: size [N x H x W]
+        output:
+            iou: size [1] (size_average=True) or [N] (size_average=False)
+    """
+    # return mask_iou_224(pred, target, eps=1e-7)
+    num_ref, T, H, W = pred.shape
+    pred = pred.view(num_ref*T, H, W)
+    target = target.view(num_ref*T, H, W)
+    assert len(pred.shape) == 3 and pred.shape == target.shape
+
+    N = pred.size(0)
+    # 像素数
+    num_pixels = pred.size(-1) * pred.size(-2)
+    # gt是否是纯黑
+    no_obj_flag = (target.sum(2).sum(1) == 0)
+
+    # 会把pred进行sigmoid变为01之间的概率
+    temp_pred = torch.sigmoid(pred)
+    # 通过阈值变成01矩阵
+    pred = (temp_pred > 0.4).int()
+    # 交集
+    inter = (pred * target).sum(2).sum(1)
+    # 并集
+    union = torch.max(pred, target).sum(2).sum(1)
+
+    inter_no_obj = ((1 - target) * (1 - pred)).sum(2).sum(1)
+    inter[no_obj_flag] = inter_no_obj[no_obj_flag]
+    union[no_obj_flag] = num_pixels
+
+    iou = torch.sum(inter / (union + eps)) / N
+
+    return iou.item()
+
+
+def _eval_pr(y_pred, y, num, device='cuda'):
+    if device.startswith('cuda'):
+        prec, recall = torch.zeros(num).to(y_pred.device), torch.zeros(num).to(y_pred.device)
+        # 0到1 生成num个阈值
+        thlist = torch.linspace(0, 1 - 1e-10, num).to(y_pred.device)
+    else:
+        prec, recall = torch.zeros(num), torch.zeros(num)
+        thlist = torch.linspace(0, 1 - 1e-10, num)
+
+    for i in range(num):
+        y_temp = (y_pred >= thlist[i]).float()
+
+        # 计算 True Positives（TP）
+        tp = (y_temp * y).sum()
+        # 一个是交集除以 预测面积 一个是除以真实面积
+        prec[i], recall[i] = tp / (y_temp.sum() + 1e-20), tp / (y.sum() + 1e-20)
+
+    return prec, recall
+
+
+def Eval_Fmeasure(pred, gt, measure_path, pr_num=255, device='cuda'):
+    r"""
+        param:
+            pred: size [N x H x W]
+            gt: size [N x H x W]
+        output:
+            iou: size [1] (size_average=True) or [N] (size_average=False)
+    """
+    num_ref, T, H, W = pred.shape
+    pred = pred.view(num_ref*T, H, W)
+    gt = gt.view(num_ref*T, H, W)
+    assert len(pred.shape) == 3
+
+
+    # sigmoid转为01之间的
+    pred = torch.sigmoid(pred) 
+    N = pred.size(0)
+    beta2 = 0.3
+    avg_f, img_num = 0.0, 0
+    score = torch.zeros(pr_num)
+
+
+    for img_id in range(N):
+        if torch.mean(gt[img_id]) == 0.0:
+            continue
+        prec, recall = _eval_pr(pred[img_id], gt[img_id], pr_num, device=device)
+        f_score = (1 + beta2) * prec * recall / (beta2 * prec + recall)
+        f_score[f_score != f_score] = 0  # for Nan
+        avg_f += f_score
+        img_num += 1
+        score = avg_f / img_num
+
+    return score.max().item()
+