89 lines
2.1 KiB
Python
89 lines
2.1 KiB
Python
import os
|
|
import sys
|
|
import pickle
|
|
import numpy as np
|
|
from torch.utils.data import DataLoader
|
|
from stl import mesh as mesh_stl
|
|
from pathlib import Path
|
|
from data.utils import Frame
|
|
|
|
|
|
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)
|
|
|
|
|
|
def create_temp_frames(folder: str, frame: Frame):
|
|
with open(Path(folder) / 'frames_temp.txt', 'w') as f:
|
|
f.writelines([f'{frame}\n' for _ in range(4)])
|
|
|
|
|
|
def pickle_load(file):
|
|
"""
|
|
Load a pickle file.
|
|
"""
|
|
with open(file, 'rb') as f:
|
|
loadout = pickle.load(f)
|
|
|
|
return loadout
|
|
|
|
|
|
def pickle_dump(loadout, file):
|
|
"""
|
|
Dump a pickle file. Create the directory if it does not exist.
|
|
"""
|
|
os.makedirs(os.path.dirname(str(file)), exist_ok=True)
|
|
|
|
with open(file, 'wb') as f:
|
|
pickle.dump(loadout, f) |