Spaces:
Sleeping
Sleeping
File size: 761 Bytes
8a6df40 6b93ef6 8a6df40 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
import torch
import os
from posenet.models.mobilenet_v1 import MobileNetV1, MOBILENET_V1_CHECKPOINTS
MODEL_DIR = 'TryYours-Virtual-Try-On/posenet_models'
DEBUG_OUTPUT = False
def load_model(model_id, output_stride=16, model_dir=MODEL_DIR):
model_path = os.path.join(model_dir, MOBILENET_V1_CHECKPOINTS[model_id] + '.pth')
if not os.path.exists(model_path):
print('Cannot find models file %s, converting from tfjs...' % model_path)
from posenet.converter.tfjs2pytorch import convert
convert(model_id, model_dir, check=False)
assert os.path.exists(model_path)
model = MobileNetV1(model_id, output_stride=output_stride)
load_dict = torch.load(model_path)
model.load_state_dict(load_dict)
return model
|