Sam2 with test log

sam2/automatic_mask_generator.py

@@ -6,7 +6,7 @@
 
 # Adapted from https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/automatic_mask_generator.py
 from typing import Any, Dict, List, Optional, Tuple
-
+from sam2.utils.Track import get_gpu_mem
 import numpy as np
 import torch
 from torchvision.ops.boxes import batched_nms, box_area  # type: ignore
@@ -222,6 +222,7 @@ def generate(self, image: np.ndarray) -> List[Dict[str, Any]]:
         return curr_anns
 
     def _generate_masks(self, image: np.ndarray) -> MaskData:
+        logs = []
         orig_size = image.shape[:2]
         crop_boxes, layer_idxs = generate_crop_boxes(
             orig_size, self.crop_n_layers, self.crop_overlap_ratio
@@ -230,9 +231,12 @@ def _generate_masks(self, image: np.ndarray) -> MaskData:
         # Iterate over image crops
         data = MaskData()
         for crop_box, layer_idx in zip(crop_boxes, layer_idxs):
+            before = get_gpu_mem()
             crop_data = self._process_crop(image, crop_box, layer_idx, orig_size)
+            after = get_gpu_mem()
             data.cat(crop_data)
-
+            logs.append((crop_box, layer_idx, before, after))
+        print(logs)
         # Remove duplicate masks between crops
         if len(crop_boxes) > 1:
             # Prefer masks from smaller crops
@@ -259,8 +263,18 @@ def _process_crop(
         x0, y0, x1, y1 = crop_box
         cropped_im = image[y0:y1, x0:x1, :]
         cropped_im_size = cropped_im.shape[:2]
+        mem_before_embed = get_gpu_mem()
+        print(
+            f"[layer {crop_layer_idx} | crop {crop_box}] "
+            f"before embed: {mem_before_embed:.1f} MB"
+        )
         self.predictor.set_image(cropped_im)
-
+        mem_after_embed = get_gpu_mem()
+        print(
+            f"[layer {crop_layer_idx} | crop {crop_box}] "
+            f"after embed:  {mem_after_embed:.1f} MB "
+            f"(+{mem_after_embed - mem_before_embed:.1f} MB)"
+        )
         # Get points for this crop
         points_scale = np.array(cropped_im_size)[None, ::-1]
         points_for_image = self.point_grids[crop_layer_idx] * points_scale
@@ -311,13 +325,19 @@ def _process_batch(
         in_labels = torch.ones(
             in_points.shape[0], dtype=torch.int, device=in_points.device
         )
+        mem_before_fwd = get_gpu_mem()
         masks, iou_preds, low_res_masks = self.predictor._predict(
             in_points[:, None, :],
             in_labels[:, None],
             multimask_output=self.multimask_output,
             return_logits=True,
         )
-
+        mem_after_fwd = get_gpu_mem()
+        print(
+            f"[batch forward] before={mem_before_fwd:.1f} MB | "
+            f"after={mem_after_fwd:.1f} MB | "
+            f"diff={mem_after_fwd - mem_before_fwd:.1f} MB"
+        )
         # Serialize predictions and store in MaskData
         data = MaskData(
             masks=masks.flatten(0, 1),

sam2/utils/Track.py

@@ -0,0 +1,19 @@
+import torch
+#Track VRAM
+def get_gpu_mem(device=None):
+    """
+    Returns current GPU memory usage (allocated) in MB.
+    Accurate and efficient.
+    """
+    if not torch.cuda.is_available():
+        return 0.0
+
+    if device is None:
+        device = torch.cuda.current_device()
+
+    torch.cuda.synchronize(device)  # ensures accurate timing
+
+    stats = torch.cuda.memory_stats(device)
+    allocated = stats["allocated_bytes.all.current"] / 1024**2
+
+    return allocated

sam2/utils/amg.py

@@ -203,39 +203,36 @@ 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.
+    Generates crop boxes only for the finest (smallest) layer.
+    The finest layer is the last one: i_layer = n_layers - 1,
+    which has (2**i_layer)**2 boxes.
     """
     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))
+    # Only finest layer (last one)
+    i_layer = n_layers - 1
+    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_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)]
+    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)
+    # Crops in XYWH format (only finest layer)
+    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)
 
     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)