def main():
"""
Main function wrapper for demo script
"""
random.seed(args["SEED"])
np.random.seed(args["SEED"])
torch.manual_seed(args["SEED"])
if torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
if args["TRAINED_WEIGHTS_FILE"] is not None:
print("Trained Weights File: %s" % (args["TRAINED_WEIGHTS_FILE"]))
print("Demo Directory: %s" % (args["DEMO_DIRECTORY"]))
model = MyNet()
model.load_state_dict(
torch.load(
args["CODE_DIRECTORY"] + args["TRAINED_WEIGHTS_FILE"],
map_location=device,
))
model.to(device)
print("Running Demo ....")
for root, dirs, files in os.walk(args["DEMO_DIRECTORY"]):
for file in files:
sampleFile = os.path.join(root, file)
preprocess_sample(sampleFile)
inp, _ = prepare_input(sampleFile)
inputBatch = torch.unsqueeze(inp, dim=0)
inputBatch = (inputBatch.float()).to(device)
model.eval()
with torch.no_grad():
outputBatch = model(inputBatch)
predictionBatch = decode(outputBatch)
pred = predictionBatch[0][:]
print("File: %s" % (file))
print("Prediction: %s" % (pred))
print("\n")
print("Demo Completed.")
else:
print("Path to trained weights file not specified.")
return
def load_model(self):
"""
:return:
"""
# TODO 1 加载模型
use_cuda = self.use_cuda
if self.o_net_path is not None:
print('=======> loading')
net = MyNet(use_cuda=False)
net.load_state_dict(torch.load(self.o_net_path))
if (use_cuda):
net.to('cpu')
net.eval()
# TODO 2 准备好数据
img_list = os.listdir(self.image_dir)
for idx, item in enumerate(img_list):
_img = Image.open(os.path.join(self.image_dir, item))
parse_result = self.parse_image_name(item)
landmark_and_format = parse_result['landmark_and_format']
name = parse_result['name']
img = self.transforms(_img)
img = img.unsqueeze(0)
pred = net(img)
pred = pred * 192
# pred = pred.detach().numpy()
print('the pred landmark is :', pred)
print("=" * 20)
# # print(pred.shape)
# # print(landmark)
#
try:
self.save_pred(_img, name, landmark_and_format,
pred.detach().numpy())
# self.visualize(_img, np.array(landmark))
# self.visualize(_img, pred.detach().numpy())
# # print(pred)
except:
print('Error:', item)