Files
reconstruction/utils/data_utils.py
T

54 lines
1.7 KiB
Python

import os
import sys
import kornia
import torch
import numpy as np
import scipy.io as sio
from math import cos, sin
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from superSMPL.IO import readPC2Frame, readOBJ
class Frame:
def __init__(self, id, frame_nr, garment):
self.id = id
self.frame_nr = frame_nr
self.garment = garment
def mask_faces(F, mask):
F_mask = np.take(mask, F.reshape(-1), axis=0)
F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
V_ind = np.zeros(mask.shape, dtype=np.int32)
V_ind[mask] = np.arange(mask.sum())
F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
return F_new
def read_frames_file(file):
frames = []
with open(file, 'r') as f:
lines = f.readlines()
for line in lines:
id, frame_nr, garment = line.split()
frames.append(Frame(id, int(frame_nr), garment))
return frames
def get_normals_from_depth(depth_image, depth_mask, camera_intrinsic):
depth_dilated = kornia.morphology.dilation(depth_image[:,None,:,:], torch.ones((3,3), device=depth_image.get_device()), engine='convolution')
depth_dilated = kornia.morphology.dilation(depth_dilated, torch.ones((3,3), device=depth_image.get_device()), engine='convolution')
depth_dilated = torch.permute(depth_dilated, [0,2,3,1])
depth_dilated[depth_mask,:] = depth_image[:,:,:,None][depth_mask,:]
depth_dilated = torch.permute(depth_dilated, [0,3,1,2])
normals = kornia.geometry.depth.depth_to_normals(depth_dilated, camera_intrinsic, normalize_points=True)
normals = torch.permute(normals, [0,2,3,1]) * depth_mask[:,:,:,None]
return normals
def center_crop_resize(img, new_center, size):
pass