ryanhlewis/Pointcloud_Alignment

1import gradio as gr2import open3d as o3d3import numpy as np4import time5import multiprocessing6import os7import tempfile8import trimesh9 10alignment_cache = {}11current_orientation = "+z"12 13def compute_alignment(files):14    """Align point clouds and cache results."""15    ck = tuple(files)16    if ck in alignment_cache: return alignment_cache[ck]17    if len(files) < 2: return (None, None)18    st = time.time()19    print(f"Aligning {len(files)} point clouds using {multiprocessing.cpu_count()} CPU cores")20    clouds, names = [], []21    for p in files:22        c = o3d.io.read_point_cloud(p)23        if not c.has_points(): return (None, f"Failed to load point cloud from {os.path.basename(p)}")24        clouds.append(c), names.append(os.path.basename(p))25    ref = clouds[0]26    trans = [np.eye(4)]27    for i, src in enumerate(clouds[1:], 1):28        vs = [0.5, 0.3, 0.1]29        t = np.eye(4)30        for v in vs:31            s_down, r_down = src.voxel_down_sample(v), ref.voxel_down_sample(v)32            rn = v * 233            s_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=rn, max_nn=30))34            r_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=rn, max_nn=30))35            d = v * 236            mi = 10 if i == 0 else 537            r = o3d.pipelines.registration.registration_icp(38                s_down, r_down, d, t,39                o3d.pipelines.registration.TransformationEstimationPointToPlane(),40                o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=mi)41            )42            t = r.transformation43        trans.append(t)44    aligned = []45    for c, t in zip(clouds, trans):46        cc = o3d.geometry.PointCloud(c)47        cc.transform(t)48        aligned.append(cc)49    ap = [np.asarray(x.points) for x in aligned]50    if not ap: return (None, "Failed to process point clouds")51    cat = np.vstack(ap)52    mn, mx = np.min(cat, 0), np.max(cat, 0)53    ctr = (mn + mx) / 254    scl = np.max(mx - mn)55    combined = o3d.geometry.PointCloud()56    for c in aligned: combined += c57    tm = time.time() - st58    msg = f"Successfully aligned {len(files)} point clouds in {tm:.2f} seconds."59    data = {60        'aligned_clouds': aligned, 'transformations': trans, 'clouds': clouds,61        'filenames': names, 'global_center': ctr, 'global_scale': scl,62        'combined_cloud': combined63    }64    alignment_cache[ck] = (data, msg)65    return data, msg66 67def visualize_alignment(files, show_colors=False, show_outlines=True, point_density=10000,68                        orientation="+z", point_size=0.005, is_initial_view=False):69    if not files or len(files) < 2: return (None, "Please upload at least 2 point cloud files (.ply)")70    res, msg = compute_alignment(files)71    if not res: return (None, msg)72    clouds = res['aligned_clouds']73    ctr, scl = res['global_center'], res['global_scale']74    cols = [[1,0,0],[0,1,0],[0,0,1],[1,1,0],[1,0,1],[0,1,1],[1,0.5,0]]75    tdir = tempfile.mkdtemp()76    global current_orientation77    current_orientation = orientation78    normed, boxes = [], []79    for i, c in enumerate(clouds):80        cp = o3d.geometry.PointCloud(c)81        if show_colors: cp.paint_uniform_color(cols[i % len(cols)])82        pts = np.asarray(cp.points)83        pts = (pts - ctr)/(scl/2)84        n = o3d.geometry.PointCloud()85        n.points = o3d.utility.Vector3dVector(pts)86        if cp.has_colors(): n.colors = cp.colors87        normed.append(n)88        if show_outlines:89            b = n.get_axis_aligned_bounding_box()90            b.color = cols[i % len(cols)]91            boxes.append(b)92    comb = o3d.geometry.PointCloud()93    for c in normed: comb += c94    if point_density < len(comb.points):95        bb = comb.get_axis_aligned_bounding_box()96        vol = np.prod(bb.get_extent())97        ppu = point_density / vol if vol > 0 else 198        vs = max(0.001, (1/ppu)**(1/3))99        comb = comb.voxel_down_sample(vs)100    vs, fs, vc = [], [], []101    pts = np.asarray(comb.points)102    if comb.has_colors(): pcols = np.asarray(comb.colors)103    else: pcols = np.tile([0.8,0.8,0.8], (len(pts),1))104    sph = trimesh.creation.icosphere(subdivisions=1, radius=point_size)105    mxp = min(point_density, len(pts))106    idxs = np.linspace(0, len(pts)-1, mxp, dtype=int)107    for i, idx in enumerate(idxs):108        p = pts[idx]109        c = pcols[idx]110        s = sph.copy()111        s.apply_translation(p)112        s.visual.vertex_colors = np.tile((c*255).astype(np.uint8), (len(s.vertices),1))113        si = len(vs)114        vs.extend(s.vertices)115        fs.extend(s.faces + si)116        vc.extend(s.visual.vertex_colors)117    cm = trimesh.Trimesh(vertices=vs, faces=fs, vertex_colors=vc)118    if show_outlines:119        box_edges = [(0,1),(1,2),(2,3),(3,0),(4,5),(5,6),(6,7),(7,4),120                     (0,4),(1,5),(2,6),(3,7)]121        for i,b in enumerate(boxes):122            mb, xb = b.min_bound, b.max_bound123            bv = [[mb[0],mb[1],mb[2]],[xb[0],mb[1],mb[2]],[xb[0],xb[1],mb[2]],[mb[0],xb[1],mb[2]],124                  [mb[0],mb[1],xb[2]],[xb[0],mb[1],xb[2]],[xb[0],xb[1],xb[2]],[mb[0],xb[1],xb[2]]]125            bc = cols[i % len(cols)]126            r = point_size*0.8127            for s0,s1 in box_edges:128                p0, p1 = bv[s0], bv[s1]129                d = np.array(p1) - np.array(p0)130                ln = np.linalg.norm(d)131                if ln<1e-6: continue132                d /= ln133                za = np.array([0,0,1])134                if abs(np.dot(d,za))>0.999: ra = np.array([1,0,0])135                else:136                    ra = np.cross(za,d)137                    ra /= np.linalg.norm(ra)138                ang = np.arccos(np.dot(za,d))139                rot = trimesh.transformations.rotation_matrix(ang,ra)140                trn = trimesh.transformations.translation_matrix(p0)141                cyl = trimesh.creation.cylinder(radius=r,height=ln,sections=8)142                cyl.apply_translation([0,0,ln/2])143                cyl.apply_transform(rot)144                cyl.apply_transform(trn)145                cyl.visual.face_colors = (np.array(bc)*255).astype(np.uint8)146                cm = trimesh.util.concatenate([cm,cyl])147    sc = trimesh.Scene(cm)148    sc.add_geometry(trimesh.creation.axis(origin_size=0.01, axis_radius=0.0025))149    if orientation == "+y":150        sc.apply_transform(trimesh.transformations.rotation_matrix(-np.pi/2,[1,0,0]))151    elif orientation == "-y":152        sc.apply_transform(trimesh.transformations.rotation_matrix(np.pi/2,[1,0,0]))153    elif orientation == "-z":154        sc.apply_transform(trimesh.transformations.rotation_matrix(np.pi,[1,0,0]))155    fo = os.path.join(tdir,"aligned_scene.glb")156    sc.export(fo)157    return (fo, msg)158 159def process_upload(files, sc, so, pd, o, ps):160    d, m = compute_alignment(files)161    if d is None: return (None, m)162    return visualize_alignment(files, sc, so, pd, o, ps, True)163 164def update_visualization_only(files, sc, so, pd, o, ps):165    if not files or len(files)<2: return None166    if tuple(files) not in alignment_cache: return None167    mo, _ = visualize_alignment(files, sc, so, pd, o, ps, False)168    return mo169 170with gr.Blocks(theme=gr.themes.Base()) as app:171    gr.Markdown("# Point Cloud Alignment Tool  -  3D ICP")172    with gr.Row():173        with gr.Column(scale=1):174            file_input = gr.File(file_count="multiple", file_types=[".ply"], label="Upload Point Cloud Files (.ply)", type="filepath")175            with gr.Row():176                show_colors = gr.Checkbox(label="Show Colored Models", value=False)177                show_outlines = gr.Checkbox(label="Show Bounding Box Outlines", value=True)178            point_density = gr.Slider(1000, 500000, 10000, step=1000, label="Point Density (fewer points = faster rendering)")179            orientation = gr.Dropdown(["+z","-z","+y","-y"], value="+z", label="Model Orientation (up direction)")180            point_size = gr.Slider(0.001, 0.05, 0.005, step=0.001, label="Point Size")181            submit_btn = gr.Button("Align Point Clouds", variant="primary")182        with gr.Column(scale=2):183            output_model = gr.Model3D(label="Aligned Point Clouds", clear_color=[0.1,0.1,0.1,1.0])184            output_text = gr.Textbox(label="Status")185    submit_btn.click(process_upload,186                     [file_input, show_colors, show_outlines, point_density, orientation, point_size],187                     [output_model, output_text])188    show_colors.change(update_visualization_only, [file_input, show_colors, show_outlines, point_density, orientation, point_size], [output_model])189    show_outlines.change(update_visualization_only, [file_input, show_colors, show_outlines, point_density, orientation, point_size], [output_model])190    point_density.change(update_visualization_only, [file_input, show_colors, show_outlines, point_density, orientation, point_size], [output_model])191    orientation.change(update_visualization_only, [file_input, show_colors, show_outlines, point_density, orientation, point_size], [output_model])192    point_size.change(update_visualization_only, [file_input, show_colors, show_outlines, point_density, orientation, point_size], [output_model])193    194if __name__ == "__main__":195    print(f"Using {multiprocessing.cpu_count()} CPU cores for parallel processing")196    app.launch(share=True)197