def main(data_file, frame_dir, segm_dir, out): data = pkl.load(open(data_file, 'rb'), encoding='iso-8859-1') segm_files = np.array( sorted( glob(os.path.join(segm_dir, '*.png')) + glob(os.path.join(segm_dir, '*.jpg')))) frame_files = np.array( sorted( glob(os.path.join(frame_dir, '*.png')) + glob(os.path.join(frame_dir, '*.jpg')))) vt = np.load('assets/basicModel_vt.npy') ft = np.load('assets/basicModel_ft.npy') f = np.load('assets/basicModel_f.npy') verts = data['vertices'] camera_c = data['camera_c'] camera_f = data['camera_f'] width = data['width'] height = data['height'] camera = ProjectPoints(t=np.zeros(3), rt=np.array([-np.pi, 0., 0.]), c=camera_c, f=camera_f, k=np.zeros(5)) visibility = VisibilityChecker(width, height, f) texture = TextureData(1000, f, vt, ft, visibility) for i, (v, frame_file, segm_file) in enumerate(tqdm(zip(verts, frame_files, segm_files))): frame = cv2.imread(frame_file) / 255. #frame = cv2.resize(frame, dsize=(1080, 1080), interpolation=cv2.INTER_LINEAR) segm = read_segmentation(segm_file) / 255. mask = np.float32(np.any(segm > 0, axis=-1)) camera.set(v=v.astype('float32')) id = os.path.splitext(os.path.basename(frame_file))[0] vis, iso, iso_segm = texture.get_data(frame, camera, mask, segm) cv2.imwrite('{}/{}_unwrap.jpg'.format(out, id), np.uint8(iso * 255)) cv2.imwrite('{}/{}_visibility.jpg'.format(out, id), np.uint8(vis * 255)) cv2.imwrite('{}/{}_segm.png'.format(out, id), np.uint8(iso_segm[:, :, ::-1] * 255))
def main(unwrap_dir, segm_template_file, gmm_file, out_file, num_iter): iso_files = np.array(sorted(glob(os.path.join(unwrap_dir, '*_unwrap.jpg')))) vis_files = np.array( sorted(glob(os.path.join(unwrap_dir, '*_visibility.jpg')))) seams = np.load('assets/basicModel_seams.npy') mask = cv2.imread('assets/tex_mask_1000.png', flags=cv2.IMREAD_GRAYSCALE) / 255. segm_template = read_segmentation(segm_template_file) gmm = pkl.load(open(gmm_file, 'rb')) num_labels = len(iso_files) texture = Texture(1000, seams, mask, segm_template, gmm) isos = [] visibilities = [] for iso_file, vis_file in zip(iso_files, vis_files): print('reading file {}...'.format(os.path.basename(iso_file))) iso = cv2.imread(iso_file) / 255. vis = cv2.imread(vis_file, flags=cv2.IMREAD_GRAYSCALE) / 255. isos.append(iso) visibilities.append(vis) texture_agg = isos[0] visibility_agg = np.array(visibilities[0]) tex, _ = texture.add_iso(texture_agg, visibility_agg, np.zeros_like(visibility_agg), inpaint=False) for i in trange(num_iter): rl = np.random.choice(num_labels) texture_agg, labels = texture.add_iso(isos[rl], visibilities[rl], rl, inpaint=i == (num_iter - 1)) cv2.imwrite(out_file.replace('.jpg', f'_{i}.jpg'), np.uint8(255 * texture_agg)) print('saving {}...'.format(os.path.basename(out_file))) cv2.imwrite(out_file, np.uint8(255 * texture_agg))
def main(data_file, frame_dir, segm_dir, out_file, num_iter): # Step 1: Make unwraps data = pkl.load(open(data_file, 'rb')) segm_files = np.array( sorted( glob(os.path.join(segm_dir, '*.png')) + glob(os.path.join(segm_dir, '*.jpg')))) frame_files = np.array( sorted( glob(os.path.join(frame_dir, '*.png')) + glob(os.path.join(frame_dir, '*.jpg')))) vt = np.load('assets/basicModel_vt.npy') ft = np.load('assets/basicModel_ft.npy') f = np.load('assets/basicModel_f.npy') verts = data['vertices'] camera_c = data['camera_c'] camera_f = data['camera_f'] width = data['width'] height = data['height'] camera = ProjectPoints(t=np.zeros(3), rt=np.array([-np.pi, 0., 0.]), c=camera_c, f=camera_f, k=np.zeros(5)) visibility = VisibilityChecker(width, height, f) texture = TextureData(1000, f, vt, ft, visibility) isos, vises, iso_segms = [], [], [] print('Unwrapping inputs...') for i, (v, frame_file, segm_file) in enumerate(zip(tqdm(verts), frame_files, segm_files)): frame = cv2.imread(frame_file) / 255. segm = read_segmentation(segm_file) / 255. mask = np.float32(np.any(segm > 0, axis=-1)) camera.set(v=v) vis, iso, iso_segm = texture.get_data(frame, camera, mask, segm) vises.append(vis) isos.append(iso) iso_segms.append(np.uint8(iso_segm * 255)) # Step 2: Segm vote gmm iso_mask = cv2.imread('assets/tex_mask_1000.png', flags=cv2.IMREAD_GRAYSCALE) / 255. iso_mask = cv2.resize(iso_mask, (1000, 1000), interpolation=cv2.INTER_NEAREST) voting = np.zeros((1000, 1000, len(LABELS_REDUCED))) gmms = {} gmm_pixels = {} for color_id in LABELS_REDUCED: gmms[color_id] = GaussianMixture(LABELS_MIXTURES[color_id]) gmm_pixels[color_id] = [] for frame, segm, vis in zip(isos, iso_segms, vises): # frame = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) / 255. tex_segm = read_segmentation(segm) tex_weights = 1 - vis tex_weights = np.sqrt(tex_weights) for i, color_id in enumerate(LABELS_REDUCED): if color_id != 'Unseen' and color_id != 'BG': where = np.all(tex_segm == LABELS_REDUCED[color_id], axis=2) voting[where, i] += tex_weights[where] gmm_pixels[color_id].extend(frame[where].tolist()) print('Fitting GMMs...') for color_id in tqdm(LABELS_REDUCED): if gmm_pixels[color_id]: gmms[color_id].fit(np.array(gmm_pixels[color_id])) for i, color_id in enumerate(LABELS_REDUCED): if color_id == 'Unseen' or color_id == 'BG': voting[:, i] = -10 voting[iso_mask == 0] = 0 voting[iso_mask == 0, 0] = 1 unaries = np.ascontiguousarray((1 - voting / len(isos)) * 10) pairwise = np.ascontiguousarray(LABEL_COMP) seams = np.load('assets/basicModel_seams.npy') edge_idx = pkl.load(open('assets/basicModel_edge_idx_1000_.pkl', 'rb')) dr_v = signal.convolve2d(iso_mask, [[-1, 1]])[:, 1:] dr_h = signal.convolve2d(iso_mask, [[-1], [1]])[1:, :] where_v = iso_mask - dr_v where_h = iso_mask - dr_h idxs = np.arange(1000**2).reshape(1000, 1000) v_edges_from = idxs[:-1, :][where_v[:-1, :] == 1].flatten() v_edges_to = idxs[1:, :][where_v[:-1, :] == 1].flatten() h_edges_from = idxs[:, :-1][where_h[:, :-1] == 1].flatten() h_edges_to = idxs[:, 1:][where_h[:, :-1] == 1].flatten() s_edges_from, s_edges_to = edges_seams(seams, 1000, edge_idx) edges_from = np.r_[v_edges_from, h_edges_from, s_edges_from] edges_to = np.r_[v_edges_to, h_edges_to, s_edges_to] edges_w = np.r_[np.ones_like(v_edges_from), np.ones_like(h_edges_from), np.ones_like(s_edges_from)] gc = gco.GCO() gc.create_general_graph(1000**2, pairwise.shape[0], True) gc.set_data_cost(unaries.reshape(1000**2, pairwise.shape[0])) gc.set_all_neighbors(edges_from, edges_to, edges_w) gc.set_smooth_cost(pairwise) gc.swap(-1) labels = gc.get_labels().reshape(1000, 1000) gc.destroy_graph() segm_colors = np.zeros((1000, 1000, 3), dtype=np.uint8) for i, color_id in enumerate(LABELS_REDUCED): segm_colors[labels == i] = LABELS_REDUCED[color_id] # Step 3: Stitch texture seams = np.load('assets/basicModel_seams.npy') mask = cv2.imread('assets/tex_mask_1000.png', flags=cv2.IMREAD_GRAYSCALE) / 255. segm_template = read_segmentation(segm_colors) num_labels = len(isos) texture = Texture(1000, seams, mask, segm_template, gmms) texture_agg = isos[0] visibility_agg = np.array(vises[0]) tex, _ = texture.add_iso(texture_agg, visibility_agg, np.zeros_like(visibility_agg), inpaint=False) print('Aggregating texture...') for i in trange(num_iter): rl = np.random.choice(num_labels) texture_agg, labels = texture.add_iso(isos[rl], vises[rl], rl, inpaint=i == (num_iter - 1)) print('saving {}...'.format(os.path.basename(out_file))) cv2.imwrite(out_file, np.uint8(255 * texture_agg))
def main(unwrap_dir, segm_out_file, gmm_out_file): iso_files = np.array(sorted(glob(os.path.join(unwrap_dir, '*_unwrap.jpg')))) segm_files = np.array(sorted(glob(os.path.join(unwrap_dir, '*_segm.png')))) vis_files = np.array( sorted(glob(os.path.join(unwrap_dir, '*_visibility.jpg')))) iso_mask = cv2.imread('assets/tex_mask_1000.png', flags=cv2.IMREAD_GRAYSCALE) / 255. iso_mask = cv2.resize(iso_mask, (1000, 1000), interpolation=cv2.INTER_NEAREST) voting = np.zeros((1000, 1000, len(LABELS_REDUCED))) gmms = {} gmm_pixels = {} for color_id in LABELS_REDUCED: gmms[color_id] = GaussianMixture(LABELS_MIXTURES[color_id]) gmm_pixels[color_id] = [] for frame_file, segm_file, vis_file in zip(iso_files, segm_files, vis_files): print('extract from {}...'.format(os.path.basename(frame_file))) frame = cv2.cvtColor(cv2.imread(frame_file), cv2.COLOR_BGR2HSV) / 255. tex_segm = read_segmentation(segm_file) tex_weights = 1 - cv2.imread(vis_file) / 255. tex_weights = np.sqrt(tex_weights) for i, color_id in enumerate(LABELS_REDUCED): if color_id != 'Unseen' and color_id != 'BG': where = np.all(tex_segm == LABELS_REDUCED[color_id], axis=2) voting[where, i] += tex_weights[where, 0] gmm_pixels[color_id].extend(frame[where].tolist()) for color_id in LABELS_REDUCED: if gmm_pixels[color_id]: print('GMM fit {}...'.format(color_id)) gmms[color_id].fit(np.array(gmm_pixels[color_id])) for i, color_id in enumerate(LABELS_REDUCED): if color_id == 'Unseen' or color_id == 'BG': voting[:, i] = -10 voting[iso_mask == 0] = 0 voting[iso_mask == 0, 0] = 1 unaries = np.ascontiguousarray((1 - voting / len(iso_files)) * 10) pairwise = np.ascontiguousarray(LABEL_COMP) seams = np.load('assets/basicModel_seams.npy') edge_idx = pkl.load(open('assets/basicModel_edge_idx_1000_.pkl', 'rb')) dr_v = signal.convolve2d(iso_mask, [[-1, 1]])[:, 1:] dr_h = signal.convolve2d(iso_mask, [[-1], [1]])[1:, :] where_v = iso_mask - dr_v where_h = iso_mask - dr_h idxs = np.arange(1000**2).reshape(1000, 1000) v_edges_from = idxs[:-1, :][where_v[:-1, :] == 1].flatten() v_edges_to = idxs[1:, :][where_v[:-1, :] == 1].flatten() h_edges_from = idxs[:, :-1][where_h[:, :-1] == 1].flatten() h_edges_to = idxs[:, 1:][where_h[:, :-1] == 1].flatten() s_edges_from, s_edges_to = edges_seams(seams, 1000, edge_idx) edges_from = np.r_[v_edges_from, h_edges_from, s_edges_from] edges_to = np.r_[v_edges_to, h_edges_to, s_edges_to] edges_w = np.r_[np.ones_like(v_edges_from), np.ones_like(h_edges_from), np.ones_like(s_edges_from)] gc = gco.GCO() gc.create_general_graph(1000**2, pairwise.shape[0], True) gc.set_data_cost(unaries.reshape(1000**2, pairwise.shape[0])) gc.set_all_neighbors(edges_from, edges_to, edges_w) gc.set_smooth_cost(pairwise) gc.swap(-1) labels = gc.get_labels().reshape(1000, 1000) gc.destroy_graph() segm_colors = np.zeros((1000, 1000, 3), dtype=np.uint8) for i, color_id in enumerate(LABELS_REDUCED): segm_colors[labels == i] = LABELS_REDUCED[color_id] cv2.imwrite('{}'.format(segm_out_file), segm_colors[:, :, ::-1]) pkl.dump(gmms, open(gmm_out_file, 'wb'))