initial commit

python_server/test_depth.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+from pathlib import Path
+import numpy as np
+import cv2
+from PIL import Image
+
+# 解除大图限制
+Image.MAX_IMAGE_PIXELS = None 
+
+from model.Depth.depth_loader import build_depth_predictor, UnifiedDepthConfig, DepthBackend
+
+# ================= 配置区 =================
+INPUT_IMAGE = "/home/dwh/Documents/毕业设计/dwh/数据集/Up the River During Qingming (detail) - Court painters.jpg"
+OUTPUT_DIR = "outputs/test_depth_v4"
+DEPTH_BACKEND = DepthBackend.DEPTH_ANYTHING_V2
+
+# 边缘捕捉参数
+# 增大这个值会让边缘更细，减小会让边缘更粗（捕捉更多微弱信息）
+EDGE_TOP_PERCENTILE = 96.0  # 选取梯度最强的前 7.0% 的像素
+# 局部增强的灵敏度，建议在 2.0 - 5.0 之间
+CLAHE_CLIP_LIMIT = 3.0 
+# 形态学核大小
+MORPH_SIZE = 5
+# =========================================
+
+def _extract_robust_edges(depth_norm: np.ndarray) -> np.ndarray:
+    """
+    通过局部增强和 Sobel 梯度提取闭合边缘
+    """
+    # 1. 转换为 8 位灰度
+    depth_u8 = (depth_norm * 255).astype(np.uint8)
+    
+    # 2. 【核心步骤】CLAHE 局部自适应对比度增强
+    # 这会强行放大古画中细微的建筑/人物深度差
+    clahe = cv2.createCLAHE(clipLimit=CLAHE_CLIP_LIMIT, tileGridSize=(16, 16))
+    enhanced_depth = clahe.apply(depth_u8)
+    
+    # 3. 高斯模糊：减少数字化噪声
+    blurred = cv2.GaussianBlur(enhanced_depth, (5, 5), 0)
+    
+    # 4. Sobel 算子计算梯度强度
+    grad_x = cv2.Sobel(blurred, cv2.CV_64F, 1, 0, ksize=3)
+    grad_y = cv2.Sobel(blurred, cv2.CV_64F, 0, 1, ksize=3)
+    grad_mag = np.sqrt(grad_x**2 + grad_y**2)
+    
+    # 5. 【统计学阈值】不再死守固定数值，而是选 Top X%
+    threshold = np.percentile(grad_mag, EDGE_TOP_PERCENTILE)
+    binary_edges = (grad_mag >= threshold).astype(np.uint8) * 255
+    
+    # 6. 形态学闭合：桥接裂缝，让线条连起来
+    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (MORPH_SIZE, MORPH_SIZE))
+    closed = cv2.morphologyEx(binary_edges, cv2.MORPH_CLOSE, kernel)
+    
+    # 7. 再次轻微膨胀，为 SAM 提供更好的引导范围
+    final_mask = cv2.dilate(closed, kernel, iterations=1)
+    
+    return final_mask
+
+def main() -> None:
+    base_dir = Path(__file__).resolve().parent
+    img_path = Path(INPUT_IMAGE)
+    out_dir = base_dir / OUTPUT_DIR
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    # 1. 初始化深度模型
+    predictor, used_backend = build_depth_predictor(UnifiedDepthConfig(backend=DEPTH_BACKEND))
+    
+    # 2. 加载图像
+    print(f"[Loading] 正在处理: {img_path.name}")
+    img_pil = Image.open(img_path).convert("RGB")
+    w_orig, h_orig = img_pil.size
+    
+    # 3. 深度预测
+    print(f"[Depth] 正在进行深度估计 (Large Image)...")
+    depth = np.asarray(predictor(img_pil), dtype=np.float32).squeeze()
+    
+    # 4. 归一化与保存
+    dmin, dmax = depth.min(), depth.max()
+    depth_norm = (depth - dmin) / (dmax - dmin + 1e-8)
+    depth_u16 = (depth_norm * 65535.0).astype(np.uint16)
+    Image.fromarray(depth_u16).save(out_dir / f"{img_path.stem}.depth.png")
+
+    # 5. 提取强鲁棒性边缘
+    print(f"[Edge] 正在应用 CLAHE + Sobel 增强算法提取边缘...")
+    edge_mask = _extract_robust_edges(depth_norm)
+    
+    # 6. 导出
+    mask_path = out_dir / f"{img_path.stem}.edge_mask_robust.png"
+    Image.fromarray(edge_mask).save(mask_path)
+
+    edge_ratio = float((edge_mask > 0).sum()) / float(edge_mask.size)
+    print("-" * 30)
+    print(f"提取完成！")
+    print(f"边缘密度: {edge_ratio:.2%} (目标通常应在 1% ~ 8% 之间)")
+    print(f"如果 Mask 依然太黑，请调低 EDGE_TOP_PERCENTILE (如 90.0)")
+    print(f"如果 Mask 太乱，请调高 EDGE_TOP_PERCENTILE (如 96.0)")
+    print("-" * 30)
+
+if __name__ == "__main__":
+    main()