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(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) # BiSeNet으로 net이리는 인스턴스 생성됨. 인자로 19 넣어서 만듦. net.cuda() # Tensor들을 GPU로 보내기 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(np.shape(img),np.shape(parsing),parsing) => torch.Size([1,3,512, 512]) / (512, 512) / [0 0 0...0 0 0]~[17 17 17... 17 17 17] # print(np.unique(parsing)) # vis_parsing_maps(image, parsing, stride=1, save_im=False, save_path=osp.join(respth, dspth)) return parsing
def evaluate(image_path='./images/img1.jpg', cp='79999_iter.pth'): # if not os.path.exists(respth): # os.makedirs(respth) n_classes = 19 net = BiSeNet(n_classes=n_classes) #net.cuda() ## uncomment this #net.load_state_dict(torch.load(cp) with map_location=torch.device('cpu')) ##net.load_state_dict(torch.load(cp)) net.load_state_dict(torch.load(cp, 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.open(image_path) image = img.resize((512, 512), Image.BILINEAR) img = to_tensor(image) img = torch.unsqueeze(img, 0) #img = img.cuda() ## uncomment this 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
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(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(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 __init__(self): sys.path.append('resources/face_parsing_pytorch/') from model import BiSeNet n_classes = 19 seg_net = BiSeNet(n_classes=n_classes).cuda() save_pth = 'resources/face_parsing_pytorch/res/cp/79999_iter.pth' seg_net.load_state_dict(torch.load(save_pth)) seg_net.eval() self.seg_net = seg_net
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 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 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 buildnet(cp='79999_iter.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, 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)), ]) return net, to_tensor
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 parser(image_path='./imgs/116.jpg', cp='cp/79999_iter.pth'): n_classes = 19 net = BiSeNet(n_classes=n_classes) device = torch.device("cpu") net.to(device) net.load_state_dict(torch.load(cp,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.open(image_path) img = to_tensor(im) img = torch.unsqueeze(img, 0) out = net(img)[0] parsing = out.squeeze(0).numpy().argmax(0) return parsing
def evaluate(face_img, cp='79999_iter.pth'): n_classes = 19 net = BiSeNet(n_classes=n_classes) # net.cuda() save_pth = osp.join('network/', cp) net.load_state_dict(torch.load(save_pth, 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 = to_tensor(face_img) img = torch.unsqueeze(img, 0) # img = img.cuda() out = net(img)[0] # print(out.shape) # torch.Size([1, 19, 512, 512]) parsing = out.squeeze(0).cpu().numpy().argmax(0) return parsing.astype(np.uint8)
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 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(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 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 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( torch.load(os.path.join(model_dir, 'cp', filename)))
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 cluster(args, outpath): # network setup n_classes = 19 net = BiSeNet(n_classes=n_classes).cuda() save_pth = 'resources/face_parsing_pytorch/res/cp/79999_iter.pth' net.load_state_dict(torch.load(save_pth)) net.eval() to_tensor = transforms.Compose([ transforms.Resize((512, 512), Image.BILINEAR), transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), ]) subsets = ['reals_easiest', 'fakes_easiest'] PatchInfo = namedtuple('PatchInfo', ['patch', 'pos', 'file', 'value']) for subset in subsets: print(subset) path = args.path cluster_dir = os.path.join(outpath, subset + '_clusters') os.makedirs(cluster_dir, exist_ok =True) with open(os.path.join(path, subset+'_files.txt')) as f: files = [line.strip() for line in f] patches = np.load(os.path.join(path, subset + '.npz')) which_model_netD = patches['which_model_netD'] fineSize = patches['finesize'] rfs = rfutil.find_rf_patches(which_model_netD, fineSize) clusters = defaultdict(list) clusters_baseline = Counter() # a counter for segmentations of random patches # assign each patch to a cluster based on segmentation for index, (patch, pos, value, file) in tqdm(enumerate( zip(patches['patch'], patches['pos'], patches['value'], files)), total=len(files)): image = Image.open(file).convert('RGB') with torch.no_grad(): tensor = to_tensor(image)[None].cuda() out = net(tensor)[0] parsing = out[0].cpu().numpy().argmax(0) interp = cv2.resize(parsing, (fineSize, fineSize), interpolation=cv2.INTER_NEAREST)[None] seg_patch = rfutil.get_patch_from_img(interp, rfs[(pos[0], pos[1])], patches['rf'], pad_value=-1) # how many pixels of each segmentation class in the patch counter_patch = Counter(seg_patch[0].ravel()) # how many pixels of each segmentation class in the full img counter_full = Counter(interp[0].ravel()) # normalize: omit padding value from normalization counter_norm = {k: counter_patch[k] / counter_full[k] for k in counter_patch.keys() if k != -1} cluster_id = max(counter_norm, key=counter_norm.get) cluster_label = cluster_assn[cluster_id] clusters[cluster_label].append(PatchInfo(patch, pos, file, value)) # pick a random patch for baseline random_patch = random.choice(list(rfs.values())) seg_patch = rfutil.get_patch_from_img(interp, random_patch, patches['rf'], pad_value=-1) counter_patch = Counter(seg_patch[0].ravel()) counter_full = Counter(interp[0].ravel()) # omit padding value from normalization counter_norm = {k: counter_patch[k] / counter_full[k] for k in counter_patch.keys() if k != -1} cluster_id = max(counter_norm, key=counter_norm.get) cluster_label = cluster_assn[cluster_id] clusters_baseline[cluster_label] += 1 # plot each cluster in a grid counts, labels = [], [] infos = [] for index, (k,v) in enumerate(sorted(clusters.items(), key=lambda item: len(item[1]))[::-1]): line = '%d: %s, %d patches' % (index, k, len(v)) print(line) infos.append(line) counts.append(len(v)) labels.append(k) cluster = np.asarray([patchinfo.patch for patchinfo in v]) files = [patchinfo.file for patchinfo in clusters[k]] normalized = (cluster[-225:] * 0.5) + 0.5 # at most 15x15 grid grid = imutil.imgrid(np.uint8(normalized * 255), pad=0, cols=int(np.ceil(np.sqrt(normalized.shape[0])))) grid_im = Image.fromarray(grid) grid_im.save(os.path.join(cluster_dir, 'cluster_%d.png' % index)) np.savez(os.path.join(cluster_dir, 'cluster_%d.npz' % index), patch=cluster, rf=patches['rf'], finesize=patches['finesize'], outsize=patches['outsize'], which_model_netD=patches['which_model_netD'], pos=np.array([patchinfo.pos for patchinfo in v]), value=np.array([patchinfo.value for patchinfo in v])) with open(os.path.join(cluster_dir, 'cluster_%d.txt' % index), 'w') as f: [f.write('%s\n' % file) for file in files] # histogram f, ax = plt.subplots(1, 1, figsize=(6, 4)) ax.bar(range(1, len(labels) + 1), counts) ax.set_xticks(range(1, len(labels) + 1)) ax.set_xticklabels(labels, rotation='vertical') ax.set_ylabel('count') f.savefig(os.path.join(cluster_dir, 'histogram.pdf'), bbox_inches='tight') # write counts to file with open(os.path.join(cluster_dir, 'counts.txt'), 'w') as f: [f.write('%s\n' % line) for line in infos] # write random patch baseline to file infos = [] for index, (k,v) in enumerate(sorted(clusters_baseline.items(), key=lambda item: item[1])[::-1]): infos.append('%d: %s, %d patches' % (index, k, v)) with open(os.path.join(cluster_dir, 'baseline.txt'), 'w') as f: [f.write('%s\n' % line) for line in infos]
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 allowed_file(filename): return '.' in filename and filename.rsplit( '.', 1)[1].lower() in ['png', 'jpeg', 'jpg'] to_tensor = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), ]) # load model device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") NCLASSES = 19 model = BiSeNet(n_classes=NCLASSES) model.load_state_dict(torch.load(model_src, map_location=torch.device(device))) model.eval() @app.route('/') def home(): return render_template("home.html") @app.route('/predict', methods=['POST']) def predict(): # get file img_file = request.files['img'] if img_file and allowed_file(img_file.filename): npimg = np.fromstring(img_file.read(), np.uint8) img = cv2.imdecode(npimg, cv2.IMREAD_COLOR)
def setup(opts): net = BiSeNet(n_classes=19) net.cuda() net.load_state_dict(torch.load(opts['checkpoint'])) net.eval() return net
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)