Update src/streamlit_app.py

src/streamlit_app.py

@@ -25,10 +25,50 @@ 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 = []
 
+# ================= YOUR RESEARCH HELPERS (DUPLICATED) =================
+@torch.no_grad()
+def _final_patch_tokens_postnorm(encoder, x):
+    """ Direct duplicate from your research code to handle ViT-L dimensioning. """
+    B, _, H, W = x.shape
+    D = int(getattr(encoder, "embed_dim"))
+    ps = getattr(encoder, "patch_embed").patch_size
+    if isinstance(ps, (tuple, list)): ps = ps[0]
+
+    pe = encoder.patch_embed(x)
+    if pe.ndim == 4:
+        if pe.shape[1] == D:
+            Hp, Wp = pe.shape[-2], pe.shape[-1]
+            tokens = pe.flatten(2).transpose(1, 2)
+        elif pe.shape[-1] == D:
+            Hp, Wp = pe.shape[1], pe.shape[2]
+            tokens = pe.view(B, Hp * Wp, D)
+        else: raise RuntimeError(f"Unexpected patch_embed output shape {pe.shape}")
+    elif pe.ndim == 3:
+        Hp, Wp = H // ps, W // ps
+        tokens = pe if pe.shape[-1] == D else pe.transpose(1, 2).contiguous()
+    else: raise RuntimeError(f"patch_embed returned {pe.ndim}D tensor.")
+
+    cls = encoder.cls_token.expand(B, 1, D)
+    tokens = torch.cat([cls, tokens], dim=1)
+
+    if hasattr(encoder, "interpolate_pos_encoding"):
+        tokens = tokens + encoder.interpolate_pos_encoding(tokens, H, W)
+    if hasattr(encoder, "pos_drop"):
+        tokens = encoder.pos_drop(tokens)
+
+    out = tokens
+    for blk in encoder.blocks: out = blk(out)
+    if hasattr(encoder, "norm") and encoder.norm is not None: out = encoder.norm(out)
+
+    out = out.squeeze(0)
+    T, N = out.shape[0], Hp * Wp
+    n_special = T - N
+    return out[n_special:, :], Hp, Wp
+
 @st.cache_resource
 def load_backbones(repo_id, dav_file, dino_file, m_size):
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    # 1. Depth Anything V2
+    # Load Depth Anything V2
     dav_path = hf_hub_download(repo_id=repo_id, filename=dav_file)
     from depth_anything_v2.dpt import DepthAnythingV2
     configs = {'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
@@ -38,17 +78,15 @@ def load_backbones(repo_id, dav_file, dino_file, m_size):
     dav_model.load_state_dict(torch.load(dav_path, map_location="cpu", weights_only=True))
     dav_model = dav_model.to(device).eval()
 
-    # 2. DINOv3 Encoder (ViT-L/16)
+    # Load DINOv3 (ViT-L/16)
     dino_path = hf_hub_download(repo_id=repo_id, filename=dino_file)
     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")
-    if "state_dict" in ckpt: ckpt = ckpt["state_dict"]
-    elif "model" in ckpt: ckpt = ckpt["model"]
-    elif "teacher" in ckpt: ckpt = ckpt["teacher"]
+    for k in ["state_dict", "model", "teacher"]:
+        if k in ckpt: ckpt = ckpt[k]
     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
 
 # ================= SIDEBAR =================
@@ -73,53 +111,36 @@ if rgb_file and prior_file:
         prior_data = src.read(1).astype(np.float32)
         h_f, w_f = prior_data.shape
 
-    # Pre-processing at specific grid resolutions
+    # 1. Generate Neural Depth (Scaled to Input Size)
     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 (DAV2)
+    with st.spinner("Processing Stage 1: Neural Depth..."):
         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 (DINOv3) - Using high-level forward_features to avoid AssertionErrors
+
+    # 2. Generate DINO Features (1024px Grid)
+    img_1024 = cv2.resize(rgb_raw.transpose(1, 2, 0), (1024, 1024))
+    with st.spinner("Processing Stage 2: Semantic Tokens..."):
         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)
-        
-        with torch.no_grad():
-            # 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 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()
-
-    # --- UI ---
-    tab_point, tab_refine = st.tabs(["🔬 Point Diagnostic (Coarse)", "🚀 BBOX Refinement (Full)"])
+        patch_tokens, Hp, Wp = _final_patch_tokens_postnorm(dino_model, dino_in)
 
+    # UI Visual Setup
     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)
 
+    # --- TABS ---
+    tab_point, tab_refine = st.tabs(["🔬 Point Diagnostic (Coarse)", "🚀 BBOX Refinement (Full)"])
+
     with tab_point:
+        st.write("Click for a **Stage 2 (Coarse)** metric estimation.")
         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))
             valid_logic = np.isfinite(prior_data)
-            feat_norm = patch_tokens / (np.linalg.norm(patch_tokens, axis=1, keepdims=True) + 1e-6)
+            feat_norm = patch_tokens.cpu().numpy()
+            feat_norm /= (np.linalg.norm(feat_norm, axis=1, keepdims=True) + 1e-6)
             tree = cKDTree(feat_norm[valid_logic.flatten()])
             target_feat = feat_norm[int(ty/h_f*Hp)*Wp + int(tx/w_f*Wp)]
             _, idx = tree.query(target_feat, k=min(100, tree.n))
@@ -128,28 +149,29 @@ 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 Elevation (Coarse)", f"{pred_h:.2f} m")
+            st.metric("Estimated Elevation", 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:
+        st.write("Draw BBOX for **Stage 3 & 4** optimization.")
         b_coords = streamlit_image_coordinates(viz_small, key="bbox_painter")
         if b_coords:
             st.session_state.bbox_points.append((int(b_coords['x']*(w_f/640)), int(b_coords['y']*(h_f/640))))
             st.rerun()
         if len(st.session_state.bbox_points) >= 2:
-            if st.button("🚀 Start 4-Stage Optimization"):
-                with st.spinner("Processing Stage 3 (Gradient) & Stage 4 (SSI)..."):
+            if st.button("🚀 Execute Full Refinement (Original Size)"):
+                with st.spinner("Optimizing..."):
                     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_prior_t = torch.from_numpy(np.nan_to_num(prior_data)).to(device).unsqueeze(0).unsqueeze(0)
+                    Anchor = torch.from_numpy((~mask).astype(np.float32)).to(device).unsqueeze(0).unsqueeze(0)
                     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()
+                    D_dav_s = ((D_dav_t - D_dav_t.mean()) / D_dav_t.std()) * D_opt.std() + D_opt.mean()
                     
                     for _ in range(300):
                         opt.zero_grad()
@@ -160,7 +182,6 @@ if rgb_file and prior_file:
                     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")
+                    st.download_button("📥 Download Full-Res DSM", buf.getvalue(), "refined.tif")
 else:
-    st.info("Upload Assets to begin.")
+    st.info("Upload imagery to begin.")