def check_models_existence() -> List:
"""
depth_edge_model_ckpt: checkpoints/edge-model.pth
depth_feat_model_ckpt: checkpoints/depth-model.pth
rgb_feat_model_ckpt: checkpoints/color-model.pth
MiDaS_model_ckpt: MiDaS/model.pt
"""
missing = []
for name, filename in MODELS.items():
if not (MODELS_DIR / filename).is_file():
missing.append(name)
continue
remote_size = get_total_size([f'{MODELS_URL_ROOT}/{MODELS[name]}'])
local_size = (MODELS_DIR / filename).lstat().st_size
if remote_size != local_size:
missing.append(name)
print(
f(
_('File {MODELS[name]} is incorrect (remote: {remote_size} != local: {local_size}). Will re-download.'
)))
if missing:
print(
f(
_('The following models are missing: {", ".join(MODELS[m] for m in missing)}! Please download them first.'
)))
print(_('You can find them there:'))
print(*[f'{MODELS_URL_ROOT}/{MODELS[m]}' for m in missing], sep='\n')
print(f(_('Please put them in {MODELS_DIR}')))
return missing
def make_video(sample, config, props, depth, normal_canvas, all_canvas):
image = imageio.imread(sample['ref_img_fi'])
mean_loc_depth = depth[depth.shape[0] // 2, depth.shape[1] // 2]
verts, colors, faces, Height, Width, hFov, vFov = props
print(f(_("Making video at {datetime.now():%Y-%m-%d %H:%M:%S.%f}")))
videos_poses, video_basename = copy.deepcopy(sample['tgts_poses']), sample['tgt_name']
top = (config.get('original_h') // 2 - sample['int_mtx'][1, 2] * config['output_h'])
left = (config.get('original_w') // 2 - sample['int_mtx'][0, 2] * config['output_w'])
down, right = top + config['output_h'], left + config['output_w']
border = [int(xx) for xx in [top, down, left, right]]
normal_canvas, all_canvas = output_3d_photo(verts.copy(), colors.copy(), faces.copy(), copy.deepcopy(Height),
copy.deepcopy(Width), copy.deepcopy(hFov), copy.deepcopy(vFov),
copy.deepcopy(sample['tgt_pose']), sample['video_postfix'],
copy.deepcopy(sample['ref_pose']),
copy.deepcopy(config['video_folder']),
image.copy(), copy.deepcopy(sample['int_mtx']), config, image,
videos_poses, video_basename, config.get('original_h'),
config.get('original_w'), border=border, depth=depth,
normal_canvas=normal_canvas, all_canvas=all_canvas,
mean_loc_depth=mean_loc_depth)
print(f(_("{len(sample['video_postfix'])} videos done in {Path(config['video_folder']).resolve()}")))
def choose(self):
'''
Call plyer filechooser API to run a filechooser Activity.
'''
path = Path(self.bnd_text_input.text)
is_valid_path = Path(
path := os.path.dirname(path)).exists() and bool(path)
properties = dict(
on_selection=self.handle_selection,
title=_("Pick a JPG file.."),
multiple=False,
path=path if is_valid_path else os.path.expanduser("~"),
filters=[(_("Image file (jpg)"), "*.jpg")])
properties.update(self.filetype)
filechooser.open_file(**properties)
def load_rgb_model(device: str, rgb_feat_model_ckpt: str) -> Inpaint_Color_Net:
print(f(_("Loading rgb model at {datetime.now():%Y-%m-%d %H:%M:%S.%f}")))
rgb_model = Inpaint_Color_Net()
rgb_feat_weight = torch.load(rgb_feat_model_ckpt, map_location=torch.device(device))
rgb_model.load_state_dict(rgb_feat_weight)
rgb_model.eval()
return rgb_model.to(device)
def load_depth_model(device: str, depth_feat_model_ckpt: str) -> Inpaint_Depth_Net:
print(f(_("Loading depth model at {datetime.now():%Y-%m-%d %H:%M:%S.%f}")))
depth_feat_model = Inpaint_Depth_Net()
depth_feat_weight = torch.load(depth_feat_model_ckpt, map_location=torch.device(device))
depth_feat_model.load_state_dict(depth_feat_weight, strict=True)
depth_feat_model = depth_feat_model.to(device)
depth_feat_model.eval()
return depth_feat_model.to(device)
def download_thread(self):
print(_('Will start downloading the models...'))
try:
download_models(
self.missing_models,
bar_total=self.bar_total,
bar_current=self.bar_current,
)
self.missing_models = []
except Exception as e:
print(e)
exc_type, exc_value, exc_traceback = sys.exc_info()
traceback.print_exception(exc_type, exc_value, exc_traceback,
file=sys.stdout)
print(_('Model download: Done (failure).'))
else:
print(_('Model download: Done (success).'))
self.revalidate()
def deep_thread(self):
print(_('Will start working...'))
try:
wrap_3d_pi_with_override(
Path(self.image_handler.bnd_image.source),
depth_handler=self.depth_handler,
bar_total=self.bar_total,
bar_current=self.bar_current,
just_depth=False,
)
except Exception as e:
print(e)
exc_type, exc_value, exc_traceback = sys.exc_info()
traceback.print_exception(exc_type, exc_value, exc_traceback,
file=sys.stdout)
print(_('Done working: failure!'))
else:
print(_('Done working: success!'))
self.tr_text = START_DEEP
self.revalidate()
class FileChoose(Button):
'''
Button that triggers 'filechooser.open_file()' and processes
the data response from filechooser Activity.
'''
selection = ListProperty([])
filetype = DictProperty({"title": _("Pick a file..."), "filters": []})
bnd_text_input = ObjectProperty(None)
bnd_image = ObjectProperty(None)
def choose(self):
'''
Call plyer filechooser API to run a filechooser Activity.
'''
path = Path(self.bnd_text_input.text)
is_valid_path = Path(
path := os.path.dirname(path)).exists() and bool(path)
properties = dict(
on_selection=self.handle_selection,
title=_("Pick a JPG file.."),
multiple=False,
path=path if is_valid_path else os.path.expanduser("~"),
filters=[(_("Image file (jpg)"), "*.jpg")])
properties.update(self.filetype)
filechooser.open_file(**properties)
def handle_selection(self, selection):
'''
Callback function for handling the selection response from Activity.
'''
self.selection = selection
def on_selection(self, *a, **k):
'''
Update TextInput.text after FileChoose.selection is changed
via FileChoose.handle_selection.
'''
self.bnd_text_input.text = str(self.selection[0])
self.bnd_image.set_source(Path(str(self.selection[0])))
def download_models(models: List,
*,
bar_total: Optional[ComplexProgressBar] = None,
bar_current: Optional[ComplexProgressBar] = None):
os.makedirs(MODELS_DIR, exist_ok=True)
to_download = []
for model in models:
if model not in MODELS:
raise ValueError(f(_("Unknown model {model}!")))
remote_url = f'{MODELS_URL_ROOT}/{MODELS[model]}'
local_path = MODELS_DIR / MODELS[model]
to_download.append((remote_url, local_path))
if bar_total is not None:
size = get_total_size([i[0] for i in to_download])
bar_total.max = size
bar_total.text = lambda *, internal_name, max_: f"{internal_name} 0 / {sizeof_fmt(max_)}"
for remote_url, local_path in to_download:
download_file(remote_url,
local_path,
bar_total=bar_total,
bar_current=bar_current)
def wrap_3d_pi_with_override(image_filename: Path, *,
depth_handler: Optional[FileChoose] = None,
bar_total: Optional[ComplexProgressBar] = None,
bar_current: Optional[ComplexProgressBar] = None,
just_depth: bool = False,
):
# with the GUI we always have exactly 1 input image file
number_input_image = 1
bar_total.reset()
bar_current.reset()
real_image_filename = image_filename.resolve()
if not real_image_filename.is_file():
raise ValueError(f(_('File {real_image_filename} does not exist!')))
chunk, ext = os.path.splitext(real_image_filename)
dirname, filename = os.path.split(chunk)
origin_conf = CONFIG_CUSTOM if CONFIG_CUSTOM.is_file() else CONFIG_ORIGIN
with open(origin_conf, 'r') as config_file:
config = yaml.load(config_file)
config['depth_folder'] = str(Path(config['depth_folder']).resolve())
config['mesh_folder'] = str(Path(config['mesh_folder']).resolve())
config['video_folder'] = str(Path(config['video_folder']).resolve())
config['src_folder'] = dirname
config['specific'] = filename
config['img_format'] = ext
for name, filename in MODELS.items():
config[name] = str(MODELS_DIR / filename)
missing_models = check_models_existence()
with open(CONFIG_CUSTOM, 'w') as config_file:
yaml.dump(config, config_file, default_flow_style=None)
if not just_depth:
from_file = Path(depth_handler.bnd_image.source).resolve()
to_file = Path(config['depth_folder'], from_file.parts[-1].rsplit('.', maxsplit=1)[0] + '.png')
if from_file.samefile(to_file):
print(f(_('Depth file {from_file} is already in the correct location. No copy needed.')))
else:
print(f(_('Copying Depth file from {from_file} to {to_file}')))
shutil.copy(str(from_file), str(to_file))
update_image_handler(image_handler=depth_handler, path=to_file)
inject_write_videofile(
num_videos=len(config['video_postfix']) * number_input_image,
total_allocated_percent=20 / 100,
bar_total=bar_total,
bar_current=bar_current,
)
if not missing_models:
bar_current.add(bar_current.max)
bar_total.add(bar_total.max * 2 / 100)
wrap_3d_photo_inpainting(
CONFIG_CUSTOM,
depth_handler=depth_handler,
bar_total=bar_total,
bar_current=bar_current,
just_depth=just_depth
)
else:
raise ValueError(_("Models have not been downloaded!"))
def load_edge_model(device: str, depth_edge_model_ckpt: str) -> Inpaint_Edge_Net:
print(f(_("Loading edge model at {datetime.now():%Y-%m-%d %H:%M:%S.%f}")))
depth_edge_model = Inpaint_Edge_Net(init_weights=True)
depth_edge_weight = torch.load(depth_edge_model_ckpt, map_location=torch.device(device))
depth_edge_model.load_state_dict(depth_edge_weight)
return depth_edge_model.to(device)
def wrap_3d_photo_inpainting(config_path, *,
depth_handler: Optional[FileChoose] = None,
bar_total: Optional[ComplexProgressBar] = None,
bar_current: Optional[ComplexProgressBar] = None,
just_depth: bool = False
):
bar_current.reset()
config = yaml.load(open(config_path, 'r'))
if config['offscreen_rendering'] is True:
vispy.use(app='egl')
init_fs(config)
sample_list = get_MiDaS_samples(config['src_folder'], config['depth_folder'], config, config['specific'])
normal_canvas, all_canvas = None, None
if isinstance(config["gpu_ids"], int) and (config["gpu_ids"] >= 0):
device = config["gpu_ids"]
else:
device = "cpu"
bar_current.add(bar_current.max)
bar_total.add(bar_total.max * 2 / 100)
print(f(_("running on device {device}")))
for idx in tqdm(range(len(sample_list))):
bar_current.reset()
depth = None
sample = sample_list[idx]
print(f(_("Current Source ==> {sample['src_pair_name']}")))
mesh_fi = os.path.join(config['mesh_folder'], sample['src_pair_name'] +'.ply')
image = imageio.imread(sample['ref_img_fi'])
print(f(_("Running depth extraction at {datetime.now():%Y-%m-%d %H:%M:%S.%f}")))
if just_depth or config['require_midas'] is True:
run_depth([sample['ref_img_fi']], config['src_folder'], config['depth_folder'],
config['MiDaS_model_ckpt'], MonoDepthNet, MiDaS_utils, target_w=640)
update_image_handler(image_handler=depth_handler,
path=Path(f"{config['depth_folder']}/{sample['src_pair_name']}.png"))
if just_depth:
bar_total.reset()
bar_total.value = bar_total.max
bar_current.reset()
bar_current.value = bar_current.max
return
bar_current.add(bar_current.max)
bar_total.add(bar_total.max * (2 / len(sample_list)) / 100)
bar_current.reset()
image = prepare_config_and_image(config=config, sample=sample, image=image)
bar_current.add(bar_current.max)
bar_total.add(bar_total.max * (2 / len(sample_list)) / 100)
bar_current.reset()
image = cv2.resize(image, (config['output_w'], config['output_h']), interpolation=cv2.INTER_AREA)
depth = read_MiDaS_depth(sample['depth_fi'], 3.0, config['output_h'], config['output_w'])
mean_loc_depth = depth[depth.shape[0]//2, depth.shape[1]//2]
bar_current.add(bar_current.max)
bar_total.add(bar_total.max * (2 / len(sample_list)) / 100)
bar_current.reset()
if not(config['load_ply'] is True and os.path.exists(mesh_fi)):
vis_photos, vis_depths = sparse_bilateral_filtering(depth.copy(), image.copy(), config, num_iter=config['sparse_iter'], spdb=False)
depth = vis_depths[-1]
model = None
torch.cuda.empty_cache()
print(_("Start Running 3D_Photo ..."))
depth_edge_model = load_edge_model(device=device, depth_edge_model_ckpt=config['depth_edge_model_ckpt'])
depth_edge_model.eval()
depth_feat_model = load_depth_model(device=device, depth_feat_model_ckpt=config['depth_feat_model_ckpt'])
rgb_model = load_rgb_model(device=device, rgb_feat_model_ckpt=config['rgb_feat_model_ckpt'])
graph = None
def up_bars(dt=None):
bar_current.add(bar_current.max * 1.5 / 100)
bar_total.add(bar_total.max * (1 / len(sample_list)) / 100)
# increase the bars every 5 sec, up to 5 min
event = schedule_interval(up_bars, 5, 60 * 5)
print(f(_("Writing depth ply (and basically doing everything) at {datetime.now():%Y-%m-%d %H:%M:%S.%f}")))
rt_info = write_ply(image,
depth,
sample['int_mtx'],
mesh_fi,
config,
rgb_model,
depth_edge_model,
depth_edge_model,
depth_feat_model)
if rt_info is False:
continue
rgb_model = None
color_feat_model = None
depth_edge_model = None
depth_feat_model = None
torch.cuda.empty_cache()
event.cancel()
bar_current.add(bar_current.max)
bar_total.value_normalized = 75 / 100
bar_current.reset()
props = read_ply(mesh_fi) if config['save_ply'] is True or config['load_ply'] is True else rt_info
make_video(
sample=sample, config=config, props=props,
depth=depth, normal_canvas=normal_canvas, all_canvas=all_canvas,
)
bar_current.value_normalized = 1
bar_total.value_normalized = 1
def tr_text(self, value):
self.text = _(value)
def __eq__(self, other):
return super(TranslatedStr, self).__eq__(_(other))
from kivy.app import App
from kivy.clock import mainthread
from kivy.properties import ListProperty, ObjectProperty
from kivy.uix.button import Button
from wrapper import wrap_3d_pi_with_override
from utilities import download_models
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from kv_classes import ComplexProgressBar, FileChoose
from kv_classes.localization import _
START_DOWNLOAD = _('Download models (required first time)')
START_DOWNLOADING = _('Downloading models...')
START_DEEP = _('Start')
START_DEEPING = _('Working, please wait...')
START_DEPTH = _('Generate Depth file')
START_DEPTHING = _('Generating Depth file...')
class TranslatedStr(str):
def __eq__(self, other):
return super(TranslatedStr, self).__eq__(_(other))
class StartButton(Button):
missing_models = ListProperty()
depth_handler: FileChoose = ObjectProperty(None)
def _find_child(self, type_: Type[T]) -> T:
for child in self.children:
if isinstance(child, type_):
return child
else:
raise AttributeError(f(_("The {self.__class__} has no {type_} child!")))
