update layout & use original method

app.py

@@ -30,59 +30,6 @@ hourglass_args = {
     },
 }
 
-def generate_mask(image, generator: SamAutomaticMaskGenerator):
-    generator.predictor.set_image(image)
-
-    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)
-        masks, iou_preds, _ = generator.predictor.predict_torch(
-            in_points[:, None, :],
-            in_labels[:, None],
-            multimask_output=True,
-            return_logits=True,
-        )
-
-        # 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(
@@ -92,8 +39,7 @@ def predict(image, speed_mode, points_per_side):
     )
     start = time.perf_counter()
     with torch.no_grad():
-        # masks = mask_generator.generate(image)
-        masks = generate_mask(image, mask_generator)
+        masks = mask_generator.generate(image)
     eta = time.perf_counter() - start
     eta_text = f"Time of generation: {eta:.2f} seconds"
 
@@ -136,7 +82,7 @@ def main():
                         label="Speed Mode",
                         multiselect=False,
                     )
-                    with gr.Row():
+                    with gr.Column():
                         run_btn = gr.Button(label="Run", id="run", value="Run")
                         clear_btn = gr.Button(label="Clear", id="clear", value="Clear")
                 with gr.Column():