update app.py

app.py

@@ -5,6 +5,8 @@ from PIL import Image
 import torchvision.transforms as transforms
 from torch import nn
 import numpy as np
+import os
+import time
 
 import gradio as gr
 
@@ -315,6 +317,7 @@ def compute_ncut(
 ):
     from ncut_pytorch import NCUT, rgb_from_tsne_3d
 
+    start = time.time()
     eigvecs, eigvals = NCUT(
         num_eig=num_eig,
         num_sample=num_sample_ncut,
@@ -322,12 +325,17 @@ def compute_ncut(
         affinity_focal_gamma=affinity_focal_gamma,
         knn=knn_ncut,
     ).fit_transform(features.reshape(-1, features.shape[-1]))
+    print(f"NCUT time: {time.time() - start:.2f}s")
+    
+    start = time.time()
     X_3d, rgb = rgb_from_tsne_3d(
         eigvecs,
         num_sample=num_sample_tsne,
         perplexity=perplexity,
         knn=knn_tsne,
     )
+    print(f"t-SNE time: {time.time() - start:.2f}s")
+    
     rgb = rgb.reshape(features.shape[:3] + (3,))
     return rgb
 
@@ -368,9 +376,13 @@ def main_fn(
         perplexity = num_sample_tsne - 1
     
     images = [image[0] for image in images]
+    
+    start = time.time()
     features = extract_features(
         images, model_name=model_name, node_type=node_type, layer=layer
     )
+    print(f"Feature extraction time: {time.time() - start:.2f}s")
+    
     rgb = compute_ncut(
         features,
         num_eig=num_eig,
@@ -391,7 +403,7 @@ demo = gr.Interface(
     main_fn,
     [
         gr.Gallery(value=default_images, label="Select images", show_label=False, elem_id="images", columns=[3], rows=[1], object_fit="contain", height="auto", type="pil"),
-        gr.Dropdown(["SAM(sam_vit_b)", "DiNO(dinov2_vitb14_reg)", "CLIP(openai/clip-vit-base-patch16"], label="Model", value="SAM(sam_vit_b)", elem_id="model_name"),
+        gr.Dropdown(["SAM(sam_vit_b)", "DiNO(dinov2_vitb14_reg)", "CLIP(openai/clip-vit-base-patch16)"], label="Model", value="SAM(sam_vit_b)", elem_id="model_name"),
         gr.Dropdown(["attn", "mlp", "block"], label="Node type", value="block", elem_id="node_type", info="attn: attention output, mlp: mlp output, block: sum of residual stream"),
         gr.Slider(0, 11, step=1, label="Layer", value=11, elem_id="layer", info="which layer of the image backbone features"),
         gr.Slider(1, 1000, step=1, label="Number of eigenvectors", value=100, elem_id="num_eig", info='increase for more object parts, decrease for whole object'),