├── screenshot.png
├── requirements.txt
├── README.md
└── main.py


/screenshot.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/ShenhanQian/nersemble-data-viewer/HEAD/screenshot.png


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | dearpygui
2 | numpy
3 | pillow
4 | pandas
5 | tyro
6 | opencv-python
7 | matplotlib


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # NeRSemble Data Viewer
 2 | 
 3 | A GUI program to inspect images in the [NeRSemble](https://tobias-kirschstein.github.io/nersemble/) dataset.
 4 | 
 5 | > [!NOTE]
 6 | > This program has been adapted to NeRSemble Dataset V2. The folder structure is no longer compatible with NeRSemble Dataset V1.
 7 | 
 8 | <div align="center">
 9 |         <image src="./screenshot.png" height=800px></image>
10 | </div>
11 | 
12 | ## Installation
13 | 
14 | ```shell
15 | pip install -r requirements.txt
16 | ```
17 | 
18 | ## Usage
19 | 
20 | Please first extract frames from the raw videos with [this script](https://github.com/ShenhanQian/VHAP/blob/main/doc/nersemble_v2.md#1-preprocess). Then you can view them with
21 | 
22 | ```shell
23 | python main.py --root_folder <root_folder>
24 | ```
25 | 
26 | ## Expected folder hierarchy
27 | 
28 | ```
29 | <root_folder>
30 | |- 017
31 |         |- sequences
32 |                 |- EXP-1-head
33 |                         |- images
34 |                                 |- cam_220700191_000000.jpg
35 |                                 |- cam_220700191_000001.jpg
36 |                                 |- ...
37 | ```
38 | 


--------------------------------------------------------------------------------
/main.py:
--------------------------------------------------------------------------------
  1 | from typing import Literal, Optional
  2 | import dearpygui.dearpygui as dpg
  3 | from pathlib import Path
  4 | import glob
  5 | import numpy as np
  6 | from PIL import Image
  7 | import pandas as pd
  8 | import tyro
  9 | from dataclasses import dataclass
 10 | import cv2
 11 | import matplotlib.pyplot as plt
 12 | 
 13 | 
 14 | region2seg_class = {
 15 |     'bg': 0,
 16 |     'skin': 1,
 17 |     'l_brow': 2,
 18 |     'r_brow': 3,
 19 |     'l_eye': 4,
 20 |     'r_eye': 5,
 21 |     'eye_g': 6,
 22 |     'l_ear': 7,
 23 |     'r_ear': 8,
 24 |     'ear_r': 9,
 25 |     'nose': 10,
 26 |     'mouth': 11,
 27 |     'u_lip': 12,
 28 |     'l_lip': 13,
 29 |     'neck': 14,
 30 |     'neck_l': 15,
 31 |     'cloth': 16,
 32 |     'hair': 17,
 33 |     'hat': 18
 34 | }
 35 | 
 36 | 
 37 | @dataclass
 38 | class NersembleDataViewerConfig:
 39 |     root_folder: Path
 40 |     width: int=1000
 41 |     height: int=1000
 42 | 
 43 | class NersembleDataViewer(object):
 44 |     def __init__(self, cfg: NersembleDataViewerConfig):
 45 |         self.root_folder = cfg.root_folder
 46 |         self.width = cfg.width
 47 |         self.height = cfg.height
 48 | 
 49 |         self.width_nav = 390
 50 |         self.height_nav = 400
 51 |         self.need_update = False
 52 | 
 53 |         # load csv
 54 |         csv_path = self.root_folder / 'metadata_sequences.csv'
 55 |         self.data = pd.read_csv(csv_path, sep=',')
 56 | 
 57 |         self.data[self.data.columns[0]] = self.data[self.data.columns[0]].astype(str)
 58 |         self.data.iloc[:, 0] = self.data.iloc[:, 0].map(lambda x: f'{int(x):03d}')
 59 |         self.data.iloc[:, 0]
 60 |         self.data.set_index("ID", inplace=True)
 61 | 
 62 |         self.data = self.data.drop(self.data.columns[0:2], axis=1)  # drop 'wears_glasses' and 'BACKGROUND'
 63 | 
 64 |         # init variables
 65 |         self.subjects = self.data.index.tolist()
 66 |         self.subjects = [x for x in self.subjects if len(list((self.root_folder / x).glob(f"sequences/*/images*/*.jpg"))) > 0]
 67 | 
 68 |         self.sequences = self.data.columns.tolist()
 69 |         self.sequences = [x for x in self.sequences if len(list((self.root_folder).glob(f"*/sequences/{x}/images*/*.jpg"))) > 0]
 70 | 
 71 |         self.reset_subject_sequence(update_items=False)
 72 | 
 73 |     def reset_subject_sequence(self, update_items=True):
 74 |         self.selected_subject = '-'
 75 |         self.selected_sequence = '-'
 76 |         self.available_subjects = self.subjects
 77 |         self.available_sequences = self.sequences
 78 | 
 79 |         if update_items:
 80 |             dpg.configure_item("combo_subject", items=['-'] + self.available_subjects, default_value=self.selected_subject)
 81 |             dpg.configure_item("combo_sequence", items=['-'] + self.available_sequences, default_value=self.selected_sequence)
 82 | 
 83 |     def reset_folder_tree(self, update_items=True):
 84 |         self.filetypes = []
 85 |         self.cameras = []
 86 |         self.timesteps = []
 87 |         self.selected_filetype = ''
 88 |         self.selected_camera = ''
 89 |         self.selected_timestep = ''
 90 |         self.selected_timestep_idx = 0
 91 |         if update_items:
 92 |             dpg.configure_item("combo_filetype", items=self.filetypes, default_value=self.selected_filetype)
 93 |             dpg.configure_item("combo_camera", items=self.cameras, default_value=self.selected_camera)
 94 |             dpg.configure_item("slider_timestep", default_value=self.selected_timestep_idx, max_value=0, min_value=0)
 95 | 
 96 |     def define_gui(self):
 97 |         dpg.create_context()
 98 | 
 99 |         with dpg.texture_registry(show=False):
100 |             dpg.add_raw_texture(width=self.width, height=self.height, default_value=np.zeros([self.height, self.width, 3]), format=dpg.mvFormat_Float_rgb, tag="texture_tag")
101 | 
102 |         # viewer window
103 |         with dpg.window(label="Viewer", pos=[0, 0], tag='viewer_tag', width=self.width, height=self.height, no_title_bar=True, no_move=True, no_bring_to_front_on_focus=True):
104 |             dpg.add_image("texture_tag", tag='image_tag', width=self.width, height=self.height)
105 | 
106 |         # navigator window
107 |         with dpg.window(label="Navigator", tag='navigator_tag', width=self.width_nav, pos=[self.width-self.width_nav-15, 0], autosize=True):
108 | 
109 |             # subject switch
110 |             with dpg.group(horizontal=True):
111 |                 def set_subject(sender, data):
112 |                     self.selected_subject = data
113 |                     if data == '-':
114 |                         self.available_sequences = self.sequences
115 |                     else:
116 |                         self.available_sequences = [x for x in self.sequences if len(list((self.root_folder).glob(f"{self.selected_subject}/sequences/{x}/images*/"))) > 0]
117 |                         
118 |                         if self.selected_sequence not in self.available_sequences:
119 |                             if len(self.available_sequences) > 0:
120 |                                 self.selected_sequence = self.available_sequences[0]
121 |                             else:
122 |                                 self.selected_sequence == '-'
123 | 
124 |                     dpg.configure_item("combo_sequence", items=['-'] + self.available_sequences, default_value=self.selected_sequence)
125 |                     self.update_folder_tree(level='sequence')
126 |                     self.need_update = True
127 |                 dpg.add_combo(['-'] + self.available_subjects, default_value=self.selected_subject, label="subject  ", height_mode=dpg.mvComboHeight_Large, callback=set_subject, tag='combo_subject')
128 |                 
129 |                 def prev_subject(sender, data):
130 |                     if self.selected_subject == '-':
131 |                         self.selected_subject = self.available_subjects[-1]
132 |                     else:
133 |                         idx = self.available_subjects.index(self.selected_subject)
134 |                         if idx > 0:
135 |                             self.selected_subject = self.available_subjects[idx-1]
136 |                         else:
137 |                             self.selected_subject = self.available_subjects[-1]
138 |                     dpg.set_value("combo_subject", value=self.selected_subject)
139 |                     set_subject(None, self.selected_subject)
140 |                 dpg.add_button(label="W", callback=prev_subject, width=19)
141 | 
142 |                 def next_subject(sender, data):
143 |                     if self.selected_subject == '-':
144 |                         self.selected_subject = self.available_subjects[0]
145 |                     else:
146 |                         idx = self.available_subjects.index(self.selected_subject)
147 |                         if idx < len(self.available_subjects)-1:
148 |                             self.selected_subject = self.available_subjects[idx+1]
149 |                         else:
150 |                             self.selected_subject = self.available_subjects[0]
151 |                     dpg.set_value("combo_subject", value=self.selected_subject)
152 |                     set_subject(None, self.selected_subject)
153 |                 dpg.add_button(label="S", callback=next_subject, width=19)
154 |             
155 |             # sequence switch
156 |             with dpg.group(horizontal=True):
157 |                 def set_sequence(sender, data):
158 |                     self.selected_sequence = data
159 |                     if data == '-':
160 |                         self.available_subjects = self.subjects
161 |                     else:
162 |                         self.available_subjects = [x for x in self.subjects if len(list((self.root_folder / x).glob(f"sequences/{self.selected_sequence}/images*/"))) > 0]
163 | 
164 |                         if self.selected_subject not in self.available_subjects:
165 |                             if len(self.available_subjects) > 0:
166 |                                 self.selected_subject = self.available_subjects[0]
167 |                             else:
168 |                                 self.selected_subject == '-'
169 |                     
170 |                     dpg.configure_item("combo_subject", items=['-'] + self.available_subjects, default_value=self.selected_subject)
171 |                     self.update_folder_tree(level='filetype')
172 |                     self.need_update = True
173 |                 dpg.add_combo(['-'] + self.available_sequences, default_value=self.selected_sequence, label="sequence ", height_mode=dpg.mvComboHeight_Large, callback=set_sequence, tag='combo_sequence')
174 | 
175 |                 def prev_sequence(sender, data):
176 |                     if self.selected_sequence == '-':
177 |                         self.selected_sequence = self.available_sequences[-1]
178 |                     else:
179 |                         idx = self.available_sequences.index(self.selected_sequence)
180 |                         if idx > 0:
181 |                             self.selected_sequence = self.available_sequences[idx-1]
182 |                         else:
183 |                             self.selected_sequence = self.available_sequences[-1]
184 |                     dpg.set_value("combo_sequence", value=self.selected_sequence)
185 |                     set_sequence(None, self.selected_sequence)
186 |                 dpg.add_button(label="A", callback=prev_sequence, width=19)
187 |                 
188 |                 def next_sequence(sender, data):
189 |                     if self.selected_sequence == '-':
190 |                         self.selected_sequence = self.available_sequences[0]
191 |                     else:
192 |                         idx = self.available_sequences.index(self.selected_sequence)
193 |                         if idx < len(self.available_sequences)-1:
194 |                             self.selected_sequence = self.available_sequences[idx+1]
195 |                         else:
196 |                             self.selected_sequence = self.available_sequences[0]
197 |                     dpg.set_value("combo_sequence", value=self.selected_sequence)
198 |                     set_sequence(None, self.selected_sequence)
199 |                 dpg.add_button(label="D", callback=next_sequence, width=19)
200 |             
201 |             # filetype switch
202 |             def set_filetype(sender, data):
203 |                 self.selected_filetype = data
204 |                 self.update_folder_tree(level='camera')
205 |                 self.need_update = True
206 |             dpg.add_combo([], label="file type", height_mode=dpg.mvComboHeight_Large, callback=set_filetype, tag='combo_filetype')
207 | 
208 | 
209 |             # camera switch
210 |             with dpg.group(horizontal=True):
211 |                 def set_camera(sender, data):
212 |                     self.selected_camera = data
213 |                     self.update_folder_tree(level='camera')
214 |                     self.need_update = True
215 |                 dpg.add_combo([], label="camera   ", height_mode=dpg.mvComboHeight_Large, callback=set_camera, tag='combo_camera')
216 | 
217 |                 def prev_camera(sender, data):
218 |                     if len(self.cameras) > 0:
219 |                         idx = self.cameras.index(self.selected_camera)
220 |                         if idx > 0:
221 |                             self.selected_camera = self.cameras[idx-1]
222 |                         else:
223 |                             self.selected_camera = self.cameras[-1]
224 |                         dpg.set_value("combo_camera", value=self.selected_camera)
225 |                         set_camera(None, self.selected_camera)
226 |                 dpg.add_button(label="", callback=prev_camera, arrow=True, direction=dpg.mvDir_Up)
227 | 
228 |                 def next_camera(sender, data):
229 |                     if len(self.cameras) > 0:
230 |                         idx = self.cameras.index(self.selected_camera)
231 |                         if idx < len(self.cameras)-1:
232 |                             self.selected_camera = self.cameras[idx+1]
233 |                         else:
234 |                             self.selected_camera = self.cameras[0]
235 |                         dpg.set_value("combo_camera", value=self.selected_camera)
236 |                         set_camera(None, self.selected_camera)
237 |                 dpg.add_button(label="Button", callback=next_camera, arrow=True, direction=dpg.mvDir_Down)
238 | 
239 |             # timestep switch
240 |             with dpg.group(horizontal=True):
241 |                 def set_timestep_slider(sender, data):
242 |                     self.selected_timestep_idx = data
243 |                     self.selected_timestep = self.timesteps[self.selected_timestep_idx]
244 |                     self.update_folder_tree(level='filetype')
245 |                     self.need_update = True
246 |                 dpg.add_slider_int(label="time step", max_value=len(self.timesteps), callback=set_timestep_slider, tag='slider_timestep')
247 | 
248 |                 def prev_timestep(sender, data):
249 |                     if len(self.timesteps) > 1:
250 |                         if self.selected_timestep_idx > 0:
251 |                             self.selected_timestep_idx -= 1
252 |                         self.selected_timestep = self.timesteps[self.selected_timestep_idx]
253 |                         dpg.set_value("slider_timestep", value=self.selected_timestep_idx)
254 |                         set_timestep_slider(None, self.selected_timestep_idx)
255 |                 dpg.add_button(label="Button", callback=prev_timestep, arrow=True, direction=dpg.mvDir_Left)
256 | 
257 |                 def next_timestep(sender, data):
258 |                     if len(self.timesteps) > 1:
259 |                         if self.selected_timestep_idx < len(self.timesteps)-1:
260 |                             self.selected_timestep_idx += 1
261 |                         self.selected_timestep = self.timesteps[self.selected_timestep_idx]
262 |                         dpg.set_value("slider_timestep", value=self.selected_timestep_idx)
263 |                         set_timestep_slider(None, self.selected_timestep_idx)            
264 |                 dpg.add_button(label="Button", callback=next_timestep, arrow=True, direction=dpg.mvDir_Right)
265 | 
266 |                 def set_timestep(sender, data):
267 |                     if self.selected_subject == '-' or self.selected_sequence == '-':
268 |                         return
269 |                     
270 |                     if sender == 'mvKey_Home':
271 |                         self.selected_timestep_idx = 0
272 |                     elif sender == 'mvKey_End':
273 |                         if len(self.timesteps) > 0:
274 |                             self.selected_timestep_idx = len(self.timesteps)-1
275 |                     self.selected_timestep = self.timesteps[self.selected_timestep_idx]
276 |                     dpg.set_value("slider_timestep", value=self.selected_timestep_idx)
277 |                     set_timestep_slider(None, self.selected_timestep_idx)
278 |             
279 |             # annotations
280 |             dpg.add_separator()
281 |             def set_annotation(sender, data):
282 |                 self.need_update = True
283 |             with dpg.group(horizontal=True):
284 |                 dpg.add_checkbox(label="landmarks (STAR)  ", callback=set_annotation, tag='checkbox_lmk_star', show=False, default_value=False)
285 |                 dpg.add_checkbox(label="landmarks (PIPnet)", callback=set_annotation, tag='checkbox_lmk_pipnet', show=False, default_value=False)
286 |             with dpg.group(horizontal=True):
287 |                 dpg.add_checkbox(label="landmarks (FA)    ", callback=set_annotation, tag='checkbox_lmk_fa', show=False, default_value=False)
288 |             with dpg.group(horizontal=True):
289 |                 dpg.add_checkbox(label="foreground        ", callback=set_annotation, tag='checkbox_fg', show=False, default_value=False)
290 |                 dpg.add_checkbox(label="segmentation      ", callback=set_annotation, tag='checkbox_seg', show=False, default_value=False)
291 |             
292 |             with dpg.collapsing_header(label="Filter regions", default_open=False, show=False, tag='collapsing_filter_regions'):
293 |                 def set_region(sender, data):
294 |                     self.need_update = True
295 | 
296 |                 n_cols = 5
297 |                 n_rows = len(region2seg_class) // n_cols + int((len(region2seg_class) % n_cols) > 0)
298 |                 
299 |                 for i in range(n_rows):
300 |                     dpg.add_group(tag=f'filter_region_group_{i}', horizontal=True)
301 | 
302 |                 for i, region in enumerate(region2seg_class.keys()):
303 |                     dpg.add_checkbox(label=f'{region:6s}', callback=set_region, tag=f'checkbox_{region}', show=True, default_value=True, parent=f'filter_region_group_{i // n_cols}')
304 | 
305 |         
306 |         # key press handlers
307 |         with dpg.handler_registry():
308 |             dpg.add_key_press_handler(dpg.mvKey_W, callback=prev_subject)
309 |             dpg.add_key_press_handler(dpg.mvKey_S, callback=next_subject)
310 |             dpg.add_key_press_handler(dpg.mvKey_A, callback=prev_sequence)
311 |             dpg.add_key_press_handler(dpg.mvKey_D, callback=next_sequence)
312 |             dpg.add_key_press_handler(dpg.mvKey_Left, callback=prev_timestep)
313 |             dpg.add_key_press_handler(dpg.mvKey_Right, callback=next_timestep)
314 |             dpg.add_key_press_handler(dpg.mvKey_Up, callback=prev_camera)
315 |             dpg.add_key_press_handler(dpg.mvKey_Down, callback=next_camera)
316 |             dpg.add_key_press_handler(dpg.mvKey_Home, callback=set_timestep, tag='mvKey_Home')
317 |             dpg.add_key_press_handler(dpg.mvKey_End, callback=set_timestep, tag='mvKey_End')
318 | 
319 |         # theme
320 |         with dpg.theme() as theme_no_padding:
321 |             with dpg.theme_component(dpg.mvAll):
322 |                 # set all padding to 0 to avoid scroll bar
323 |                 dpg.add_theme_style(dpg.mvStyleVar_WindowPadding, 0, 0, category=dpg.mvThemeCat_Core)
324 |                 dpg.add_theme_style(dpg.mvStyleVar_FramePadding, 0, 0, category=dpg.mvThemeCat_Core)
325 |                 dpg.add_theme_style(dpg.mvStyleVar_CellPadding, 0, 0, category=dpg.mvThemeCat_Core)
326 |         dpg.bind_item_theme("viewer_tag", theme_no_padding)
327 | 
328 |     def resize_windows(self):
329 |         dpg.configure_item('viewer_tag', width=self.width, height=self.height)
330 |         dpg.configure_item('navigator_tag', width=self.width_nav, height=self.height_nav, pos=[self.width-self.width_nav-15, 0])
331 | 
332 |         dpg.delete_item('texture_tag')
333 |         dpg.delete_item('image_tag')
334 |         with dpg.texture_registry(show=False):
335 |             dpg.add_raw_texture(width=self.width, height=self.height, default_value=np.zeros([self.height, self.width, 3]), format=dpg.mvFormat_Float_rgb, tag="texture_tag")
336 |         dpg.add_image("texture_tag", tag='image_tag', parent='viewer_tag')
337 |         self.need_update = True
338 | 
339 |     def run(self):
340 |         self.reset_folder_tree(update_items=False)
341 |         self.define_gui()
342 |         dpg.create_viewport(title='NersembleData Viewer', width=self.width, height=self.height, resizable=True)
343 |         dpg.setup_dearpygui()
344 |         dpg.show_viewport()
345 | 
346 |         while dpg.is_dearpygui_running():
347 |             if self.width != dpg.get_viewport_width() or self.height != dpg.get_viewport_height():
348 |                 self.width = dpg.get_viewport_width()
349 |                 self.height = dpg.get_viewport_height()
350 |                 self.resize_windows()
351 | 
352 |             if self.need_update:
353 |                 self.update_viewer()
354 |                 self.need_update = False
355 |             dpg.render_dearpygui_frame()
356 |         dpg.destroy_context()
357 |     
358 |     def iterdir(self, subject=None, sequence=None, filetype=None):
359 |         if filetype is not None:
360 |             return [x.name for x in (self.root_folder / subject / 'sequences' / sequence / filetype).iterdir()]
361 |         elif sequence is not None:
362 |             return [x.name for x in (self.root_folder / subject / 'sequences' / sequence).iterdir() if x.is_dir()]
363 |         elif subject is not None:
364 |             return [x.name for x in (self.root_folder / subject / 'sequences').iterdir()]
365 |         else:
366 |             raise ValueError("Invalid arguments")
367 |     
368 |     def update_annotations(self):
369 |         lmk_star_path = self.root_folder / self.selected_subject / 'sequences' / self.selected_sequence / 'landmark2d' / 'STAR' / f"{self.selected_camera.replace('cam_', '')}.npz"
370 |         if lmk_star_path.exists():
371 |             dpg.configure_item('checkbox_lmk_star', show=True)
372 |         else:
373 |             dpg.configure_item('checkbox_lmk_star', show=False)
374 | 
375 |         lmk_pipnet_path = self.root_folder / self.selected_subject / 'sequences'/ self.selected_sequence / 'landmark2d' / 'PIPnet' / f"{self.selected_camera.replace('cam_', '')}.npy"
376 |         if lmk_pipnet_path.exists():
377 |             dpg.configure_item('checkbox_lmk_pipnet', show=True)
378 |         else:
379 |             dpg.configure_item('checkbox_lmk_pipnet', show=False)
380 | 
381 |         lmk_fa_path = self.root_folder / self.selected_subject / 'sequences'/ self.selected_sequence / 'landmark2d' / 'face-alignment' / f"{self.selected_camera.replace('cam_', '')}.npz"
382 |         if lmk_fa_path.exists():
383 |             dpg.configure_item('checkbox_lmk_fa', show=True)
384 |         else:
385 |             dpg.configure_item('checkbox_lmk_fa', show=False)
386 |         
387 |         fg = self.root_folder / self.selected_subject / 'sequences'/ self.selected_sequence / 'alpha_maps' / f'{self.selected_camera}_{self.selected_timestep}.jpg'
388 |         if fg.exists():
389 |             dpg.configure_item('checkbox_fg', show=True)
390 |         else:
391 |             dpg.configure_item('checkbox_fg', show=False)
392 |         
393 |         seg = self.root_folder / self.selected_subject / 'sequences'/ self.selected_sequence / 'bisenet_segmentation_masks' / f'{self.selected_camera}_{self.selected_timestep}.jpg'
394 |         if seg.exists():
395 |             dpg.configure_item('checkbox_seg', show=True)
396 |             dpg.configure_item('collapsing_filter_regions', show=True)
397 |         else:
398 |             dpg.configure_item('checkbox_seg', show=False)
399 |             dpg.configure_item('collapsing_filter_regions', show=False)
400 | 
401 |     def update_folder_tree(self, level=Optional[Literal['sequence', 'filetype', 'camera', 'timestep']]):
402 |         if self.selected_sequence == '-' or self.selected_subject == '-':
403 |             self.reset_folder_tree()
404 |             return
405 |         
406 |         update_sequence = level in ['sequence']
407 |         update_filetype = level in ['sequence', 'filetype']
408 |         update_camera = level in ['sequence', 'filetype', 'camera']
409 |         update_timestep = level in ['sequence', 'filetype', 'camera', 'timestep']
410 | 
411 |         if update_sequence:
412 |             self.available_sequences = [
413 |                 x for x in self.iterdir(self.selected_subject) \
414 |                     if len(list((self.root_folder / self.selected_subject).glob(f'sequences/{x}/images/*.jpg'))) > 0
415 |             ]
416 |             if self.selected_sequence not in self.available_sequences:
417 |                 self.selected_sequence = self.available_sequences[0]
418 |             dpg.configure_item("combo_sequence", items=self.available_sequences, default_value=self.selected_sequence)
419 | 
420 |         if update_filetype:
421 |             self.filetypes = self.iterdir(self.selected_subject, self.selected_sequence)
422 |             self.filetypes_images = sorted([x for x in self.filetypes if 'image' in x], reverse=True)
423 |             if self.selected_filetype not in self.filetypes:
424 |                 self.selected_filetype = self.filetypes_images[0] if len(self.filetypes_images) > 0 else self.filetypes[0]
425 |             dpg.configure_item("combo_filetype", items=self.filetypes, default_value=self.selected_filetype)
426 |             
427 |         if update_camera:
428 |             fnames = self.iterdir(self.selected_subject, self.selected_sequence, self.selected_filetype)
429 |             self.cameras = sorted(set(['_'.join(f.split('.')[0].split('_')[:-1]) for f in fnames if '.mp4' not in f]))
430 |             if self.selected_camera not in self.cameras:
431 |                 self.selected_camera = self.cameras[0]
432 |             dpg.configure_item("combo_camera", items=self.cameras, default_value=self.selected_camera)
433 | 
434 |         if update_timestep:
435 |             fnames = self.iterdir(self.selected_subject, self.selected_sequence, self.selected_filetype)
436 |             self.timesteps = sorted(set([f.split('.')[0].split('_')[-1] for f in fnames]))
437 |             if self.selected_timestep not in self.timesteps:
438 |                 self.selected_timestep = self.timesteps[0]
439 |             self.selected_timestep_idx = self.timesteps.index(self.selected_timestep)
440 |             dpg.configure_item("slider_timestep", max_value=len(self.timesteps)-1, default_value=self.selected_timestep_idx)
441 |         
442 |         self.update_annotations()
443 | 
444 |     def update_viewer(self):
445 |         if self.selected_sequence != '-' and self.selected_subject != '-':
446 |             
447 |             path = self.root_folder / self.selected_subject / 'sequences' / self.selected_sequence / self.selected_filetype
448 |             path = glob.glob(f'{str(path)}/{self.selected_camera}_{self.selected_timestep}*')
449 |             if len(path) == 0:
450 |                 return
451 |             
452 |             path = path[0]
453 |             if 'jpg' not in path.lower() and 'png' not in path.lower():
454 |                 return
455 |             
456 |             if self.selected_filetype == 'bisenet_segmentation_masks':
457 |                 # load as uint8 and get float32 after applying colormap
458 |                 img = self.load_image(path, Image.NEAREST)
459 |                 cm = plt.get_cmap('tab20c')
460 |                 img = cm(img[:, :, 0])[:, :, :3].astype(np.float32)
461 |             else:
462 |                 # directly load as float32
463 |                 img = self.load_image(path)
464 |                 img = img.astype(np.float32) / 255
465 |             
466 |             if dpg.get_item_configuration("checkbox_lmk_star")['show'] and dpg.get_value("checkbox_lmk_star"):
467 |                 npz_path = self.root_folder / self.selected_subject / 'sequences'/ self.selected_sequence / 'landmark2d' / 'STAR' / f"{self.selected_camera.replace('cam_', '')}.npz"
468 |                 npz = np.load(npz_path)
469 |                 lmk_star = npz['face_landmark_2d']
470 |                 bbox_star = npz['bounding_box']
471 | 
472 |                 color = (0, 255, 0)
473 |                 for lmk in lmk_star[self.selected_timestep_idx]:
474 |                     x = int(lmk[0] * img.shape[1])
475 |                     y = int(lmk[1] * img.shape[0])
476 |                     cv2.circle(img, (x, y), 2, color, -1)
477 | 
478 |                 # x1, y1, x2, y2 = bbox_star[self.selected_timestep_idx][:4]
479 |                 # x1 = int(x1 * img.shape[1])
480 |                 # y1 = int(y1 * img.shape[0])
481 |                 # x2 = int(x2 * img.shape[1])
482 |                 # y2 = int(y2 * img.shape[0])
483 |                 # cv2.rectangle(img, (x1, y1), (x2, y2), color, 1)
484 | 
485 |             if dpg.get_item_configuration("checkbox_lmk_pipnet")['show'] and dpg.get_value("checkbox_lmk_pipnet"):
486 |                 npy_path = self.root_folder / self.selected_subject / 'sequences'/ self.selected_sequence / 'landmark2d' / 'PIPnet' / f"{self.selected_camera.replace('cam_', '')}.npy"
487 |                 lmk_pipnet = np.load(npy_path)
488 | 
489 |                 color = (0, 0, 255)
490 |                 for lmk in lmk_pipnet[self.selected_timestep_idx]:
491 |                     x = int(lmk[0] * img.shape[1])
492 |                     y = int(lmk[1] * img.shape[0])
493 |                     cv2.circle(img, (x, y), 2, color, -1)
494 |             
495 |             if dpg.get_item_configuration("checkbox_lmk_fa")['show'] and dpg.get_value("checkbox_lmk_fa"):
496 |                 npz_path = self.root_folder / self.selected_subject / 'sequences'/ self.selected_sequence / 'landmark2d' / 'face-alignment' / f"{self.selected_camera.replace('cam_', '')}.npz"
497 |                 npz = np.load(npz_path)
498 |                 lmk_fa = npz['face_landmark_2d']
499 |                 bbox_fa = npz['bounding_box']
500 | 
501 |                 color = (255, 0, 0)
502 |                 for lmk in lmk_fa[self.selected_timestep_idx]:
503 |                     x = int(lmk[0] * img.shape[1])
504 |                     y = int(lmk[1] * img.shape[0])
505 |                     cv2.circle(img, (x, y), 2, color, -1)
506 | 
507 |                 # x1, y1, x2, y2 = bbox_fa[self.selected_timestep_idx][:4]
508 |                 # x1 = int(x1 * img.shape[1])
509 |                 # y1 = int(y1 * img.shape[0])
510 |                 # x2 = int(x2 * img.shape[1])
511 |                 # y2 = int(y2 * img.shape[0])
512 |                 # cv2.rectangle(img, (x1, y1), (x2, y2), color, 1)
513 |             
514 |             if dpg.get_item_configuration("checkbox_fg")['show'] and dpg.get_value("checkbox_fg"):
515 |                 fg_path = self.root_folder / self.selected_subject / 'sequences'/ self.selected_sequence / 'alpha_maps' / f'{self.selected_camera}_{self.selected_timestep}.jpg'
516 |                 fg_alpha = self.load_image(fg_path).astype(np.float32)[..., :3] / 255
517 |                 img = img * (fg_alpha) + np.ones_like(img) * (1 - fg_alpha)
518 |             
519 |             if dpg.get_item_configuration("checkbox_seg")['show'] and dpg.get_value("checkbox_seg"):
520 |                 seg_path = self.root_folder / self.selected_subject / 'sequences'/ self.selected_sequence / 'bisenet_segmentation_masks' / f'{self.selected_camera}_{self.selected_timestep}.jpg'
521 |                 seg = self.load_image(seg_path, Image.NEAREST)
522 |                 cm = plt.get_cmap('tab20c')
523 |                 seg = cm(seg[:, :, 0])[:, :, :3].astype(np.float32)
524 |                 img = img * 0.5 + seg * 0.5
525 |             
526 |             if dpg.get_item_configuration("collapsing_filter_regions")['show']:
527 |                 seg_path = self.root_folder / self.selected_subject / 'sequences'/ self.selected_sequence / 'bisenet_segmentation_masks' / f'{self.selected_camera}_{self.selected_timestep}.jpg'
528 |                 seg = self.load_image(seg_path, Image.NEAREST)
529 |                 for region, seg_class in region2seg_class.items():
530 |                     if not dpg.get_value(f"checkbox_{region}"):
531 |                         mask = np.ones_like(img)
532 |                         mask[seg == seg_class] = 0
533 |                         img = img * mask + np.ones_like(img) * (1 - mask)
534 |             
535 |             img = np.pad(img, ((0, self.height - img.shape[0]), (0, self.width - img.shape[1]), (0, 0)), mode='constant', constant_values=0)
536 | 
537 |             dpg.set_value("texture_tag", img)
538 | 
539 |     def load_image(self, path, resample=Image.BILINEAR):
540 |         img = Image.open(path)
541 |         scale = min(self.height / img.height, self.width / img.width)
542 |         img = img.resize((int(img.width * scale), int(img.height * scale)), resample)
543 |         img = np.asarray(img)
544 |         if img.ndim == 2:
545 |             img = np.repeat(img[:, :, np.newaxis], 3, axis=2)
546 |         return img
547 | 
548 | 
549 | if __name__ == '__main__':
550 |     cfg = tyro.cli(NersembleDataViewerConfig)
551 |     app = NersembleDataViewer(cfg)
552 |     app.run()
553 | 


--------------------------------------------------------------------------------