class GUI(tk.Tk): MODE_LOW = 1 MODE_HIGH = 2 MODE_FLOW = 3 CHANNEL_MASK = 1 CHANNEL_NORMAL = 2 CHANNEL_DEPTH = 3 CHANNEL_AO = 4 CHANNEL_COLOR = 5 def __init__(self): tk.Tk.__init__(self) self.title('Dataset Viewer') self.input_name = None # members to be filled later self.pilImage = None self.tikImage = None self.current_high = None self.current_low = None self.current_flow = None self.dataset_folder = None self.last_folder = INITIAL_DIR self.entries = [] self.num_frames = 0 self.selected_time = 0 # root self.root_panel = tk.Frame(self) self.root_panel.pack(side="bottom", fill="both", expand="yes") self.panel = tk.Label(self.root_panel) self.black = np.zeros((512, 512, 3)) self.setImage(self.black) self.panel.pack(side="left", fill="both", expand="yes") options1 = tk.Label(self.root_panel) # Input inputFrame = ttk.LabelFrame(options1, text="Input", relief=tk.RIDGE) tk.Button(inputFrame, text="Open Folder", command=lambda: self.openFolder()).pack() self.dataset_entry = 0 self.dataset_entry_slider = tk.Scale( inputFrame, from_=0, to=0, orient=tk.HORIZONTAL, resolution=1, label="Entry", showvalue=0, command=lambda e: self.setEntry(int(e))) self.dataset_entry_slider.pack(anchor=tk.W, fill=tk.X) self.dataset_time = 0 self.dataset_time_slider = tk.Scale( inputFrame, from_=0, to=0, orient=tk.HORIZONTAL, resolution=1, label="Time", showvalue=0, command=lambda e: self.setTime(int(e))) self.dataset_time_slider.pack(anchor=tk.W, fill=tk.X) inputFrame.pack(fill=tk.X) # Mode modeFrame = ttk.LabelFrame(options1, text="Mode", relief=tk.RIDGE) self.mode = tk.IntVar() self.mode.set(GUI.MODE_LOW) self.mode.trace_add('write', lambda a, b, c: self.updateImage()) tk.Radiobutton(modeFrame, text="Low", variable=self.mode, value=GUI.MODE_LOW).pack(anchor=tk.W) tk.Radiobutton(modeFrame, text="High", variable=self.mode, value=GUI.MODE_HIGH).pack(anchor=tk.W) tk.Radiobutton(modeFrame, text="Flow", variable=self.mode, value=GUI.MODE_FLOW).pack(anchor=tk.W) modeFrame.pack(fill=tk.X) # Channel channelsFrame = ttk.LabelFrame(options1, text="Channel", relief=tk.RIDGE) self.channel_mode = tk.IntVar() self.channel_mode.set(GUI.CHANNEL_COLOR) self.channel_mode.trace_add('write', lambda a, b, c: self.updateImage()) tk.Radiobutton(channelsFrame, text="Mask", variable=self.channel_mode, value=GUI.CHANNEL_MASK).pack(anchor=tk.W) tk.Radiobutton(channelsFrame, text="Normal", variable=self.channel_mode, value=GUI.CHANNEL_NORMAL).pack(anchor=tk.W) tk.Radiobutton(channelsFrame, text="Depth", variable=self.channel_mode, value=GUI.CHANNEL_DEPTH).pack(anchor=tk.W) tk.Radiobutton(channelsFrame, text="AO", variable=self.channel_mode, value=GUI.CHANNEL_AO).pack(anchor=tk.W) tk.Radiobutton(channelsFrame, text="Color", variable=self.channel_mode, value=GUI.CHANNEL_COLOR).pack(anchor=tk.W) channelsFrame.pack(fill=tk.X) # Shading self.shading = ScreenSpaceShading('cpu') self.shading.fov(30) self.ambient_light_color = np.array([0.1, 0.1, 0.1]) self.shading.ambient_light_color(self.ambient_light_color) self.diffuse_light_color = np.array([0.8, 0.8, 0.8]) self.shading.diffuse_light_color(self.diffuse_light_color) self.specular_light_color = np.array([0.02, 0.02, 0.02]) self.shading.specular_light_color(self.specular_light_color) self.shading.specular_exponent(16) self.shading.light_direction(np.array([0.1, 0.1, 1.0])) self.material_color = np.array([1.0, 1.0, 1.0]) #[1.0, 0.3, 0.0]) self.shading.material_color(self.material_color) self.shading.ambient_occlusion(0.5) self.shading.background(np.array([1.0, 1.0, 1.0])) # Save tk.Button(options1, text="Save Image", command=lambda: self.saveImage()).pack() self.saveFolder = "/" options1.pack(side="left") def openFolder(self): dataset_folder = filedialog.askdirectory(initialdir=self.last_folder) print(dataset_folder) if dataset_folder is not None: self.dataset_folder = str(dataset_folder) print("New folder selected:", self.dataset_folder) self.last_folder = self.dataset_folder # find number of entries self.current_low = None self.current_high = None self.current_flow = None entries = [] for i in range(0, 10000): file_low = os.path.join(self.dataset_folder, "low_%05d.npy" % i) if os.path.isfile(file_low): entries.append((file_low, os.path.join(self.dataset_folder, "high_%05d.npy" % i), os.path.join(self.dataset_folder, "flow_%05d.npy" % i))) else: break print("Number of entries found:", len(entries)) self.entries = entries self.dataset_entry = 0 self.dataset_entry_slider.config(to=len(entries)) self.setEntry(0) else: print("No folder selected") def setEntry(self, entry): entry = min(entry, len(self.entries) - 1) self.dataset_entry_slider.config(label='Entry: %d' % int(entry)) if len(self.entries) == 0: self.current_low = None self.current_high = None self.current_flow = None else: self.current_low = np.load(self.entries[entry][0]) self.current_high = np.load(self.entries[entry][1]) self.current_flow = np.load(self.entries[entry][2]) self.num_frames = self.current_low.shape[0] self.dataset_time_slider.config(to=self.num_frames) print("Entry loaded") self.setTime(0) def setTime(self, entry): entry = min(entry, self.num_frames - 1) self.dataset_time_slider.config(label='Time: %d' % int(entry)) self.selected_time = entry self.updateImage() def setImage(self, img): self.pilImage = pl.Image.fromarray( np.clip((img * 255).transpose((1, 0, 2)), 0, 255).astype(np.uint8)) self.tikImage = ImageTk.PhotoImage(self.pilImage) self.panel.configure(image=self.tikImage) def saveImage(self): filename = filedialog.asksaveasfilename( initialdir=self.saveFolder, title="Save as", filetypes=(("jpeg files", "*.jpg"), ("png files", "*.png"), ("all files", "*.*"))) if len(filename) == 0: return if len(os.path.splitext(filename)[1]) == 0: filename = filename + ".jpg" self.pilImage.save(filename) self.saveFolder = os.path.dirname(filename) def updateImage(self): if self.current_low is None: # no image loaded self.setImage(self.black) return def selectChannel(img): if self.channel_mode.get() == GUI.CHANNEL_MASK: mask = img[0:1, :, :] * 0.5 + 0.5 return np.concatenate((mask, mask, mask)) elif self.channel_mode.get() == GUI.CHANNEL_NORMAL: return img[1:4, :, :] * 0.5 + 0.5 elif self.channel_mode.get() == GUI.CHANNEL_DEPTH: return np.concatenate( (img[4:5, :, :], img[4:5, :, :], img[4:5, :, :])) elif self.channel_mode.get() == GUI.CHANNEL_AO: if img.shape[0] == 6: return np.concatenate( (img[5:6, :, :], img[5:6, :, :], img[5:6, :, :])) else: return np.zeros((3, img.shape[1], img.shape[2]), dtype=np.float32) elif self.channel_mode.get() == GUI.CHANNEL_COLOR: shading_input = torch.unsqueeze(torch.from_numpy(img), 0) shading_output = self.shading(shading_input)[0] return torch.clamp(shading_output, 0, 1).cpu().numpy() if self.mode.get() == GUI.MODE_LOW: img = self.current_low[self.selected_time, :, :, :] img = selectChannel(img) img = img.transpose((2, 1, 0)) img = cv.resize(img, dsize=None, fx=UPSCALE, fy=UPSCALE, interpolation=cv.INTER_NEAREST) self.setImage(img) elif self.mode.get() == GUI.MODE_HIGH: img = self.current_high[self.selected_time, :, :, :] img = selectChannel(img) img = img.transpose((2, 1, 0)) self.setImage(img) elif self.mode.get() == GUI.MODE_FLOW: #img = np.stack(( # cv.inpaint(self.current_flow[self.selected_time,0,:,:], np.uint8(self.current_low[self.selected_time,0,:,:]==0), 3, cv.INPAINT_NS), # cv.inpaint(self.current_flow[self.selected_time,1,:,:], np.uint8(self.current_low[self.selected_time,0,:,:]==0), 3, cv.INPAINT_NS), # np.zeros((self.current_flow.shape[2], self.current_flow.shape[3]))), axis=0).astype(np.float32) img = np.concatenate( (self.current_flow[self.selected_time, 0:2, :, :], np.zeros((1, self.current_flow.shape[2], self.current_flow.shape[3]), dtype=np.float32)), ) img = (img * 10 + 0.5) img = img.transpose((2, 1, 0)) img = cv.resize(img, dsize=None, fx=UPSCALE, fy=UPSCALE, interpolation=cv.INTER_NEAREST) self.setImage(img)
datasetfile = os.path.join(data_dir, DATASETS[i]['file']) print('Open', datasetfile) renderer = inference.Renderer(renderer_path, datasetfile, material, camera) time.sleep(5) renderer.send_command("aoradius=%5.3f\n" % float(AO_RADIUS)) # create shading shading = ScreenSpaceShading(torch.device('cpu')) shading.fov(30) shading.light_direction(np.array([0.1, 0.1, 1.0])) shading.ambient_light_color(np.array(DATASETS[i]['ambient']) / 255.0) shading.diffuse_light_color(np.array(DATASETS[i]['diffuse']) / 255.0) shading.specular_light_color(np.array(DATASETS[i]['specular']) / 255.0) shading.specular_exponent(SPECULAR_EXPONENT) shading.material_color(np.array(DATASETS[i]['material']) / 255.0) shading.ambient_occlusion(AO_STRENGTH) shading.background(np.array(BACKGROUND)) # render each model for k, m in enumerate(MODELS): print('Render', m['name']) p = m['path'] outputName = os.path.join( OUTPUT_FOLDER, "%s.%s.mp4" % (DATASETS[i]['name'], m['name'])) writer = imageio.get_writer(outputName, fps=FPS) camera.currentYaw = 0 previous_image = None if p == MODEL_GROUND_TRUTH: renderer.send_command("aosamples=%d\n"% \ int(AO_SAMPLES if m['ao'] else 0))
class GUI(tk.Tk): CHANNEL_MASK = 1 CHANNEL_NORMAL = 2 CHANNEL_DEPTH = 3 CHANNEL_AO = 4 CHANNEL_COLOR = 5 CHANNEL_FLOW = 6 def __init__(self): tk.Tk.__init__(self) self.title('Dataset Viewer') self.input_name = None # members to be filled later self.pilImage = None self.tikImage = None self.num_frames = 0 self.selected_entry = 0 self.selected_time = 0 self.last_folder = None self.dataset_file = None self.hdf5_file = None self.dset_keys = [] self.dset = None self.mode = None # root self.root_panel = tk.Frame(self) self.root_panel.pack(side="bottom", fill="both", expand="yes") self.panel = tk.Label(self.root_panel) self.black = np.zeros((512, 512, 3)) self.setImage(self.black) self.panel.pack(side="left", fill="both", expand="yes") options1 = tk.Label(self.root_panel) # Input inputFrame = ttk.LabelFrame(options1, text="Input", relief=tk.RIDGE) tk.Button(inputFrame, text="Open HDF5", command=lambda: self.openHDF5()).pack() listbox_frame = tk.Frame(inputFrame) self.dset_listbox_scrollbar = tk.Scrollbar(listbox_frame, orient=tk.VERTICAL) self.dset_listbox = tk.Listbox( listbox_frame, selectmode=tk.SINGLE, yscrollcommand=self.dset_listbox_scrollbar.set) self.dset_listbox_scrollbar.config(command=self.dset_listbox.yview) self.dset_listbox_scrollbar.pack(side=tk.RIGHT, fill=tk.Y) self.dset_listbox.pack(side=tk.LEFT, anchor=tk.W, fill=tk.X, expand=1) self.dset_listbox.bind("<Double-Button-1>", self.setDsetCallback) listbox_frame.pack(anchor=tk.W, fill=tk.X) self.selected_dset = None self.dataset_entry_slider = tk.Scale( inputFrame, from_=0, to=0, orient=tk.HORIZONTAL, resolution=1, label="Entry", showvalue=0, command=lambda e: self.setEntry(int(e))) self.dataset_entry_slider.pack(anchor=tk.W, fill=tk.X) self.dataset_time = 0 self.dataset_time_slider = tk.Scale( inputFrame, from_=0, to=0, orient=tk.HORIZONTAL, resolution=1, label="Time", showvalue=0, command=lambda e: self.setTime(int(e))) self.dataset_time_slider.pack(anchor=tk.W, fill=tk.X) inputFrame.pack(fill=tk.X) # Channel channelsFrame = ttk.LabelFrame(options1, text="Channel", relief=tk.RIDGE) self.channel_mode = tk.IntVar() self.channel_mode.set(GUI.CHANNEL_COLOR) self.channel_mode.trace_add('write', lambda a, b, c: self.updateImage()) tk.Radiobutton(channelsFrame, text="Mask", variable=self.channel_mode, value=GUI.CHANNEL_MASK).pack(anchor=tk.W) tk.Radiobutton(channelsFrame, text="Normal", variable=self.channel_mode, value=GUI.CHANNEL_NORMAL).pack(anchor=tk.W) tk.Radiobutton(channelsFrame, text="Depth", variable=self.channel_mode, value=GUI.CHANNEL_DEPTH).pack(anchor=tk.W) tk.Radiobutton(channelsFrame, text="AO", variable=self.channel_mode, value=GUI.CHANNEL_AO).pack(anchor=tk.W) tk.Radiobutton(channelsFrame, text="Color", variable=self.channel_mode, value=GUI.CHANNEL_COLOR).pack(anchor=tk.W) tk.Radiobutton(channelsFrame, text="Flow", variable=self.channel_mode, value=GUI.CHANNEL_FLOW).pack(anchor=tk.W) channelsFrame.pack(fill=tk.X) # Shading self.shading = ScreenSpaceShading('cpu') self.shading.fov(30) self.ambient_light_color = np.array([0.1, 0.1, 0.1]) self.shading.ambient_light_color(self.ambient_light_color) self.diffuse_light_color = np.array([0.8, 0.8, 0.8]) self.shading.diffuse_light_color(self.diffuse_light_color) self.specular_light_color = np.array([0.02, 0.02, 0.02]) self.shading.specular_light_color(self.specular_light_color) self.shading.specular_exponent(16) self.shading.light_direction(np.array([0.1, 0.1, 1.0])) self.material_color = np.array([1.0, 1.0, 1.0]) #[1.0, 0.3, 0.0]) self.shading.material_color(self.material_color) self.shading.ambient_occlusion(0.5) self.shading.background(np.array([1.0, 1.0, 1.0])) # Save tk.Button(options1, text="Save Image", command=lambda: self.saveImage()).pack() self.saveFolder = "/" options1.pack(side="left") def openHDF5(self): dataset_file = filedialog.askopenfilename(initialdir=self.last_folder, title="Select HDF5 file", filetypes=(("HDF5 files", "*.hdf5"), )) print(dataset_file) if dataset_file is not None: self.dataset_file = str(dataset_file) print("New hdf5 selected:", self.dataset_file) self.title(self.dataset_file) self.last_folder = os.path.dirname(self.dataset_file) # load hdf5 file self.hdf5_file = h5py.File(dataset_file, "r") # list datasets self.dset_listbox.delete(0, tk.END) self.dset_keys = list(self.hdf5_file.keys()) for key in self.dset_keys: self.dset_listbox.insert(tk.END, key) self.dset_listbox.selection_set(first=0) self.setDset(self.dset_keys[0]) else: print("No folder selected") def setDsetCallback(self, *args): items = self.dset_listbox.curselection() items = [self.dset_keys[int(item)] for item in items] self.setDset(items[0]) def setDset(self, name): self.selected_dset = name print("Select dataset '%s'" % (self.selected_dset)) self.dset = self.hdf5_file[self.selected_dset] self.mode = self.dset.attrs.get("Mode", "IsoUnshaded") print("Mode:", self.mode) # find number of entries entries = self.dset.shape[0] num_frames = self.dset.shape[1] print("Number of entries found:", entries, "with", num_frames, "timesteps") print("Image size:", self.dset.shape[3], "x", self.dset.shape[4]) self.entries = entries self.num_frames = num_frames self.dataset_entry_slider.config(to=entries - 1) self.dataset_time_slider.config(to=num_frames - 1) self.setEntry( self.selected_entry if self.selected_entry < entries else 0) def setEntry(self, entry): entry = min(entry, self.entries - 1) self.dataset_entry_slider.config(label='Entry: %d' % int(entry)) self.selected_entry = entry self.updateImage() def setTime(self, entry): entry = min(entry, self.num_frames - 1) self.dataset_time_slider.config(label='Time: %d' % int(entry)) self.selected_time = entry self.updateImage() def setImage(self, img): if img.dtype == np.uint8: self.pilImage = pl.Image.fromarray(img.transpose((1, 0, 2))) else: self.pilImage = pl.Image.fromarray( np.clip((img * 255).transpose((1, 0, 2)), 0, 255).astype(np.uint8)) if self.pilImage.size[0] <= 256: self.pilImage = self.pilImage.resize( (self.pilImage.size[0] * 2, self.pilImage.size[1] * 2), pl.Image.NEAREST) self.tikImage = ImageTk.PhotoImage(self.pilImage) self.panel.configure(image=self.tikImage) def saveImage(self): filename = filedialog.asksaveasfilename( initialdir=self.saveFolder, title="Save as", filetypes=(("png files", "*.png"), ("jpeg files", "*.jpg"), ("all files", "*.*"))) if len(filename) == 0: return if len(os.path.splitext(filename)[1]) == 0: filename = filename + ".png" self.pilImage.save(filename) self.saveFolder = os.path.dirname(filename) def updateImage(self): if self.dset is None: # no image loaded self.setImage(self.black) return def selectChannel(img): if self.mode == "DVR": mask = img[3:4, :, :] if self.channel_mode.get() == GUI.CHANNEL_MASK: return np.concatenate((mask, mask, mask)) elif self.channel_mode.get() == GUI.CHANNEL_COLOR: return img[0:3, :, :] elif self.channel_mode.get() == GUI.CHANNEL_NORMAL: return (img[4:7, :, :] * 0.5 + 0.5) * mask elif self.channel_mode.get() == GUI.CHANNEL_DEPTH: return np.concatenate([img[7:8, :, :]] * 3, axis=0) elif self.channel_mode.get() == GUI.CHANNEL_FLOW: return np.concatenate( (img[8:9, :, :] * 10, img[9:10, :, :] * 10, np.zeros_like(img[6:7, :, :]))) + 0.5 else: return self.black else: # IsoUnshaded if self.channel_mode.get() == GUI.CHANNEL_MASK: if img.dtype == np.uint8: mask = img[0:1, :, :] else: mask = img[0:1, :, :] * 0.5 + 0.5 return np.concatenate((mask, mask, mask)) elif self.channel_mode.get() == GUI.CHANNEL_NORMAL: if img.dtype == np.uint8: return img[1:4, :, :] else: return img[1:4, :, :] * 0.5 + 0.5 elif self.channel_mode.get() == GUI.CHANNEL_DEPTH: return np.concatenate( (img[4:5, :, :], img[4:5, :, :], img[4:5, :, :])) elif self.channel_mode.get() == GUI.CHANNEL_AO: return np.concatenate( (img[5:6, :, :], img[5:6, :, :], img[5:6, :, :])) elif self.channel_mode.get() == GUI.CHANNEL_FLOW: return np.concatenate( (img[6:7, :, :] * 10, img[7:8, :, :] * 10, np.zeros_like(img[6:7, :, :]))) + 0.5 elif self.channel_mode.get() == GUI.CHANNEL_COLOR: if img.dtype == np.uint8: shading_input = torch.unsqueeze( torch.from_numpy(img.astype(np.float32) / 255.0), 0) shading_input[:, 1: 4, :, :] = shading_input[:, 1: 4, :, :] * 2 - 1 else: shading_input = torch.unsqueeze( torch.from_numpy(img), 0) shading_output = self.shading(shading_input)[0] return torch.clamp(shading_output, 0, 1).cpu().numpy() img = self.dset[self.selected_entry, self.selected_time, :, :, :] img = selectChannel(img) img = img.transpose((2, 1, 0)) self.setImage(img)