pmd-python/main.py

import cv2
import os
import sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import time
import glob
import json
import argparse
import pickle
import matplotlib.pyplot as plt

from datetime import datetime
import numpy as np
import scipy.io as sio
from scipy.io import savemat, loadmat
np.random.seed(42)

from src.utils import binarize, get_meshgrid, get_world_points, get_camera_points, get_screen_points, write_point_cloud, \
              get_white_mask, get_world_points_from_mask, get_meshgrid_contour, post_process
from src.phase import extract_phase, unwrap_phase
from src.recons import reconstruction_cumsum
from src.calibration import calibrate_world, calibrate_screen, map_screen_to_world
from src.vis import plot_coords
from src.eval import get_eval_result, find_notch

def list_directories(path):
    # List all items in the given directory and filter only directories
    return [item for item in os.listdir(path) if os.path.isdir(os.path.join(path, item))]

def parse_args():

    parser = argparse.ArgumentParser(description="")
    parser.add_argument("--img_path", type=str, default='D:\\file\\20240913-data\\20240913105234292')  
    parser.add_argument("--camera_path", type=str, default='D:\\file\\标定数据\\标定数据0913\\calibration_0913')
    parser.add_argument("--screen_path", type=str, default='D:\\file\\screen0920')
    parser.add_argument("--cfg", type=str, default="config/cfg_3freq_wafer.json")
    #parser.add_argument("--cfg", type=str, default="D:\code\pmdrecons-python\config\cfg_3freq_wafer.json")
    parser.add_argument("--save_path", type=str, default="debug")
    parser.add_argument("--grid_spacing", type=int, default=1)  
    parser.add_argument("--debug", action='store_true')
    parser.add_argument("--smooth", action='store_true')
    parser.add_argument("--align", action='store_true')
    parser.add_argument("--denoise", action='store_true')
    args = parser.parse_args()
    return args 


def main(cfg, img_folder, camera_path, screen_path, save_path, grid_spacing, debug):
    n_cam = 4
    num_freq = cfg['num_freq']
    start_time = time.time()
    #debug = True
    print(f"开始执行时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("\n1. 相机标定")
    preprocess_start = time.time()
    
    camera_subdir = list_directories(camera_path)
    camera_subdir.sort()
    assert len(camera_subdir) == 4, f"found {len(camera_subdir)} cameras, should be 4"
    cam_para_path = os.path.join(camera_path, "cam_params.pkl")

    if os.path.exists(cam_para_path):
        with open(cam_para_path, 'rb') as pkl_file:
            cam_params = pickle.load(pkl_file)
    else:
        cam_params = []
        for i in range(n_cam):
            cam_img_path = glob.glob(os.path.join(camera_path, camera_subdir[i], "*.bmp"))
            cam_img_path.sort()
            cam_param_raw = calibrate_world(cam_img_path, i, cfg['world_chessboard_size'], cfg['world_square_size'], debug=False)
            cam_params.append(cam_param_raw)
        with open(cam_para_path, 'wb') as pkl_file:
           pickle.dump(cam_params, pkl_file)
   
    preprocess_end = time.time()
    print(f"   完成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"   耗时: {preprocess_end - preprocess_start:.2f} 秒") 
    # import pdb; pdb.set_trace()
 
    print("\n2. 屏幕标定")
    screen_cal_start = time.time()
    screen_img_path = glob.glob(os.path.join(screen_path, "*.bmp"))
    screen_para_path = os.path.join(screen_path, "screen_params.pkl")    
    if os.path.exists(screen_para_path):
        with open(screen_para_path, 'rb') as pkl_file:
            screen_params = pickle.load(pkl_file)[0]
    else:
        screen_params = calibrate_screen(screen_img_path, cam_params[0]['camera_matrix'], cam_params[0]['distortion_coefficients'], cfg['screen_chessboard_size'], cfg['screen_square_size'], debug=True)
        with open(screen_para_path, 'wb') as pkl_file:
            pickle.dump([screen_params], pkl_file)

    screen_to_world = map_screen_to_world(screen_params, cam_params[0])  
    screen_cal_end = time.time()
    print(f"   完成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"   耗时: {screen_cal_end - screen_cal_start:.2f} 秒")
   
    print("\n3. 相位提取，相位展开")
    phase_start = time.time()
    x_uns, y_uns = [], []
    binary_masks = []
    for cam_id in range(n_cam):
        white_path = os.path.join(img_folder, f'{cam_id}_frame_24.bmp')
        binary = get_white_mask(white_path)
        binary_masks.append(binary)
        phases = extract_phase(img_folder, cam_id, binary, cam_params[cam_id]['camera_matrix'], cam_params[cam_id]['distortion_coefficients'], num_freq=num_freq)
        x_un, y_un = unwrap_phase(phases, save_path, num_freq, debug=False)
        x_uns.append(x_un)
        y_uns.append(y_un)
    
    phase_end = time.time()
    print(f"   完成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"   耗时: {phase_end - phase_start:.2f} 秒")
 
    print("\n4. 获得不同坐标系下点的位置")
    get_point_start = time.time()
    total_cloud_point = np.empty((0, 3))
    for i in range(n_cam):
        if i > 5:
            continue
        contours_point = get_meshgrid_contour(binary_masks[i], grid_spacing, save_path, debug=False)
        world_points = get_world_points(contours_point, cam_params[i], grid_spacing, cfg['d'], save_path, debug=debug)
        camera_points, u_p, v_p = get_camera_points(world_points, cam_params[i], save_path, i, debug=debug)
        
        point_data = {'x_w': world_points[:, 0], 'y_w': world_points[:, 1], 'z_w': world_points[:, 2], 
                    'x_c': camera_points[:, 0], 'y_c': camera_points[:, 1], 'z_c': camera_points[:, 2],  
                    'u_p': u_p, 'v_p': v_p}
        
        screen_points = get_screen_points(point_data, x_uns[i], y_uns[i], screen_params, screen_to_world, cfg, save_path, i, debug=debug)
        #plot_coords(world_points, camera_points, screen_points)
        z, smoothed, aligned, denoise = reconstruction_cumsum(world_points, camera_points, screen_points, save_path, i, debug=debug, smooth=args.smooth, align=args.align, denoise=args.denoise)
        write_point_cloud(os.path.join(img_folder, str(i) + '_cloudpoint.txt'), world_points[:, 0], world_points[:, 1], 1000 * denoise[:,2])

        total_cloud_point = np.vstack([total_cloud_point, np.column_stack((denoise[:, 0], denoise[:, 1], 1000 * denoise[:,2]))])
        print('point cloud has been written in file')

    write_point_cloud(os.path.join(img_folder, 'cloudpoint.txt'), total_cloud_point[:, 0], total_cloud_point[:, 1], total_cloud_point[:,2])
    
    if debug:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')

        # 提取 x, y, z 坐标
        x_vals = total_cloud_point[:, 0]
        y_vals = total_cloud_point[:, 1]
        z_vals = total_cloud_point[:, 2]

        # 绘制3D点云
        ax.scatter(x_vals, y_vals, z_vals, c=z_vals, cmap='viridis', marker='o')

        # 设置轴标签和标题
        ax.set_xlabel('X (mm)')
        ax.set_ylabel('Y (mm)')
        ax.set_zlabel('Z (mm)')
        ax.set_title('3D Point Cloud Visualization')

        plt.show()
    get_point_end = time.time()
    print(f"   完成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"   耗时: {get_point_end - get_point_start:.2f} 秒")
    
    print("\n5. 后处理")
    post_process_start = time.time()
    post_process(args.img_path, debug)
    post_process_end = time.time()
    print(f"   完成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"   耗时: {post_process_end - post_process_start:.2f} 秒")

    print("\n6. 评估")
    eval_start = time.time()
    point_cloud_path = os.path.join(img_folder, str(i) + '_cloudpoint.txt')
    json_path = os.path.join(img_folder, str(i) + 'result.json')
    theta_notch = 0
    get_eval_result(point_cloud_path, json_path, theta_notch)
    post_process(args.img_path, debug)
    eval_end = time.time()
    print(f"   完成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"   耗时: {eval_end - eval_start:.2f} 秒")


    return True
        

if __name__ == '__main__': 

    start_time = time.time()
    
    args = parse_args()
    img_folder = args.img_path
    
    
    cfg = json.load(open(args.cfg, 'r'))
    if not os.path.exists(args.save_path):
        os.makedirs(args.save_path)
    
    args.smooth = True
    args.align = True
    args.denoise = True
    args.debug = 0

    main(cfg, img_folder, args.camera_path, args.screen_path, args.save_path, args.grid_spacing, args.debug)

    end_time = time.time()
    print(f"总运行时间: {end_time - start_time:.2f} 秒")