Update src/streamlit_app.py

src/streamlit_app.py

@@ -28,10 +28,19 @@ if current_dir not in sys.path:
 def load_backbones(repo_id, dav_file, dino_file, m_size):
     device = "cuda" if torch.cuda.is_available() else "cpu"
     
-    # 1. Load Depth Anything V2
+    # 1. Load Depth Anything V2 with error handling for config
     dav_path = hf_hub_download(repo_id=repo_id, filename=dav_file)
-    from depth_anything_v2 import build_backbone
-    dav_model = build_backbone(depth_size=m_size)
+    from depth_anything_v2.dpt import DepthAnythingV2
+    
+    # Manual config matching the expected ViT architecture
+    model_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]},
+    }
+    
+    # Initialize without the problematic 'encoder_cond_dim'
+    dav_model = DepthAnythingV2(**model_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()
 
@@ -66,147 +75,120 @@ def extract_dino_features(encoder, img_tensor, device):
         for blk in encoder.blocks: tokens = blk(tokens)
         if hasattr(encoder, "norm"): tokens = encoder.norm(tokens)
         
-        patch_tokens = tokens[:, 1:, :].squeeze(0).cpu().numpy() # [N, 1024]
+        patch_tokens = tokens[:, 1:, :].squeeze(0).cpu().numpy()
         
     pca = PCA(n_components=3, whiten=True)
     Z = pca.fit_transform(patch_tokens)
     Z_img = Z.reshape(Hp, Wp, 3)
-    # sigmoid(2*z) notebook style
     Z_rgb = 1.0 / (1.0 + np.exp(-2.0 * Z_img))
     
     norm_features = patch_tokens / (np.linalg.norm(patch_tokens, axis=1, keepdims=True) + 1e-6)
     return norm_features, Z_rgb.transpose(2, 0, 1), (Hp, Wp)
 
 # ================= UI DESIGN =================
-st.set_page_config(layout="wide", page_title="Prior2DSM Research", page_icon="🛰️")
+st.set_page_config(layout="wide", page_title="Prior2DSM Pro", page_icon="🛰️")
 
-# Custom CSS
 st.markdown("""
     <style>
     .stTabs [data-baseweb="tab"] { background-color: #f0f2f6; border-radius: 4px; padding: 10px; }
     .stTabs [aria-selected="true"] { background-color: #2e7d32; color: white; }
-    .stButton>button { border-radius: 8px; background-color: #2e7d32; color: white; width: 100%; }
+    .stButton>button { border-radius: 8px; background-color: #2e7d32; color: white; width: 100%; height: 3.5em; }
     </style>
     """, unsafe_allow_html=True)
 
-st.title("🛰️ Prior2DSM: Precision Refinement & Research")
+st.title("🛰️ Prior2DSM: Precision Terrain Refinement")
 
 if 'points' not in st.session_state:
     st.session_state.points = []
 
 with st.sidebar:
     st.header("📂 Data Ingestion")
-    rgb_file = st.file_uploader("Source Image (TIF, PNG, JPG)", type=["tif", "tiff", "png", "jpg", "jpeg"])
+    rgb_file = st.file_uploader("Source (TIF, PNG, JPG)", type=["tif", "tiff", "png", "jpg", "jpeg"])
     prior_file = st.file_uploader("Prior DSM (TIF)", type=["tif", "tiff"])
-    mask_file = st.file_uploader("Optional: Change Mask (TIF)", type=["tif", "tiff"])
     
     st.divider()
-    st.header("⚙️ Model Configuration")
-    repo = st.text_input("HF Repo", "osherr/Prior2DSM")
+    st.header("⚙️ Configuration")
+    m_size = st.selectbox("DAV2 Size", ["vitl", "vitb", "vits"])
+    repo = st.text_input("Repo", "osherr/Prior2DSM")
     dav_ckpt = st.text_input("DAV2 Weights", "depth_anything_v2_vitl.pth")
     dino_ckpt = st.text_input("DINO Weights", "dinov3_vitl16_pretrain.pth")
-    m_size = st.selectbox("DAV2 Size", ["vitl", "vitb", "vits"])
 
 if rgb_file and prior_file:
-    # 1. Load Assets
     dav_model, dino_model, device = load_backbones(repo, dav_ckpt, dino_ckpt, m_size)
     
     with rasterio.open(BytesIO(rgb_file.read())) as src:
         n_bands = src.count
-        sel_bands = st.sidebar.multiselect("Select RGB Bands", range(1, n_bands+1), default=[1, 2, 3] if n_bands >=3 else [1])
+        sel_bands = st.sidebar.multiselect("RGB Bands", range(1, n_bands+1), default=[1, 2, 3] if n_bands >=3 else [1])
         rgb_raw = src.read(sel_bands)
     
     with rasterio.open(BytesIO(prior_file.read())) as src:
         prior_data = src.read(1).astype(np.float32)
-        prior_meta = src.profile
-        h_full, w_full = prior_meta['height'], prior_meta['width']
+        p_meta = src.profile
+        h_f, w_f = p_meta['height'], p_meta['width']
 
-    # 2. Extract Features
+    # Inference Prep
     img_448 = cv2.resize(rgb_raw[:3].transpose(1, 2, 0), (448, 448))
     img_1024 = cv2.resize(rgb_raw[:3].transpose(1, 2, 0), (1024, 1024))
     
-    with st.spinner("Processing Neural Backbones..."):
-        # DAV2 Relative Depth
+    with st.spinner("Processing Backbones..."):
         dav_in = torch.from_numpy(img_448).transpose(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_full, w_full), mode="bilinear").squeeze().cpu().numpy()
+            rel_depth = F.interpolate(rel_depth.unsqueeze(1), size=(h_f, w_f), mode="bilinear").squeeze().cpu().numpy()
         
-        # DINO Semantic PCA
         dino_in = torch.from_numpy(img_1024).transpose(2, 0, 1).unsqueeze(0).to(device).float() / 255.0
-        dino_feats, pca_map, (Hp, Wp) = extract_dino_features(dino_model, dino_in, device)
-        pca_viz = F.interpolate(torch.from_numpy(pca_map).unsqueeze(0), size=(h_full, w_full), mode="bilinear").squeeze().numpy()
+        _, pca_raw, _ = extract_dino_features(dino_model, dino_in, device)
+        pca_viz = F.interpolate(torch.from_numpy(pca_raw).unsqueeze(0), size=(h_f, w_f), mode="bilinear").squeeze().numpy()
 
-    # 3. GUI Layout
+    # UI Tabs
     tab_mask, tab_diag, tab_refine = st.tabs(["🖌️ Masking", "🔬 Diagnostics", "🏗️ Refinement"])
 
     with tab_mask:
-        st.subheader("Interactive Change Masking")
         viz_rgb = np.clip(rgb_raw[:3].transpose(1, 2, 0) / (np.percentile(rgb_raw[:3], 98) + 1e-6), 0, 1)
         viz_small = (cv2.resize(viz_rgb, (640, 640)) * 255).astype(np.uint8)
-        
         col_m1, col_m2 = st.columns([2, 1])
         with col_m1:
             coords = streamlit_image_coordinates(viz_small, key="mask_painter")
             if coords:
-                scale_h, scale_w = h_full/640, w_full/640
-                st.session_state.points.append((int(coords['x']*scale_w), int(coords['y']*scale_h)))
+                st.session_state.points.append((int(coords['x']*(w_f/640)), int(coords['y']*(h_f/640))))
                 st.rerun()
         with col_m2:
             if st.button("🗑️ Clear Mask"):
                 st.session_state.points = []
                 st.rerun()
-            st.write(f"Points captured: {len(st.session_state.points)}")
+            st.write(f"Points: {len(st.session_state.points)}")
 
-    # Merge Masks
-    final_mask = np.zeros((h_full, w_full), dtype=bool)
+    mask = np.zeros((h_f, w_f), dtype=bool)
     if len(st.session_state.points) >= 2:
         for i in range(0, len(st.session_state.points)-1, 2):
             p1, p2 = st.session_state.points[i], st.session_state.points[i+1]
-            final_mask[min(p1[1], p2[1]):max(p1[1], p2[1]), min(p1[0], p2[0]):max(p1[0], p2[0])] = True
-
-    with tab_diag:
-        st.subheader("DINO-Powered Semantic Analysis")
-        col_d1, col_d2 = st.columns([2, 1])
-        with col_d1:
-            d_coords = streamlit_image_coordinates(viz_small, key="diag_selector")
-        
-        if d_coords:
-            tx, ty = int(d_coords['x'] * (w_full/640)), int(d_coords['y'] * (h_full/640))
-            st.info(f"Analyzing location: {tx}, {ty}")
-            # [Research Diagnostic Logic: cKDTree search in DINO features...]
-            st.write("Target Elevation:", prior_data[ty, tx])
+            mask[min(p1[1], p2[1]):max(p1[1], p2[1]), min(p1[0], p2[0]):max(p1[0], p2[0])] = True
 
     with tab_refine:
-        st.subheader("Fine Surface Refinement")
-        if st.button("🚀 Run Two-Stage Optimization"):
-            with st.spinner("Executing Fine Refinement..."):
+        if st.button("🚀 Start Fine Refinement"):
+            with st.spinner("Optimizing Surface..."):
                 D_dav = torch.from_numpy(rel_depth)[None, None].to(device)
-                D_prior = torch.from_numpy(prior_data.astype(np.float32))[None, None].to(device)
-                Anchor = torch.from_numpy((~final_mask).astype(np.float32))[None, None].to(device)
-                
-                # Step 1: Global Trend Alignment
+                D_prior = torch.from_numpy(prior_data)[None, None].to(device)
+                Anchor = torch.from_numpy((~mask).astype(np.float32))[None, None].to(device)
                 D_opt = D_prior.clone().detach().requires_grad_(True)
-                opt1 = torch.optim.Adam([D_opt], lr=1e-2)
+                opt = torch.optim.Adam([D_opt], lr=8e-3)
                 D_dav_s = ((D_dav - D_dav.mean()) / D_dav.std()) * D_prior.std() + D_prior.mean()
-                
-                # Step 2: Fine Detail Injection (Fine Refinement)
-                for _ in range(300):
-                    opt1.zero_grad()
+                for _ in range(250):
+                    opt.zero_grad()
                     l_anchor = F.mse_loss(D_opt * Anchor, D_prior * Anchor)
                     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))
                     (1200 * l_anchor + 180 * l_geom).backward()
-                    opt1.step()
-
+                    opt.step()
+                
                 final = D_opt.detach().squeeze().cpu().numpy()
-                prior_meta.update(count=1, dtype="float32")
                 out_buf = BytesIO()
-                with rasterio.open(out_buf, "w", **prior_meta) as dst:
+                p_meta.update(count=1, dtype="float32")
+                with rasterio.open(out_buf, "w", **p_meta) as dst:
                     dst.write(final, 1)
-                
                 st.success("Refinement Complete!")
-                st.download_button("📥 Download Refined TIF", out_buf.getvalue(), "refined_448_fine.tif")
+                st.download_button("📥 Download Refined DSM", out_buf.getvalue(), "refined.tif")
+
 else:
-    st.info("Upload TIF assets and configure models in the sidebar.")
+    st.info("Upload TIF/PNG assets to begin.")