stylegan3/viz/pickle_widget.py

171 lines
6.6 KiB
Python
Raw Normal View History

# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2021-10-07 11:55:26 +02:00
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
import glob
import os
import re
import dnnlib
import imgui
import numpy as np
from gui_utils import imgui_utils
from . import renderer
#----------------------------------------------------------------------------
def _locate_results(pattern):
return pattern
#----------------------------------------------------------------------------
class PickleWidget:
def __init__(self, viz):
self.viz = viz
self.search_dirs = []
self.cur_pkl = None
self.user_pkl = ''
self.recent_pkls = []
self.browse_cache = dict() # {tuple(path, ...): [dnnlib.EasyDict(), ...], ...}
self.browse_refocus = False
self.load('', ignore_errors=True)
def add_recent(self, pkl, ignore_errors=False):
try:
resolved = self.resolve_pkl(pkl)
if resolved not in self.recent_pkls:
self.recent_pkls.append(resolved)
except:
if not ignore_errors:
raise
def load(self, pkl, ignore_errors=False):
viz = self.viz
viz.clear_result()
viz.skip_frame() # The input field will change on next frame.
try:
resolved = self.resolve_pkl(pkl)
name = resolved.replace('\\', '/').split('/')[-1]
self.cur_pkl = resolved
self.user_pkl = resolved
viz.result.message = f'Loading {name}...'
viz.defer_rendering()
if resolved in self.recent_pkls:
self.recent_pkls.remove(resolved)
self.recent_pkls.insert(0, resolved)
except:
self.cur_pkl = None
self.user_pkl = pkl
if pkl == '':
viz.result = dnnlib.EasyDict(message='No network pickle loaded')
else:
viz.result = dnnlib.EasyDict(error=renderer.CapturedException())
if not ignore_errors:
raise
@imgui_utils.scoped_by_object_id
def __call__(self, show=True):
viz = self.viz
recent_pkls = [pkl for pkl in self.recent_pkls if pkl != self.user_pkl]
if show:
imgui.text('Pickle')
imgui.same_line(viz.label_w)
changed, self.user_pkl = imgui_utils.input_text('##pkl', self.user_pkl, 1024,
flags=(imgui.INPUT_TEXT_AUTO_SELECT_ALL | imgui.INPUT_TEXT_ENTER_RETURNS_TRUE),
width=(-1 - viz.button_w * 2 - viz.spacing * 2),
help_text='<PATH> | <URL> | <RUN_DIR> | <RUN_ID> | <RUN_ID>/<KIMG>.pkl')
if changed:
self.load(self.user_pkl, ignore_errors=True)
if imgui.is_item_hovered() and not imgui.is_item_active() and self.user_pkl != '':
imgui.set_tooltip(self.user_pkl)
imgui.same_line()
if imgui_utils.button('Recent...', width=viz.button_w, enabled=(len(recent_pkls) != 0)):
imgui.open_popup('recent_pkls_popup')
imgui.same_line()
if imgui_utils.button('Browse...', enabled=len(self.search_dirs) > 0, width=-1):
imgui.open_popup('browse_pkls_popup')
self.browse_cache.clear()
self.browse_refocus = True
if imgui.begin_popup('recent_pkls_popup'):
for pkl in recent_pkls:
clicked, _state = imgui.menu_item(pkl)
if clicked:
self.load(pkl, ignore_errors=True)
imgui.end_popup()
if imgui.begin_popup('browse_pkls_popup'):
def recurse(parents):
key = tuple(parents)
items = self.browse_cache.get(key, None)
if items is None:
items = self.list_runs_and_pkls(parents)
self.browse_cache[key] = items
for item in items:
if item.type == 'run' and imgui.begin_menu(item.name):
recurse([item.path])
imgui.end_menu()
if item.type == 'pkl':
clicked, _state = imgui.menu_item(item.name)
if clicked:
self.load(item.path, ignore_errors=True)
if len(items) == 0:
with imgui_utils.grayed_out():
imgui.menu_item('No results found')
recurse(self.search_dirs)
if self.browse_refocus:
imgui.set_scroll_here()
viz.skip_frame() # Focus will change on next frame.
self.browse_refocus = False
imgui.end_popup()
paths = viz.pop_drag_and_drop_paths()
if paths is not None and len(paths) >= 1:
self.load(paths[0], ignore_errors=True)
viz.args.pkl = self.cur_pkl
def list_runs_and_pkls(self, parents):
items = []
run_regex = re.compile(r'\d+-.*')
pkl_regex = re.compile(r'network-snapshot-\d+\.pkl')
for parent in set(parents):
if os.path.isdir(parent):
for entry in os.scandir(parent):
if entry.is_dir() and run_regex.fullmatch(entry.name):
items.append(dnnlib.EasyDict(type='run', name=entry.name, path=os.path.join(parent, entry.name)))
if entry.is_file() and pkl_regex.fullmatch(entry.name):
items.append(dnnlib.EasyDict(type='pkl', name=entry.name, path=os.path.join(parent, entry.name)))
items = sorted(items, key=lambda item: (item.name.replace('_', ' '), item.path))
return items
def resolve_pkl(self, pattern):
assert isinstance(pattern, str)
assert pattern != ''
# URL => return as is.
if dnnlib.util.is_url(pattern):
return pattern
# Short-hand pattern => locate.
path = _locate_results(pattern)
# Run dir => pick the last saved snapshot.
if os.path.isdir(path):
pkl_files = sorted(glob.glob(os.path.join(path, 'network-snapshot-*.pkl')))
if len(pkl_files) == 0:
raise IOError(f'No network pickle found in "{path}"')
path = pkl_files[-1]
# Normalize.
path = os.path.abspath(path)
return path
#----------------------------------------------------------------------------