STL转12视图

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

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='train.txt')