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)