62 lines
1.5 KiB
Python
62 lines
1.5 KiB
Python
import os
|
|
import sys
|
|
import numpy as np
|
|
from torch.utils.data import DataLoader
|
|
from stl import mesh as mesh_stl
|
|
|
|
|
|
|
|
class Accumulator():
|
|
def __init__(self):
|
|
self.dict = dict()
|
|
|
|
def update(self, dict):
|
|
for key in dict.keys():
|
|
if key in self.dict.keys():
|
|
self.dict[key].append(dict[key])
|
|
else:
|
|
self.dict[key] = [dict[key]]
|
|
|
|
def get(self):
|
|
result = dict()
|
|
for key in self.dict.keys():
|
|
result[key] = np.mean(self.dict[key])
|
|
return result
|
|
|
|
def reset(self):
|
|
self.dict = dict()
|
|
|
|
|
|
class InfiniteDataLoader(DataLoader):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
# Initialize an iterator over the dataset.
|
|
self.dataset_iterator = super().__iter__()
|
|
|
|
def __iter__(self):
|
|
return self
|
|
|
|
def __next__(self):
|
|
try:
|
|
batch = next(self.dataset_iterator)
|
|
except StopIteration:
|
|
# Dataset exhausted, use a new fresh iterator.
|
|
self.dataset_iterator = super().__iter__()
|
|
batch = next(self.dataset_iterator)
|
|
return batch
|
|
|
|
def text_to_html(lines):
|
|
html = ''
|
|
for line in lines:
|
|
html += "<pre>" + line + "</pre>"
|
|
return html
|
|
|
|
|
|
|
|
def save_stl(file, mesh, faces):
|
|
obj = mesh_stl.Mesh(np.zeros(faces.shape[0], dtype=mesh_stl.Mesh.dtype))
|
|
for i, f in enumerate(faces):
|
|
for j in range(3):
|
|
obj.vectors[i][j] = mesh[f[j],:]
|
|
|
|
obj.save(file) |