File size: 698 Bytes
bcf59c3 |
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 |
import os
from collections import OrderedDict
import torch
from .u2net_cloth_segm import U2NET
def load_cloth_segm_model(device, checkpoint_path, in_ch=3, out_ch=1):
if not os.path.exists(checkpoint_path):
print("Invalid path")
return
model = U2NET(in_ch=in_ch, out_ch=out_ch)
model_state_dict = torch.load(checkpoint_path, map_location=device)
new_state_dict = OrderedDict()
for k, v in model_state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
model = model.to(device=device)
print("Checkpoints loaded from path: {}".format(checkpoint_path))
return model
|