class RandomizedExtractor(object): def __init__(self, dataset): self.aggregator = Aggregator() self.dataset = dataset self.target_index, target = self.__choose_random_sample(dataset) self.target = copy.deepcopy(target) #self.target.set_mesh_and_parts_colour([1, 0, 0]) self.set_colour_coded_sources( dataset) ############# ADD TO RESET TARGET?? ############## def __choose_random_sample(self, dataset): data_size = len(dataset) sample_idx = random.randint(0, data_size - 1) sample = dataset[sample_idx] return sample_idx, sample def reset_target(self, dataset): idx = self.target_index while idx == self.target_index: self.target_index, target = self.__choose_random_sample(dataset) self.target = copy.deepcopy(target) #self.set_colour_coded_sources(self.dataset) def get_target_labels_with_bounding_box(self): assert (self.target is not None) bboxes_dict = {} for part in self.target.parts: if part.label not in bboxes_dict: bboxes_dict[ part.label] = self.aggregator.get_super_bounding_box( self.target, part.label) return bboxes_dict def get_target_labels_with_parts(self): assert (self.target is not None) bboxes_dict = {} for part in self.target.parts: if part.label not in bboxes_dict: bboxes_dict[ part.label] = self.aggregator.get_all_parts_by_label( self.target, part.label) return bboxes_dict def get_target_parts_by_label(self, label): return self.aggregator.get_all_parts_by_label(self.target, label) def make_target_parts_invisible_by_label(self, label): assert (self.target is not None) for part in self.target.parts: if part.label == label: part.render = False def add_parts_to_target(self, parts): assert (self.target is not None) self.target.parts.extend(parts) def replace_target_parts_by_label(self, label, parts): assert (self.target is not None) target_parts = [] for part in self.target.parts: if part.label != label: target_parts.append(part) target_parts.extend(parts) self.target.parts = target_parts def find_source_part(self, dataset, label, iter_limit=10): bbox = None idx = self.target_index iter_count = 0 while bbox is None or idx == self.target_index: idx, source_mesh = self.__choose_random_sample(dataset) parts, bbox = self.aggregator.get_parts_with_bounding_box( source_mesh, label) iter_count += 1 if iter_count == iter_limit: break return parts, bbox def __reindex_faces(part_faces, new_offset=0): part_faces_np = np.asarray(part_faces, dtype=np.int32) part_faces_np = part_faces_np - np.amin(part_faces_np) + new_offset return part_faces_np def get_triangle_part_meshes(self, parts): triangle_meshes = [] for part in parts: tri_mesh = o3d.geometry.TriangleMesh() tri_mesh.vertices = o3d.utility.Vector3dVector( np.asarray(part.vertices, dtype=np.float64)) tri_mesh.triangles = o3d.utility.Vector3iVector( np.asarray(part.faces, dtype=np.int32)) triangle_meshes.append(tri_mesh) return triangle_meshes def smooth_part_meshes(self, parts): triangle_meshes = self.get_triangle_part_meshes(parts) for part, mesh in zip(parts, triangle_meshes): mesh.filter_smooth_laplacian(number_of_iterations=3) part.vertices = np.asarray(mesh.vertices, dtype=np.float64).tolist() part.faces = np.asarray(mesh.triangles, dtype=np.int32).tolist() return parts def set_colour_coded_sources(self, dataset): data_size = len(dataset) cols = [[1, 0, 0], [0.2, 0.2, 0.6], [0.2, 0.7, 0.3], [1, 0.7, 0.2], [0.5, 0.1, 0.5], [1, 0.4, 0.4], [0.8, 0.5, 0.1], [1, 0, 0.5], [0.5, 0.6, 0.2], [0.3, 0.6, 0.9]] # List of colours # 1: Dark Blue, 2: Green, 3: Yellow, 4: Purple i = 0 for idx in range(data_size): if idx == self.target_index: # For target self.target.set_mesh_and_parts_colour(cols[i]) #print("Setting colour for target with colour red at idx ", idx) #continue mesh = dataset[idx] print("Setting colour for ", idx, " with colour number ", i) mesh.set_mesh_and_parts_colour(cols[i]) i = i + 1 def get_source_meshes(self): data_size = len(self.dataset) meshes = [] i = 0 for idx in range(data_size): if idx == self.target_index: continue meshes.append(self.dataset[idx]) return meshes