| import sys | |
| sys.path.append("..") | |
| from models.Mask3D.mask3d import get_model, load_mesh, prepare_data, map_output_to_pointcloud, save_colorized_mesh | |
| import torch | |
| class Network_3D(): | |
| def __init__(self, config): | |
| self.model = get_model(config["network3d"]["pretrained_path"]) | |
| self.model.eval() | |
| self.device = torch.device("cuda:0") | |
| self.model.to(self.device) | |
| def get_class_agnostic_masks(self, pointcloud_file, point2segment=None): | |
| data, points, colors, features, unique_map, inverse_map, point2segment, point2segment_full = prepare_data(pointcloud_file, self.device) | |
| with torch.no_grad(): | |
| outputs = self.model(data, raw_coordinates=features, point2segment=[point2segment] if point2segment is not None else None) | |
| return map_output_to_pointcloud(outputs, inverse_map, point2segment, point2segment_full) | |