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(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(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(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) # 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(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/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 __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 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 __init__(self, model_path, csv_path): # retrieve label info self.label_info = get_label_info(csv_path) # build model and load weight self.model = BiSeNet(12, 'resnet18') self.device = torch.device( 'cuda:0' if torch.cuda.is_available() else 'cpu') self.model.load_state_dict(torch.load(model_path)) self.model.to(self.device).eval() self.transform = transforms.Compose([ transforms.ToPILImage(), transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) # retrieve person color self.person_color = self.label_info['Pedestrian'][:-1]
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 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 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 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 train(fintune_model, data_root, respth): # dataset n_classes = 19 n_img_per_gpu = 16 n_workers = 8 cropsize = [448, 448] 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=True, num_workers=n_workers, pin_memory=True, drop_last=True) # model ignore_idx = -100 use_cuda = torch.cuda.is_available() device = torch.device("cuda:0" if use_cuda else "cpu") net = BiSeNet(n_classes=n_classes) net = net.to(device) if os.access(fintune_model, os.F_OK) and (fintune_model is not None): # checkpoint chkpt = torch.load(fintune_model, map_location=device) net.load_state_dict(chkpt) print('load fintune model : {}'.format(fintune_model)) else: print('no fintune model') 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_epoch = 1000 optim = Optimizer.SGD(net.parameters(), lr=lr_start, momentum=momentum, weight_decay=weight_decay) ## train loop msg_iter = 50 loss_avg = [] st = glob_st = time.time() # diter = iter(dl) epoch = 0 flag_change_lr_cnt = 0 # 学习率更新计数器 init_lr = lr_start # 学习率 best_loss = np.inf loss_mean = 0. # 损失均值 loss_idx = 0. # 损失计算计数器 print('start training ~') it = 0 for epoch in range(max_epoch): net.train() # 学习率更新策略 if loss_mean != 0.: if best_loss > (loss_mean / loss_idx): flag_change_lr_cnt = 0 best_loss = (loss_mean / loss_idx) else: flag_change_lr_cnt += 1 if flag_change_lr_cnt > 30: init_lr = init_lr * 0.1 set_learning_rate(optimizer, init_lr) flag_change_lr_cnt = 0 loss_mean = 0. # 损失均值 loss_idx = 0. # 损失计算计数器 for i, (im, lb) in enumerate(dl): 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_mean += loss.item() loss_idx += 1. loss.backward() optim.step() if it % msg_iter == 0: print('epoch <{}/{}> -->> <{}/{}> -> iter {} : loss {:.5f}, loss_mean :{:.5f}, best_loss :{:.5f},lr :{:.6f},batch_size : {}'.\ format(epoch,max_epoch,i,int(ds.__len__()/n_img_per_gpu),it,loss.item(),loss_mean/loss_idx,best_loss,init_lr,n_img_per_gpu)) # print(msg) if (it) % 500 == 0: state = net.module.state_dict() if hasattr( net, 'module') else net.state_dict() torch.save(state, respth + '/model/face_parse_latest.pth') # evaluate(dspth='./images', cp='{}_iter.pth'.format(it)) it += 1 torch.save(state, respth + '/model/face_parse_epoch_{}.pth'.format(epoch))
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):
import torch import torchvision import torchvision.transforms.functional as F import time from PIL import Image from model import FCN, UNet, SegNet, PSPNet, ENet, LinkNet, ICNet, BiSeNet, DFANet # model = torch.load('/home/froven/桌面/semantic-segmentation model/model_parameter/BiSeNet.pth').cuda() # model = torch.load('/home/froven/桌面/semantic-segmentation model/model_parameter/BiSeNet.pth').cuda() model = BiSeNet(num_class=19).cuda() model.eval() # data = torch.rand((1, 3, 1024, 1024)).cuda() image = Image.open('/home/froven/桌面/newdata/test/image/IMG_2047_img.png') image = F.crop(image, 200, 500, 512, 512) transform = torchvision.transforms.Compose( [ torchvision.transforms.ToTensor(), torchvision.transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) image = transform(image).float() data = (image.unsqueeze(0)).cuda() epoch_start_time = time.time() predict = model(data) predict = predict.max(1)[1].squeeze().cpu().data.numpy() print((time.time() - epoch_start_time))
def setup(opts): net = BiSeNet(n_classes=19) net.cuda() net.load_state_dict(torch.load(opts['checkpoint'])) net.eval() return net