Update src/streamlit_app.py

src/streamlit_app.py

@@ -28,19 +28,15 @@ 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 with error handling for config
+    # 1. Load Depth Anything V2 (Manual Init)
     dav_path = hf_hub_download(repo_id=repo_id, filename=dav_file)
     from depth_anything_v2.dpt import DepthAnythingV2
-    
-    # Manual config matching the expected ViT architecture
-    model_configs = {
+    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 = 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()
 
@@ -60,10 +56,12 @@ def load_backbones(repo_id, dav_file, dino_file, m_size):
     return dav_model, dino_model, device
 
 def extract_dino_features(encoder, img_tensor, device):
-    """Extracts patch tokens and computes PCA map [3, H, W] using sigmoid(2z) style."""
+    """CHW input tensor -> normalized features and PCA map."""
     B, C, H, W = img_tensor.shape
-    x = (img_tensor - torch.tensor(IMAGENET_MEAN, device=device).view(1,3,1,1)) / \
-         torch.tensor(IMAGENET_STD, device=device).view(1,3,1,1)
+    # Standardize
+    mean = torch.tensor(IMAGENET_MEAN, device=device).view(1, 3, 1, 1)
+    std = torch.tensor(IMAGENET_STD, device=device).view(1, 3, 1, 1)
+    x = (img_tensor - mean) / std
     
     with torch.no_grad():
         pe = encoder.patch_embed(x)
@@ -74,7 +72,6 @@ def extract_dino_features(encoder, img_tensor, device):
         tokens = tokens + encoder.interpolate_pos_encoding(tokens, H, W)
         for blk in encoder.blocks: tokens = blk(tokens)
         if hasattr(encoder, "norm"): tokens = encoder.norm(tokens)
-        
         patch_tokens = tokens[:, 1:, :].squeeze(0).cpu().numpy()
         
     pca = PCA(n_components=3, whiten=True)
@@ -85,12 +82,12 @@ def extract_dino_features(encoder, img_tensor, device):
     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 Pro", page_icon="🛰️")
+# ================= UI =================
+st.set_page_config(layout="wide", page_title="Prior2DSM Pro")
 
 st.markdown("""
     <style>
-    .stTabs [data-baseweb="tab"] { background-color: #f0f2f6; border-radius: 4px; padding: 10px; }
+    .stTabs [data-baseweb="tab"] { background-color: #f0f2f6; border-radius: 4px; padding: 10px; font-weight: bold; }
     .stTabs [aria-selected="true"] { background-color: #2e7d32; color: white; }
     .stButton>button { border-radius: 8px; background-color: #2e7d32; color: white; width: 100%; height: 3.5em; }
     </style>
@@ -107,18 +104,18 @@ with st.sidebar:
     prior_file = st.file_uploader("Prior DSM (TIF)", type=["tif", "tiff"])
     
     st.divider()
-    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_sat493m-eadcf0ff.pth")
+    st.header("⚙️ Settings")
+    repo = st.text_input("Repo ID", "osherr/Prior2DSM")
+    dav_ckpt = st.text_input("DAV2 Filename", "depth_anything_v2_vitl.pth")
+    dino_ckpt = st.text_input("DINO Filename", "dinov3_vitl16_pretrain_sat493m-eadcf0ff.pth")
+    m_size = st.selectbox("DAV2 Encoder", ["vitl", "vitb", "vits"])
 
 if rgb_file and prior_file:
     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("RGB Bands", range(1, n_bands+1), default=[1, 2, 3] if n_bands >=3 else [1])
+        sel_bands = st.sidebar.multiselect("Select 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:
@@ -126,69 +123,80 @@ if rgb_file and prior_file:
         p_meta = src.profile
         h_f, w_f = p_meta['height'], p_meta['width']
 
-    # Inference Prep
+    # --- PREPARE INPUTS ---
+    # Convert HWC (OpenCV/NumPy) to CHW (PyTorch) using .permute() instead of .transpose()
     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 Backbones..."):
-        dav_in = torch.from_numpy(img_448).transpose(2, 0, 1).unsqueeze(0).to(device).float() / 255.0
+    with st.spinner("Processing Neural Backbones..."):
+        # Fix: Use .permute(2, 0, 1) for reordering 3 dimensions
+        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()
         
-        dino_in = torch.from_numpy(img_1024).transpose(2, 0, 1).unsqueeze(0).to(device).float() / 255.0
-        _, pca_raw, _ = extract_dino_features(dino_model, dino_in, device)
+        dino_in = torch.from_numpy(img_1024).permute(2, 0, 1).unsqueeze(0).to(device).float() / 255.0
+        dino_feats, 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()
 
-    # UI Tabs
-    tab_mask, tab_diag, tab_refine = st.tabs(["🖌️ Masking", "🔬 Diagnostics", "🏗️ Refinement"])
+    tab_mask, tab_diag, tab_refine = st.tabs(["🖌️ Interactive Mask", "🔬 Diagnostics", "🏗️ Refinement"])
 
     with tab_mask:
+        col_m1, col_m2 = st.columns([2, 1])
         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:
+            st.write("Draw rectangles to mask change areas.")
             coords = streamlit_image_coordinates(viz_small, key="mask_painter")
             if coords:
                 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"):
+            if st.button("🗑️ Reset Mask"):
                 st.session_state.points = []
                 st.rerun()
-            st.write(f"Points: {len(st.session_state.points)}")
+            st.metric("Points", len(st.session_state.points))
 
+    # Mask construction
     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]
             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:
+        if st.button("Generate Diagnostic PCA View"):
+            fig, ax = plt.subplots()
+            pca_rgb_view = 1.0 / (1.0 + np.exp(-2.0 * pca_viz.transpose(1, 2, 0)))
+            ax.imshow(pca_rgb_view)
+            st.pyplot(fig)
+
     with tab_refine:
-        if st.button("🚀 Start Fine Refinement"):
-            with st.spinner("Injecting Fine Laplacian Details..."):
-                # Move to device once
-                D_dav = torch.from_numpy(rel_depth).to(device).view(1, 1, h_f, w_f)
-                D_prior = torch.from_numpy(prior_data).to(device).view(1, 1, h_f, w_f)
-                Anchor = torch.from_numpy((~mask).astype(np.float32)).to(device).view(1, 1, h_f, w_f)
+        if st.button("🚀 Run Fine Refinement"):
+            with st.spinner("Refining Terrain..."):
+                D_dav = torch.from_numpy(rel_depth)[None, None].to(device)
+                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)
                 optimizer = torch.optim.Adam([D_opt], lr=8e-3)
-                
-                # Align Global Scale
                 D_dav_s = ((D_dav - D_dav.mean()) / D_dav.std()) * D_prior.std() + D_prior.mean()
                 
                 for _ in range(300):
                     optimizer.zero_grad()
-                    # Loss 1: Anchor constraint (don't move valid ground)
                     l_anchor = F.mse_loss(D_opt * Anchor, D_prior * Anchor)
-                    # Loss 2: Laplacian Detail (transfer high-frequency features)
                     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 + 250 * l_geom).backward()
+                    (1200 * l_anchor + 180 * l_geom).backward()
                     optimizer.step()
-
+                
+                res = D_opt.detach().squeeze().cpu().numpy()
+                buf = BytesIO()
+                p_meta.update(count=1, dtype="float32")
+                with rasterio.open(buf, "w", **p_meta) as dst: dst.write(res, 1)
+                st.success("Refinement Complete!")
+                st.download_button("📥 Download Refined DSM", buf.getvalue(), "refined_fine.tif")
 else:
-    st.info("Upload TIF/PNG assets to begin.")
+    st.info("Upload your assets to begin.")