简介:数据处理中有scale的过程,而计算指标的时候需要进行inverse scale。因此保存scale的参数到json。如果已有json文件,那么可以通过指定json文件路径完成追加。
输入:

  1. in_dir:被计算scale的文件夹
  2. out_dir:保存scale info的文件夹
  3. padding:scale为3Dmesh到-1~1的空隙大小
  4. scale_info_file:是否有现有的json文件

import os
import argparse
import numpy as np
import trimesh
from tqdm import tqdm
import json

class Scale:
    """
    Scales a bunch of meshes.
    """

    def __init__(self):
        """
        Constructor.
        """

        parser = self.get_parser()
        self.options = parser.parse_args()

    def get_parser(self):
        """
        Get parser of tool.

        :return: parser
        """

        parser = argparse.ArgumentParser(description='Scale a set of meshes stored as PLY files.')
        parser.add_argument('--in_dir', type=str, help='Path to input directory.')
        parser.add_argument('--out_dir', type=str, help='Path to output directory; files within are overwritten!')
        parser.add_argument('--padding', type=float, default=0.1, help='Relative padding applied on each side.')
        parser.add_argument('--scale_info_file', type=str, default=None, help='Path to existing scale info JSON file to append to.')
        return parser

    def read_directory(self, directory):
        """
        Read directory.

        :param directory: path to directory
        :return: list of files
        """

        files = []
        for filename in os.listdir(directory):
            files.append(os.path.normpath(os.path.join(directory, filename)))

        return files

    def run(self):
        """
        Run the tool, i.e. scale all found PLY files.
        """

        assert os.path.exists(self.options.in_dir)
        assert os.path.exists(os.path.join(self.options.in_dir, 'tooth_crown'))
        assert os.path.exists(os.path.join(self.options.in_dir, 'jaw'))
         
        os.makedirs(self.options.out_dir, exist_ok=True)
        files = self.read_directory(os.path.join(self.options.in_dir, 'tooth_crown'))

        scale_info = []
        if self.options.scale_info_file and os.path.exists(self.options.scale_info_file):
            with open(self.options.scale_info_file, 'r') as f:
                scale_info = json.load(f)

        for filepath in tqdm(files, desc='Processing scaling: ', total=len(files)):
            if filepath[-4:] != '.ply':
                continue
            mesh = trimesh.load(filepath)
            mesh_upper_jaw = trimesh.load(os.path.join(self.options.in_dir, "jaw", "upper_" + filepath.split('/')[-1]))
            mesh_lower_jaw = trimesh.load(os.path.join(self.options.in_dir, "jaw", "lower_" + filepath.split('/')[-1]))

            # Get extents of model.
            min = mesh.vertices.min(axis=0)
            max = mesh.vertices.max(axis=0)

            # Set the center (although this should usually be the origin already).
            centers = (
                (min[0] + max[0]) / 2,
                (min[1] + max[1]) / 2,
                (min[2] + max[2]) / 2
            )

            sizes = max - min
            logest_size = sizes.max()

            translation = (
                -centers[0],
                -centers[1],
                -centers[2]
            )

            scales = (
                1 / (logest_size + self.options.padding * 2 * logest_size)
            )

            mesh.vertices += translation
            mesh.vertices *= scales

            mesh_upper_jaw.vertices += translation
            mesh_upper_jaw.vertices *= scales

            mesh_lower_jaw.vertices += translation
            mesh_lower_jaw.vertices *= scales

            mesh_name = os.path.splitext(os.path.basename(filepath))[0]
            print('[Data] %s extents before %f - %f, %f - %f, %f - %f' % (mesh_name, min[0], max[0], min[1], max[1], min[2], max[2]))
            min = mesh.vertices.min(axis=0)
            max = mesh.vertices.max(axis=0)
            print('[Data] %s extents after %f - %f, %f - %f, %f - %f' % (mesh_name, min[0], max[0], min[1], max[1], min[2], max[2]))

            scale_info.append({
                "mesh_name": mesh_name,
                "translation": translation,
                "scales": scales
            })

        output_scale_info_file = os.path.join(self.options.out_dir, 'scale_info.json')
        with open(output_scale_info_file, 'w') as f:
            json.dump(scale_info, f, indent=4)

if __name__ == '__main__':
    app = Scale()
    app.run()

最后修改:2025 年 02 月 17 日
如果觉得我的文章对你有用,请随意赞赏