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)
Esempio n. 2
0
    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))
Esempio n. 3
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)