File size: 1,351 Bytes
42759e6 fc034ff 42759e6 588b256 8439fac 87a5a37 8439fac 7258769 fc034ff 588b256 8439fac 588b256 8439fac 588b256 a83935b 42759e6 a83935b a99680e 5ee3f67 588b256 a83935b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
import io
from PIL import Image as PImage
import numpy as np
from collections import defaultdict
import cv2
from typing import Tuple, List
from scipy.spatial.distance import cdist
from hoho.read_write_colmap import read_cameras_binary, read_images_binary, read_points3D_binary
from hoho.color_mappings import gestalt_color_mapping, ade20k_color_mapping
from geom_solver import GeomSolver, my_empty_solution, fully_connected_solution
def predict(entry, visualize=False) -> Tuple[np.ndarray, List[int]]:
vertices0, edges0 = my_empty_solution()
try:
vertices, edges = GeomSolver().solve(entry)
except:
print('ERROR')
vertices, edges = fully_connected_solution()
if (len(edges) < 1) and (len(vertices) >= 2):
# print("Added only edges")
edges = edges0
if (len(vertices) < 2) or (len(edges) < 1):
# print("Added empty solution")
vertices, edges = vertices0, edges0
if visualize:
from hoho.viz3d import plot_estimate_and_gt
plot_estimate_and_gt( vertices,
edges,
entry['wf_vertices'],
entry['wf_edges'])
if vertices.shape[-1] != 3:
print("Wrong size")
vertices, edges = vertices0, edges0
return entry['__key__'], vertices, edges
|