Clone/DIS
1
0
Fork 0
mirror of https://github.com/xuebinqin/DIS.git synced 2024-11-29 18:23:08 +01:00
Code Issues Actions Packages Projects Releases Wiki Activity

Inference.py modified

Browse Source
This commit is contained in:
PiggyJerry 2022-08-20 19:10:44 +04:00
parent ff195223e5
commit 9a37251d6b
1 changed files with 5 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
IS-Net/Inference.py
View File

@ -28,7 +28,7 @@ if __name__ == "__main__":
net=net.cuda() net=net.cuda()
else: else:
net.load_state_dict(torch.load(model_path,map_location="cpu")) net.load_state_dict(torch.load(model_path,map_location="cpu"))
net.eval()
im_list = glob(dataset_path+"/*.jpg") im_list = glob(dataset_path+"/*.jpg")
for i, im_path in tqdm(enumerate(im_list), total=len(im_list)): for i, im_path in tqdm(enumerate(im_list), total=len(im_list)):
print("im_path: ", im_path) print("im_path: ", im_path)
@ -37,11 +37,12 @@ if __name__ == "__main__":
im = im[:, :, np.newaxis] im = im[:, :, np.newaxis]
im_shp=im.shape[0:2] im_shp=im.shape[0:2]
im_tensor = torch.tensor(im, dtype=torch.float32).permute(2,0,1) im_tensor = torch.tensor(im, dtype=torch.float32).permute(2,0,1)
im_tensor = F.upsample(torch.unsqueeze(im_tensor,0), input_size, mode="bilinear") im_tensor = F.upsample(torch.unsqueeze(im_tensor,0), input_size, mode="bilinear").type(torch.uint8)
image = normalize(im_tensor,[0.5,0.5,0.5],[1.0,1.0,1.0]).type(torch.uint8) image = torch.divide(im_tensor,255.0)
image = normalize(image,[0.5,0.5,0.5],[1.0,1.0,1.0])
if torch.cuda.is_available(): if torch.cuda.is_available():
image=image.cuda() image=image.cuda()
image = torch.divide(image,255.0)
result=net(image) result=net(image)
result=torch.squeeze(F.upsample(result[0][0],im_shp,mode='bilinear'),0) result=torch.squeeze(F.upsample(result[0][0],im_shp,mode='bilinear'),0)
ma = torch.max(result) ma = torch.max(result)