CADNET数据集

环境配置

conda create -n transform python=3.6
pip install numpy-stl
pip install matplotlib==2.2.3
pip install Pillow

File_Searcher.py

• 将CADNET_3317数据集的所有模型的路径存到两个文件夹中

from os import walk, path


class File_Searcher:
    def __init__(self, file_extension=None, search_folder='Model_File_Folder', output_file='file_list', limit=2000):
        """
        :param file_extension: tuple of strings of file extensions which need to search in the search_folder. by default
          file_extension = ['.STL', '.stl']
        :param search_folder: Folder name where files will be searched
        :param output_file: .txt file name which will have all the file path which have file extension == file_extension
        """
        assert path.isdir(search_folder)
        self.output_file = output_file
        self.limit = limit
        self.search_folder = search_folder
        if file_extension is None:
            file_extension = ['.STL', '.stl']
        self.file_extension = file_extension
        self.generate_file()

    def generate_file(self):
        output_file_count = 1
        text_file = open(f"{self.output_file}_{output_file_count}.txt", 'w')
        count = 0
        for r, d, files in walk(self.search_folder):
            for file in files:
                _, extension = path.splitext(file)
                if extension in self.file_extension:
                    text_file.write(path.join(r, file) + '\n')
                    count += 1
                    if count >= self.limit:
                        text_file.close()
                        output_file_count += 1
                        text_file = open(f"{self.output_file}_{output_file_count}.txt", 'w')
                        count = 0


if __name__ == "__main__":
    # lis = ['test', 'train']
    # for name in lis:
    File_Searcher(search_folder='E:\\data\\弄好的\\CADNET\\CADNET_3317', output_file='file_list',
                  limit=2000)

Mesh_to_Image.py

• 将其三维模型转化成20视图

from os import walk, path


class File_Searcher:
    def __init__(self, file_extension=None, search_folder='Model_File_Folder', output_file='file_list', limit=2000):
        """
        :param file_extension: tuple of strings of file extensions which need to search in the search_folder. by default
          file_extension = ['.STL', '.stl']
        :param search_folder: Folder name where files will be searched
        :param output_file: .txt file name which will have all the file path which have file extension == file_extension
        """
        assert path.isdir(search_folder)
        self.output_file = output_file
        self.limit = limit
        self.search_folder = search_folder
        if file_extension is None:
            file_extension = ['.STL', '.stl']
        self.file_extension = file_extension
        self.generate_file()

    def generate_file(self):
        output_file_count = 1
        text_file = open(f"{self.output_file}_{output_file_count}.txt", 'w')
        count = 0
        for r, d, files in walk(self.search_folder):
            for file in files:
                _, extension = path.splitext(file)
                if extension in self.file_extension:
                    text_file.write(path.join(r, file) + '\n')
                    count += 1
                    if count >= self.limit:
                        text_file.close()
                        output_file_count += 1
                        text_file = open(f"{self.output_file}_{output_file_count}.txt", 'w')
                        count = 0


if __name__ == "__main__":
    # lis = ['test', 'train']
    # for name in lis:
    File_Searcher(search_folder='E:\\data\\弄好的\\CADNET\\CADNET_3317', output_file='file_list',
                  limit=2000)

demo.py

• 将其三维模型转化成12视图。

import numpy as np
from PIL import Image, ImageChops
from os import path, makedirs
from stl import mesh
from mpl_toolkits import mplot3d
from matplotlib import pyplot as plt

def get_trim_image(data):
    def trim(im, border):
        bg = Image.new(im.mode, im.size, border)
        diff = ImageChops.difference(im, bg)
        bbox = diff.getbbox()
        if bbox:
            return im.crop(bbox)

    image = Image.fromarray(data)
    image.thumbnail((500, 500), Image.ANTIALIAS)
    return trim(image, (255, 255, 255))

class Mesh_to_Image:
    def __init__(self, image_folder='data_train', files_list='file_list.txt'):
        if not path.exists(image_folder):
            makedirs(image_folder)
        self.image_folder = image_folder
        self.files_list = files_list
        self.generate_Image()

    def generate_Image(self):
        file_no = 0
        files = open(self.files_list, 'r')
        for file in files:
            file_no += 1
            file_path = file.strip()
            self.file_name, exe = path.splitext(file_path)
            print(f"{file_no} {self.file_name}{exe}")

            figure = plt.figure(figsize=(10, 10))
            axes = mplot3d.Axes3D(figure)

            Mesh = mesh.Mesh.from_file(file_path)
            axes.add_collection3d(mplot3d.art3d.Poly3DCollection(Mesh.vectors, edgecolor='k'))

            scale = Mesh.points.flatten('C')
            axes.auto_scale_xyz(scale, scale, scale)
            axes.set_axis_off()

            for angle_no, angle in enumerate(range(0, 360, 30), start=1):
                axes.view_init(30, angle)

                file_base_name, _ = path.splitext(path.basename(file_path))
                folder_path = path.join(self.image_folder, file_base_name)
                file_name = f"{angle_no}.png"
                image_file_path = path.join(folder_path, file_name)

                if not path.exists(folder_path):
                    makedirs(folder_path)

                figure.canvas.draw()
                data = np.fromstring(figure.canvas.tostring_rgb(), dtype=np.uint8, sep='')
                data = data.reshape(figure.canvas.get_width_height()[::-1] + (3,))
                image = get_trim_image(data)
                image.save(image_file_path)
                print(f"{file_no}, {angle_no} >> {image_file_path}  Done!")

if __name__ == "__main__":
    Mesh_to_Image(image_folder='Data_lfd', files_list='file_list_1.txt')

split_train_test.py

• 将转化后的数据集分成训练集和测试集。

import os, random, shutil


def split_train_test(dir, test_split=5):


    if not os.path.exists(dir):
        raise Exception("dir Not exist!")

    output_dir = dir + "_split_train_test"

    class_dir_list = os.listdir(dir)

    for class_dir in class_dir_list:
        class_dir_path = os.path.join(dir, class_dir)

        model_dir_list = os.listdir(class_dir_path)
        file_no = len(model_dir_list)
        test_file_no = file_no//test_split
        for _ in range(test_file_no):
            test_dir = random.choice(model_dir_list)
            model_dir_list.remove(test_dir)
            source_dir = os.path.join(class_dir_path, test_dir)
            destination_dir = os.path.join(output_dir, "test_dir", class_dir, test_dir)

            #print("source_dir: ", source_dir)
            #print("destination_dir: ", destination_dir)


            shutil.copytree(source_dir, destination_dir,  symlinks=True, ignore=None)

        for train_dir in model_dir_list:
            source_dir = os.path.join(class_dir_path, train_dir)
            destination_dir = os.path.join(output_dir, "train_dir", class_dir, train_dir)

            shutil.copytree(source_dir, destination_dir,  symlinks=False, ignore=None)


if __name__ == "__main__":
    split_train_test('C:\\Users\\su\\Desktop\\Data_lfd')
    # dir = 'Image_Folder'
    # for folder in os.listdir(dir):
    #     split_train_test(os.path.join(dir,folder))

change.py

• 将指定文件夹下的所有子文件下的所有文件依次遍历提取并重新的命名整合到一块。

import os
import shutil

# 定义原始文件夹路径和目标文件夹路径
original_folder = r'C:\Users\su\Desktop\Data_lfd_split_train_test\test_dir\U-shaped_parts'
target_folder = r'C:\Users\su\Desktop\CADNETv12\U-shaped_parts\test'

# 遍历原始文件夹及其所有子文件夹
folder_index = 0
for root, dirs, files in os.walk(original_folder):
    for file in files:
        # 检查文件是否为图片文件
        if file.endswith('.png') or file.endswith('.jpg') or file.endswith('.jpeg'):
            # 构建原始文件的完整路径
            original_file_path = os.path.join(root, file)
            # 构建目标文件的完整路径
            # 修改目标文件的名称，添加文件夹索引和图片索引作为文件名的一部分，并添加前导零
            folder_name = os.path.basename(root)
            target_file_name = f"U-shaped_parts_{folder_index:04d}_{files.index(file)+1:03d}.png"
            target_file_path = os.path.join(target_folder, target_file_name)
            # 复制文件到目标文件夹
            shutil.copy(original_file_path, target_file_path)
            print(f"复制文件：{original_file_path} 到 {target_file_path}")
    folder_index += 1

print("文件复制完成！")