CalvinBackup/neuralchen-SimSwap
1
0
Fork 0
Code Issues 288 Pull Requests 16 Actions Packages Projects Releases 2 Wiki Activity

Update test_one_image.py

Browse Source
This commit is contained in:
Smiril
2022-07-12 12:19:20 +02:00
committed by GitHub
parent cd3d89638d
commit 8ee255f35f
1 changed files with 25 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test_one_image.py
+25 -10
View File
@@ -49,16 +49,31 @@ if __name__ == '__main__':
img_att = img_b.view(-1, img_b.shape[0], img_b.shape[1], img_b.shape[2]) img_att = img_b.view(-1, img_b.shape[0], img_b.shape[1], img_b.shape[2])
# convert numpy to tensor # convert numpy to tensor
img_id = img_id.cuda() if not torch.backends.mps.is_available():
img_att = img_att.cuda() if not torch.backends.mps.is_built():
img_id = img_id.cuda()
img_att = img_att.cuda()
else:
print("ERROR")
else:
img_id = img_id.has_mps()
img_att = img_att.has_mps()
#create latent id #create latent id
img_id_downsample = F.interpolate(img_id, size=(112,112)) if not torch.backends.mps.is_available():
latend_id = model.netArc(img_id_downsample) if not torch.backends.mps.is_built():
latend_id = latend_id.detach().to('cpu') img_id_downsample = F.interpolate(img_id, size=(112,112))
latend_id = latend_id/np.linalg.norm(latend_id,axis=1,keepdims=True) latend_id = model.netArc(img_id_downsample)
latend_id = latend_id.to('cuda') latend_id = latend_id.detach().to('cpu')
latend_id = latend_id/np.linalg.norm(latend_id,axis=1,keepdims=True)
latend_id = latend_id.to('cuda')
else:
print("ERROR")
else:
img_id_downsample = F.interpolate(img_id, size=(112,112))
latend_id = model.netArc(img_id_downsample)
latend_id = latend_id.detach().to('cpu')
latend_id = latend_id/np.linalg.norm(latend_id,axis=1,keepdims=True)
latend_id = latend_id.to('has_mps')
############## Forward Pass ###################### ############## Forward Pass ######################
img_fake = model(img_id, img_att, latend_id, latend_id, True) img_fake = model(img_id, img_att, latend_id, latend_id, True)
@@ -83,4 +98,4 @@ if __name__ == '__main__':
output = output*255 output = output*255
cv2.imwrite(opt.output_path + 'result.jpg',output) cv2.imwrite(opt.output_path + 'result.jpg',output)