pmd-python/pmd.py

#encoding=utf-8
import time
import os
import sys
import pandas as pd


sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import time
import glob
import numpy as np
np.random.seed(42)
from datetime import datetime
import json
from src.utils import get_meshgrid, get_world_points, get_camera_points, get_screen_points, write_point_cloud, get_white_mask, get_meshgrid_contour, post_process, find_notch, post_process_with_grad
from src.phase import extract_phase, unwrap_phase
from src.recons import reconstruction_cumsum, poisson_recons_with_smoothed_gradients
from src.pcl_postproc import smooth_pcl, align2ref
import matplotlib.pyplot as plt
from src.calibration import calibrate_world, calibrate_screen, map_screen_to_world
import argparse
from src.vis import plot_coords
import cv2
from src.eval import get_eval_result
import pickle
from collections import defaultdict
from scipy.io import loadmat, savemat



def pmdstart(config_path, img_folder):
    start_time = time.time()
    print(f"config_path: {config_path}")
    #time.sleep(15)
    main(config_path, img_folder)
    print(f"img_folder: {img_folder}")
    print('test pass')
    end_time = time.time()
    print(f"Time taken: {end_time - start_time} seconds")
    return True


def main(config_path, img_folder):
    current_dir = os.path.dirname(os.path.abspath(__file__))
    os.chdir(current_dir)

    cfg = json.load(open(config_path, 'r'))
    n_cam = cfg['cam_num']
    num_freq = cfg['num_freq']
    save_path = 'debug'
    debug = False
    grid_spacing = cfg['grid_spacing']
    num_freq = cfg['num_freq']
    smooth = False
    align = False
    denoise = False
    #cammera_img_path = 'D:\\data\\four_cam\\calibrate\\calibrate-1008'
    screen_img_path = 'D:\\data\\four_cam\\calibrate\\cam3-screen-1008'
    cammera_img_path = 'D:\\data\\four_cam\\calibrate\\calibrate-1016'
    #screen_img_path = 'D:\\data\\four_cam\\calibrate\\screen0920'

    print(f"开始执行时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("\n1. 相机标定")
    preprocess_start = time.time()

    cam_para_path = os.path.join(current_dir, 'config', cfg['cam_params'])
    print('cam_para_path = ', cam_para_path)
    print('current_dir = ', current_dir)
    #cam_para_path = 'D:\\code\\pmd-python\\config\\cam_params.pkl'
    if os.path.exists(cam_para_path):
    #if False:
        with open(cam_para_path, 'rb') as pkl_file:
            cam_params = pickle.load(pkl_file)
    else:
        cam_params = []
        camera_subdir = [item for item in os.listdir(cammera_img_path) if os.path.isdir(os.path.join(cammera_img_path, item))]
        camera_subdir.sort()
        assert len(camera_subdir) == 4, f"found {len(camera_subdir)} cameras, should be 4"
        for i in range(n_cam):
            cam_img_path = glob.glob(os.path.join(cammera_img_path, camera_subdir[i], "*.bmp"))
            cam_img_path.sort()
            #print('cam_img_path = ', cam_img_path)
            cam_param_raw = calibrate_world(cam_img_path, i, cfg['world_chessboard_size'], cfg['world_square_size'], debug=0)
            cam_params.append(cam_param_raw)
        # with open(cam_para_path, 'wb') as pkl_file:
        #    pickle.dump(cam_params, pkl_file) 

    print("\n2. 屏幕标定")
    screen_cal_start = time.time()

    screen_img_path = glob.glob(os.path.join(screen_img_path, "*.bmp"))
    screen_para_path = os.path.join('config', cfg['screen_params'])   
    if os.path.exists(screen_para_path):
    #if False:
        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[3]['camera_matrix'], cam_params[3]['distortion_coefficients'], cfg['screen_chessboard_size'], cfg['screen_square_size'], debug=0)
        # with open(screen_para_path, 'wb') as pkl_file:
        #     pickle.dump([screen_params], pkl_file)
    
    
    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):
        print('cam_id = ', cam_id)
        #white_path = os.path.join(img_folder, f'{cam_id}_frame_24.bmp') 
        white_path = 'D:\\code\\code of PMD\\code of PMD\\picture\\white.bmp'
        if n_cam == 3:
            binary = get_white_mask(white_path, bin_thresh=12, debug=0)
            
        elif n_cam == 1:
            #angle_rad, binary = find_notch(white_path, n_cam, debug=0)
            binary = get_white_mask(white_path, bin_thresh=12, debug=0)

        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=0)
        
        #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} 秒")
    
    # matlab align
    #cam_params = loadmat('D:\\code\\code of PMD\\code of PMD\\calibration\\calibrationSessionworld.mat')
    #screen_params = loadmat('D:\\code\\code of PMD\\code of PMD\\calibration\\calibrationSessionscreen.mat')
   
    #print('cam_params = ', cam_params)
    #print('screen_params = ', screen_params)

    cam_params = []
    screen_params = []
    mtx = np.array([[6944.89018564351, 0, 2709.44699784446], [0, 6942.91962497637, 1882.05677580185], [0,0,1]])
    R = cv2.Rodrigues(np.array([0.0732148059333282,-0.265812028310130,-0.0532640604086260]).reshape(3,1))[0]
    T = np.array([-15.179977,-42.247126,246.92182]).reshape(3,1)
    
    cam_calibration_data = {
        'camera_matrix': mtx,
        'distortion_coefficients': 0,
        'rotation_matrix': R,
        'translation_vector': T,
        'error': 0
    }

    mtx = np.array([[6944.89018564351, 0, 2709.44699784446], [0, 6942.91962497637, 1882.05677580185], [0,0,1]])
    R = np.array([[0.96514996,-0.042578806,0.25821037],[0.029285983,0.99805061,0.055111798],[-0.26005361,-0.045629206,0.96451547]])
    T = np.array([30.1970,-49.3108,507.4424]).reshape(3,1)
    
    screen_calibration_data = {
        'screen_matrix': mtx,
        'screen_distortion': 0,
        'screen_rotation_matrix': R,
        'screen_translation_vector': T,
        'screen_error': 0
    }

    
    cam_params.append(cam_calibration_data)
    #screen_params.append(screen_calibration_data)
    screen_params = screen_calibration_data

    x_un = loadmat('D:\\code\\code of PMD\\code of PMD\\phase_information\\x_phase_unwrapped.mat')
    y_un = loadmat('D:\\code\\code of PMD\\code of PMD\\phase_information\\y_phase_unwrapped.mat')

    x_un = x_un['x_phase_unwrapped']
    y_un = y_un['y_phase_unwrapped']

    x_uns, y_uns = [], []
    x_uns.append(x_un)
    y_uns.append(y_un)
   
    #fig, axes = plt.subplots(1, 2, figsize=(12, 6))


        # 第一个子图
    # cax0 = axes[0].imshow(x_un)
    # axes[0].set_title('x_phase_unwrapped')
    # axes[0].set_xlabel('X Axis')
    # axes[0].set_ylabel('Y Axis')
    # fig.colorbar(cax0, ax=axes[0])

    #     # 第二个子图
    # cax1 = axes[1].imshow(y_un)
    # axes[1].set_title('y_phase_unwrapped')
    # axes[1].set_xlabel('X Axis')
    # axes[1].set_ylabel('Y Axis')
    # fig.colorbar(cax0, ax=axes[1])

    #     # 调整子图之间的间距
    # plt.tight_layout()
        #plt.savefig(os.path.join(save_path, "phase_unwrapped.png"))
    #plt.show()
    
    #return 0
    print('screen_params = ', screen_params)

    if n_cam == 1:
        screen_to_world = map_screen_to_world(screen_params, cam_params[0]) 
    elif n_cam == 4:
        screen_to_world = map_screen_to_world(screen_params, cam_params[3])  
    else:
        print('camera number should be 1 or 4')
        return 0
    

    print("\n4. 获得不同坐标系下点的位置")
    get_point_start = time.time()
    total_cloud_point = np.empty((0, 3))
    total_gradient = np.empty((0, 4))
    total_boundary_point = np.empty((0, 3))
    for i in range(n_cam):
        print('cam_id = ', i)
        contours_point = get_meshgrid_contour(binary_masks[i], save_path, debug=False)
        world_points, world_points_boundary, world_points_boundary_3 = get_world_points(contours_point, cam_params[i], i, grid_spacing, cfg['d'], erosion_pixels=2, debug=0)
        camera_points, u_p, v_p = get_camera_points(world_points, cam_params[i], save_path, i, debug=0)
        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_raw,  gradient_xy = reconstruction_cumsum(world_points, camera_points, screen_points, debug=0, smooth=smooth, align=align, denoise=denoise)
        
        z_raw_xy = np.round(z_raw[:, :2]).astype(int)
        
        # # 创建布尔掩码，初始为 True
        # mask = np.ones(len(z_raw_xy), dtype=bool)

        # # 遍历每个边界点，标记它们在 aligned 中的位置
        # for boundary_point in world_points_boundary:
        #     # 标记与当前边界点相同 xy 坐标的行
        #     mask &= ~np.all(z_raw_xy == boundary_point[:2], axis=1)
        
        # # 使用掩码过滤出非边界点
        # non_boundary_points = z_raw[mask]
        # non_boundary_aligned, rotation_matrix = align2ref(non_boundary_points)
        
        # 创建布尔掩码，初始为 True
        mask = np.ones(len(z_raw_xy), dtype=bool)

        # 遍历每个边界点，标记它们在 aligned 中的位置
        for boundary_point in world_points_boundary_3:
            # 标记与当前边界点相同 xy 坐标的行
            mask &= ~np.all(z_raw_xy == boundary_point[:2], axis=1)
        
        # 使用掩码过滤出非边界点
        non_boundary_points = z_raw[mask] 
        # z_raw_aligned = non_boundary_points @ rotation_matrix.T
        # z_raw_aligned[:,2] = z_raw_aligned[:,2] - np.mean(z_raw_aligned[:, 2])

        #z_raw_aligned = non_boundary_points
        z_raw_aligned, _  = align2ref(non_boundary_points)

        #non_boundary_points = smoothed
        write_point_cloud(os.path.join(img_folder, str(i) + '_cloudpoint.txt'), np.round(z_raw_aligned[:, 0]), np.round(z_raw_aligned[:, 1]),  z_raw_aligned[:, 2])
        np.savetxt(os.path.join(img_folder, str(i) + '_gradient.txt'), gradient_xy, fmt='%.10f', delimiter=',')
        total_cloud_point = np.vstack([total_cloud_point, np.column_stack((z_raw_aligned[:, 0], z_raw_aligned[:, 1], z_raw_aligned[:, 2]))])
        total_gradient = np.vstack([total_gradient, np.column_stack((gradient_xy[:, 0], gradient_xy[:, 1], gradient_xy[:, 2], gradient_xy[:, 3]))])
    
    if 0:
        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('z_raw 3D Point Cloud Visualization gradient')
        plt.show()

        # fig = plt.figure()
        # ax = fig.add_subplot(111, projection='3d')
        # smoothed_total = smooth_pcl(total_cloud_point, 3)
        #  # 提取 x, y, z 坐标
        # x_vals = smoothed_total[:, 0]
        # y_vals = smoothed_total[:, 1] 
        # z_vals = smoothed_total[:, 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('smoothed 3D Point Cloud Visualization')
       

    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()
    total_cloud_point[:,0] = np.round(total_cloud_point[:,0])
    total_cloud_point[:,1] = np.round(total_cloud_point[:,1])
    #fitted_points = post_process(total_cloud_point, debug=0)
    fitted_points = post_process_with_grad(img_folder, n_cam, 1)
    #fitted_points = total_cloud_point
    #align_fitted, _ = align2ref(fitted_points)
    align_fitted = fitted_points
    
    write_point_cloud(os.path.join(img_folder, 'cloudpoint.txt'), np.round(align_fitted[:, 0]-np.mean(align_fitted[:, 0])), np.round(align_fitted[:, 1]-np.mean(align_fitted[:, 1])), align_fitted[:,2]-np.min(align_fitted[:,2]))
    if 1:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')

        # 提取 x, y, z 坐标
        x_vals = np.round(fitted_points[:, 0]-np.mean(fitted_points[:, 0]))
        y_vals = np.round(fitted_points[:, 1]-np.mean(fitted_points[:, 1]))
        z_vals = align_fitted[:,2]-np.min(align_fitted[:,2])

        x_vals = fitted_points[:, 0]
        y_vals = fitted_points[:, 1]
        z_vals = fitted_points[:, 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('post 3D Point Cloud Visualization')
        plt.show()
    
    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, 'cloudpoint.txt')
    json_path = os.path.join(img_folder, 'result.json')
    theta_notch = 0
    #get_eval_result(point_cloud_path, json_path, theta_notch, 0)
    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__':
    config_path = 'config\\cfg_3freq_wafer_matlab.json'
    #img_folder = 'D:\\data\\four_cam\\1008_storage\\pingjing_20241017092315228\\'   #'D:\\data\\four_cam\\betone_1011\\20241011142348762-1'
   # img_folder = 'D:\\huchao\\inspect_server_202409241013_py\\storage\\20241023193453708\\'
    img_folder = 'D:\\data\\one_cam\\pingjing_20241024154405250\\'
    #img_folder = 'D:\\data\\one_cam\\betone-20241025095352783\\'
    #img_folder = 'D:\\data\\one_cam\\20241025113857041-neg\\'
    #'
    json_path = os.path.join(img_folder, 'result.json')

    # 20241011142348762-1  20241011142901251-2   20241011143925746-3   20241011144821292-4 
    #     0.60                0.295                0.221                  0.346

    # x max =  653.5925
    # y max =  692.1735648
    # x max =  251.0775
    # y max =  293.05856482
    # x max =  184.7275
    # y max =  239.00919669
    # x max =  342.2525
    # y max =  386.92087185

    pmdstart(config_path, img_folder)
    #fitted_points = post_process_with_grad(img_folder, 1)