def main(input_size, model_path): demo = Predictor(model_path, input_size) while True: glob_pattern = os.path.join("/app/data/lsun_room/images/", '*.jpg') img_fn = random.choice(glob(glob_pattern)) img = np.array(Image.open(img_fn).resize(input_size)) output, label_map = demo.process(img) output_layout = output.copy() output_layout_edges = output.copy() for label in np.unique(label_map): label_map = discard_smallest_blobs(label_map, label, discard_value=-1, plot_diff=False) pts = np.flip(np.array(np.where(label_map == label)).T, axis=1) hull = ConvexHull(pts, qhull_options='Qt') #rot_angle, area, width, height, center_point, corner_points = minBoundingRect(pts[hull.vertices]) corner_points = pts[hull.vertices] #corner_points = minimum_bounding_rectangle(pts) corner_points = corner_points.astype(np.int32) polygon_pts = corner_points #polygon_pts_rdp = rdp(corner_points, epsilon=10) cv2.fillPoly(output_layout, pts=[polygon_pts], color=COLORS[label]) cv2.polylines(output_layout_edges, [polygon_pts], True, COLORS[label]) f, axarr = plt.subplots(1, 4) axarr[0].imshow(img, interpolation='bicubic') axarr[1].imshow(output.astype('float32'), interpolation='bicubic') axarr[2].imshow(output_layout.astype('uint8'), interpolation='bicubic') axarr[3].imshow(output_layout_edges.astype('uint8'), interpolation='bicubic') plt.title(img_fn) plt.show() alpha = 0.9 output = cv2.addWeighted(output, alpha, output_layout, 1 - alpha, 0) scipy.misc.imsave('output/super_res_output.jpg', output)
def main(input_size, model_path): demo = Predictor(model_path, input_size) img = Image.open('/app/hubstairs.jpg').resize(input_size) output, label_map = demo.process(img) lines, output = gen_linear_layout_map(label_map, output) minmax = lambda x, y: (min(320, max(0, x)), min(320, max(0, y))) for pt1, pt2 in lines: cv2.circle(output, minmax(*pt1), 10, [255,0,0], thickness=1, lineType=8, shift=0) cv2.circle(output, minmax(*pt2), 10, [255,0,0], thickness=1, lineType=8, shift=0) for pt3, pt4 in lines: if pt1 != pt3 and pt2 != pt4: x, y = get_intersect(pt1, pt2, pt3, pt4) cv2.circle(output, (int(x), int(y)), 5, [255,0,0], thickness=1, lineType=8, shift=0) plt.imshow(output, cmap = 'gray', interpolation = 'bicubic') plt.xticks([]), plt.yticks([]) # to hide tick values on X and Y axis plt.show() scipy.misc.imsave('output/super_res_output.jpg', output)