File size: 2,897 Bytes
8b79d57 |
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 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
import torch
from model import MattingNetwork
from torch.utils.data import DataLoader
from torch.utils.data.dataset import Dataset
import glob
import os
import cv2
import pdb
import argparse
class ItwDataset(Dataset):
def __init__(self, input_pth, step, rotate):
self.input_pth_list = glob.glob(os.path.join(input_pth, '*.png')) + \
glob.glob(os.path.join(input_pth, '*.jpg'))
self.input_pth_list.sort()
self.input_pth_list = self.input_pth_list[::step]
self.rotate = rotate
# pdb.set_trace()
def __len__(self):
return len(self.input_pth_list)
def __getitem__(self, index):
render_path = self.input_pth_list[index]
# pdb.set_trace()
img = cv2.imread(render_path)
if self.rotate == '+90':
img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
elif self.rotate == '-90':
img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
elif self.rotate == '180':
img = cv2.rotate(img, cv2.ROTATE_180)
img = torch.from_numpy(img)
img = img.permute(2,0,1)/255.
img = img.unsqueeze(0)
# img = torch.flip(img, dims = [0])
# print(img.shape)
# img = img[::-1,...]
# img = img.unsqueeze(0)
return {
'img': img,
'file_name': os.path.basename(render_path)[:-4]
}
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--input_pth', type = str)
parser.add_argument('--output_pth', type = str)
parser.add_argument('--device', type = str, default = 'cpu')
parser.add_argument('--step', type = int, default = 1)
parser.add_argument('--rotate', type = str, default = '')
args = parser.parse_args()
device = torch.device(f'cuda:{args.device}')
downsample_ratio = 0.4
model = MattingNetwork(variant='mobilenetv3').eval().to(device) # Or variant="resnet50"
model.load_state_dict(torch.load('./checkpoint/rvm_mobilenetv3.pth'))
rec = [None] * 4 # Initial recurrent states are None
frame_dataset = ItwDataset(args.input_pth, args.step, args.rotate)
# pdb.set_trace()
if not os.path.exists(args.output_pth):
os.makedirs(args.output_pth)
for data in frame_dataset:
save_img_pth = os.path.join(args.output_pth, data['file_name'] + '.png')
if os.path.exists(save_img_pth):
print(save_img_pth + ' exists!')
continue
# print('in')
with torch.no_grad():
fgr, pha, *rec = model(data['img'].to(device), *rec, downsample_ratio)
# pdb.set_trace()
mask_infer = torch.round(pha.repeat(1,3,1,1))*255
mask_infer = mask_infer.squeeze(0).permute(1,2,0).detach().cpu().numpy()
# pdb.set_trace()
cv2.imwrite(save_img_pth, mask_infer)
print(data['file_name']) |