fix bugs and add sam time box

app.py

@@ -1,10 +1,18 @@
 import os
+import time
 import torch
 import numpy as np
 
 import gradio as gr
 
 from segment_anything import build_sam, SamAutomaticMaskGenerator
+from segment_anything.utils.amg import (
+    batch_iterator,
+    MaskData,
+    calculate_stability_score,
+    batched_mask_to_box,
+    is_box_near_crop_edge,
+)
 
 os.system(r'python -m wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth')
 
@@ -22,14 +30,76 @@ hourglass_args = {
     },
 }
 
+def generate_mask(image, generator: SamAutomaticMaskGenerator):
+    start = time.perf_counter()
+    generator.predictor.set_image(image)
+    eta1 = time.perf_counter() - start
+
+    image_size = image.shape[:2]
+    points_scale = np.array(image_size)[None, ::-1]
+    points_for_image = generator.point_grids[0] * points_scale
+    for (points,) in batch_iterator(generator.points_per_batch, points_for_image):
+        transformed_points = generator.predictor.transform.apply_coords(points, image_size)
+        in_points = torch.as_tensor(transformed_points, device=generator.predictor.device)
+        in_labels = torch.ones(in_points.shape[0], dtype=torch.int, device=in_points.device)
+        start = time.perf_counter()
+        masks, iou_preds, _ = generator.predictor.predict_torch(
+            in_points[:, None, :],
+            in_labels[:, None],
+            multimask_output=True,
+            return_logits=True,
+        )
+        eta2 = time.perf_counter() - start
+
+        # Serialize predictions and store in MaskData
+        data = MaskData(
+            masks=masks.flatten(0, 1),
+            iou_preds=iou_preds.flatten(0, 1),
+            points=torch.as_tensor(points.repeat(masks.shape[1], axis=0)),
+        )
+        del masks
+
+        # Filter by predicted IoU
+        if generator.pred_iou_thresh > 0.0:
+            keep_mask = data["iou_preds"] > generator.pred_iou_thresh
+            data.filter(keep_mask)
+
+        # Calculate stability score
+        data["stability_score"] = calculate_stability_score(
+            data["masks"], generator.predictor.model.mask_threshold, generator.stability_score_offset
+        )
+        if generator.stability_score_thresh > 0.0:
+            keep_mask = data["stability_score"] >= generator.stability_score_thresh
+            data.filter(keep_mask)
+
+        # Threshold masks and calculate boxes
+        data["masks"] = data["masks"] > generator.predictor.model.mask_threshold
+    
+    # Write mask records
+    curr_anns = []
+    for idx in range(len(data["masks"])):
+        ann = {
+            "segmentation": data["masks"][idx].numpy(),
+            "area": data["masks"][idx].sum().item(),
+        }
+        curr_anns.append(ann)
+
+    return curr_anns
+
+
 def predict(image, speed_mode, points_per_side):
     points_per_side = int(points_per_side)
     mask_generator = SamAutomaticMaskGenerator(
-        build_sam(checkpoint="sam_vit_h_4b8939.pth", hourglass_kwargs=hourglass_args[speed_mode]), 
+        build_sam(checkpoint="sam_vit_h_4b8939.pth", **hourglass_args[speed_mode]), 
         points_per_side=points_per_side,
         points_per_batch=64 if points_per_side > 12 else points_per_side * points_per_side
     )
-    masks = mask_generator.generate(image)
+    start = time.perf_counter()
+    with torch.no_grad():
+        # masks = mask_generator.generate(image)
+        masks = generate_mask(image, mask_generator)
+    eta = time.perf_counter() - start
+    eta_text = f"Time of generation: {eta:.2f} seconds"
 
     if len(masks) == 0:
         return image
@@ -41,7 +111,7 @@ def predict(image, speed_mode, points_per_side):
         img = img * (1 - m[..., None]) + color_mask * m[..., None]
 
     image = ((image + img * 255) / 2).astype(np.uint8)
-    return image
+    return image, eta_text
 
 description = """
 #  <center>Expedit-SAM (Expedite Segment Anything Model without any training)</center>
@@ -73,7 +143,9 @@ def main():
                     with gr.Row():
                         run_btn = gr.Button(label="Run", id="run", value="Run")
                         clear_btn = gr.Button(label="Clear", id="clear", value="Clear")
-                output_image = gr.Image(label="Output Image")
+                with gr.Column():
+                    output_image = gr.Image(label="Output Image")
+                    eta_label = gr.Label(label="ETA")
             gr.Examples(
                 examples=[
                     ["./notebooks/images/dog.jpg"],
@@ -88,7 +160,7 @@ def main():
         run_btn.click(
             fn=predict, 
             inputs=[input_image, speed_mode, points_per_side], 
-            outputs=output_image
+            outputs=[output_image, eta_label]
         )
         clear_btn.click(
             fn=lambda: [None, None], 

segment_anything/modeling/hourglass_image_encoder.py

@@ -203,7 +203,7 @@ class TokenReconstructionBlock(UnpoolingBase):
             mink = torch.min(topk, dim=-1).values
             mink = mink.unsqueeze(-1).repeat(1, 1, weight.shape[-1])
             mask = torch.ge(weight, mink)
-            zero = Variable(torch.zeros_like(weight)).cuda()
+            zero = Variable(torch.zeros_like(weight)).to(weight.device)
             attention = torch.where(mask, weight, zero)
         attention = F.normalize(attention, dim=2)
         ret = torch.einsum("bnm, bmc -> bnc", attention, x)
@@ -233,10 +233,10 @@ class HourglassImageEncoderViT(ImageEncoderViT):
         global_attn_indexes: Tuple[int, ...] = (),
         hourglass_clustering_location: int = -1,
         hourglass_num_cluster: int = None,
-        hourglass_cluster_iters: int = 3,
-        hourglass_temperture: float = 0.1,
-        hourglass_cluster_window_size: int = 12,
-        hourglass_reconstruction_k: int = 36,
+        hourglass_cluster_iters: int = 5,
+        hourglass_temperture: float = 0.01,
+        hourglass_cluster_window_size: int = 5,
+        hourglass_reconstruction_k: int = 20,
     ) -> None:
         """
         Args: