Update src/streamlit_app.py

src/streamlit_app.py

@@ -18,12 +18,9 @@ from PIL import Image
 # --- RESEARCH CONFIG ---
 IMAGENET_MEAN = (0.485, 0.456, 0.406)
 IMAGENET_STD  = (0.229, 0.224, 0.225)
-current_dir = os.path.dirname(os.path.abspath(__file__))
-if current_dir not in sys.path:
-    sys.path.append(current_dir)
 
 # ================= PAGE CONFIG & STATE INIT =================
-st.set_page_config(layout="wide", page_title="Prior2DSM Pro")
+st.set_page_config(layout="wide", page_title="Prior2DSM Pro", page_icon="🛰️")
 
 if 'bbox_points' not in st.session_state:
     st.session_state.bbox_points = []
@@ -46,71 +43,26 @@ def load_backbones(repo_id, dav_file, dino_file, m_size):
     from dinov3.models.vision_transformer import DinoVisionTransformer
     dino_model = DinoVisionTransformer(img_size=1024, patch_size=16, embed_dim=1024, depth=24, num_heads=16, qkv_bias=True).to(device).eval()
     ckpt = torch.load(dino_path, map_location="cpu")
-    for key in ["state_dict", "model", "teacher"]:
-        if key in ckpt: ckpt = ckpt[key]
+    if "state_dict" in ckpt: ckpt = ckpt["state_dict"]
+    elif "model" in ckpt: ckpt = ckpt["model"]
+    elif "teacher" in ckpt: ckpt = ckpt["teacher"]
     clean_ckpt = {k.replace("module.", "").replace("backbone.", "").replace("teacher.backbone.", ""): v for k, v in ckpt.items()}
     dino_model.load_state_dict(clean_ckpt, strict=False)
-    return dav_model, dino_model, device
-
-@torch.no_grad()
-def extract_dino_features_safe(encoder, x):
-    """Ultimate robustness check for CLS/Patch token concatenation."""
-    B, _, H, W = x.shape
-    D = int(getattr(encoder, "embed_dim")) # 1024
-    
-    # Get patch embeddings
-    pe = encoder.patch_embed(x)
-    
-    # Robust shape handling
-    if pe.ndim == 4: # (B, C_out, Hp, Wp)
-        Hp, Wp = pe.shape[-2], pe.shape[-1]
-        tokens = pe.flatten(2).transpose(1, 2)
-    else: # (B, N, C_out)
-        Hp, Wp = H // 16, W // 16
-        tokens = pe
-
-    # CRITICAL FIX: If patch_embed didn't project to 1024 yet (e.g., returned 64)
-    if tokens.shape[-1] != D:
-        # Check for a projection layer often found in custom ViTs
-        if hasattr(encoder.patch_embed, 'proj'):
-            # If proj exists but wasn't used, we skip and use blocks which handle it
-            pass 
-        # Emergency reshape/pad or block pass
-        # In DINO ViT-L, the first block usually expects 1024. 
-        # If we have 64, it's because we are looking at raw conv outputs.
-        # Let's use the internal forward_features if available or mimic it:
-        tokens = encoder.prepare_tokens_with_masks(x) if hasattr(encoder, "prepare_tokens_with_masks") else None
-
-    if tokens is None:
-        # Manual fallback
-        cls_token = encoder.cls_token.expand(B, -1, -1)
-        # Re-derive tokens correctly if first attempt failed
-        pe = encoder.patch_embed(x)
-        tokens = pe.flatten(2).transpose(1, 2)
-        # If still 64, we must use the encoder's own logic to get to 1024
-        # Most DINOv3 models use a linear layer inside patch_embed
-        tokens = torch.cat((cls_token, tokens), dim=1) # This is where it fails if tokens is (N, 64)
     
-    # --- The standard research path from your notebook ---
-    tokens = encoder.prepare_tokens_with_masks(x) # This is the safest way in DINOv3
-    for blk in encoder.blocks:
-        tokens = blk(tokens)
-    if hasattr(encoder, "norm"):
-        tokens = encoder.norm(tokens)
-    
-    patch_tokens = tokens[:, 1:, :].squeeze(0).cpu().numpy()
-    return patch_tokens, H // 16, W // 16
+    return dav_model, dino_model, device
 
-# ================= MAIN APP =================
+# ================= SIDEBAR =================
 with st.sidebar:
     st.header("📂 Data Assets")
     rgb_file = st.file_uploader("Source RGB (TIF)", type=["tif", "tiff"])
     prior_file = st.file_uploader("Prior DSM (TIF)", type=["tif", "tiff"])
+    st.divider()
     repo = st.text_input("Repo ID", "osherr/Prior2DSM")
     dav_ckpt = st.text_input("DAV2 Weights", "depth_anything_v2_vitl.pth")
     dino_ckpt = st.text_input("DINO Weights", "dinov3_vitl16_pretrain_sat493m-eadcf0ff.pth")
-    m_size = st.selectbox("DAV2 Encoder", ["vitl", "vitb", "vits"])
+    m_size = st.selectbox("DAV2 Size", ["vitl", "vitb", "vits"])
 
+# ================= MAIN APP =================
 if rgb_file and prior_file:
     dav_model, dino_model, device = load_backbones(repo, dav_ckpt, dino_ckpt, m_size)
     
@@ -121,42 +73,48 @@ if rgb_file and prior_file:
         prior_data = src.read(1).astype(np.float32)
         h_f, w_f = prior_data.shape
 
-    # PRE-PROCESSING
+    # Pre-processing at specific grid resolutions
     img_448 = cv2.resize(rgb_raw.transpose(1, 2, 0), (448, 448))
     img_1024 = cv2.resize(rgb_raw.transpose(1, 2, 0), (1024, 1024))
     
     with st.spinner("Processing Backbones..."):
-        # 1. Depth (448px)
+        # 1. Depth (DAV2)
         dav_in = torch.from_numpy(img_448).permute(2, 0, 1).unsqueeze(0).to(device).float() / 255.0
         with torch.no_grad():
             rel_depth = dav_model(dav_in)
             if isinstance(rel_depth, (list, tuple)): rel_depth = rel_depth[-1]
             rel_depth = F.interpolate(rel_depth.unsqueeze(1), size=(h_f, w_f), mode="bilinear").squeeze().cpu().numpy()
         
-        # 2. DINO Features (1024px)
+        # 2. DINO Features (DINOv3) - Using high-level forward_features to avoid AssertionErrors
         dino_in = (torch.from_numpy(img_1024).permute(2, 0, 1).unsqueeze(0).to(device).float() / 255.0 - torch.tensor(IMAGENET_MEAN, device=device).view(1,3,1,1)) / torch.tensor(IMAGENET_STD, device=device).view(1,3,1,1)
         
-        # USE THE ENCODER'S INTERNAL METHOD TO AVOID CONCAT ERRORS
         with torch.no_grad():
-            # DINOv3 standard: forward_features or prepare_tokens_with_masks
-            tokens = dino_model.prepare_tokens_with_masks(dino_in)
-            for blk in dino_model.blocks: tokens = blk(tokens)
-            if hasattr(dino_model, "norm"): tokens = dino_model.norm(tokens)
+            # forward_features handles RoPE and masks correctly
+            features = dino_model.forward_features(dino_in)
+            if isinstance(features, (list, tuple)):
+                tokens = features[0] # Typical DINOv3 returns (tokens, metadata)
+            else:
+                tokens = features
+            
+            # Extract only the patch tokens (ignoring CLS/Registers)
+            # DINOv3 often uses 1 CLS + N patches
             patch_tokens = tokens[:, 1:, :].squeeze(0).cpu().numpy()
             Hp, Wp = 1024 // 16, 1024 // 16
 
-        # PCA Style from Notebook
+        # PCA Generation
         pca = PCA(n_components=3, whiten=True)
         Z = pca.fit_transform(patch_tokens).reshape(Hp, Wp, 3)
+        # Notebook Sigmoid 2*Z Style
         Z_rgb = 1.0 / (1.0 + np.exp(-2.0 * Z))
         pca_full = F.interpolate(torch.from_numpy(Z_rgb.transpose(2,0,1)).unsqueeze(0), size=(h_f, w_f), mode="bilinear").squeeze().numpy()
 
-    # --- TABS ---
-    tab_point, tab_refine = st.tabs(["🔬 Point Diagnostic", "🚀 BBOX Refinement"])
+    # --- UI ---
+    tab_point, tab_refine = st.tabs(["🔬 Point Diagnostic (Coarse)", "🚀 BBOX Refinement (Full)"])
+
+    viz_rgb = np.clip(rgb_raw.transpose(1, 2, 0) / (np.percentile(rgb_raw, 98) + 1e-6), 0, 1)
+    viz_small = (cv2.resize(viz_rgb, (640, 640)) * 255).astype(np.uint8)
 
     with tab_point:
-        viz_rgb = np.clip(rgb_raw.transpose(1, 2, 0) / (np.percentile(rgb_raw, 98) + 1e-6), 0, 1)
-        viz_small = (cv2.resize(viz_rgb, (640, 640)) * 255).astype(np.uint8)
         p_diag = streamlit_image_coordinates(viz_small, key="point_diag")
         if p_diag:
             tx, ty = int(p_diag['x']*(w_f/640)), int(p_diag['y']*(h_f/640))
@@ -170,7 +128,7 @@ if rgb_file and prior_file:
             k_rel, k_elev = c_rel[np.argsort(spatial_dists)[:60]], c_elev[np.argsort(spatial_dists)[:60]]
             reg = LinearRegression().fit(k_rel.reshape(-1, 1), k_elev)
             pred_h = reg.predict([[rel_depth[ty, tx]]])[0]
-            st.metric("Estimated Local Elevation", f"{pred_h:.2f} m")
+            st.metric("Estimated Elevation (Coarse)", f"{pred_h:.2f} m")
             fig, ax = plt.subplots(); ax.scatter(k_rel, k_elev, color='lawngreen'); ax.plot(k_rel, reg.predict(k_rel.reshape(-1,1)), 'b-'); st.pyplot(fig)
 
     with tab_refine:
@@ -180,28 +138,29 @@ if rgb_file and prior_file:
             st.rerun()
         if len(st.session_state.bbox_points) >= 2:
             if st.button("🚀 Start 4-Stage Optimization"):
-                mask = np.zeros((h_f, w_f), dtype=bool)
-                for i in range(0, len(st.session_state.bbox_points)-1, 2):
-                    p1, p2 = st.session_state.bbox_points[i], st.session_state.bbox_points[i+1]
-                    mask[min(p1[1], p2[1]):max(p1[1], p2[1]), min(p1[0], p2[0]):max(p1[0], p2[0])] = True
-                
-                # STAGE 3 & 4: Optimization
-                D_dav_t = torch.from_numpy(rel_depth)[None, None].to(device)
-                D_prior_t = torch.from_numpy(np.nan_to_num(prior_data))[None, None].to(device)
-                Anchor = torch.from_numpy((~mask).astype(np.float32))[None, None].to(device)
-                D_opt = D_prior_t.clone().detach().requires_grad_(True)
-                opt = torch.optim.Adam([D_opt], lr=8e-3)
-                D_dav_s = ((D_dav_t - D_dav_t.mean()) / D_dav_t.std()) * D_prior_t.std() + D_prior_t.mean()
-                
-                for _ in range(300):
-                    opt.zero_grad()
-                    l_a = F.mse_loss(D_opt * Anchor, D_prior_t * Anchor)
-                    l_l = F.l1_loss(D_opt - F.avg_pool2d(D_opt, 5, 1, 2), D_dav_s - F.avg_pool2d(D_dav_s, 5, 1, 2))
-                    (1200 * l_a + 180 * l_l).backward(); opt.step()
-                
-                st.success("Optimization Complete.")
-                buf = BytesIO(); meta_full.update(count=1, dtype="float32")
-                with rasterio.open(buf, "w", **meta_full) as dst: dst.write(D_opt.detach().squeeze().cpu().numpy(), 1)
-                st.download_button("📥 Download Refined DSM", buf.getvalue(), "refined.tif")
+                with st.spinner("Processing Stage 3 (Gradient) & Stage 4 (SSI)..."):
+                    mask = np.zeros((h_f, w_f), dtype=bool)
+                    for i in range(0, len(st.session_state.bbox_points)-1, 2):
+                        p1, p2 = st.session_state.bbox_points[i], st.session_state.bbox_points[i+1]
+                        mask[min(p1[1], p2[1]):max(p1[1], p2[1]), min(p1[0], p2[0]):max(p1[0], p2[0])] = True
+                    
+                    D_dav_t = torch.from_numpy(rel_depth)[None, None].to(device)
+                    D_prior_t = torch.from_numpy(np.nan_to_num(prior_data))[None, None].to(device)
+                    Anchor = torch.from_numpy((~mask).astype(np.float32))[None, None].to(device)
+                    D_opt = D_prior_t.clone().detach().requires_grad_(True)
+                    opt = torch.optim.Adam([D_opt], lr=8e-3)
+                    D_dav_s = ((D_dav_t - D_dav_t.mean()) / D_dav_t.std()) * D_prior_t.std() + D_prior_t.mean()
+                    
+                    for _ in range(300):
+                        opt.zero_grad()
+                        l_a = F.mse_loss(D_opt * Anchor, D_prior_t * Anchor)
+                        l_l = F.l1_loss(D_opt - F.avg_pool2d(D_opt, 5, 1, 2), D_dav_s - F.avg_pool2d(D_dav_s, 5, 1, 2))
+                        (1500 * l_a + 250 * l_l).backward(); opt.step()
+                    
+                    res = D_opt.detach().squeeze().cpu().numpy()
+                    buf = BytesIO(); meta_full.update(count=1, dtype="float32")
+                    with rasterio.open(buf, "w", **meta_full) as dst: dst.write(res, 1)
+                    st.success("DSM Refinement Complete!")
+                    st.download_button("📥 Download Final DSM", buf.getvalue(), "refined_full_res.tif")
 else:
     st.info("Upload Assets to begin.")