我需要 彩色的,点云重建得好的。代码:import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
import argparse
import trimesh
from tqdm import tqdm
from typing import Tuple
import time
def normalize_disparity_map(disparity_map):
"""Normalize disparity map for visualization, clip negative values."""
disparity_valid = disparity_map[disparity_map > 0]
if len(disparity_valid) == 0:
return np.zeros_like(disparity_map)
d_min, d_max = disparity_valid.min(), disparity_valid.max()
normalized = np.zeros_like(disparity_map, dtype=np.float32)
if d_max > d_min:
normalized = (disparity_map - d_min) / (d_max - d_min)
normalized[disparity_map <= 0] = 0
return normalized
def visualize_disparity_map(disparity_map, gt_map, save_path=None):
"""
Visualize or save the estimated and ground truth disparity maps side-by-side.
Use JET colormap for better visual interpretation (colorful).
"""
# Normalize both maps
disp_norm = normalize_disparity_map(disparity_map)
gt_norm = normalize_disparity_map(gt_map)
# Apply jet colormap -> (H, W, 3)
disp_color = plt.cm.jet(disp_norm)[:, :, :3] # Remove alpha
gt_color = plt.cm.jet(gt_norm)[:, :, :3]
# Concatenate horizontally: [result | GT]
concat_color = np.hstack([disp_color, gt_color])
if save_path is None:
plt.figure(figsize=(12, 5))
plt.imshow(concat_color)
plt.axis('off')
plt.title("Estimated Disparity (left) vs Ground Truth (right)")
plt.show()
else:
os.makedirs(os.path.dirname(save_path), exist_ok=True)
plt.imsave(save_path, concat_color, cmap='jet') # Save as colorful PNG
print(f"Disparity map saved to {save_path}")
def task1_compute_disparity_map_simple(
ref_img: np.ndarray,
sec_img: np.ndarray,
window_size: int,
disparity_range: Tuple[int, int],
matching_function: str
):
"""
Compute disparity map using simple window-based matching.
Uses OpenCV's boxFilter for fast cost aggregation.
Returns a single-channel float32 disparity map.
"""
ref = ref_img.astype(np.float32)
sec = sec_img.astype(np.float32)
H, W = ref.shape
dmin, dmax = disparity_range
if dmin < 0:
raise ValueError("min_disparity should be >= 0.")
disparities = list(range(dmin, dmax + 1))
ksize = (window_size, window_size)
N = float(window_size * window_size)
eps = 1e-6
cost_volume = np.empty((len(disparities), H, W), dtype=np.float32)
for idx, d in enumerate(disparities):
shifted_sec = np.zeros_like(sec)
if d == 0:
shifted_sec = sec.copy()
else:
shifted_sec[:, d:] = sec[:, :-d]
if matching_function == "SSD":
diff = ref - shifted_sec
cost = cv2.boxFilter(diff * diff, -1, ksize, normalize=False, borderType=cv2.BORDER_REFLECT101)
elif matching_function == "SAD":
diff = np.abs(ref - shifted_sec)
cost = cv2.boxFilter(diff, -1, ksize, normalize=False, borderType=cv2.BORDER_REFLECT101)
elif matching_function == "normalized_correlation":
sumI = cv2.boxFilter(ref, -1, ksize, normalize=False, borderType=cv2.BORDER_REFLECT101)
sumJ = cv2.boxFilter(shifted_sec, -1, ksize, normalize=False, borderType=cv2.BORDER_REFLECT101)
sumI2 = cv2.boxFilter(ref * ref, -1, ksize, normalize=False, borderType=cv2.BORDER_REFLECT101)
sumJ2 = cv2.boxFilter(shifted_sec * shifted_sec, -1, ksize, normalize=False, borderType=cv2.BORDER_REFLECT101)
sumIJ = cv2.boxFilter(ref * shifted_sec, -1, ksize, normalize=False, borderType=cv2.BORDER_REFLECT101)
num = sumIJ - (sumI * sumJ) / N
den = np.sqrt(np.maximum(sumI2 - sumI*sumI/N, 0)) * np.sqrt(np.maximum(sumJ2 - sumJ*sumJ/N, 0)) + eps
ncc = num / den
cost = -ncc # Maximize NCC => Minimize (-NCC)
else:
raise ValueError(f"Unknown matching function: {matching_function}")
if d > 0:
cost[:, :d] = np.inf
cost_volume[idx] = cost
best_idx = np.argmin(cost_volume, axis=0)
disparity_map = np.take(np.array(disparities, dtype=np.float32), best_idx)
best_cost = np.min(cost_volume, axis=0)
disparity_map[~np.isfinite(best_cost)] = 0.0
return disparity_map
def task1_simple_disparity(ref_img, sec_img, gt_map, img_name='tsukuba'):
"""
Run Task 1: test different configurations of window size, disparity range, and matching function.
Save all results as colorful disparity images and log runtime.
"""
window_sizes = [5, 9, 15]
disparity_range = (0, 64)
matching_functions = ['SSD', 'SAD', 'normalized_correlation']
disparity_maps = []
print(f"🚀 Starting Task 1: Testing multiple configurations on '{img_name}'...")
for window_size in window_sizes:
for matching_function in matching_functions:
start_time = time.time()
print(f"⚙️ Computing: ws={window_size}, func={matching_function}")
disparity_map = task1_compute_disparity_map_simple(
ref_img, sec_img, window_size, disparity_range, matching_function)
runtime = time.time() - start_time
print(f"⏱️ Done in {runtime:.2f}s")
disparity_maps.append((disparity_map.copy(), window_size, matching_function, disparity_range))
# Save path with full info
dmin, dmax = disparity_range
save_path = f"output/task1_{img_name}_{window_size}_{dmin}_{dmax}_{matching_function}.png"
visualize_disparity_map(disparity_map, gt_map, save_path=save_path)
# Log runtime for report
with open("output/runtime_log.txt", "a") as f:
f.write(f"Task1,{img_name},{window_size},{dmin},{dmax},{matching_function},{runtime:.4f}\n")
return disparity_maps
def task2_compute_depth_map(disparity_map, baseline=0.2, focal_length=615):
"""
Convert disparity to depth using z = fB / d
Ignore zero/negative disparities.
"""
depth_map = np.zeros_like(disparity_map, dtype=np.float32)
valid = disparity_map > 0
depth_map[valid] = (focal_length * baseline) / (disparity_map[valid] + 1e-8)
return depth_map
def task2_visualize_pointcloud(
ref_img_bgr: np.ndarray,
disparity_map: np.ndarray,
save_path: str = 'output/task2_tsukuba.ply'
):
"""
Generate a FULLY COLORFUL 3D point cloud from disparity.
- X,Y: pixel coordinates
- Z: depth derived from disparity
- Color: true RGB from reference image (convert BGR → RGB)
Filters out invalid depths and extreme outliers.
"""
# Calibration parameters (Tsukuba dataset typical values)
baseline = 0.2 # meters
focal_length = 615 # pixels
depth_map = task2_compute_depth_map(disparity_map, baseline, focal_length)
# Remove infinite/nan and clip top 1%
valid = np.isfinite(depth_map) & (depth_map > 0)
if np.sum(valid) == 0:
print("⚠️ No valid depth points found!")
return
# Clip far outliers
max_depth = np.percentile(depth_map[valid], 99)
valid &= (depth_map <= max_depth)
H, W = depth_map.shape
xs, ys = np.meshgrid(np.arange(W), np.arange(H), indexing='xy')
points = np.stack([xs[valid], ys[valid], depth_map[valid]], axis=1)
# Convert BGR to RGB and extract colors
ref_rgb = cv2.cvtColor(ref_img_bgr, cv2.COLOR_BGR2RGB)
colors = ref_rgb[valid].reshape(-1, 3).astype(np.uint8)
# Create and save PLY
pc = trimesh.PointCloud(vertices=points, colors=colors)
os.makedirs(os.path.dirname(save_path), exist_ok=True)
pc.export(save_path, file_type='ply')
print(f"🎨 Colored point cloud saved to {save_path}")
def task3_compute_disparity_map_dp(ref_img, sec_img):
"""
Dynamic Programming based stereo matching along horizontal scanlines.
Handles occlusions via penalty terms.
"""
ref = ref_img.astype(np.float32)
sec = sec_img.astype(np.float32)
H, W = ref.shape
occlusion_penalty = 20.0
max_disparity = 64
disparity_map = np.zeros((H, W), dtype=np.float32)
t0 = time.perf_counter()
for r in tqdm(range(H), desc="DP Scanline Processing"):
L = ref[r]
R = sec[r]
dp = np.full((W+1, W+1), np.inf, dtype=np.float32)
move = np.zeros((W+1, W+1), dtype=np.uint8) # 0=match, 1=occL, 2=occR
dp[0, 0] = 0.0
for i in range(1, W+1):
dp[i, 0] = dp[i-1, 0] + occlusion_penalty
move[i, 0] = 1
for i in range(1, W+1):
li = L[i-1]
for j in range(1, i+1):
# Match only if within max_disparity
disp = i - j
if disp <= max_disparity:
c_match = dp[i-1, j-1] + abs(li - R[j-1])
else:
c_match = np.inf
c_occL = dp[i-1, j] + occlusion_penalty
c_occR = dp[i, j-1] + occlusion_penalty
if c_match <= c_occL and c_match <= c_occR:
dp[i, j] = c_match
move[i, j] = 0
elif c_occL <= c_occR:
dp[i, j] = c_occL
move[i, j] = 1
else:
dp[i, j] = c_occR
move[i, j] = 2
# Backtrack
i, j = W, W
while i > 0 or j > 0:
m = move[i, j]
if i > 0 and j > 0 and m == 0:
disparity_map[r, i-1] = float(i - j)
i -= 1
j -= 1
elif i > 0 and (j == 0 or m == 1):
disparity_map[r, i-1] = 0.0
i -= 1
else:
j -= 1
t1 = time.perf_counter()
print(f"[Task3] DP runtime: {t1 - t0:.3f}s")
return disparity_map
def main(tasks):
# Load images
try:
moebius_img1 = cv2.imread("data/moebius1.png")
moebius_img1_gray = cv2.cvtColor(moebius_img1, cv2.COLOR_BGR2GRAY).astype(np.float32)
moebius_img2 = cv2.imread("data/moebius2.png")
moebius_img2_gray = cv2.cvtColor(moebius_img2, cv2.COLOR_BGR2GRAY).astype(np.float32)
moebius_gt = cv2.imread("data/moebius_gt.png", cv2.IMREAD_GRAYSCALE).astype(np.float32)
except Exception as e:
print("Moebius data not available:", e)
tsukuba_img1 = cv2.imread("data/tsukuba1.jpg")
tsukuba_img1_gray = cv2.cvtColor(tsukuba_img1, cv2.COLOR_BGR2GRAY).astype(np.float32)
tsukuba_img2 = cv2.imread("data/tsukuba2.jpg")
tsukuba_img2_gray = cv2.cvtColor(tsukuba_img2, cv2.COLOR_BGR2GRAY).astype(np.float32)
tsukuba_gt = cv2.imread("data/tsukuba_gt.jpg", cv2.IMREAD_GRAYSCALE).astype(np.float32)
# Ensure output directory exists
os.makedirs("output", exist_ok=True)
# Clear or init runtime log
with open("output/runtime_log.txt", "a") as f:
f.write("Method,Image,WindowSize,dMin,dMax,MatchingFunction,Runtime(s)\n")
# Task 0: OpenCV Baseline
if '0' in tasks:
print('🔧 Running Task 0: OpenCV StereoBM baseline...')
stereo = cv2.StereoBM.create(numDisparities=64, blockSize=15)
tsukuba_disparity_cv2 = stereo.compute(tsukuba_img1_gray.astype(np.uint8),
tsukuba_img2_gray.astype(np.uint8)).astype(np.float32)
tsukuba_disparity_cv2[tsukuba_disparity_cv2 < 0] = 0
visualize_disparity_map(tsukuba_disparity_cv2, tsukuba_gt,
save_path="output/task0_tsukuba_colormap.png")
if '2' in tasks:
task2_visualize_pointcloud(tsukuba_img1, tsukuba_disparity_cv2,
save_path='output/task2_tsukuba_cv2.ply')
# Task 1: Simple Matching
if '1' in tasks:
print('🔍 Running Task 1: Window-based Matching with Multiple Settings...')
start_time = time.time()
disparity_maps = task1_simple_disparity(tsukuba_img1_gray, tsukuba_img2_gray, tsukuba_gt, img_name='tsukuba')
total_time = time.time() - start_time
print(f"🏁 Task 1 completed in {total_time:.2f}s")
if '2' in tasks:
print('🎨 Generating colored point clouds for each Task 1 result...')
for dm, ws, mf, dr in disparity_maps:
dmin, dmax = dr
path = f'output/task2_tsukuba_{ws}_{dmin}_{dmax}_{mf}.ply'
task2_visualize_pointcloud(tsukuba_img1, dm, save_path=path)
# Task 3: DP Matching
if '3' in tasks:
print('⚡ Running Task 3: Dynamic Programming Matching...')
start_time = time.time()
tsukuba_disparity_dp = task3_compute_disparity_map_dp(tsukuba_img1_gray, tsukuba_img2_gray)
runtime = time.time() - start_time
print(f"⏱️ Task 3 runtime: {runtime:.2f}s")
with open("output/runtime_log.txt", "a") as f:
f.write(f"Task3,tsukuba,DP,0,64,DP,{runtime:.4f}\n")
visualize_disparity_map(tsukuba_disparity_dp, tsukuba_gt,
save_path='output/task3_tsukuba_colormap.png')
if '2' in tasks:
task2_visualize_pointcloud(tsukuba_img1, tsukuba_disparity_dp,
save_path='output/task2_tsukuba_dp.ply')
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Homework 4: Binocular Stereo (FULL COLOR VERSION)')
parser.add_argument('--tasks', type=str, default='0123', help='Tasks to run: e.g., 0, 12, 023')
args = parser.parse_args()
main(args.tasks)
。 严格按照题目要求,不要修改原代码