| import copy |
| import math |
| import numpy as np |
| import torch |
| import json |
|
|
| class WireframeGraph: |
| def __init__(self, |
| vertices: torch.Tensor, |
| v_confidences: torch.Tensor, |
| edges: torch.Tensor, |
| edge_weights: torch.Tensor, |
| frame_width: int, |
| frame_height: int): |
| self.vertices = vertices |
| self.v_confidences = v_confidences |
| self.edges = edges |
| self.weights = edge_weights |
| self.frame_width = frame_width |
| self.frame_height = frame_height |
|
|
| @classmethod |
| def xyxy2indices(cls,junctions, lines): |
| |
| |
| |
| dist1 = torch.norm(junctions[None,:,:]-lines[:,None,:2],dim=-1) |
| dist2 = torch.norm(junctions[None,:,:]-lines[:,None,2:],dim=-1) |
| idx1 = torch.argmin(dist1,dim=-1) |
| idx2 = torch.argmin(dist2,dim=-1) |
| return torch.stack((idx1,idx2),dim=-1) |
| @classmethod |
| def load_json(cls, fname): |
| with open(fname,'r') as f: |
| data = json.load(f) |
|
|
| |
| vertices = torch.tensor(data['vertices']) |
| v_confidences = torch.tensor(data['vertices-score']) |
| edges = torch.tensor(data['edges']) |
| edge_weights = torch.tensor(data['edges-weights']) |
| height = data['height'] |
| width = data['width'] |
|
|
| return WireframeGraph(vertices,v_confidences,edges,edge_weights,width,height) |
|
|
| @property |
| def is_empty(self): |
| for key, val in self.__dict__.items(): |
| if val is None: |
| return True |
| return False |
|
|
| @property |
| def num_vertices(self): |
| if self.is_empty: |
| return 0 |
| return self.vertices.shape[0] |
| |
| @property |
| def num_edges(self): |
| if self.is_empty: |
| return 0 |
| return self.edges.shape[0] |
|
|
|
|
| def line_segments(self, threshold = 0.05, device=None, to_np=False): |
| is_valid = self.weights>threshold |
| p1 = self.vertices[self.edges[is_valid,0]] |
| p2 = self.vertices[self.edges[is_valid,1]] |
| ps = self.weights[is_valid] |
|
|
| lines = torch.cat((p1,p2,ps[:,None]),dim=-1) |
| if device is not None: |
| lines = lines.to(device) |
| if to_np: |
| lines = lines.cpu().numpy() |
|
|
| return lines |
| |
| |
| def rescale(self, image_width, image_height): |
| scale_x = float(image_width)/float(self.frame_width) |
| scale_y = float(image_height)/float(self.frame_height) |
|
|
| self.vertices[:,0] *= scale_x |
| self.vertices[:,1] *= scale_y |
| self.frame_width = image_width |
| self.frame_height = image_height |
|
|
| def jsonize(self): |
| return { |
| 'vertices': self.vertices.cpu().tolist(), |
| 'vertices-score': self.v_confidences.cpu().tolist(), |
| 'edges': self.edges.cpu().tolist(), |
| 'edges-weights': self.weights.cpu().tolist(), |
| 'height': self.frame_height, |
| 'width': self.frame_width, |
| } |
| def __repr__(self) -> str: |
| return "WireframeGraph\n"+\ |
| "Vertices: {}\n".format(self.num_vertices)+\ |
| "Edges: {}\n".format(self.num_edges,) + \ |
| "Frame size (HxW): {}x{}".format(self.frame_height,self.frame_width) |
|
|
| |
| if __name__ == "__main__": |
| graph = WireframeGraph.load_json('NeuS/public_data/bmvs_clock/hawp/000.json') |
| print(graph) |
| |
