Spaces:
Build error
Build error
File size: 1,942 Bytes
fd5e0f7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
'''
Dust3R reconstrucion
GeoWizard Estimation
Smooth Projection
'''
import torch
import PIL,cv2
import numpy as np
from PIL import Image
from ops.gs.basic import Frame
from ops.utils import *
from ops.depth_pro import Depth_Pro_Tool
from ops.connect import Smooth_Connect_Tool
class Reconstruct_Tool():
def __init__(self,cfg) -> None:
self.cfg = cfg
self._load_model()
self.connector = Smooth_Connect_Tool()
def _load_model(self):
self.pro_dpt = Depth_Pro_Tool(ckpt=self.cfg.model.mde.dpt_pro.ckpt,device='cpu')
def _ProDpt_(self, rgb, intrinsic=None):
# conduct reconstruction
print('Pro_dpt[1/3] Move Pro_dpt.model to GPU...')
self.pro_dpt.to('cuda')
print('Pro_dpt[2/3] Pro_dpt Estimation...')
f_px = intrinsic[0,0] if intrinsic is not None else None
metric_dpt,intrinsic = self.pro_dpt(rgb,f_px)
print('Pro_dpt[3/3] Move Pro_dpt.model to GPU...')
self.pro_dpt.to('cpu')
torch.cuda.empty_cache()
edge_mask = edge_filter(metric_dpt,times=0.05)
return metric_dpt, intrinsic, edge_mask
def _Guide_ProDpt_(self, rgb, intrinsic=None, refer_dpt=None, refer_msk=None):
# conduct reconstruction
print('Pro_dpt[1/3] Move Pro_dpt.model to GPU...')
self.pro_dpt.to('cuda')
print('Pro_dpt[2/3] Pro_dpt Estimation...')
f_px = intrinsic[0,0] if intrinsic is not None else None
metric_dpt,intrinsic = self.pro_dpt(rgb,f_px=f_px)
metric_dpt_connect = self.connector._affine_dpt_to_GS(refer_dpt,metric_dpt,~refer_msk)
print('Pro_dpt[3/3] Move Pro_dpt.model to GPU...')
self.pro_dpt.to('cpu')
torch.cuda.empty_cache()
edge_mask = edge_filter(metric_dpt_connect,times=0.05)
return metric_dpt_connect, metric_dpt, intrinsic, edge_mask
# ------------- TODO: Metricv2 + Guide-GeoWizard ------------------ #
|