hololens/stable-diffusion-webui-depthmap-script
1
1from pathlib import Path
2from PIL import Image
3
4try:
5 from tqdm import trange
6except:
7 from builtins import range as trange
8
9import torch, gc
10import cv2
11import os.path
12import numpy as np
13import copy
14import platform
15import math
16
17# Our code
18from src.misc import *
19from src.common_constants import GenerationOptions as go
20from src.common_constants import *
21from src.stereoimage_generation import create_stereoimages
22from src.normalmap_generation import create_normalmap
23from src.depthmap_generation import ModelHolder
24from src import backbone
25
26try:
27 # 3d-photo-inpainting imports
28 from inpaint.mesh import write_mesh, read_mesh, output_3d_photo
29 from inpaint.networks import Inpaint_Color_Net, Inpaint_Depth_Net, Inpaint_Edge_Net
30 from inpaint.utils import path_planning
31 from inpaint.bilateral_filtering import sparse_bilateral_filtering
32except Exception as e:
33 print('Impaint import failed. Impaint will not work.')
34 import traceback
35 traceback.print_exc()
36
37global video_mesh_data, video_mesh_fn
38video_mesh_data = None
39video_mesh_fn = None
40
41model_holder = ModelHolder()
42
43
44def convert_to_i16(arr):
45 # Single channel, 16 bit image. This loses some precision!
46 # uint16 conversion uses round-down, therefore values should be [0; 2**16)
47 numbytes = 2
48 max_val = (2 ** (8 * numbytes))
49 out = np.clip(arr * max_val + 0.0001, 0, max_val - 0.1) # -0.1 from above is needed to avoid overflowing
50 return out.astype("uint16")
51
52def convert_i16_to_rgb(image, like):
53 # three channel, 8 bits per channel image
54 output = np.zeros_like(like)
55 output[:, :, 0] = image / 256.0
56 output[:, :, 1] = image / 256.0
57 output[:, :, 2] = image / 256.0
58 return output
59
60
61class CoreGenerationFunnelInp:
62 """This class takes a dictionary and creates a core_generation_funnel inp.
63 Non-applicable parameters are silently discarded (no error)"""
64 def __init__(self, values):
65 if isinstance(values, CoreGenerationFunnelInp):
66 values = values.values
67 values = {(k.name if isinstance(k, GenerationOptions) else k).lower(): v for k, v in values.items()}
68
69 self.values = {}
70 for setting in GenerationOptions:
71 name = setting.name.lower()
72 self.values[name] = values[name] if name in values else setting.df
73
74 def __getitem__(self, item):
75 if isinstance(item, GenerationOptions):
76 return self.values[item.name.lower()]
77 return self.values[item]
78
79 def __getattr__(self, item):
80 return self[item]
81
82
83def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp, ops=None):
84 if len(inputimages) == 0 or inputimages[0] is None:
85 return
86 if inputdepthmaps is None or len(inputdepthmaps) == 0:
87 inputdepthmaps: list[Image] = [None for _ in range(len(inputimages))]
88 inputdepthmaps_complete = all([x is not None for x in inputdepthmaps])
89
90 inp = CoreGenerationFunnelInp(inp)
91
92 if ops is None:
93 ops = backbone.gather_ops()
94 model_holder.update_settings(**ops)
95
96 # TODO: ideally, run_depthmap should not save meshes - that makes the function not pure
97 print(SCRIPT_FULL_NAME)
98 print(f'Backbone: {backbone.USED_BACKBONE.name}')
99
100 backbone.unload_sd_model()
101
102 # TODO: this still should not be here
103 background_removed_images = []
104 # remove on base image before depth calculation
105 if inp[go.GEN_REMBG]:
106 if inp[go.PRE_DEPTH_BACKGROUND_REMOVAL]:
107 inputimages = batched_background_removal(inputimages, inp[go.REMBG_MODEL])
108 background_removed_images = inputimages
109 else:
110 background_removed_images = batched_background_removal(inputimages, inp[go.REMBG_MODEL])
111
112 # init torch device
113 if inp[go.COMPUTE_DEVICE] == 'GPU':
114 if torch.cuda.is_available():
115 device = torch.device("cuda")
116 else:
117 print('WARNING: Cuda device was not found, cpu will be used')
118 device = torch.device("cpu")
119 else:
120 device = torch.device("cpu")
121 print("device: %s" % device)
122
123 # TODO: This should not be here
124 inpaint_imgs = []
125 inpaint_depths = []
126
127 try:
128 if not inputdepthmaps_complete:
129 print("Loading model(s) ..")
130 model_holder.ensure_models(inp[go.MODEL_TYPE], device, inp[go.BOOST], inp[go.TILING_MODE])
131 print("Computing output(s) ..")
132 # iterate over input images
133 for count in trange(0, len(inputimages)):
134 # Convert single channel input (PIL) images to rgb
135 if inputimages[count].mode == 'I':
136 inputimages[count].point(lambda p: p * 0.0039063096, mode='RGB')
137 inputimages[count] = inputimages[count].convert('RGB')
138
139 raw_prediction = None
140 """Raw prediction, as returned by a model. None if input depthmap is used."""
141 raw_prediction_invert = False
142 """True if near=dark on raw_prediction"""
143 out = None
144
145 if inputdepthmaps is not None and inputdepthmaps[count] is not None:
146 # use custom depthmap
147 dp = inputdepthmaps[count]
148 if isinstance(dp, Image.Image):
149 if dp.width != inputimages[count].width or dp.height != inputimages[count].height:
150 try: # LANCZOS may fail on some formats
151 dp = dp.resize((inputimages[count].width, inputimages[count].height), Image.Resampling.LANCZOS)
152 except:
153 dp = dp.resize((inputimages[count].width, inputimages[count].height))
154 # Trying desperately to rescale image to [0;1) without actually normalizing it
155 # Normalizing is avoided, because we want to preserve the scale of the original depthmaps
156 # (batch mode, video mode).
157 if len(dp.getbands()) == 1:
158 out = np.asarray(dp, dtype="float")
159 out_max = out.max()
160 if out_max < 256:
161 bit_depth = 8
162 elif out_max < 65536:
163 bit_depth = 16
164 else:
165 bit_depth = 32
166 out /= 2.0 ** bit_depth
167 else:
168 out = np.asarray(dp, dtype="float")[:, :, 0]
169 out /= 256.0
170 else:
171 # Should be in interval [0; 1], values outside of this range will be clipped.
172 out = np.asarray(dp, dtype="float")
173 assert inputimages[count].height == out.shape[0], "Custom depthmap height mismatch"
174 assert inputimages[count].width == out.shape[1], "Custom depthmap width mismatch"
175 else:
176 # override net size (size may be different for different images)
177 if inp[go.NET_SIZE_MATCH]:
178 # Round up to a multiple of 32 to avoid potential issues
179 # TODO: buggs for Depth Anything
180 net_width = (inputimages[count].width + 31) // 32 * 32
181 net_height = (inputimages[count].height + 31) // 32 * 32
182 else:
183 net_width = inp[go.NET_WIDTH]
184 net_height = inp[go.NET_HEIGHT]
185 raw_prediction, raw_prediction_invert = \
186 model_holder.get_raw_prediction(inputimages[count], net_width, net_height)
187
188 # output
189 if abs(raw_prediction.max() - raw_prediction.min()) > np.finfo("float").eps:
190 out = np.copy(raw_prediction)
191 # TODO: some models may output negative values, maybe these should be clamped to zero.
192 if raw_prediction_invert:
193 out *= -1
194 if inp[go.DO_OUTPUT_DEPTH_PREDICTION]:
195 yield count, 'depth_prediction', np.copy(out)
196 if inp[go.CLIPDEPTH]:
197 if inp[go.CLIPDEPTH_MODE] == 'Range':
198 out = (out - out.min()) / (out.max() - out.min()) # normalize to [0; 1]
199 out = np.clip(out, inp[go.CLIPDEPTH_FAR], inp[go.CLIPDEPTH_NEAR])
200 elif inp[go.CLIPDEPTH_MODE] == 'Outliers':
201 fb, nb = np.percentile(out, [inp[go.CLIPDEPTH_FAR] * 100.0, inp[go.CLIPDEPTH_NEAR] * 100.0])
202 out = np.clip(out, fb, nb)
203 out = (out - out.min()) / (out.max() - out.min()) # normalize to [0; 1]
204 else:
205 # Regretfully, the depthmap is broken and will be replaced with a black image
206 out = np.zeros(raw_prediction.shape)
207
208 # Maybe we should not use img_output for everything, since we get better accuracy from
209 # the raw_prediction. However, it is not always supported. We maybe would like to achieve
210 # reproducibility, so depthmap of the image should be the same as generating the depthmap one more time.
211 img_output = convert_to_i16(out)
212 """Depthmap (near=bright), as uint16"""
213
214 # if 3dinpainting, store maps for processing in second pass
215 if inp[go.GEN_INPAINTED_MESH]:
216 inpaint_imgs.append(inputimages[count])
217 inpaint_depths.append(img_output)
218
219 # applying background masks after depth
220 if inp[go.GEN_REMBG]:
221 print('applying background masks')
222 background_removed_image = background_removed_images[count]
223 # maybe a threshold cut would be better on the line below.
224 background_removed_array = np.array(background_removed_image)
225 bg_mask = (background_removed_array[:, :, 0] == 0) & (background_removed_array[:, :, 1] == 0) & (
226 background_removed_array[:, :, 2] == 0) & (background_removed_array[:, :, 3] <= 0.2)
227 img_output[bg_mask] = 0 # far value
228
229 yield count, 'background_removed', background_removed_image
230
231 if inp[go.SAVE_BACKGROUND_REMOVAL_MASKS]:
232 bg_array = (1 - bg_mask.astype('int8')) * 255
233 mask_array = np.stack((bg_array, bg_array, bg_array, bg_array), axis=2)
234 mask_image = Image.fromarray(mask_array.astype(np.uint8))
235
236 yield count, 'foreground_mask', mask_image
237
238 # A weird quirk: if user tries to save depthmap, whereas custom depthmap is used,
239 # custom depthmap will be outputed
240 if inp[go.DO_OUTPUT_DEPTH]:
241 img_depth = cv2.bitwise_not(img_output) if inp[go.OUTPUT_DEPTH_INVERT] else img_output
242 if inp[go.OUTPUT_DEPTH_COMBINE]:
243 axis = 1 if inp[go.OUTPUT_DEPTH_COMBINE_AXIS] == 'Horizontal' else 0
244 img_concat = Image.fromarray(np.concatenate(
245 (inputimages[count], convert_i16_to_rgb(img_depth, inputimages[count])),
246 axis=axis))
247 yield count, 'concat_depth', img_concat
248 else:
249 yield count, 'depth', Image.fromarray(img_depth)
250
251 if inp[go.GEN_STEREO]:
252 # print("Generating stereoscopic image(s)..")
253 stereoimages = create_stereoimages(
254 inputimages[count], img_output,
255 inp[go.STEREO_DIVERGENCE], inp[go.STEREO_SEPARATION],
256 inp[go.STEREO_MODES],
257 inp[go.STEREO_BALANCE], inp[go.STEREO_OFFSET_EXPONENT], inp[go.STEREO_FILL_ALGO])
258 for c in range(0, len(stereoimages)):
259 yield count, inp[go.STEREO_MODES][c], stereoimages[c]
260
261 if inp[go.GEN_NORMALMAP]:
262 normalmap = create_normalmap(
263 img_output,
264 inp[go.NORMALMAP_PRE_BLUR_KERNEL] if inp[go.NORMALMAP_PRE_BLUR] else None,
265 inp[go.NORMALMAP_SOBEL_KERNEL] if inp[go.NORMALMAP_SOBEL] else None,
266 inp[go.NORMALMAP_POST_BLUR_KERNEL] if inp[go.NORMALMAP_POST_BLUR] else None,
267 inp[go.NORMALMAP_INVERT]
268 )
269 yield count, 'normalmap', normalmap
270
271 if inp[go.GEN_HEATMAP]:
272 from dzoedepth.utils.misc import colorize
273 heatmap = Image.fromarray(colorize(img_output, cmap='inferno'))
274 yield count, 'heatmap', heatmap
275
276 # gen mesh
277 if inp[go.GEN_SIMPLE_MESH]:
278 print(f"\nGenerating (occluded) mesh ..")
279 basename = 'depthmap'
280 meshsimple_fi = get_uniquefn(outpath, basename, 'obj', 'simple')
281
282 depthi = raw_prediction if raw_prediction is not None else out
283 depthi_min, depthi_max = depthi.min(), depthi.max()
284 # try to map output to sensible values for non zoedepth models, boost, or custom maps
285 if inp[go.MODEL_TYPE] not in [7, 8, 9] or inp[go.BOOST] or inputdepthmaps[count] is not None:
286 # invert if midas
287 if inp[go.MODEL_TYPE] > 0 or inputdepthmaps[count] is not None: # TODO: Weird
288 depthi = depthi_max - depthi + depthi_min
289 depth_max = depthi.max()
290 depth_min = depthi.min()
291 # make positive
292 if depthi_min < 0:
293 depthi = depthi - depthi_min
294 depth_max = depthi.max()
295 depth_min = depthi.min()
296 # scale down
297 if depthi.max() > 10.0:
298 depthi = 4.0 * (depthi - depthi_min) / (depthi_max - depthi_min)
299 # offset
300 depthi = depthi + 1.0
301
302 mesh = create_mesh(inputimages[count], depthi, keep_edges=not inp[go.SIMPLE_MESH_OCCLUDE],
303 spherical=(inp[go.SIMPLE_MESH_SPHERICAL]))
304 mesh.export(meshsimple_fi)
305 yield count, 'simple_mesh', meshsimple_fi
306
307 print("Computing output(s) done.")
308 except Exception as e:
309 import traceback
310 if 'out of memory' in str(e).lower():
311 print(str(e))
312 suggestion = "out of GPU memory, could not generate depthmap! " \
313 "Here are some suggestions to work around this issue:\n"
314 if inp[go.BOOST]:
315 suggestion += " * Disable BOOST (generation will be faster, but the depthmap will be less detailed)\n"
316 if backbone.USED_BACKBONE != backbone.BackboneType.STANDALONE:
317 suggestion += " * Run DepthMap in the standalone mode - without launching the SD WebUI\n"
318 if device != torch.device("cpu"):
319 suggestion += " * Select CPU as the processing device (this will be slower)\n"
320 if inp[go.MODEL_TYPE] != 6:
321 suggestion +=\
322 " * Use a different model (generally, more memory-consuming models produce better depthmaps)\n"
323 if not inp[go.BOOST]:
324 suggestion += " * Reduce net size (this could reduce quality)\n"
325 print('Fail.\n')
326 raise Exception(suggestion)
327 else:
328 print('Fail.\n')
329 raise e
330 finally:
331 if backbone.get_opt('depthmap_script_keepmodels', True):
332 model_holder.offload() # Swap to CPU memory
333 else:
334 model_holder.unload_models()
335 gc.collect()
336 backbone.torch_gc()
337
338 # TODO: This should not be here
339 if inp[go.GEN_INPAINTED_MESH]:
340 try:
341 mesh_fi = run_3dphoto(device, inpaint_imgs, inpaint_depths, inputnames, outpath,
342 inp[go.GEN_INPAINTED_MESH_DEMOS],
343 1, "mp4")
344 yield 0, 'inpainted_mesh', mesh_fi
345 except Exception as e:
346 print(f'{str(e)}, some issue with generating inpainted mesh')
347
348 backbone.reload_sd_model()
349 print("All done.\n")
350
351
352def get_uniquefn(outpath, basename, ext, suffix=''):
353 basecount = backbone.get_next_sequence_number(outpath, basename)
354 if basecount > 0:
355 basecount -= 1
356 if suffix != '':
357 suffix = f'-{suffix}' # Dash is important for selecting unique filenames (see get_next_sequence_number)
358 for i in range(500):
359 fullfn = os.path.join(outpath, f"{basename}-{basecount + i:04}{suffix}.{ext}")
360 if not os.path.exists(fullfn):
361 return fullfn
362 return f"{basename}-99999{suffix}.{ext}" # Failback, should never be executed
363
364
365def run_3dphoto(device, img_rgb, img_depth, inputnames, outpath, gen_inpainted_mesh_demos, vid_ssaa, vid_format):
366 mesh_fi = ''
367 try:
368 print("Running 3D Photo Inpainting .. ")
369 edgemodel_path = './models/3dphoto/edge_model.pth'
370 depthmodel_path = './models/3dphoto/depth_model.pth'
371 colormodel_path = './models/3dphoto/color_model.pth'
372 # create paths to model if not present
373 os.makedirs('./models/3dphoto/', exist_ok=True)
374
375 ensure_file_downloaded(
376 edgemodel_path,
377 ["https://huggingface.co/spaces/Epoching/3D_Photo_Inpainting/resolve/e389e564fd2a55cf/checkpoints/edge-model.pth",
378 "https://filebox.ece.vt.edu/~jbhuang/project/3DPhoto/model/edge-model.pth"],
379 "b1d768bd008ad5fe9f540004f870b8c3d355e4939b2009aa4db493fd313217c9")
380 ensure_file_downloaded(
381 depthmodel_path,
382 ["https://huggingface.co/spaces/Epoching/3D_Photo_Inpainting/resolve/e389e564fd2a55cf/checkpoints/depth-model.pth",
383 "https://filebox.ece.vt.edu/~jbhuang/project/3DPhoto/model/depth-model.pth"],
384 "2d0e63e89a22762ddfa8bc8c9f8c992e5532b140123274ffc6e4171baa1b76f8")
385 ensure_file_downloaded(
386 colormodel_path,
387 ["https://huggingface.co/spaces/Epoching/3D_Photo_Inpainting/resolve/e389e564fd2a55cf/checkpoints/color-model.pth",
388 "https://filebox.ece.vt.edu/~jbhuang/project/3DPhoto/model/color-model.pth"],
389 "383c9b1db70097907a6f9c8abb0303e7056f50d5456a36f34ab784592b8b2c20"
390 )
391
392 print("Loading edge model ..")
393 depth_edge_model = Inpaint_Edge_Net(init_weights=True)
394 depth_edge_weight = torch.load(edgemodel_path, map_location=torch.device(device))
395 depth_edge_model.load_state_dict(depth_edge_weight)
396 depth_edge_model = depth_edge_model.to(device)
397 depth_edge_model.eval()
398 print("Loading depth model ..")
399 depth_feat_model = Inpaint_Depth_Net()
400 depth_feat_weight = torch.load(depthmodel_path, map_location=torch.device(device))
401 depth_feat_model.load_state_dict(depth_feat_weight, strict=True)
402 depth_feat_model = depth_feat_model.to(device)
403 depth_feat_model.eval()
404 depth_feat_model = depth_feat_model.to(device)
405 print("Loading rgb model ..")
406 rgb_model = Inpaint_Color_Net()
407 rgb_feat_weight = torch.load(colormodel_path, map_location=torch.device(device))
408 rgb_model.load_state_dict(rgb_feat_weight)
409 rgb_model.eval()
410 rgb_model = rgb_model.to(device)
411
412 config = {}
413 config["gpu_ids"] = 0
414 config['extrapolation_thickness'] = 60
415 config['extrapolate_border'] = True
416 config['depth_threshold'] = 0.04
417 config['redundant_number'] = 12
418 config['ext_edge_threshold'] = 0.002
419 config['background_thickness'] = 70
420 config['context_thickness'] = 140
421 config['background_thickness_2'] = 70
422 config['context_thickness_2'] = 70
423 config['log_depth'] = True
424 config['depth_edge_dilate'] = 10
425 config['depth_edge_dilate_2'] = 5
426 config['largest_size'] = 512
427 config['repeat_inpaint_edge'] = True
428 config['ply_fmt'] = "bin"
429
430 config['save_ply'] = backbone.get_opt('depthmap_script_save_ply', False)
431 config['save_obj'] = True
432
433 if device == torch.device("cpu"):
434 config["gpu_ids"] = -1
435
436 for count in trange(0, len(img_rgb)):
437 basename = 'depthmap'
438 if inputnames is not None:
439 if inputnames[count] is not None:
440 p = Path(inputnames[count])
441 basename = p.stem
442
443 mesh_fi = get_uniquefn(outpath, basename, 'obj')
444
445 print(f"\nGenerating inpainted mesh .. (go make some coffee) ..")
446
447 # from inpaint.utils.get_MiDaS_samples
448 W = img_rgb[count].width
449 H = img_rgb[count].height
450 int_mtx = np.array([[max(H, W), 0, W // 2], [0, max(H, W), H // 2], [0, 0, 1]]).astype(np.float32)
451 if int_mtx.max() > 1:
452 int_mtx[0, :] = int_mtx[0, :] / float(W)
453 int_mtx[1, :] = int_mtx[1, :] / float(H)
454
455 # how inpaint.utils.read_MiDaS_depth() imports depthmap
456 disp = img_depth[count].astype(np.float32)
457 disp = disp - disp.min()
458 disp = cv2.blur(disp / disp.max(), ksize=(3, 3)) * disp.max()
459 disp = (disp / disp.max()) * 3.0
460 depth = 1. / np.maximum(disp, 0.05)
461
462 # rgb input
463 img = np.asarray(img_rgb[count])
464 if len(img.shape) > 2 and img.shape[2] == 4:
465 # convert the image from RGBA2RGB
466 img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
467
468 # run sparse bilateral filter
469 config['sparse_iter'] = 5
470 config['filter_size'] = [7, 7, 5, 5, 5]
471 config['sigma_s'] = 4.0
472 config['sigma_r'] = 0.5
473 vis_photos, vis_depths = sparse_bilateral_filtering(depth.copy(), img.copy(), config,
474 num_iter=config['sparse_iter'], spdb=False)
475 depth = vis_depths[-1]
476
477 # bilat_fn = os.path.join(outpath, basename +'_bilatdepth.png')
478 # cv2.imwrite(bilat_fn, depth)
479
480 rt_info = write_mesh(img,
481 depth,
482 int_mtx,
483 mesh_fi,
484 config,
485 rgb_model,
486 depth_edge_model,
487 depth_edge_model,
488 depth_feat_model)
489
490 if rt_info is not False and gen_inpainted_mesh_demos:
491 run_3dphoto_videos(mesh_fi, basename, outpath, 300, 40,
492 [0.03, 0.03, 0.05, 0.03],
493 ['double-straight-line', 'double-straight-line', 'circle', 'circle'],
494 [0.00, 0.00, -0.015, -0.015],
495 [0.00, 0.00, -0.015, -0.00],
496 [-0.05, -0.05, -0.05, -0.05],
497 ['dolly-zoom-in', 'zoom-in', 'circle', 'swing'], False, vid_format, vid_ssaa)
498
499 backbone.torch_gc()
500
501 finally:
502 del rgb_model
503 rgb_model = None
504 del depth_edge_model
505 depth_edge_model = None
506 del depth_feat_model
507 depth_feat_model = None
508 backbone.torch_gc()
509
510 return mesh_fi
511
512
513def run_3dphoto_videos(mesh_fi, basename, outpath, num_frames, fps, crop_border, traj_types, x_shift_range,
514 y_shift_range, z_shift_range, video_postfix, vid_dolly, vid_format, vid_ssaa):
515 import vispy
516 try:
517 if platform.system() == 'Windows':
518 vispy.use(app='PyQt5')
519 elif platform.system() == 'Darwin':
520 vispy.use('PyQt6')
521 else:
522 vispy.use(app='egl')
523 except:
524 import traceback
525 print(traceback.format_exc())
526 print('Trying an alternative...')
527 for u in ['PyQt5', 'PyQt6', 'egl']:
528 try:
529 vispy.use(app=u)
530 break
531 except:
532 print(f'On {u}')
533 print(traceback.format_exc())
534 # Honestly, I don't know if it actually helps at all
535
536 # read ply
537 global video_mesh_data, video_mesh_fn
538 if video_mesh_fn is None or video_mesh_fn != mesh_fi:
539 try:
540 del video_mesh_data
541 except:
542 print("del video_mesh_data failed")
543 video_mesh_fn = mesh_fi
544 video_mesh_data = read_mesh(mesh_fi)
545
546 verts, colors, faces, Height, Width, hFov, vFov, mean_loc_depth = video_mesh_data
547
548 original_w = output_w = W = Width
549 original_h = output_h = H = Height
550 int_mtx = np.array([[max(H, W), 0, W // 2], [0, max(H, W), H // 2], [0, 0, 1]]).astype(np.float32)
551 if int_mtx.max() > 1:
552 int_mtx[0, :] = int_mtx[0, :] / float(W)
553 int_mtx[1, :] = int_mtx[1, :] / float(H)
554
555 config = {}
556 config['video_folder'] = outpath
557 config['num_frames'] = num_frames
558 config['fps'] = fps
559 config['crop_border'] = crop_border
560 config['traj_types'] = traj_types
561 config['x_shift_range'] = x_shift_range
562 config['y_shift_range'] = y_shift_range
563 config['z_shift_range'] = z_shift_range
564 config['video_postfix'] = video_postfix
565 config['ssaa'] = vid_ssaa
566
567 # from inpaint.utils.get_MiDaS_samples
568 generic_pose = np.eye(4)
569 assert len(config['traj_types']) == len(config['x_shift_range']) == \
570 len(config['y_shift_range']) == len(config['z_shift_range']) == len(config['video_postfix']), \
571 "The number of elements in 'traj_types', 'x_shift_range', 'y_shift_range', 'z_shift_range' and \
572 'video_postfix' should be equal."
573 tgt_pose = [[generic_pose * 1]]
574 tgts_poses = []
575 for traj_idx in range(len(config['traj_types'])):
576 tgt_poses = []
577 sx, sy, sz = path_planning(config['num_frames'], config['x_shift_range'][traj_idx],
578 config['y_shift_range'][traj_idx],
579 config['z_shift_range'][traj_idx], path_type=config['traj_types'][traj_idx])
580 for xx, yy, zz in zip(sx, sy, sz):
581 tgt_poses.append(generic_pose * 1.)
582 tgt_poses[-1][:3, -1] = np.array([xx, yy, zz])
583 tgts_poses += [tgt_poses]
584 tgt_pose = generic_pose * 1
585
586 # seems we only need the depthmap to calc mean_loc_depth, which is only used when doing 'dolly'
587 # width and height are already in the ply file in the comments ..
588 # might try to add the mean_loc_depth to it too
589 # did just that
590 # mean_loc_depth = img_depth[img_depth.shape[0]//2, img_depth.shape[1]//2]
591
592 print("Generating videos ..")
593
594 normal_canvas, all_canvas = None, None
595 videos_poses, video_basename = copy.deepcopy(tgts_poses), basename
596 top = (original_h // 2 - int_mtx[1, 2] * output_h)
597 left = (original_w // 2 - int_mtx[0, 2] * output_w)
598 down, right = top + output_h, left + output_w
599 border = [int(xx) for xx in [top, down, left, right]]
600 normal_canvas, all_canvas, fn_saved = output_3d_photo(verts.copy(), colors.copy(), faces.copy(),
601 copy.deepcopy(Height), copy.deepcopy(Width),
602 copy.deepcopy(hFov), copy.deepcopy(vFov),
603 copy.deepcopy(tgt_pose), config['video_postfix'],
604 copy.deepcopy(generic_pose),
605 copy.deepcopy(config['video_folder']),
606 None, copy.deepcopy(int_mtx), config, None,
607 videos_poses, video_basename, original_h, original_w,
608 border=border, depth=None, normal_canvas=normal_canvas,
609 all_canvas=all_canvas,
610 mean_loc_depth=mean_loc_depth, dolly=vid_dolly,
611 fnExt=vid_format)
612 return fn_saved
613
614def run_makevideo(fn_mesh, vid_numframes, vid_fps, vid_traj, vid_shift, vid_border, dolly, vid_format, vid_ssaa,
615 outpath=None, basename=None):
616 if len(fn_mesh) == 0 or not os.path.exists(fn_mesh):
617 raise Exception("Could not open mesh.")
618
619 vid_ssaa = int(vid_ssaa)
620
621 # traj type
622 if vid_traj == 0:
623 vid_traj = ['straight-line']
624 elif vid_traj == 1:
625 vid_traj = ['double-straight-line']
626 elif vid_traj == 2:
627 vid_traj = ['circle']
628
629 num_fps = int(vid_fps)
630 num_frames = int(vid_numframes)
631 shifts = vid_shift.split(',')
632 if len(shifts) != 3:
633 raise Exception("Translate requires 3 elements.")
634 x_shift_range = [float(shifts[0])]
635 y_shift_range = [float(shifts[1])]
636 z_shift_range = [float(shifts[2])]
637
638 borders = vid_border.split(',')
639 if len(borders) != 4:
640 raise Exception("Crop Border requires 4 elements.")
641 crop_border = [float(borders[0]), float(borders[1]), float(borders[2]), float(borders[3])]
642
643 if not outpath:
644 outpath = backbone.get_outpath()
645
646 if not basename:
647 # output path and filename mess ..
648 basename = Path(fn_mesh).stem
649
650 # unique filename
651 basecount = backbone.get_next_sequence_number(outpath, basename)
652 if basecount > 0: basecount = basecount - 1
653 fullfn = None
654 for i in range(500):
655 fn = f"{basecount + i:05}" if basename == '' else f"{basename}-{basecount + i:04}"
656 fullfn = os.path.join(outpath, f"{fn}_." + vid_format)
657 if not os.path.exists(fullfn):
658 break
659 basename = Path(fullfn).stem
660 basename = basename[:-1]
661
662 print("Loading mesh ..")
663
664 fn_saved = run_3dphoto_videos(fn_mesh, basename, outpath, num_frames, num_fps, crop_border, vid_traj, x_shift_range,
665 y_shift_range, z_shift_range, [''], dolly, vid_format, vid_ssaa)
666
667 return fn_saved[-1], fn_saved[-1], ''
668
669def unload_models():
670 model_holder.unload_models()
671
672
673# TODO: code borrowed from the internet to be marked as such and to reside in separate files
674
675def batched_background_removal(inimages, model_name):
676 from rembg import new_session, remove
677 print('creating background masks')
678 outimages = []
679
680 # model path and name
681 bg_model_dir = Path.joinpath(Path().resolve(), "models/rem_bg")
682 os.makedirs(bg_model_dir, exist_ok=True)
683 os.environ["U2NET_HOME"] = str(bg_model_dir)
684
685 # starting a session
686 background_removal_session = new_session(model_name)
687 for count in range(0, len(inimages)):
688 bg_remove_img = np.array(remove(inimages[count], session=background_removal_session))
689 outimages.append(Image.fromarray(bg_remove_img))
690 # The line below might be redundant
691 del background_removal_session
692 return outimages
693
694
695def pano_depth_to_world_points(depth):
696 """
697 360 depth to world points
698 given 2D depth is an equirectangular projection of a spherical image
699 Treat depth as radius
700 longitude : -pi to pi
701 latitude : -pi/2 to pi/2
702 """
703
704 # Convert depth to radius
705 radius = depth.flatten()
706
707 lon = np.linspace(-np.pi, np.pi, depth.shape[1])
708 lat = np.linspace(-np.pi / 2, np.pi / 2, depth.shape[0])
709
710 lon, lat = np.meshgrid(lon, lat)
711 lon = lon.flatten()
712 lat = lat.flatten()
713
714 # Convert to cartesian coordinates
715 x = radius * np.cos(lat) * np.cos(lon)
716 y = radius * np.cos(lat) * np.sin(lon)
717 z = radius * np.sin(lat)
718
719 pts3d = np.stack([x, y, z], axis=1)
720
721 return pts3d
722
723
724def depth_edges_mask(depth):
725 """Returns a mask of edges in the depth map.
726 Args:
727 depth: 2D numpy array of shape (H, W) with dtype float32.
728 Returns:
729 mask: 2D numpy array of shape (H, W) with dtype bool.
730 """
731 # Compute the x and y gradients of the depth map.
732 depth_dx, depth_dy = np.gradient(depth)
733 # Compute the gradient magnitude.
734 depth_grad = np.sqrt(depth_dx ** 2 + depth_dy ** 2)
735 # Compute the edge mask.
736 mask = depth_grad > 0.05
737 return mask
738
739
740def create_mesh(image, depth, keep_edges=False, spherical=False):
741 import trimesh
742 from dzoedepth.utils.geometry import depth_to_points, create_triangles
743 maxsize = backbone.get_opt('depthmap_script_mesh_maxsize', 2048)
744
745 # limit the size of the input image
746 image.thumbnail((maxsize, maxsize))
747
748 if not spherical:
749 pts3d = depth_to_points(depth[None])
750 else:
751 pts3d = pano_depth_to_world_points(depth)
752
753 pts3d = pts3d.reshape(-1, 3)
754
755 verts = pts3d.reshape(-1, 3)
756 image = np.array(image)
757 if keep_edges:
758 triangles = create_triangles(image.shape[0], image.shape[1])
759 else:
760 triangles = create_triangles(image.shape[0], image.shape[1], mask=~depth_edges_mask(depth))
761 colors = image.reshape(-1, 3)
762
763 mesh = trimesh.Trimesh(vertices=verts, faces=triangles, vertex_colors=colors)
764
765 # rotate 90deg over X when spherical
766 if spherical:
767 angle = math.pi / 2
768 direction = [1, 0, 0]
769 center = [0, 0, 0]
770 rot_matrix = trimesh.transformations.rotation_matrix(angle, direction, center)
771 mesh.apply_transform(rot_matrix)
772
773 return mesh
774