Update src/streamlit_app.py

src/streamlit_app.py

@@ -14,20 +14,26 @@ from scipy.spatial import cKDTree
 from huggingface_hub import hf_hub_download
 from streamlit_image_coordinates import streamlit_image_coordinates
 
-# --- קונפיגורציה ---
+# --- CONFIG & INITIALIZATION ---
 IMAGENET_MEAN = (0.485, 0.456, 0.406)
 IMAGENET_STD  = (0.229, 0.224, 0.225)
 
-st.set_page_config(layout="wide", page_title="Prior2DSM Research Pro")
+st.set_page_config(layout="wide", page_title="Prior2DSM Research Suite")
 
+if 'bbox_coords' not in st.session_state:
+    st.session_state.bbox_coords = []
+
+# --- MODEL LOADING (Safe Init) ---
 @st.cache_resource
 def load_models(repo_id, dav_file, dino_file, m_size):
     device = "cuda" if torch.cuda.is_available() else "cpu"
     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]},
-               'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
-               'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]}}
+    configs = {
+        'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
+        'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
+        'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]}
+    }
     dav_model = DepthAnythingV2(**configs[m_size])
     dav_model.load_state_dict(torch.load(dav_path, map_location="cpu", weights_only=True))
     dav_model = dav_model.to(device).eval()
@@ -39,11 +45,12 @@ def load_models(repo_id, dav_file, dino_file, m_size):
     if "state_dict" in ckpt: ckpt = ckpt["state_dict"]
     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
 
 @st.cache_data(show_spinner=False)
-def run_backbone_inference(_dav, _dino, rgb_raw, h_f, w_f, device):
-    # 1. Depth Inference (Scaling to Input Size)
+def run_full_inference(_dav, _dino, rgb_raw, h_f, w_f, device):
+    # 1. DAV2 Depth
     img_448 = cv2.resize(rgb_raw.transpose(1, 2, 0), (448, 448))
     dav_in = torch.from_numpy(img_448).permute(2, 0, 1).unsqueeze(0).to(device).float() / 255.0
     with torch.no_grad():
@@ -51,142 +58,170 @@ def run_backbone_inference(_dav, _dino, rgb_raw, h_f, w_f, device):
         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 & PCA
+    # 2. DINO PCA
     img_1024 = cv2.resize(rgb_raw.transpose(1, 2, 0), (1024, 1024))
     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():
         features = _dino.forward_features(dino_in)
-        # תיקון לחילוץ טוקנים אם מוחזר מילון (dict)
         tokens = features.get('x_norm_patchtokens', features.get('main')) if isinstance(features, dict) else features[0]
-        if tokens.shape[1] > 4096: # If CLS is present
-            patch_tokens = tokens[:, 1:, :].squeeze(0).cpu().numpy()
-        else:
-            patch_tokens = tokens.squeeze(0).cpu().numpy()
+        patch_tokens = tokens[:, 1:, :].squeeze(0).cpu().numpy()
         
-        Hp, Wp = 1024 // 16, 1024 // 16
-
-    # Generate PCA for visualization
     pca_obj = PCA(n_components=3, whiten=True)
-    Z = pca_obj.fit_transform(patch_tokens).reshape(Hp, Wp, 3)
-    Z_rgb = 1.0 / (1.0 + np.exp(-2.0 * Z)) # Sigmoid PCA
-    pca_viz = F.interpolate(torch.from_numpy(Z_rgb.transpose(2,0,1)).unsqueeze(0).float(), size=(h_f, w_f), mode="bilinear").squeeze().numpy()
+    Z = pca_obj.fit_transform(patch_tokens).reshape(64, 64, 3).transpose(2, 0, 1)
+    pca_full = F.interpolate(torch.from_numpy(Z).unsqueeze(0).float(), size=(h_f, w_f), mode="bilinear").squeeze(0)
+    pca_map = F.normalize(pca_full, p=2, dim=0).numpy()
 
-    return rel_depth, patch_tokens, pca_viz, Hp, Wp
+    return rel_depth, pca_map
 
-# ================= UI Sidebar =================
+# --- SIDEBAR ---
 with st.sidebar:
-    st.header("📂 Data Ingestion")
-    rgb_file = st.file_uploader("Source RGB (TIF)", type=["tif", "tiff"])
-    prior_file = st.file_uploader("Prior DSM (TIF)", type=["tif", "tiff"])
-    mask_file = st.file_uploader("Change Mask (TIF)", type=["tif", "tiff"])
-    mask_band = st.selectbox("Mask Band Index", [1, 2, 3], index=1)
+    st.header("📂 Research Data")
+    rgb_file = st.file_uploader("Source RGB", type=["tif", "tiff"])
+    prior_file = st.file_uploader("Prior DSM", type=["tif", "tiff"])
+    mask_file = st.file_uploader("Change Mask", type=["tif", "tiff"])
+    mask_band = st.selectbox("Mask Band", [1, 2, 3], index=1)
     st.divider()
+    K_val = st.slider("K Neighbors (Diagnostic)", 5, 50, 15)
     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"])
 
-# ================= Main App =================
+# --- MAIN LOGIC ---
 if rgb_file and prior_file and mask_file:
-    dav_m, dino_m, dev = load_models(repo, dav_ckpt, dino_ckpt, m_size)
+    dav_m, dino_m, dev = load_models(repo, dav_ckpt, dino_ckpt, "vitl")
     
     with rasterio.open(BytesIO(rgb_file.read())) as src:
         rgb_raw = src.read([1, 2, 3])
+        h_f, w_f = src.height, src.width
         meta_full = src.profile
     with rasterio.open(BytesIO(prior_file.read())) as src:
         prior_data = src.read(1).astype(np.float32)
-        h_f, w_f = prior_data.shape
     with rasterio.open(BytesIO(mask_file.getvalue())) as src:
         change_mask = src.read(mask_band) > 0
 
-    # Running Stage 1 & 2
-    rel_depth, patch_tokens, pca_viz, Hp, Wp = run_backbone_inference(dav_m, dino_m, rgb_raw, h_f, w_f, dev)
+    rel_depth, pca_map = run_full_inference(dav_m, dino_m, rgb_raw, h_f, w_f, dev)
 
-    # --- DIAGNOSTIC TABS ---
-    st.header("🔬 Research Diagnostics")
-    tab_vis, tab_knn, tab_refine = st.tabs(["🖼️ Visual Maps", "🎯 KNN Analysis", "🚀 SSI Refinement"])
+    tab_diag, tab_bbox, tab_ssi = st.tabs(["🔬 Point Diagnostic", "⊞ Grid/BBOX Update", "🚀 Final SSI"])
 
-    with tab_vis:
-        col1, col2, col3 = st.columns(3)
+    # --- TAB 1: INTERACTIVE DIAGNOSTIC (YOUR LOGIC) ---
+    with tab_diag:
+        st.subheader("Spatial + Semantic Metric Diagnostic")
         viz_rgb = np.clip(rgb_raw.transpose(1, 2, 0) / (np.percentile(rgb_raw, 98) + 1e-6), 0, 1)
-        
-        with col1:
-            st.subheader("Semantic PCA (DINO)")
-            st.image(pca_viz.transpose(1, 2, 0), use_container_width=True)
-        with col2:
-            st.subheader("Relative Depth (DAV2)")
-            fig_d, ax_d = plt.subplots()
-            ax_d.imshow(rel_depth, cmap='magma')
-            ax_d.axis('off')
-            st.pyplot(fig_d)
-        with col3:
-            st.subheader("Change Mask")
-            st.image(change_mask.astype(float), use_container_width=True)
-
-    with tab_knn:
         viz_small = (cv2.resize(viz_rgb, (640, 640)) * 255).astype(np.uint8)
-        col_c, col_r = st.columns([1, 1])
-        with col_c:
-            st.write("Click a point to see KNN consistency:")
+        
+        col_map, col_res = st.columns([1, 1])
+        with col_map:
             p_click = streamlit_image_coordinates(viz_small, key="diag_click")
         
         if p_click:
             tx, ty = int(p_click['x']*(w_f/640)), int(p_click['y']*(h_f/640))
-            valid_logic = np.isfinite(prior_data) & (~change_mask)
-            # Match PCA map to full resolution for retrieval
-            pca_full = F.normalize(torch.from_numpy(pca_viz), p=2, dim=0).numpy()
-            valid_pca = pca_full[:, valid_logic].T
-            tree = cKDTree(valid_pca)
             
-            target_pca = pca_full[:, ty, tx]
-            _, idx = tree.query(target_pca, k=100)
+            # Logic: valid anchors
+            valid_logic = np.isfinite(prior_data) & (~change_mask)
+            valid_coords = np.argwhere(valid_logic)
+            valid_pca = pca_map[:, valid_logic].T
+            pca_tree = cKDTree(valid_pca)
             
-            c_rel, c_elev = rel_depth[valid_logic][idx], prior_data[valid_logic][idx]
-            reg = LinearRegression().fit(c_rel.reshape(-1, 1), c_elev)
-            pred_h = reg.predict([[rel_depth[ty, tx]]])[0]
+            # 1. Search candidates by PCA (Semantic match)
+            target_pca = pca_map[:, ty, tx]
+            dist, idx = pca_tree.query(target_pca, k=100) # Start with 100
             
-            with col_r:
-                st.metric("KNN Predicted Elevation", f"{pred_h:.2f} m")
-                fig_knn, ax_knn = plt.subplots()
-                ax_knn.scatter(c_rel, c_elev, alpha=0.3, label='Semantic Anchors')
-                ax_knn.plot(c_rel, reg.predict(c_rel.reshape(-1,1)), 'r-', label='Local Fit')
-                ax_knn.scatter([rel_depth[ty, tx]], [pred_h], color='cyan', marker='*', s=200, label='Prediction')
-                ax_knn.legend()
-                st.pyplot(fig_knn)
-
-    with tab_refine:
-        if st.button("🚀 Run SSI Final Optimization"):
-            with st.spinner("Injecting Fine Details..."):
-                # המרה לטנזורים וניתוק מהגרף (DETACH) למניעת השגיאה
-                D_dav_t = torch.from_numpy(rel_depth).to(dev).unsqueeze(0).unsqueeze(0).detach()
-                D_prior_t = torch.from_numpy(np.nan_to_num(prior_data)).to(dev).unsqueeze(0).unsqueeze(0).detach()
-                Anchor = torch.from_numpy((~change_mask).astype(np.float32)).to(dev).unsqueeze(0).unsqueeze(0).detach()
+            c_pred = rel_depth[valid_logic][idx]
+            c_metric = prior_data[valid_logic][idx]
+            c_coords = valid_coords[idx]
+
+            # 2. Filter by Physical distance (Spatial match) - Use 20
+            spatial_dists = np.sqrt((c_coords[:, 0] - ty)**2 + (c_coords[:, 1] - tx)**2)
+            spatial_indices = np.argsort(spatial_dists)[:20]
+            s_pred, s_metric, s_coords = c_pred[spatial_indices], c_metric[spatial_indices], c_coords[spatial_indices]
+
+            # 3. Final Geometric filter - Use 10 (K=10)
+            target_rel = rel_depth[ty, tx]
+            depth_diff = np.abs(s_pred - target_rel)
+            final_indices = np.argsort(depth_diff)[:10]
+            rejected_indices = np.argsort(depth_diff)[10:]
+
+            k_pred, k_metric, k_coords = s_pred[final_indices], s_metric[final_indices], s_coords[final_indices]
+            r_coords = s_coords[rejected_indices]
+
+            # Local Fit
+            local_model = LinearRegression().fit(k_pred.reshape(-1, 1), k_metric)
+            pred_h = local_model.predict([[target_rel]])[0]
+
+            with col_res:
+                st.metric("Estimated Local Elevation", f"{pred_h:.2f} m")
+                fig, ax = plt.subplots(1, 2, figsize=(12, 5))
+                # Map view with points
+                ax[0].imshow(viz_rgb)
+                ax[0].scatter(r_coords[:, 1], r_coords[:, 0], color='red', marker='x', s=10, label='Rejected')
+                ax[0].scatter(k_coords[:, 1], k_coords[:, 0], color='lawngreen', s=25, label='Anchors')
+                ax[0].scatter([tx], [ty], color='cyan', marker='P', s=80, label='Target')
+                ax[0].set_title("Anchor Spatial Selection")
+                ax[0].legend()
+                
+                # Scatter view
+                ax[1].scatter(s_pred, s_metric, color='red', marker='x', alpha=0.3)
+                ax[1].scatter(k_pred, k_metric, color='lawngreen', s=50, edgecolors='black')
+                xr = np.linspace(s_pred.min(), s_pred.max(), 100).reshape(-1, 1)
+                ax[1].plot(xr, local_model.predict(xr), 'b-', label='Local Fit')
+                ax[1].scatter([target_rel], [pred_h], color='cyan', marker='*', s=150)
+                ax[1].set_title("Metric Regression")
+                st.pyplot(fig)
+
+    # --- TAB 2: BBOX / GRID UPDATE ---
+    with tab_bbox:
+        st.subheader("Select Coordinate Range for Update")
+        # Draw prior with grid
+        fig_g, ax_g = plt.subplots(figsize=(10, 10))
+        ax_g.imshow(prior_data, cmap='terrain')
+        # Add grid
+        grid_space = 64
+        for x in range(0, w_f, grid_space): ax_g.axvline(x, color='white', alpha=0.2, lw=0.5)
+        for y in range(0, h_f, grid_space): ax_g.axhline(y, color='white', alpha=0.2, lw=0.5)
+        
+        viz_grid = (cv2.resize(prior_data, (640, 640)) - prior_data.min()) / (prior_data.max() - prior_data.min())
+        p_bbox = streamlit_image_coordinates((viz_grid * 255).astype(np.uint8), key="bbox_sel")
+        
+        if p_bbox:
+            bx, by = int(p_bbox['x']*(w_f/640)), int(p_bbox['y']*(h_f/640))
+            st.session_state.bbox_coords.append((bx, by))
+        
+        if st.session_state.bbox_coords:
+            st.write(f"Points Selected: {st.session_state.bbox_coords}")
+            if st.button("Reset BBOX"): 
+                st.session_state.bbox_coords = []
+                st.rerun()
+
+    # --- TAB 3: SSI OPTIMIZATION ---
+    with tab_ssi:
+        if st.button("🚀 Run SSI Final Refinement"):
+            with st.spinner("Optimizing Surface..."):
+                D_dav_t = torch.from_numpy(rel_depth)[None, None].to(dev)
+                D_prior_t = torch.from_numpy(np.nan_to_num(prior_data)).to(dev).unsqueeze(0).unsqueeze(0)
                 
-                # המשתנה היחיד שדורש גרדיאנט
-                D_opt = D_prior_t.clone().requires_grad_(True)
-                optimizer = torch.optim.Adam([D_opt], lr=8e-3)
+                # If BBOX selected, update mask
+                final_mask = change_mask.copy()
+                if len(st.session_state.bbox_coords) >= 2:
+                    c1, c2 = st.session_state.bbox_coords[-2], st.session_state.bbox_coords[-1]
+                    final_mask[min(c1[1], c2[1]):max(c1[1], c2[1]), min(c1[0], c2[0]):max(c1[0], c2[0])] = True
                 
-                # נרמול עומק יחסי למרחב המטרי (Alignment)
+                Anchor = torch.from_numpy((~final_mask).astype(np.float32)).to(dev).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() + 1e-6)) * D_prior_t.std() + D_prior_t.mean()
-                D_dav_s = D_dav_s.detach() # חשוב מאוד למניעת Backward שגיאה
                 
                 for _ in range(350):
-                    optimizer.zero_grad()
-                    # Loss 1: Anchor (נעילת אזורים יציבים)
-                    l_anchor = F.mse_loss(D_opt * Anchor, D_prior_t * Anchor)
-                    # Loss 2: Laplacian Detail (הזרקת טקסטורה)
-                    l_geom = 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_anchor + 200 * l_geom).backward()
-                    optimizer.step()
+                    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.detach() - F.avg_pool2d(D_dav_s.detach(), 5, 1, 2))
+                    (1500 * l_a + 200 * 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("Optimization Complete!")
-                st.download_button("📥 Download Refined DSM", buf.getvalue(), "refined_final.tif")
+                st.success("Refinement Complete!")
+                st.download_button("Download Refined DSM", buf.getvalue(), "refined.tif")
 else:
-    st.info("אנא העלה קבצים כדי להתחיל.")
+    st.info("Upload Assets to begin.")