sort_and_merge.py

import os
from pathlib import Path
import random

from tqdm import tqdm
import argparse
import numpy as np
import shutil
import ray
import glob

from eval.check_valid import check_step_valid_soild
from diffusion.utils import *
from OCC.Core.BRepLProp import BRepLProp_SLProps
from OCC.Core.BRepGProp import brepgprop
from OCC.Core.TopAbs import TopAbs_SOLID
import trimesh

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"


def random_rgba():
    return tuple(np.random.randint(100, 256, 3).tolist() + [255])


def idx_to_rgba(idx):
    r = idx // 100
    g = (idx // 10) % 10
    b = idx % 10
    return (r + 100, g + 100, b + 100, 255)


def rgb_to_idx(rgb):
    r, g, b = rgb
    return (r - 100) * 100 + (g - 100) * 10 + (b - 100)


def normalize_mesh(mesh):
    bounds = mesh.bounds
    scale = 2.0 / (bounds[1] - bounds[0]).max()
    mesh.apply_scale(scale)
    mesh.apply_translation(-mesh.centroid)
    return mesh


def arrange_meshes(files, out_path, intervals=0.5, color_mode="index", is_normalized=False):
    assert color_mode in ["random", "index"]
    if type(files[0]) == str:
        if is_normalized:
            meshes = [normalize_mesh(trimesh.load(file)) for file in files]
        else:
            meshes = [trimesh.load(file) for file in files]
    elif type(files[0]) == trimesh.Trimesh:
        if is_normalized:
            meshes = [normalize_mesh(mesh) for mesh in files]
        else:
            meshes = [mesh for mesh in files]
    else:
        raise ValueError("Invalid input type")
    num_meshes = len(meshes)

    grid_size = int(np.ceil(np.sqrt(num_meshes)))
    combined = []
    for idx, mesh in enumerate(meshes):
        row = idx // grid_size
        col = idx % grid_size
        translation = [col * (2 + intervals), -row * (2 + intervals), 0]
        if color_mode == "index":
            mesh.visual.face_colors = np.array([idx_to_rgba(idx)] * len(mesh.faces))
        elif color_mode == "random":
            mesh.visual.face_colors = np.array([random_rgba()] * len(mesh.faces))
        else:
            raise ValueError("Invalid color mode")
        mesh.apply_translation(translation)
        combined.append(mesh)

    combined_mesh = trimesh.util.concatenate(combined)
    combined_mesh.export(out_path)


def arrange_meshes_row(file_paths, out_path, intervals=0.5, color_mode="random"):
    assert color_mode in ["random", "index"]
    meshes = [normalize_mesh(trimesh.load(file)) for file in file_paths]
    num_meshes = len(meshes)

    combined = []
    for idx, mesh in enumerate(meshes):
        translation = [idx * (2 + intervals), 0, 0]
        if color_mode == "index":
            mesh.visual.face_colors = np.array([idx_to_rgba(idx)] * len(mesh.faces))
        elif color_mode == "random":
            mesh.visual.face_colors = np.array([random_rgba()] * len(mesh.faces))
        else:
            raise ValueError("Invalid color mode")
        mesh.apply_translation(translation)
        combined.append(mesh)

    combined_mesh = trimesh.util.concatenate(combined)
    combined_mesh.export(out_path)


def explore_primitive(shape, primitive):
    primitive_list = []
    explorer = TopExp_Explorer(shape, primitive)
    while explorer.More():
        primitive_list.append(explorer.Current())
        explorer.Next()

    return primitive_list


def compute_solid_complexity(file_path, num_samples=4):
    try:
        shape = read_step_file(file_path, as_compound=False, verbosity=False)
    except:
        return {"is_valid_solid": False, "mean_curvature": -1, "num_faces": -1, "num_edges": -1, "num_vertices": -1}

    is_valid = check_step_valid_soild(file_path)

    sample_point_curvature = []

    face_list = explore_primitive(shape, TopAbs_FACE)
    edge_list = explore_primitive(shape, TopAbs_EDGE)
    vetex_list = explore_primitive(shape, TopAbs_VERTEX)

    for face in face_list:
        surf_adaptor = BRepAdaptor_Surface(face)
        u_min, u_max, v_min, v_max = (surf_adaptor.FirstUParameter(), surf_adaptor.LastUParameter(), surf_adaptor.FirstVParameter(),
                                      surf_adaptor.LastVParameter())

        u_samples = np.linspace(u_min, u_max, int(np.sqrt(num_samples)))
        v_samples = np.linspace(v_min, v_max, int(np.sqrt(num_samples)))

        for u in u_samples:
            for v in v_samples:
                props = BRepLProp_SLProps(surf_adaptor, u, v, 2, 1e-6)
                if props.IsCurvatureDefined():
                    mean_curvature = props.MeanCurvature()
                    sample_point_curvature.append(abs(mean_curvature))

    if len(sample_point_curvature) > 0:
        mean_curvature = np.mean(sample_point_curvature)
    else:
        mean_curvature = 0.0

    if mean_curvature < 1e-4:
        mean_curvature = 0.0

    volume_props = GProp_GProps()
    surface_props = GProp_GProps()
    brepgprop.VolumeProperties(shape, volume_props)
    brepgprop.SurfaceProperties(shape, surface_props)

    complexity = len(face_list) + surface_props.Mass() / volume_props.Mass() + mean_curvature

    return {"is_valid_solid": is_valid, "mean_curvature": mean_curvature,
            "num_faces"     : len(face_list), "num_edges": len(edge_list), "num_vertices": len(vetex_list),
            "complexity"    : complexity}


compute_solid_complexity_remote = ray.remote(compute_solid_complexity)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_root", type=str, required=True)
    parser.add_argument("--out_root", type=str, required=False)
    parser.add_argument("--src_root", type=str, required=False)
    parser.add_argument("--random", action='store_true')
    parser.add_argument("--sort", action='store_true')
    parser.add_argument("--sample_num", type=int, default=100)
    parser.add_argument("--use_ray", action='store_true')
    parser.add_argument("--valid", action='store_true')
    parser.add_argument("--index", action='store_true')
    args = parser.parse_args()
    data_root = args.data_root
    out_root = args.out_root
    src_root = None if not args.src_root else Path(args.src_root)
    random = args.random
    sort = args.sort
    sample_num = args.sample_num
    use_ray = args.use_ray
    onlyvalid = args.valid

    if not out_root:
        out_root = data_root + "_choose"

    # can only choose one of random, sort, seg
    if sum([random, sort]) != 1:
        raise ValueError("Only and must set one of random, sort")

    if os.path.exists(out_root):
        shutil.rmtree(out_root)
    os.makedirs(out_root)
    folder_names = [f for f in os.listdir(data_root) if os.path.isdir(os.path.join(data_root, f))]
    # folder_names = folder_names[:10]
    if random:
        sample_folder = random.sample(folder_names, sample_num)
        for folder in tqdm(sample_folder):
            shutil.copytree(str(os.path.join(data_root, folder)), str(os.path.join(out_root, folder)))
        exit(0)

    # accumulate the data(mean_curvature, num_faces, num_edges, num_vertices) of each sample
    folder_names.sort()
    folder_scores = {}
    if not use_ray:
        pbar = tqdm(folder_names)
        for folder in pbar:
            pbar.set_description(f"Processing {folder}")
            file_path = glob.glob(os.path.join(data_root, folder, "*.step"))
            file_path.sort()
            if len(file_path) == 0:
                continue
            file_path = file_path[0]
            score = compute_solid_complexity(file_path)
            folder_scores[folder] = score
    else:
        ray.init(
                local_mode=False,
        )
        futures = []
        futures_folder_names = []
        for folder in tqdm(folder_names):
            file_path = glob.glob(os.path.join(data_root, folder, "*.step"))
            if len(file_path) == 0:
                continue
            file_path = file_path[0]
            futures.append(compute_solid_complexity_remote.remote(file_path))
            futures_folder_names.append(folder)

        for idx in tqdm(range(len(futures))):
            folder_name = futures_folder_names[idx]
            result = ray.get(futures[idx])
            folder_scores[folder_name] = result

    if onlyvalid:
        valid_folder_scores = {}
        for folder, score in folder_scores.items():
            if score['is_valid_solid']:
                valid_folder_scores[folder] = score
        folder_scores = valid_folder_scores

    if sort:
        # sort the folders based on the mean_curvature, num_faces, num_edges, num_vertices
        sorted_folders = sorted(folder_scores.items(),
                                key=lambda x: (x[1]["num_faces"], x[1]["mean_curvature"],),
                                reverse=True)

        for idx, folder in enumerate(tqdm(sorted_folders)):
            shutil.copytree(str(os.path.join(data_root, folder[0])), str(os.path.join(out_root, f"{idx:05d}_{folder[0]}")))
            if src_root is not None:
                for file in (src_root/folder[0]).glob("*"):
                    if "_pc.ply" in file.name or "_txt.txt" in file.name or ".png" in file.name:
                        shutil.copy(str(file), str(os.path.join(out_root, f"{idx:05d}_{folder[0]}")))
        ray.shutdown()
        mesh_path_list = glob.glob(os.path.join(out_root, "**", "*.stl"), recursive=True)
        mesh_path_list.sort()
        prefix = Path(data_root).name
        name = os.path.join(out_root, "{}.ply".format(prefix))
        if args.index:
            arrange_meshes(mesh_path_list, name, color_mode="index")
        else:
            arrange_meshes(mesh_path_list, name, color_mode="random")
        print(f"arranged mesh is saved to {name}")

    print("Done")