init + inference.py патч
This commit is contained in:
151
data/MBD/infer.py
Normal file
151
data/MBD/infer.py
Normal file
@ -0,0 +1,151 @@
|
||||
import torch
|
||||
import argparse
|
||||
import numpy as np
|
||||
import torch.nn.functional as F
|
||||
import glob
|
||||
import cv2
|
||||
from tqdm import tqdm
|
||||
|
||||
import time
|
||||
import os
|
||||
from model.deep_lab_model.deeplab import *
|
||||
from MBD import mask_base_dewarper
|
||||
import time
|
||||
|
||||
from utils import cvimg2torch,torch2cvimg
|
||||
|
||||
|
||||
|
||||
def net1_net2_infer(model,img_paths,args):
|
||||
|
||||
### validate on the real datasets
|
||||
seg_model=model
|
||||
seg_model.eval()
|
||||
for img_path in tqdm(img_paths):
|
||||
if os.path.exists(img_path.replace('_origin','_capture')):
|
||||
continue
|
||||
t1 = time.time()
|
||||
### segmentation mask predict
|
||||
img_org = cv2.imread(img_path)
|
||||
h_org,w_org = img_org.shape[:2]
|
||||
img = cv2.resize(img_org,(448, 448))
|
||||
img = cv2.GaussianBlur(img,(15,15),0,0)
|
||||
img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
|
||||
img = cvimg2torch(img)
|
||||
|
||||
with torch.no_grad():
|
||||
pred = seg_model(img.cuda())
|
||||
mask_pred = pred[:,0,:,:].unsqueeze(1)
|
||||
mask_pred = F.interpolate(mask_pred,(h_org,w_org))
|
||||
mask_pred = mask_pred.squeeze(0).squeeze(0).cpu().numpy()
|
||||
mask_pred = (mask_pred*255).astype(np.uint8)
|
||||
kernel = np.ones((3,3))
|
||||
mask_pred = cv2.dilate(mask_pred,kernel,iterations=3)
|
||||
mask_pred = cv2.erode(mask_pred,kernel,iterations=3)
|
||||
mask_pred[mask_pred>100] = 255
|
||||
mask_pred[mask_pred<100] = 0
|
||||
### tps transform base on the mask
|
||||
# dewarp, grid = mask_base_dewarper(img_org,mask_pred)
|
||||
try:
|
||||
dewarp, grid = mask_base_dewarper(img_org,mask_pred)
|
||||
except:
|
||||
print('fail')
|
||||
grid = np.meshgrid(np.arange(w_org),np.arange(h_org))/np.array([w_org,h_org]).reshape(2,1,1)
|
||||
grid = torch.from_numpy((grid-0.5)*2).float().unsqueeze(0).permute(0,2,3,1)
|
||||
dewarp = torch2cvimg(F.grid_sample(cvimg2torch(img_org),grid))[0]
|
||||
grid = grid[0].numpy()
|
||||
# cv2.imshow('in',cv2.resize(img_org,(512,512)))
|
||||
# cv2.imshow('out',cv2.resize(dewarp,(512,512)))
|
||||
# cv2.waitKey(0)
|
||||
cv2.imwrite(img_path.replace('_origin','_capture'),dewarp)
|
||||
cv2.imwrite(img_path.replace('_origin','_mask_new'),mask_pred)
|
||||
|
||||
grid0 = cv2.resize(grid[:,:,0],(128,128))
|
||||
grid1 = cv2.resize(grid[:,:,1],(128,128))
|
||||
grid = np.stack((grid0,grid1),axis=-1)
|
||||
np.save(img_path.replace('_origin','_grid1'),grid)
|
||||
|
||||
|
||||
def net1_net2_infer_single_im(img,model_path):
|
||||
seg_model = DeepLab(num_classes=1,
|
||||
backbone='resnet',
|
||||
output_stride=16,
|
||||
sync_bn=None,
|
||||
freeze_bn=False)
|
||||
seg_model = torch.nn.DataParallel(seg_model, device_ids=range(torch.cuda.device_count()))
|
||||
seg_model.cuda()
|
||||
checkpoint = torch.load(model_path)
|
||||
seg_model.load_state_dict(checkpoint['model_state'])
|
||||
### validate on the real datasets
|
||||
seg_model.eval()
|
||||
### segmentation mask predict
|
||||
img_org = img
|
||||
h_org,w_org = img_org.shape[:2]
|
||||
img = cv2.resize(img_org,(448, 448))
|
||||
img = cv2.GaussianBlur(img,(15,15),0,0)
|
||||
img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
|
||||
img = cvimg2torch(img)
|
||||
|
||||
with torch.no_grad():
|
||||
# from torchtoolbox.tools import summary
|
||||
# print(summary(seg_model,torch.rand((1, 3, 448, 448)).cuda())) 59.4M 135.6G
|
||||
|
||||
pred = seg_model(img.cuda())
|
||||
mask_pred = pred[:,0,:,:].unsqueeze(1)
|
||||
mask_pred = F.interpolate(mask_pred,(h_org,w_org))
|
||||
mask_pred = mask_pred.squeeze(0).squeeze(0).cpu().numpy()
|
||||
mask_pred = (mask_pred*255).astype(np.uint8)
|
||||
kernel = np.ones((3,3))
|
||||
mask_pred = cv2.dilate(mask_pred,kernel,iterations=3)
|
||||
mask_pred = cv2.erode(mask_pred,kernel,iterations=3)
|
||||
mask_pred[mask_pred>100] = 255
|
||||
mask_pred[mask_pred<100] = 0
|
||||
### tps transform base on the mask
|
||||
# dewarp, grid = mask_base_dewarper(img_org,mask_pred)
|
||||
# try:
|
||||
# dewarp, grid = mask_base_dewarper(img_org,mask_pred)
|
||||
# except:
|
||||
# print('fail')
|
||||
# grid = np.meshgrid(np.arange(w_org),np.arange(h_org))/np.array([w_org,h_org]).reshape(2,1,1)
|
||||
# grid = torch.from_numpy((grid-0.5)*2).float().unsqueeze(0).permute(0,2,3,1)
|
||||
# dewarp = torch2cvimg(F.grid_sample(cvimg2torch(img_org),grid))[0]
|
||||
# grid = grid[0].numpy()
|
||||
# cv2.imshow('in',cv2.resize(img_org,(512,512)))
|
||||
# cv2.imshow('out',cv2.resize(dewarp,(512,512)))
|
||||
# cv2.waitKey(0)
|
||||
# cv2.imwrite(img_path.replace('_origin','_capture'),dewarp)
|
||||
# cv2.imwrite(img_path.replace('_origin','_mask_new'),mask_pred)
|
||||
|
||||
# grid0 = cv2.resize(grid[:,:,0],(128,128))
|
||||
# grid1 = cv2.resize(grid[:,:,1],(128,128))
|
||||
# grid = np.stack((grid0,grid1),axis=-1)
|
||||
# np.save(img_path.replace('_origin','_grid1'),grid)
|
||||
return mask_pred
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser(description='Hyperparams')
|
||||
parser.add_argument('--img_folder', nargs='?', type=str, default='./all_data',help='Data path to load data')
|
||||
parser.add_argument('--img_rows', nargs='?', type=int, default=448,
|
||||
help='Height of the input image')
|
||||
parser.add_argument('--img_cols', nargs='?', type=int, default=448,
|
||||
help='Width of the input image')
|
||||
parser.add_argument('--seg_model_path', nargs='?', type=str, default='checkpoints/mbd.pkl',
|
||||
help='Path to previous saved model to restart from')
|
||||
args = parser.parse_args()
|
||||
|
||||
seg_model = DeepLab(num_classes=1,
|
||||
backbone='resnet',
|
||||
output_stride=16,
|
||||
sync_bn=None,
|
||||
freeze_bn=False)
|
||||
seg_model = torch.nn.DataParallel(seg_model, device_ids=range(torch.cuda.device_count()))
|
||||
seg_model.cuda()
|
||||
checkpoint = torch.load(args.seg_model_path)
|
||||
seg_model.load_state_dict(checkpoint['model_state'])
|
||||
|
||||
im_paths = glob.glob(os.path.join(args.img_folder,'*_origin.*'))
|
||||
|
||||
net1_net2_infer(seg_model,im_paths,args)
|
||||
|
||||
Reference in New Issue
Block a user