Upload app.py

app.py

@@ -52,18 +52,6 @@ def resize(img, target_res=224, resize=True, to_pil=True, edge=False, sampling_f
         canvas = Image.fromarray(canvas)
     return canvas
 
-# ─── Configuration ───────────────────────────────────────────────
-num_patches = 30
-target_res = num_patches * 14
-ckpt_file = "ckpts/dino_spair_0300.pth"
-
-# ─── Model setup ─────────────────────────────────────────────────
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-aggre_net = AggregationNetwork(feature_dims=[768], projection_dim=768, device=device)
-aggre_net.load_pretrained_weights(torch.load(ckpt_file, map_location=device))
-aggre_net_dummy  = DummyAggregationNetwork()
-extractor_vit = ViTExtractor('dinov2_vitb14', stride=14, device=device)
-
 # ─── Feature extraction ──────────────────────────────────────────
 def get_processed_features_dino(num_patches, img,use_dummy):
     batch = extractor_vit.preprocess_pil(img)
@@ -85,7 +73,7 @@ def get_sim(
     coord: tuple[int,int],
     feat1: torch.Tensor,
     feat2: torch.Tensor,
-    img_size: int = target_res
+    img_size: int = 420
 ) -> np.ndarray:
     """
     Upsamples the DINO features to `img_size`, then computes cosine‐similarity
@@ -182,6 +170,18 @@ def reload_img(
     
 
 
+# ─── Configuration ───────────────────────────────────────────────
+num_patches = 30
+target_res = num_patches * 14
+ckpt_file = "ckpts/dino_spair_0300.pth"
+
+# ─── Model setup ─────────────────────────────────────────────────
+device = 'cpu' #'cuda' if torch.cuda.is_available() else 'cpu'
+aggre_net = AggregationNetwork(feature_dims=[768], projection_dim=768, device=device)
+aggre_net.load_pretrained_weights(torch.load(ckpt_file, map_location=device))
+aggre_net_dummy  = DummyAggregationNetwork()
+extractor_vit = ViTExtractor('dinov2_vitb14', stride=14, device=device)
+
 # ─── Build Gradio UI ──────────────────────────────────────────────
 with gr.Blocks() as demo:
     # Hidden states to hold features
@@ -194,7 +194,7 @@ with gr.Blocks() as demo:
     intro_text = gr.Markdown("""
     ## Do It Yourself: Learning Semantic Correspondence from Pseudo-Labels
     [Project Page](https://example.com) | [GitHub Repository](https://github.com/example/repo)
-                             
+                            
     Welcome to the DIY-SC demo!
     Upload two images and select a keypoint in the source image. This demo will compute and visualize the feature similarity map and a corresponding point in the target image.
     You can choose between the DIY-SC (DINOv2) or the DINOv2 feature extractor.
@@ -240,4 +240,5 @@ with gr.Blocks() as demo:
         outputs=[src,tgt]
     )
 
-demo.launch(share=True)
+if __name__ == "__main__":
+    demo.launch(share=True)