def evaluate(cfg,
respth='./res/test_res',
dspth='./data',
cp='model_final_diss.pth'):
if not os.path.exists(respth):
os.makedirs(respth)
n_classes = 19
net = BiSeNet(n_classes=n_classes)
if cfg["GPU"]["enable"]:
net.cuda(cfg["GPU"]["name"])
net.load_state_dict(torch.load(cp))
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():
for image_path in os.listdir(dspth):
img = Image.open(osp.join(dspth, image_path))
image = img.resize((512, 512), Image.BILINEAR)
img = to_tensor(image)
img = torch.unsqueeze(img, 0)
img = img.cuda(cfg["GPU"]["name"]) if cfg["GPU"]["enable"] else img
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
parsing = \
torch.nn.functional.interpolate(torch.tensor(np.expand_dims(parsing, axis=(0, 1))).type(torch.tensor),
size=(cfg["mask_size"], cfg["mask_size"]))[0][0].numpy()
hkl.dump(parsing, join(respth, image_path[:-4]))
def evaluate(image_path="./imgs/116.jpg", cp="cp/79999_iter.pth"):
# if not os.path.exists(respth):
# os.makedirs(respth)
n_classes = 19
net = BiSeNet(n_classes=n_classes)
if CUDA:
net.cuda()
net.load_state_dict(torch.load(cp, map_location=torch.device("cpu")))
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 = Image.fromarray(image_path)
img = Image.open(image_path)
image = img.resize((512, 512), Image.BILINEAR)
img = to_tensor(image)
img = torch.unsqueeze(img, 0)
if CUDA:
img = img.cuda()
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
print(parsing)
print(np.unique(parsing))
vis_parsing_maps(image, parsing, stride=1)
return parsing
def evaluate(image_path='./imgs/116.jpg', cp='cp/79999_iter.pth'):
# if not os.path.exists(respth):
# os.makedirs(respth)
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
net.load_state_dict(torch.load(cp))
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 = Image.open(image_path)
image = img.resize((512, 512), Image.BILINEAR)
img = to_tensor(image)
img = torch.unsqueeze(img, 0)
img = img.cuda()
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
# print(parsing)
# print(np.unique(parsing))
# vis_parsing_maps(image, parsing, stride=1, save_im=False, save_path=osp.join(respth, dspth))
return parsing
class FaceSegementer():
def __init__(self):
self.model = BiSeNet(n_classes=19)
self.model.cuda()
weights_filename = 'face-parsing-02dd3f6f.pth'
url = 'https://rewriting.csail.mit.edu/data/models/' + weights_filename
try:
sd = torch.hub.load_state_dict_from_url(url) # pytorch 1.1+
except Exception:
sd = torch.hub.model_zoo.load_url(url) # pytorch 1.0
self.model.load_state_dict(sd)
self.model.eval()
def segment_batch(self, xs):
results = []
og_size = xs.shape[2:]
xs = F.interpolate(xs, size=(512, 512))
for x in xs:
x = x.unsqueeze(dim=0)
out = self.model(x)[0]
mask = torch.from_numpy(
out.squeeze(0).cpu().numpy().argmax(0)).unsqueeze(dim=0)
results.append(mask)
masks = torch.stack(results).float()
masks = F.interpolate(masks, size=og_size).long()
return masks
def evaluate(dspth='./', cp=''):
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
save_pth = cp
print('model : {}'.format(save_pth))
net.load_state_dict(torch.load(save_pth))
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():
idx = 0
for image_path in os.listdir(dspth):
img = Image.open(osp.join(dspth, image_path))
image = img.resize((512, 512), Image.BILINEAR)
img = to_tensor(image)
img = torch.unsqueeze(img, 0)
img = img.cuda()
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
# print(parsing)
idx += 1
print('<{}> image : '.format(idx), np.unique(parsing))
test_result = './test_result'
if not osp.exists(test_result):
os.makedirs(test_result)
vis_parsing_maps(image,
parsing,
stride=1,
save_im=True,
save_path=osp.join(test_result, image_path))
def evaluate(respth='./res', dspth='./data'):
## logger
logger = logging.getLogger()
## model
logger.info('\n')
logger.info('====' * 20)
logger.info('evaluating the model ...\n')
logger.info('setup and restore model')
n_classes = 19
net = BiSeNet(n_classes=n_classes)
save_pth = osp.join(respth, 'model_final.pth')
net.load_state_dict(torch.load(save_pth))
net.cuda()
net.eval()
## dataset
batchsize = 5
n_workers = 2
dsval = CityScapes(dspth, mode='val')
dl = DataLoader(dsval,
batch_size=batchsize,
shuffle=False,
num_workers=n_workers,
drop_last=False)
## evaluator
logger.info('compute the mIOU')
evaluator = MscEval(net, dl)
## eval
mIOU = evaluator.evaluate()
logger.info('mIOU is: {:.6f}'.format(mIOU))
def evaluate(respth='./res/test_res',
dspth='./data',
cp='model_final_diss.pth'):
if not os.path.exists(respth):
os.makedirs(respth)
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
save_pth = osp.join('res/cp', cp)
net.load_state_dict(torch.load(save_pth))
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():
for image_path in os.listdir(dspth):
img = Image.open(osp.join(dspth, image_path))
image = img.resize((512, 512), Image.BILINEAR)
img = to_tensor(image)
img = torch.unsqueeze(img, 0)
img = img.cuda()
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
# print(parsing)
print(np.unique(parsing))
vis_parsing_maps(image,
parsing,
stride=1,
save_im=True,
save_path=osp.join(respth, image_path))
def evaluate(respth='./results/data_src',
dspth='../data',
cp='79999_iter.pth'):
respth = osp.join(os.path.abspath(os.path.dirname(__file__)), respth)
if not os.path.exists(respth):
os.makedirs(respth)
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
model_path = osp.join(os.path.abspath(os.path.dirname(__file__)), cp)
data_path = osp.join(os.path.abspath(os.path.dirname(__file__)), dspth)
net.load_state_dict(torch.load(model_path))
net.eval()
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
face_detector = FaceDetector(device="cuda", verbose=False)
with torch.no_grad():
for image_path in os.listdir(data_path):
img = Image.open(osp.join(data_path, image_path))
image = img.resize((512, 512), Image.BILINEAR)
# --------------------------------------------------------------------------------------
bgr_image = cv2.imread(osp.join(data_path, image_path))
detected_faces = face_detector.detect_from_image(
cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB))
lx, ly, rx, ry = int(detected_faces[0][0]), int(
detected_faces[0][1]), int(detected_faces[0][2]), int(
detected_faces[0][3])
face = bgr_image[ly:ry, lx:rx, :].copy()
# upscale the image
face = cv2.resize(face, (0, 0), fx=5, fy=5)
# cv2.imwrite(osp.join(respth, image_path), face)
face_rgb = cv2.cvtColor(face, cv2.COLOR_BGR2RGB)
# --------------------------------------------------------------------------------------
img = to_tensor(face_rgb)
img = torch.unsqueeze(img, 0)
img = img.cuda()
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
print(np.unique(parsing))
out_face_mask = vis_parsing_maps(face_rgb,
parsing,
stride=1,
save_im=True,
save_path=osp.join(
respth, image_path))
origin_face_mask = cv2.resize(out_face_mask, (rx - lx, ry - ly))
image_mask = np.zeros(np.shape(bgr_image)[:2], dtype=np.uint8)
image_mask[ly:ry, lx:rx] = origin_face_mask
result = merge(bgr_image, image_mask)
cv2.imwrite(osp.join(respth, image_path)[:-4] + '.png', result)
def create_faceparse_model(path_model):
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
print('model : {}'.format(path_model))
net.load_state_dict(torch.load(path_model))
net.eval()
return net
def setup_face_parsing():
global parser
parser_file = downloads.download_from_gdrive(
gdrive_fileid='154JgKpzCPW82qINcVieuPH3fZ2e0P812',
output_path='face-parsing/79999_iter.pth')
n_classes = 19
parser = BiSeNet(n_classes=n_classes)
parser.cuda()
parser.load_state_dict(torch.load(parser_file))
parser.eval()
def main(haar_cascade_filepath, state_dict):
app = QtWidgets.QApplication(sys.argv)
n_classes = 19
net = BiSeNet(
n_classes=n_classes) # BiSeNet으로 net이리는 인스턴스 생성됨. 인자로 19 넣어서 만듦.
net.cuda() # Tensor들을 GPU로 보내기
net.load_state_dict(state_dict)
net.eval()
main_window = MainWindow()
main_widget = MainWidget(haar_cascade_filepath, net, main_window)
main_window.setCentralWidget(main_widget)
main_window.show()
sys.exit(app.exec_())
def label_images(images: torch.tensor, cp='model_final_diss.pth'):
"""
Labels images
Args:
images: a (n x H x W x C) tensor with the images
Returns:
a (n x H x W) numpy array with the class labels
"""
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
save_pth = os.path.join('res/cp', cp)
net.load_state_dict(torch.load(save_pth))
# need to move color channels dimension to the second dimension due to weird PIL transform issues
images = images.transpose(1, 3).transpose(2, 3)
transform_images = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((512, 512)),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
net.eval()
with torch.no_grad():
labels = []
for image in tqdm.tqdm(images, desc="segmenting frames"):
img = transform_images(image)
img = torch.unsqueeze(img, 0)
img = img.cuda()
labels_probs = net(img)[0]
labels.append(labels_probs.squeeze(0).detach().cpu().argmax(0))
return torch.stack(labels, dim=0)
def evaluate(image_path, cp):
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
net.load_state_dict(torch.load(cp))
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 = Image.open(image_path)
image = img.resize((512, 512), Image.BILINEAR)
img = to_tensor(image)
img = torch.unsqueeze(img, 0)
img = img.cuda()
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
vis_parsing_maps(image, parsing, stride=1)
return parsing
def getface(input_image):
save_pth = baseLoc+'face-parsing.PyTorch/79999_iter.pth'
input_image = Image.fromarray(input_image)
n_classes = 19
net = BiSeNet(n_classes=n_classes)
if torch.cuda.is_available():
net.cuda()
net.load_state_dict(torch.load(save_pth))
else:
net.load_state_dict(torch.load(save_pth, map_location=lambda storage, loc: storage))
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 = input_image.resize((512, 512), Image.BILINEAR)
img = to_tensor(img)
img = torch.unsqueeze(img, 0)
if torch.cuda.is_available():
img = img.cuda()
out = net(img)[0]
if torch.cuda.is_available():
parsing = out.squeeze(0).cpu().numpy().argmax(0)
else:
parsing = out.squeeze(0).numpy().argmax(0)
parsing = Image.fromarray(np.uint8(parsing))
parsing = parsing.resize(input_image.size)
parsing = np.array(parsing)
return parsing
def evaluate(respth='./res', dspth='./data'):
## logger
logger = logging.getLogger()
## model
logger.info('\n')
logger.info('====' * 20)
logger.info('evaluating the model ...\n')
logger.info('setup and restore model')
n_classes = 19
net = BiSeNet(n_classes=n_classes)
save_pth = osp.join(respth, 'model_final.pth')
net.load_state_dict(torch.load(save_pth))
net.cuda()
net.eval()
## dataset
batchsize = 2
n_workers = 1
dsval = CityScapes(dspth, mode='val')
dl = DataLoader(dsval,
batch_size=batchsize,
shuffle=False,
num_workers=n_workers,
drop_last=False)
## evaluator
logger.info('compute the mIOU')
evaluator = MscEval(net, dl)
# ## export to ONNX
# dummy_input = Variable(torch.randn(batchsize, 3, 1024, 1024)).cuda()
# onnx.export(net, dummy_input, "bisenet.proto", verbose=True)
## eval
mIOU = evaluator.evaluate()
logger.info('mIOU is: {:.6f}'.format(mIOU))
def exportImgAPI(
img,
style,
dspth='/home/KLTN_TheFaceOfArtFaceParsing/Updates/face_parsing/test-img',
cp='model_final_diss.pth'):
# img = Image.open(osp.join(dspth, image_path))
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
save_pth = osp.join(
'/home/KLTN_TheFaceOfArtFaceParsing/Updates/face_parsing/res/cp', cp)
net.load_state_dict(torch.load(save_pth, map_location='cpu'))
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 = Image.fromarray(img)
image = img.resize((512, 512), Image.BILINEAR)
img = to_tensor(image)
img = torch.unsqueeze(img, 0)
img = img.cuda()
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
# print(parsing)
print(np.unique(parsing))
return vis_parsing_maps(image,
style,
parsing,
stride=1,
save_im=True,
save_path='')
import os
import numpy as np
import torch
import math
from torch.utils.data import DataLoader
from model import BiSeNet
from face_dataset import FaceMask
from sklearn.metrics import f1_score
if __name__ == '__main__':
model_dir = '/home/jihyun/workspace/face_parsing/face-parsing.PyTorch/res/exp'
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
net.eval()
data_root = '/home/jihyun/workspace/face_parsing/dataset/CelebAMask-HQ/'
cropsize = [448, 448]
n_img_per_gpu = 16
ds = FaceMask(data_root, cropsize=cropsize, mode='val')
dl = DataLoader(ds, batch_size=16, shuffle=False, drop_last=True)
f = open(os.path.join(model_dir, 'f_measure.txt'), 'w+')
for filename in os.listdir(os.path.join(model_dir, 'cp')):
if filename.endswith('70000_iter.pth'):
net.load_state_dict(
def testval(model_path, sv_dir='res', sv_pred=False):
n_classes = 17
confusion_matrix = np.zeros(
(n_classes, n_classes)) # num_classes x num_classes
cropsize = [448, 448]
data_root = '/home/data2/DATASET/vschallenge'
batch_size = 1
ds = FaceMask(data_root, cropsize=cropsize, mode='val')
dl = DataLoader(ds, batch_size=batch_size, shuffle=False, num_workers=0)
with torch.no_grad():
net = BiSeNet(n_classes=n_classes)
net.cuda()
net.load_state_dict(torch.load(model_path))
net.eval()
for iter, data in enumerate(dl):
im = data['img']
lb = data['label']
impth = data['impth']
lb = torch.squeeze(lb, 1)
im = im.cuda()
lb = lb.cuda()
out = net(im)[0]
size = lb.size()
pred = out.cpu().numpy().argmax(1)
gt = lb.cpu().numpy()
vis_parsing_maps(im.cpu(),
pred,
stride=1,
save_im=True,
save_path='res',
imspth=impth)
# vis_parsing_maps(im.cpu(), gt, stride=1, save_im=True, save_path='res_gt', imspth=impth)
confusion_matrix += get_confusion_matrix(
lb, out, size, n_classes, ignore=-1) # [16, 19, 448, 448]
if sv_pred:
sv_path = os.path.join(sv_dir, 'test_results')
if not os.path.exists(sv_path):
os.mkdir(sv_path)
# cv2.imwrite(sv_path+'/', ) add imname
if iter % 5 == 0:
pos = confusion_matrix.sum(1)
res = confusion_matrix.sum(0)
tp = np.diag(confusion_matrix)
pixel_acc = tp.sum() / pos.sum()
mean_acc = (tp / np.maximum(1.0, pos)).mean()
IoU_array = (tp / np.maximum(1.0, pos + res - tp))
mean_IoU = IoU_array.mean()
print(
'index/allimg {}/{}. mean_IoU: {:1.5f}. pixel_acc: {:1.5f}. mean_acc: {:1.5f}.'
.format(iter * batch_size,
len(dl) * batch_size, mean_IoU, pixel_acc,
mean_acc))
pos = confusion_matrix.sum(1)
res = confusion_matrix.sum(0)
tp = np.diag(confusion_matrix)
pixel_acc = tp.sum() / pos.sum()
mean_acc = (tp / np.maximum(1.0, pos)).mean()
IoU_array = (tp / np.maximum(1.0, pos + res - tp))
mean_IoU = IoU_array.mean()
print(
'mean_IoU: {:1.5f}. pixel_acc: {:1.5f}. mean_acc: {:1.5f}.'.format(
mean_IoU, pixel_acc, mean_acc))
return mean_IoU, IoU_array, pixel_acc, mean_acc
def train(opt):
# saving setting
opt.saved_path = opt.saved_path + opt.project
opt.log_path = os.path.join(opt.saved_path, 'tensorboard')
os.makedirs(opt.log_path, exist_ok=True)
os.makedirs(opt.saved_path, exist_ok=True)
# gpu setting
os.environ["CUDA_VISIBLE_DEVICES"] = '2, 3, 4, 5, 6'
gpu_number = torch.cuda.device_count()
# dataset setting
n_classes = 17
n_img_all_gpu = opt.batch_size * gpu_number
cropsize = [448, 448]
data_root = '/home/data2/DATASET/vschallenge'
num_workers = opt.num_workers
ds = FaceMask(data_root, cropsize=cropsize, mode='train')
dl = DataLoader(ds,
batch_size=n_img_all_gpu,
shuffle=True,
num_workers=num_workers,
drop_last=True)
ds_eval = FaceMask(data_root, cropsize=cropsize, mode='val')
dl_eval = DataLoader(ds_eval,
batch_size=n_img_all_gpu,
shuffle=True,
num_workers=num_workers,
drop_last=True)
ignore_idx = -100
net = BiSeNet(n_classes=n_classes)
# load last weights
if opt.load_weights is not None:
if opt.load_weights.endswith('.pth'):
weights_path = opt.load_weights
else:
weights_path = get_last_weights(opt.saved_path)
try:
last_step = int(
os.path.basename(weights_path).split('_')[-1].split('.')[0])
except:
last_step = 0
try:
ret = net.load_state_dict(torch.load(weights_path), strict=False)
except RuntimeError as e:
print(f'[Warning] Ignoring {e}')
print(
'[Warning] Don\'t panic if you see this, this might be because you load a pretrained weights '
'with different number of classes. The rest of the weights should be loaded already.'
)
print(
f'[Info] loaded weights: {os.path.basename(weights_path)}, resuming checkpoint from step: {last_step}'
)
else:
last_step = 0
print('[Info] initializing weights...')
writer = SummaryWriter(
opt.log_path +
f'/{datetime.datetime.now().strftime("%Y%m%d-%H%M%S")}/')
net = net.cuda()
net = nn.DataParallel(net)
score_thres = 0.7
n_min = n_img_all_gpu * cropsize[0] * cropsize[1] // opt.batch_size
LossP = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
Loss2 = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
Loss3 = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
# optimizer
momentum = 0.9
weight_decay = 5e-4
lr_start = opt.lr
max_iter = 80000
power = 0.9
warmup_steps = 1000
warmup_start_lr = 1e-5
optim = Optimizer(model=net.module,
lr0=lr_start,
momentum=momentum,
wd=weight_decay,
warmup_steps=warmup_steps,
warmup_start_lr=warmup_start_lr,
max_iter=max_iter,
power=power)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optim.optim,
patience=3,
verbose=True)
# train loop
loss_avg = []
step = max(0, last_step)
max_iter = len(dl)
best_epoch = 0
epoch = 0
best_loss = 1e5
net.train()
try:
for epoch in range(opt.num_epochs):
last_epoch = step // max_iter
if epoch < last_epoch:
continue
epoch_loss = []
progress_bar = tqdm(dl)
for iter, data in enumerate(progress_bar):
if iter < step - last_epoch * max_iter:
progress_bar.update()
continue
try:
im = data['img']
lb = data['label']
lb = torch.squeeze(lb, 1)
im = im.cuda()
lb = lb.cuda()
optim.zero_grad()
out, out16, out32 = net(im)
lossp = LossP(out, lb)
loss2 = Loss2(out16, lb)
loss3 = Loss3(out32, lb)
loss = lossp + loss2 + loss3
if loss == 0 or not torch.isfinite(loss):
continue
loss.backward()
optim.step()
loss_avg.append(loss.item())
# print training log message
# progress_bar.set_description(
# 'Epoch: {}/{}. Iteration: {}/{}. p_loss: {:.5f}. 2_loss: {:.5f}. 3_loss: {:.5f}. loss_avg: {:.5f}'.format(
# epoch, opt.num_epochs, iter + 1, max_iter, lossp.item(),
# loss2.item(), loss3.item(), loss.item()))
print(
'p_loss: {:.5f}. 2_loss: {:.5f}. 3_loss: {:.5f}. loss_avg: {:.5f}'
.format(lossp.item(), loss2.item(), loss3.item(),
loss.item()))
writer.add_scalars('Lossp', {'train': lossp}, step)
writer.add_scalars('loss2', {'train': loss2}, step)
writer.add_scalars('loss3', {'train': loss3}, step)
writer.add_scalars('loss_avg', {'train': loss}, step)
# log learning_rate
lr = optim.lr
writer.add_scalar('learning_rate', lr, step)
step += 1
if step % opt.save_interval == 0 and step > 0:
save_checkpoint(net, f'Bisenet_{epoch}_{step}.pth')
print('checkpoint...')
except Exception as e:
print('[Erro]', traceback.format_exc())
print(e)
continue
scheduler.step(np.mean(epoch_loss))
if epoch % opt.val_interval == 0:
net.eval()
loss_p = []
loss_2 = []
loss_3 = []
for iter, data in enumerate(dl_eval):
with torch.no_grad():
im = data['img']
lb = data['label']
lb = torch.squeeze(lb, 1)
im = im.cuda()
lb = lb.cuda()
out, out16, out32 = net(im)
lossp = LossP(out, lb)
loss2 = Loss2(out16, lb)
loss3 = Loss3(out32, lb)
loss = lossp + loss2 + loss3
if loss == 0 or not torch.isfinite(loss):
continue
loss_p.append(lossp.item())
loss_2.append(loss2.item())
loss_3.append(loss3.item())
lossp = np.mean(loss_p)
loss2 = np.mean(loss_2)
loss3 = np.mean(loss_3)
loss = lossp + loss2 + loss3
print(
'Val. Epoch: {}/{}. p_loss: {:1.5f}. 2_loss: {:1.5f}. 3_loss: {:1.5f}. Total_loss: {:1.5f}'
.format(epoch, opt.num_epochs, lossp, loss2, loss3, loss))
writer.add_scalars('Total_loss', {'val': loss}, step)
writer.add_scalars('p_loss', {'val': lossp}, step)
writer.add_scalars('2_loss', {'val': loss2}, step)
writer.add_scalars('3_loss', {'val': loss3}, step)
if loss + opt.es_min_delta < best_loss:
best_loss = loss
best_epoch = epoch
save_checkpoint(net, f'Bisenet_{epoch}_{step}.pth')
net.train() # ??
# Early stopping
if epoch - best_epoch > opt.es_patience > 0:
print(
'[Info] Stop training at epoch {}. The lowest loss achieved is {}'
.format(epoch, loss))
break
except KeyboardInterrupt:
save_checkpoint(net, f'Bisenet_{epoch}_{step}.pth')
writer.close()
writer.close()
def evaluate(respth='./res/test_res',
dspth='./data',
cp='model_final_diss.pth'):
if not os.path.exists(respth):
os.makedirs(respth)
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
save_pth = osp.join(respth, 'cp', cp)
net.load_state_dict(torch.load(save_pth))
net.eval()
no_iter = str(int(cp.split('_')[0]))
org_respth = respth[:]
respth = os.path.join(respth, no_iter)
if not os.path.exists(respth):
os.makedirs(respth)
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
''' added '''
cropsize = [448, 448]
n_img_per_gpu = 16
data_root = '/home/jihyun/workspace/face_parsing/dataset/CelebAMask-HQ/'
ds = FaceMask(data_root, cropsize=cropsize, mode='val')
dl = DataLoader(ds, batch_size=16, shuffle=False, drop_last=True)
n_min = n_img_per_gpu * cropsize[0] * cropsize[1] // 16
score_thres = 0.7
ignore_idx = -100
loss_avg = []
LossP = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
Loss2 = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
Loss3 = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
with torch.no_grad():
for i, sample in enumerate(dl):
im, lb = sample
im = im.cuda()
lb = lb.cuda()
lb = torch.squeeze(lb, 1)
out, out16, out32 = net(im)
lossp = LossP(out, lb)
loss2 = Loss2(out16, lb)
loss3 = Loss3(out32, lb)
loss = lossp + loss2 + loss3
loss_avg.append(loss.item())
loss_avg = sum(loss_avg) / len(loss_avg)
f = open(osp.join(org_respth, 'loss.log'), 'a')
f.write(' eval_loss: ' + str(loss_avg) + '\n')
f.close()
for image_path in os.listdir(dspth):
img = Image.open(osp.join(dspth, image_path))
image = img.resize((512, 512), Image.BILINEAR)
img = to_tensor(image)
img = torch.unsqueeze(img, 0)
img = img.cuda()
out, out16, out32 = net(img)
parsing = out.squeeze(0).cpu().numpy().argmax(0)
vis_parsing_maps(image,
parsing,
stride=1,
save_im=True,
save_path=osp.join(respth, image_path))
def train():
args = parse_args()
torch.cuda.set_device(args.local_rank)
dist.init_process_group(backend='nccl',
init_method='tcp://127.0.0.1:33241',
world_size=torch.cuda.device_count(),
rank=args.local_rank)
setup_logger(respth)
# dataset
n_classes = 19
n_img_per_gpu = 16
n_workers = 8
cropsize = [448, 448]
data_root = '/home/data2/DATASET/CelebAMask-HQ/'
ds = FaceMask(data_root, cropsize=cropsize, mode='train')
sampler = torch.utils.data.distributed.DistributedSampler(ds)
dl = DataLoader(ds,
batch_size=n_img_per_gpu,
shuffle=False,
sampler=sampler,
num_workers=n_workers,
pin_memory=True,
drop_last=True)
# model
ignore_idx = -100
net = BiSeNet(n_classes=n_classes)
net.cuda()
net.train()
net = nn.parallel.DistributedDataParallel(net,
device_ids=[
args.local_rank,
],
output_device=args.local_rank)
score_thres = 0.7
n_min = n_img_per_gpu * cropsize[0] * cropsize[1] // 16
LossP = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
Loss2 = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
Loss3 = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
## optimizer
momentum = 0.9
weight_decay = 5e-4
lr_start = 1e-2
max_iter = 80000
power = 0.9
warmup_steps = 1000
warmup_start_lr = 1e-5
optim = Optimizer(model=net.module,
lr0=lr_start,
momentum=momentum,
wd=weight_decay,
warmup_steps=warmup_steps,
warmup_start_lr=warmup_start_lr,
max_iter=max_iter,
power=power)
## train loop
msg_iter = 50
loss_avg = []
st = glob_st = time.time()
diter = iter(dl)
epoch = 0
for it in range(max_iter):
try:
im, lb = next(diter)
if not im.size()[0] == n_img_per_gpu:
raise StopIteration
except StopIteration:
epoch += 1
sampler.set_epoch(epoch)
diter = iter(dl)
im, lb = next(diter)
im = im.cuda()
lb = lb.cuda()
H, W = im.size()[2:]
lb = torch.squeeze(lb, 1)
optim.zero_grad()
out, out16, out32 = net(im)
lossp = LossP(out, lb)
loss2 = Loss2(out16, lb)
loss3 = Loss3(out32, lb)
loss = lossp + loss2 + loss3
loss.backward()
optim.step()
loss_avg.append(loss.item())
# print training log message
if (it + 1) % msg_iter == 0:
loss_avg = sum(loss_avg) / len(loss_avg)
lr = optim.lr
ed = time.time()
t_intv, glob_t_intv = ed - st, ed - glob_st
eta = int((max_iter - it) * (glob_t_intv / it))
eta = str(datetime.timedelta(seconds=eta))
msg = ', '.join([
'it: {it}/{max_it}',
'lr: {lr:4f}',
'loss: {loss:.4f}',
'eta: {eta}',
'time: {time:.4f}',
]).format(it=it + 1,
max_it=max_iter,
lr=lr,
loss=loss_avg,
time=t_intv,
eta=eta)
logger.info(msg)
loss_avg = []
st = ed
if dist.get_rank() == 0:
if (it + 1) % 5000 == 0:
state = net.module.state_dict() if hasattr(
net, 'module') else net.state_dict()
if dist.get_rank() == 0:
torch.save(state, './res/cp/{}_iter.pth'.format(it))
evaluate(dspth='/home/data2/DATASET/CelebAMask-HQ/test-img',
cp='{}_iter.pth'.format(it))
# dump the final model
save_pth = osp.join(respth, 'model_final_diss.pth')
# net.cpu()
state = net.module.state_dict() if hasattr(net,
'module') else net.state_dict()
if dist.get_rank() == 0:
torch.save(state, save_pth)
logger.info('training done, model saved to: {}'.format(save_pth))
def train():
setup_logger(respth)
# dataset
gpu_number = torch.cuda.device_count()
n_classes = 19
n_img_all_gpu = 16 * gpu_number
cropsize = [448, 448]
data_root = '/home/data2/DATASET/CelebAMask-HQ/'
ds = FaceMask(data_root, cropsize=cropsize, mode='train')
dl = DataLoader(
ds,
batch_size=n_img_all_gpu,
shuffle=True,
)
# model
ignore_idx = -100
net = BiSeNet(n_classes=n_classes)
net.cuda()
net.train()
net = nn.DataParallel(net)
net = net.cuda()
score_thres = 0.7
n_min = n_img_all_gpu * cropsize[0] * cropsize[1] // 16
LossP = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
Loss2 = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
Loss3 = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
# optimizer
momentum = 0.9
weight_decay = 5e-4
lr_start = 1e-2
max_iter = 80000
power = 0.9
warmup_steps = 1000
warmup_start_lr = 1e-5
optim = Optimizer(model=net.module,
lr0=lr_start,
momentum=momentum,
wd=weight_decay,
warmup_steps=warmup_steps,
warmup_start_lr=warmup_start_lr,
max_iter=max_iter,
power=power)
# train loop
msg_iter = 2
loss_avg = []
st = glob_st = time.time()
diter = iter(dl)
for it in range(max_iter):
# try:
im, lb = next(diter)
im = im.cuda()
lb = lb.cuda()
# H, W = im.size()[2:]
lb = torch.squeeze(lb, 1)
optim.zero_grad()
out, out16, out32 = net(im)
lossp = LossP(out, lb)
loss2 = Loss2(out16, lb)
loss3 = Loss3(out32, lb)
loss = lossp + loss2 + loss3
loss.backward()
optim.step()
loss_avg.append(loss.item())
# print training log message
if (it + 1) % msg_iter == 0:
loss_avg = sum(loss_avg) / len(loss_avg)
lr = optim.lr
ed = time.time()
t_intv, glob_t_intv = ed - st, ed - glob_st
eta = int((max_iter - it) * (glob_t_intv / it))
eta = str(datetime.timedelta(seconds=eta))
msg = ', '.join([
'it: {it}/{max_it}',
'lr: {lr:4f}',
'loss: {loss:.4f}',
'eta: {eta}',
'time: {time:.4f}',
]).format(it=it + 1,
max_it=max_iter,
lr=lr,
loss=loss_avg,
time=t_intv,
eta=eta)
logger.info(msg)
loss_avg = []
st = ed
# dump the final model
save_pth = osp.join(respth, 'model_final_diss.pth')
# net.cpu()
state = net.module.state_dict() if hasattr(net,
'module') else net.state_dict()
if dist.get_rank() == 0:
torch.save(state, save_pth)
logger.info('training done, model saved to: {}'.format(save_pth))
def setup(opts):
net = BiSeNet(n_classes=19)
net.cuda()
net.load_state_dict(torch.load(opts['checkpoint']))
net.eval()
return net
def train():
os.environ["CUDA_VISIBLE_DEVICES"] = '6, 7'
args = parse_args()
# 1. 初始化horovod
hvd.init()
# 2. 给当前进程分配对应的gpu，local_rank()返回的是当前是第几个进程
torch.cuda.set_device(hvd.local_rank())
# torch.cuda.set_device(args.local_rank)
# dist.init_process_group(
# backend = 'nccl',
# init_method = 'tcp://127.0.0.1:33271',
# world_size = 2,
# world_size = torch.cuda.device_count(),
# rank=args.local_rank
# )
setup_logger(respth)
## dataset
n_classes = 19
n_img_per_gpu = 8
n_workers = 4
cropsize = [1024, 1024]
ds = CityScapes('/dataset/cityscapes/leftImg8bit_trainvaltest',
cropsize=cropsize,
mode='train')
sampler = torch.utils.data.distributed.DistributedSampler(
ds, num_replicas=hvd.size(), rank=hvd.rank())
dl = DataLoader(ds,
batch_size=n_img_per_gpu,
shuffle=False,
sampler=sampler,
num_workers=n_workers,
pin_memory=True,
drop_last=True)
## model
ignore_idx = 255
net = BiSeNet(n_classes=n_classes)
net.cuda()
# 5. 初始化的时候广播参数，这个是为了在一开始的时候同步各个gpu之间的参数
hvd.broadcast_parameters(net.state_dict(), root_rank=0)
net.train()
# net = nn.parallel.DistributedDataParallel(net,
# device_ids = [args.local_rank, ],
# output_device = args.local_rank
# )
score_thres = 0.7
n_min = n_img_per_gpu * cropsize[0] * cropsize[1] // 16
criteria_p = OhemCELoss(thresh=score_thres,
n_min=n_min,
ignore_lb=ignore_idx)
criteria_16 = OhemCELoss(thresh=score_thres,
n_min=n_min,
ignore_lb=ignore_idx)
criteria_32 = OhemCELoss(thresh=score_thres,
n_min=n_min,
ignore_lb=ignore_idx)
## optimizer
momentum = 0.9
weight_decay = 5e-4
lr_start = 1e-2
max_iter = 80000
power = 0.9
warmup_steps = 1000
warmup_start_lr = 1e-5
optim = Optimizer(model=net,
lr0=lr_start,
momentum=momentum,
wd=weight_decay,
warmup_steps=warmup_steps,
warmup_start_lr=warmup_start_lr,
max_iter=max_iter,
power=power)
hvd.broadcast_optimizer_state(optim.optim, root_rank=0)
optim = hvd.DistributedOptimizer(optim.optim,
named_parameters=net.named_parameters())
## train loop
msg_iter = 50
loss_avg = []
st = glob_st = time.time()
diter = iter(dl)
epoch = 0
for it in range(max_iter):
try:
im, lb = next(diter)
if not im.size()[0] == n_img_per_gpu: raise StopIteration
except StopIteration:
epoch += 1
sampler.set_epoch(epoch)
diter = iter(dl)
im, lb = next(diter)
im = im.cuda()
lb = lb.cuda()
H, W = im.size()[2:]
lb = torch.squeeze(lb, 1)
optim.zero_grad()
out, out16, out32 = net(im)
lossp = criteria_p(out, lb)
loss2 = criteria_16(out16, lb)
loss3 = criteria_32(out32, lb)
loss = lossp + loss2 + loss3
loss.backward()
optim.step()
loss_avg.append(loss.item())
## print training log message
if (it + 1) % msg_iter == 0:
loss_avg = sum(loss_avg) / len(loss_avg)
lr = optim.lr
ed = time.time()
t_intv, glob_t_intv = ed - st, ed - glob_st
eta = int((max_iter - it) * (glob_t_intv / it))
eta = str(datetime.timedelta(seconds=eta))
msg = ', '.join([
'it: {it}/{max_it}',
'lr: {lr:4f}',
'loss: {loss:.4f}',
'eta: {eta}',
'time: {time:.4f}',
]).format(it=it + 1,
max_it=max_iter,
lr=lr,
loss=loss_avg,
time=t_intv,
eta=eta)
logger.info(msg)
loss_avg = []
st = ed
## dump the final model
save_pth = osp.join(respth, 'model_final.pth')
net.cpu()
state = net.module.state_dict() if hasattr(net,
'module') else net.state_dict()
if dist.get_rank() == 0: torch.save(state, save_pth)
logger.info('training done, model saved to: {}'.format(save_pth))
def evaluate(source='./data', dest='./res/test_res', cp='evaluate.pth'):
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
save_pth = osp.join('res/cp', cp)
net.load_state_dict(torch.load(save_pth))
net.eval()
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
if path.isfile(source):
full_path = True
#dspth = np.sort(np.loadtxt(source, dtype=np.str))
dspth = np.sort(
np.asarray(pd.read_csv(source, delimiter=' ',
header=None)).flatten())
else:
full_path = False
dspth = listdir(source)
with torch.no_grad():
for image_name in tqdm(dspth):
if not full_path:
image_path = path.join(source, image_name)
else:
image_path = image_name
image_name = path.split(image_name)[1]
img = Image.open(image_path)
if full_path:
sub_folder = path.basename(
path.normpath(path.split(image_path)[0]))
dest_path = path.join(dest, sub_folder)
if not path.exists(dest_path):
makedirs(dest_path)
else:
dest_path = dest
save_path = osp.join(dest_path, image_name)
if os.path.isfile(save_path[:-4] + '_mask.png'):
print('Skipping...')
continue
image = img.resize((512, 512), Image.BILINEAR)
img = to_tensor(image)
img = torch.unsqueeze(img, 0)
img = img.cuda()
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
# print(parsing)
# print(np.unique(parsing))
vis_parsing_maps(image,
parsing,
stride=1,
save_im=True,
save_path=save_path)
