def loading_model():
# init model
resume = "stage1_sad_54.4.pth"
model = net.VGG16(stage = 1)
ckpt = torch.load(resume, map_location=torch.device('cpu'))
model.load_state_dict(ckpt['state_dict'], strict=True)
return model
def predictImg(image,trimap,row,col):
result_dir = os.getcwd()+"/pred"
if not os.path.exists(result_dir):
os.makedirs(result_dir)
# parameters setting
parser = argparse.ArgumentParser()
args = parser.parse_args()
args.cuda = False
args.resume = "models/deep_image_matting/stage_sad.pth"
args.stage = 1
args.crop_or_resize = "whole"
args.max_size = 1600
# init model
model = net.VGG16(args)
ckpt = torch.load(args.resume, map_location='cpu')
model.load_state_dict(ckpt['state_dict'], strict=True)
# model = model.cuda()
# infer one by one
torch.cuda.empty_cache()
with torch.no_grad():
pred_mattes = inference_img_whole(args, model, image, trimap)
pred_mattes = (pred_mattes * 255).astype(np.uint8)
pred_mattes[trimap == 255] = 255
pred_mattes[trimap == 0] = 0
# target = image * pred_mattes[:,:,np.newaxis]
# cv2.imshow('target', cv2.cvtColor(target,cv2.COLOR_BGR2RGB))
# cv2.waitKey(0)
# bg = cv2.imread(os.getcwd() + '/composebg/input_image.jpg')
bg = create_bgimage(row,col)
# print('00000000')
# cv2.imshow('img33',image)
im, bg = composite4(image, bg, pred_mattes, col, row)
timestamp = str(int(time.time()))
cv2.imwrite(os.getcwd()+'/pred/test'+timestamp+'.png', pred_mattes)
fileNameAndSuffix=os.path.split(image_path)[1].split('.')
newfileNameAndSuffix=os.getcwd()+'/pred/'+fileNameAndSuffix[0]+'_replace_bg'+timestamp+'.'+fileNameAndSuffix[1]
print('=============:'+newfileNameAndSuffix)
cv2.imwrite(newfileNameAndSuffix,im)
if isDebug:
cv2.waitKey(0)
return newfileNameAndSuffix
def getnewalpha(image, mask, cFlag):
if image.shape[2] == 4: # get rid of alpha channel
image = image[:, :, 0:3]
if mask.shape[2] == 4: # get rid of alpha channel
mask = mask[:, :, 0:3]
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
trimap = mask[:, :, 0]
cudaFlag = False
if torch.cuda.is_available() and not cFlag:
cudaFlag = True
args = Namespace(crop_or_resize='whole',
cuda=cudaFlag,
max_size=1600,
resume=baseLoc +
'weights/deepmatting/stage1_sad_57.1.pth',
stage=1)
model = net.VGG16(args)
if cudaFlag:
ckpt = torch.load(args.resume)
else:
ckpt = torch.load(args.resume, map_location=torch.device("cpu"))
model.load_state_dict(ckpt['state_dict'], strict=True)
if cudaFlag:
model = model.cuda()
# ckpt = torch.load(args.resume)
# model.load_state_dict(ckpt['state_dict'], strict=True)
# model = model.cuda()
torch.cuda.empty_cache()
with torch.no_grad():
pred_mattes = inference_img_whole(args, model, image, trimap)
pred_mattes = (pred_mattes * 255).astype(np.uint8)
pred_mattes[trimap == 255] = 255
pred_mattes[trimap == 0] = 0
# pred_mattes = np.repeat(pred_mattes[:, :, np.newaxis], 3, axis=2)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
pred_mattes = np.dstack((image, pred_mattes))
return pred_mattes
def build_model(args):
model = net.VGG16(args)
model.apply(weight_init)
start_epoch = 1
if args.pretrain and os.path.isfile(args.pretrain):
print("=> loading pretrain '{}'".format(args.pretrain))
ckpt = torch.load(args.pretrain)
model.load_state_dict(ckpt['state_dict'],strict=False)
print("=> loaded pretrain '{}' (epoch {})".format(args.pretrain, ckpt['epoch']))
if args.resume and os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
ckpt = torch.load(args.resume)
start_epoch = ckpt['epoch']
model.load_state_dict(ckpt['state_dict'],strict=True)
print("=> loaded checkpoint '{}' (epoch {})".format(args.resume, ckpt['epoch']))
return start_epoch, model
def build_model(args, logger):
model = net.VGG16(args)
model.apply(weight_init)
start_epoch = 1
best_sad = 100000000.
if args.pretrain and os.path.isfile(args.pretrain):
logger.info("loading pretrain '{}'".format(args.pretrain))
ckpt = torch.load(args.pretrain)
model.load_state_dict(ckpt['state_dict'], strict=False)
logger.info("loaded pretrain '{}' (epoch {})".format(
args.pretrain, ckpt['epoch']))
if args.resume and os.path.isfile(args.resume):
logger.info("=> loading checkpoint '{}'".format(args.resume))
ckpt = torch.load(args.resume)
start_epoch = ckpt['epoch']
best_sad = ckpt['best_sad']
model.load_state_dict(ckpt['state_dict'], strict=True)
logger.info(
"=> loaded checkpoint '{}' (epoch {} bestSAD {:.3f})".format(
args.resume, ckpt['epoch'], ckpt['best_sad']))
return start_epoch, model, best_sad
def main():
print("===> Loading args")
args = get_args()
print("===> Environment init")
#os.environ["CUDA_VISIBLE_DEVICES"] = "0"
if args.cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
model = net.VGG16(args)
ckpt = torch.load(args.resume)
if args.not_strict:
model.load_state_dict(ckpt['state_dict'], strict=False)
else:
model.load_state_dict(ckpt['state_dict'], strict=True)
if args.cuda:
model = model.cuda()
print("===> Load dataset")
dataset = gen_dataset(args.imgDir, args.trimapDir)
mse_diffs = 0.
sad_diffs = 0.
cnt = len(dataset)
cur = 0
t0 = time.time()
for img_path, trimap_path in dataset:
img = cv2.imread(img_path)
trimap = cv2.imread(trimap_path)[:, :, 0]
assert(img.shape[:2] == trimap.shape[:2])
img_info = (img_path.split('/')[-1], img.shape[0], img.shape[1])
cur += 1
print('[{}/{}] {}'.format(cur, cnt, img_info[0]))
with torch.no_grad():
torch.cuda.empty_cache()
if args.crop_or_resize == "whole":
origin_pred_mattes = inference_img_whole(args, model, img, trimap)
elif args.crop_or_resize == "crop":
origin_pred_mattes = inference_img_by_crop(args, model, img, trimap)
else:
origin_pred_mattes = inference_img_by_resize(args, model, img, trimap)
# only attention unknown region
origin_pred_mattes[trimap == 255] = 1.
origin_pred_mattes[trimap == 0 ] = 0.
# origin trimap
pixel = float((trimap == 128).sum())
# eval if gt alpha is given
if args.alphaDir != '':
alpha_name = os.path.join(args.alphaDir, img_info[0])
assert(os.path.exists(alpha_name))
alpha = cv2.imread(alpha_name)[:, :, 0] / 255.
assert(alpha.shape == origin_pred_mattes.shape)
mse_diff = ((origin_pred_mattes - alpha) ** 2).sum() / pixel
sad_diff = np.abs(origin_pred_mattes - alpha).sum()
mse_diffs += mse_diff
sad_diffs += sad_diff
print("sad:{} mse:{}".format(sad_diff, mse_diff))
origin_pred_mattes = (origin_pred_mattes * 255).astype(np.uint8)
res = origin_pred_mattes.copy()
# only attention unknown region
res[trimap == 255] = 255
res[trimap == 0 ] = 0
if not os.path.exists(args.saveDir):
os.makedirs(args.saveDir)
cv2.imwrite(os.path.join(args.saveDir, img_info[0]), res)
print("Avg-Cost: {} s/image".format((time.time() - t0) / cnt))
if args.alphaDir != '':
print("Eval-MSE: {}".format(mse_diffs / cur))
print("Eval-SAD: {}".format(sad_diffs / cur))
configFile = "models/face_detector.pbtxt"
face_detector = cv2.dnn.readNetFromTensorflow(modelFile, configFile)
n_classes = 19
seg_net = BiSeNet(n_classes=n_classes)
seg_net.load_state_dict(torch.load('models/segmentation.pth', map_location='cpu'))
parser = argparse.ArgumentParser()
args = parser.parse_args()
args.resume = "models/matting.pth"
args.stage = 1
args.crop_or_resize = "whole"
args.max_size = 512
args.cuda = False
# init model
matting_model = net.VGG16(args)
ckpt = torch.load(args.resume, map_location='cpu')
matting_model.load_state_dict(ckpt['state_dict'], strict=True)
def segmentation(image):
seg_net.eval()
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
with torch.no_grad():
img = cv2.resize(image, (512, 512), cv2.INTER_LINEAR)
img = to_tensor(img)
img = torch.unsqueeze(img, 0)
from deploy import inference_img_whole
# input file list
image_path = "boy-1518482_1920_12_img.png"
trimap_path = "boy-1518482_1920_12.png"
image = cv2.imread(image_path)
trimap = cv2.imread(trimap_path)
# print(trimap.shape)
trimap = trimap[:, :, 0]
# init model
args = Namespace(crop_or_resize='whole',
cuda=True,
max_size=1600,
resume='model/stage1_sad_57.1.pth',
stage=1)
model = net.VGG16(args)
ckpt = torch.load(args.resume)
model.load_state_dict(ckpt['state_dict'], strict=True)
model = model.cuda()
torch.cuda.empty_cache()
with torch.no_grad():
pred_mattes = inference_img_whole(args, model, image, trimap)
pred_mattes = (pred_mattes * 255).astype(np.uint8)
pred_mattes[trimap == 255] = 255
pred_mattes[trimap == 0] = 0
# print(pred_mattes)
# cv2.imwrite('out.png', pred_mattes)
# import matplotlib.pyplot as plt
# plt.imshow(image)
def main():
global args, best_prec1
args = parser.parse_args()
print(args)
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
torch.cuda.set_device(int(args.gpu))
# model = net.VGG16_GCN(args.threshold)
model = net.VGG16()
model.cuda()
cudnn.benchmark = True
train_trans = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomCrop(32, 4),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
val_trans = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
train_loader = torch.utils.data.DataLoader(datasets.CIFAR10(
root=args.data, train=True, transform=train_trans, download=True),
batch_size=args.batch_size,
shuffle=True,
num_workers=2,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(datasets.CIFAR10(
root=args.data, train=False, transform=val_trans),
batch_size=args.batch_size,
shuffle=False,
num_workers=2,
pin_memory=True)
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
decreasing_lr = list(map(int, args.decreasing_lr.split(',')))
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
milestones=decreasing_lr,
gamma=0.1)
# scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, float(args.epochs))
if not os.path.exists(args.save_dir):
os.mkdir(args.save_dir)
for epoch in range(args.epochs):
print("The learning rate is {}".format(
optimizer.param_groups[0]['lr']))
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion)
scheduler.step()
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
if is_best:
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer,
},
filename=os.path.join(args.save_dir, 'best_model.pt'))
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer,
},
filename=os.path.join(args.save_dir, 'checkpoint.pt'))