diff --git a/README.md b/README.md index ca4247f..60a5d60 100644 --- a/README.md +++ b/README.md @@ -97,6 +97,20 @@ path_female_star = '/mypath/female/model.npz' path_neutral_star = '/mypath/neutral/model.npz' ``` +Alternatively, place a file called `star.ini` in your user home directory, i.e. +- Linux: `/home/username/star.ini` +- Windows: `C:\Users\username\star.ini` +- MacOs: `/Users/username/star.ini` + +The `star.ini` file should specify the paths, e.g. +``` +[DEFAULT] +path_male_star = /Users/username/star_1_1/male/model.npz +path_female_star = /Users/username/star_1_1/female/model.npz +path_neutral_star = /Users/username/star_1_1/neutral/model.npz +data_type = float32 +``` + 7. Install with pip ``` pip install . diff --git a/demo/load_torch.py b/demo/load_torch.py index 29d8e24..05999ea 100644 --- a/demo/load_torch.py +++ b/demo/load_torch.py @@ -34,11 +34,13 @@ batch_size=1 m = STAR(gender='male',num_betas=num_betas) +device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') + # Zero pose -poses = torch.cuda.FloatTensor(np.zeros((batch_size,72))) -betas = torch.cuda.FloatTensor(betas) +poses = torch.FloatTensor(np.zeros((batch_size,72)), device=device) +betas = torch.FloatTensor(betas, device=device) -trans = torch.cuda.FloatTensor(np.zeros((batch_size,3))) +trans = torch.FloatTensor(np.zeros((batch_size,3)), device=device) model = star.forward(poses, betas,trans) shaped = model.v_shaped[-1, :, :] diff --git a/star/config.py b/star/config.py index 22a7851..2665124 100644 --- a/star/config.py +++ b/star/config.py @@ -19,13 +19,25 @@ # Code Developed by: # Ahmed A. A. Osman -import os +from configparser import ConfigParser +from pathlib import Path + path_male_star = '' path_female_star = '' path_neutral_star = '' - data_type = 'float32' +# allow users to specify paths in a config file outside the source +config_path = Path.home() / 'star.ini' +if config_path.exists(): + config = ConfigParser() + config.read(config_path) + + path_male_star = config['DEFAULT'].get('path_male_star', '') + path_female_star = config['DEFAULT'].get('path_female_star', '') + path_neutral_star = config['DEFAULT'].get('path_neutral_star', '') + data_type = config['DEFAULT'].get('data_type', 'float32') + if data_type not in ['float16','float32','float64']: raise RuntimeError('Invalid data type %s'%(data_type)) diff --git a/star/pytorch/star.py b/star/pytorch/star.py index 0c89407..e26c187 100644 --- a/star/pytorch/star.py +++ b/star/pytorch/star.py @@ -33,7 +33,7 @@ class STAR(nn.Module): - def __init__(self,gender='female',num_betas=10): + def __init__(self,gender='female',num_betas=10,device=None): super(STAR, self).__init__() if gender not in ['male','female','neutral']: @@ -55,18 +55,21 @@ def __init__(self,gender='female',num_betas=10): rows,cols = np.where(J_regressor!=0) vals = J_regressor[rows,cols] self.num_betas = num_betas + + if device: + device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') # Model sparse joints regressor, regresses joints location from a mesh - self.register_buffer('J_regressor', torch.cuda.FloatTensor(J_regressor)) + self.register_buffer('J_regressor', torch.FloatTensor(J_regressor).to(device)) # Model skinning weights - self.register_buffer('weights', torch.cuda.FloatTensor(star_model['weights'])) + self.register_buffer('weights', torch.FloatTensor(star_model['weights']).to(device)) # Model pose corrective blend shapes - self.register_buffer('posedirs', torch.cuda.FloatTensor(star_model['posedirs'].reshape((-1,93)))) + self.register_buffer('posedirs', torch.FloatTensor(star_model['posedirs'].reshape((-1,93))).to(device)) # Mean Shape - self.register_buffer('v_template', torch.cuda.FloatTensor(star_model['v_template'])) + self.register_buffer('v_template', torch.FloatTensor(star_model['v_template']).to(device)) # Shape corrective blend shapes - self.register_buffer('shapedirs', torch.cuda.FloatTensor(np.array(star_model['shapedirs'][:,:,:num_betas]))) + self.register_buffer('shapedirs', torch.FloatTensor(np.array(star_model['shapedirs'][:,:,:num_betas])).to(device)) # Mesh traingles self.register_buffer('faces', torch.from_numpy(star_model['f'].astype(np.int64))) self.f = star_model['f'] diff --git a/star/pytorch/utils.py b/star/pytorch/utils.py index d8f161f..1bfd279 100644 --- a/star/pytorch/utils.py +++ b/star/pytorch/utils.py @@ -80,8 +80,9 @@ def with_zeros(input): :param input: A tensor of dimensions batch size x 3 x 4 :return: A tensor batch size x 4 x 4 (appended with 0,0,0,1) ''' + device = input.device batch_size = input.shape[0] - row_append = torch.cuda.FloatTensor(([0.0, 0.0, 0.0, 1.0])) + row_append = torch.FloatTensor(([0.0, 0.0, 0.0, 1.0])).to(device) row_append.requires_grad = False padded_tensor = torch.cat([input, row_append.view(1, 1, 4).repeat(batch_size, 1, 1)], 1) return padded_tensor