Rename app_new.py to app.py

app_new.py → app.py

@@ -1,4 +1,4 @@
-# Code credit: [EdgeSAM Demo](https://huggingface.co/spaces/yunyangx/EfficientSAM).
+# Code credit: [EfficientSAM Demo](https://huggingface.co/spaces/yunyangx/EfficientSAM).
 
 import copy
 import os  # noqa
@@ -15,11 +15,12 @@ from tinysam import sam_model_registry, SamPredictor
 from huggingface_hub import snapshot_download
 
 
-snapshot_download("merve/tinysam", local_dir="tinysam")
+#snapshot_download("merve/tinysam", local_dir="tinysam")
 
 model_type = "vit_t"
-sam = sam_model_registry[model_type](checkpoint="./tinysam/tinysam.pth")
+sam = sam_model_registry[model_type](checkpoint="./tinysam.pth")
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(device)
 sam.to(device=device)
 sam.eval()
 predictor = SamPredictor(sam)
@@ -101,51 +102,23 @@ def segment_with_boxs(
     nd_image = np.array(image)
     img_tensor = ToTensor()(nd_image)
 
-    coord_np = np.array(session_state['coord_list'])
-    label_np = np.array(session_state['label_list'])
+    #coord_np = np.array(session_state['coord_list'])
+    #label_np = np.array(session_state['label_list'])
+    print(scaled_points, scaled_point_label)
+    predictor.set_image(np.array(image))
+    input_box = scaled_points.reshape([4])
+    print('box', input_box)
     masks, scores, logits = predictor.predict(
-        point_coords=coord_np,
-        point_labels=label_np,
+        point_coords=None, #scaled_points,
+        point_labels=None, #scaled_point_label,
+        box=input_box[None, :]
     )
     print(f'scores: {scores}')
     area = masks.sum(axis=(1, 2))
     print(f'area: {area}')
     annotations = np.expand_dims(masks[scores.argmax()], axis=0)
     
-    print(img_tensor.shape)
-    pts_sampled = torch.reshape(torch.tensor(scaled_points), [1, 1, -1, 2])
-    pts_sampled = pts_sampled[:, :, :2, :]
-    pts_labels = torch.reshape(torch.tensor([2, 3]), [1, 1, 2])
-
-    predicted_logits, predicted_iou = model(
-        img_tensor[None, ...].to(device),
-        pts_sampled.to(device),
-        pts_labels.to(device),
-    )
-    predicted_logits = predicted_logits.cpu()
-    all_masks = torch.ge(torch.sigmoid(predicted_logits[0, 0, :, :, :]), 0.5).numpy()
-    predicted_iou = predicted_iou[0, 0, ...].cpu().detach().numpy()
-
-
-    max_predicted_iou = -1
-    selected_mask_using_predicted_iou = None
-    selected_predicted_iou = None
-
-    for m in range(all_masks.shape[0]):
-        curr_predicted_iou = predicted_iou[m]
-        if (
-            curr_predicted_iou > max_predicted_iou
-            or selected_mask_using_predicted_iou is None
-        ):
-            max_predicted_iou = curr_predicted_iou
-            selected_mask_using_predicted_iou = all_masks[m:m+1]
-            selected_predicted_iou = predicted_iou[m:m+1]
-
-    results = format_results(selected_mask_using_predicted_iou, selected_predicted_iou, predicted_logits, 0)
-
-    annotations = results[0]["segmentation"]
-    annotations = np.array([annotations])
-    print(scaled_points.shape)
+    print(annotations)
     fig = fast_process(
         annotations=annotations,
         image=image,
@@ -365,7 +338,7 @@ with gr.Blocks(css=css, title="TinySAM") as demo:
                 gr.Examples(
                     examples=examples,
                     inputs=[cond_img_p],
-                    examples_per_page=4,
+                    examples_per_page=6,
                 )
 
             with gr.Column():
@@ -396,7 +369,7 @@ with gr.Blocks(css=css, title="TinySAM") as demo:
                     examples=examples,
                     inputs=[cond_img_b],
 
-                    examples_per_page=4,
+                    examples_per_page=6,
                 )
 
             with gr.Column():