def main(model='res_unet.ResUNet', res=(256, )): from face_tran.utils.obj_factory import obj_factory model = obj_factory(model) if len(res) == 1: img = torch.rand(1, model.in_nc, res, res) pred = model(img) print(pred.shape) else: img = [] for i in range(1, len(res) + 1): img.append(torch.rand(1, model.in_nc, res[-i], res[-i])) pred = model(img) print(pred.shape)
def main(root_path, img_list, pil_transforms): import cv2 import torchvision.transforms as transforms from face_tran.utils.obj_factory import obj_factory import face_tran.utils as utils pil_transforms = [obj_factory(t) for t in pil_transforms ] if pil_transforms is not None else [] pil_transforms = transforms.Compose(pil_transforms) dataset = ImageListDataset(root_path, img_list, transform=pil_transforms) for img, label in dataset: img = np.array(img)[:, :, ::-1].copy() cv2.imshow('img', img) print('label = ' + str(label)) cv2.waitKey(0)
def main(root_path, split, pil_transforms): import cv2 import torchvision.transforms as transforms import numpy as np from face_tran.utils.obj_factory import obj_factory pil_transforms = [obj_factory(t) for t in pil_transforms ] if pil_transforms is not None else [] pil_transforms = transforms.Compose(pil_transforms) dataset = CelebaHQ(root=root_path, split=split, pil_transforms=pil_transforms) for img, label in dataset: img = np.array(img)[:, :, ::-1].copy() cv2.imshow('img', img) print('label = ' + str(label)) cv2.waitKey(0)
def __init__(self, det_model_path=None, lms_model_path=None, seg_model_path=None, gpus=None, cpu_only=None, max_size=640, batch_size=8, conf_threshold=0.5, nms_threshold=0.4, verbose=0): self.max_size = max_size self.batch_size = batch_size self.conf_threshold = conf_threshold self.nms_threshold = nms_threshold self.verbose = verbose self.device = tran_utils.set_device(gpus, cpu_only) # Load face detection model if det_model_path is not None: print('Loading face detection model: "' + os.path.basename(det_model_path) + '"...') self.detection_net = torch.jit.load(det_model_path, map_location=self.device) if self.detection_net is None: raise RuntimeError('Failed to load face detection model!') # Load face landmarks model if lms_model_path is not None: print('Loading face landmarks model: "' + os.path.basename(lms_model_path) + '"...') self.landmarks_net = torch.jit.load(lms_model_path, map_location=self.device) if self.landmarks_net is None: raise RuntimeError('Failed to load face landmarks model!') # Load face segmentation model if seg_model_path is not None: print('Loading face segmentation model: "' + os.path.basename(seg_model_path) + '"...') if seg_model_path.endswith('.pth'): checkpoint = torch.load(seg_model_path) self.segmentation_net = obj_factory(checkpoint['arch']).to(self.device) self.segmentation_net.load_state_dict(checkpoint['state_dict']) else: self.segmentation_net = torch.jit.load(seg_model_path, map_location=self.device) if self.segmentation_net is None: raise RuntimeError('Failed to load face segmentation model!') self.segmentation_net.eval()
def main(exp_dir='/data/experiments', pretrained_model=None, train_dir=None, val_dir=None, workers=4, iterations=None, epochs=180, start_epoch=0, batch_size=64, learning_rate=0.1, momentum=0.9, weight_decay=1e-4, resume_dir=None, seed=None, gpus=None, tensorboard=False, train_dataset=None, val_dataset=None, optimizer='optim.SGD(lr=0.1,momentum=0.9,weight_decay=1e-4)', scheduler='lr_scheduler.StepLR(step_size=30,gamma=0.1)', log_freq=20, pil_transforms=None, tensor_transforms=None, arch='resnet18'): best_prec1 = 0 # Validation if not os.path.isdir(exp_dir): raise RuntimeError('Experiment directory was not found: \'' + exp_dir + '\'') # Check CUDA device availability use_cuda = torch.cuda.is_available() if use_cuda: gpus = list(range(torch.cuda.device_count())) if not gpus else gpus print('=> using GPU devices: {}'.format(', '.join(map(str, gpus)))) else: gpus = None print('=> using CPU device') device = torch.device('cuda:{}'.format(gpus[0])) if gpus else torch.device('cpu') # Initialize loggers logger = SummaryWriter(log_dir=exp_dir) if tensorboard else None # Initialize datasets pil_transforms = [obj_factory(t) for t in pil_transforms] if pil_transforms is not None else [] tensor_transforms = [obj_factory(t) for t in tensor_transforms] if tensor_transforms is not None else [] if not tensor_transforms: tensor_transforms = [transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])] img_transforms = transforms.Compose(pil_transforms + tensor_transforms) val_dataset = train_dataset if val_dataset is None else val_dataset train_dataset = obj_factory(train_dataset, train_dir, transform=img_transforms) if val_dir: val_dataset = obj_factory(val_dataset, val_dir, transform=img_transforms) # Initialize data loaders sampler = None if iterations is not None: sampler = tutils.data.sampler.SubsetRandomSampler(np.random.choice(len(train_dataset), iterations)) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=batch_size, sampler=sampler, num_workers=workers, pin_memory=True, drop_last=True, shuffle=not sampler) if val_dir: if iterations is not None: sampler = tutils.data.sampler.SubsetRandomSampler(np.random.choice(len(val_dataset), iterations)) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=batch_size, sampler=sampler, num_workers=workers, pin_memory=True, drop_last=True, shuffle=not sampler) # Create model model = obj_factory(arch) model.apply(utils.init_weights) if pretrained_model is not None: pretrained_weights = torch.load(pretrained_model) if arch.startswith('vgg'): model.features.load_state_dict(pretrained_weights, strict=False) else: model.load_state_dict(pretrained_weights, strict=False) # Optimizer and scheduler optimizer = obj_factory(optimizer, model.parameters()) scheduler = obj_factory(scheduler, optimizer) # Optionally resume from a checkpoint checkpoint_dir = exp_dir if resume_dir is None else resume_dir model_path = os.path.join(checkpoint_dir, 'model_latest.pth') if os.path.isfile(model_path): print("=> loading checkpoint from '{}'".format(checkpoint_dir)) checkpoint = torch.load(model_path) best_prec1 = checkpoint['best_prec1'] start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) scheduler.load_state_dict(checkpoint['scheduler']) else: print("=> no checkpoint found at '{}'".format(checkpoint_dir)) # Support multiple GPUs if gpus and len(gpus) > 1: model = nn.DataParallel(model, gpus).to(device) else: model.to(device) # define loss function (criterion) criterion = nn.CrossEntropyLoss().to(device) cudnn.benchmark = True for epoch in range(start_epoch, epochs): if not isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): scheduler.step() # train for one epoch train(train_loader, model, criterion, optimizer, epoch, epochs, device, logger, log_freq) # evaluate on validation set val_loss, prec1 = validate(val_loader, model, criterion, epoch, epochs, device, logger, log_freq) # remember best prec@1 and save checkpoint is_best = prec1 > best_prec1 best_prec1 = max(prec1, best_prec1) utils.save_checkpoint(exp_dir, 'model', { 'epoch': epoch + 1, 'arch': arch, 'state_dict': model.module.state_dict() if gpus and len(gpus) > 1 else model.state_dict(), 'best_prec1': prec1, 'optimizer': optimizer.state_dict(), 'scheduler': scheduler.state_dict(), }, is_best) if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): scheduler.step(val_loss)
def main(exp_dir='/data/experiments', pretrained_dir=None, train_dir=None, val_dir=None, workers=4, iterations=None, epochs=90, start_epoch=0, batch_size=1, resume_dir=None, seed=None, gpus=None, tensorboard=False, train_dataset=None, val_dataset=None, optimizer='optim.SGD(lr=0.1,momentum=0.9,weight_decay=1e-4)', scheduler='lr_scheduler.StepLR(step_size=30,gamma=0.1)', log_freq=20, pil_transforms=None, tensor_transforms=None, visual_augmentations=None, geo_augmentations=None, arch='resnet18', pretrained=False, cudnn_benchmark=True): best_iou = 0 # Validation if not os.path.isdir(exp_dir): raise RuntimeError('Experiment directory was not found: \'' + exp_dir + '\'') # Seed if seed is not None: random.seed(seed) torch.manual_seed(seed) cudnn.deterministic = True warnings.warn('You have chosen to seed training. ' 'This will turn on the CUDNN deterministic setting, ' 'which can slow down your training considerably! ' 'You may see unexpected behavior when restarting ' 'from checkpoints.') # Check CUDA device availability use_cuda = torch.cuda.is_available() if use_cuda: gpus = list(range(torch.cuda.device_count())) if not gpus else gpus print('=> using GPU devices: {}'.format(', '.join(map(str, gpus)))) else: gpus = None print('=> using CPU device') device = torch.device('cuda:{}'.format( gpus[0])) if gpus else torch.device('cpu') # Initialize loggers logger = SummaryWriter(log_dir=exp_dir) if tensorboard else None # Initialize transforms pil_transforms = [obj_factory(t) for t in pil_transforms ] if type(pil_transforms) == list else pil_transforms pil_transforms = seg_transforms.ComposePair(pil_transforms) visual_augmentations = [obj_factory(t) for t in visual_augmentations ] if visual_augmentations else [] if visual_augmentations: visual_augmentations = transforms.Compose(visual_augmentations) # geo_augmentations = [obj_factory(t) for t in geo_augmentations] if visual_augmentations else [] # if geo_augmentations: # geo_augmentations = transforms.Compose(geo_augmentations) tensor_transforms = [obj_factory(t) for t in tensor_transforms ] if tensor_transforms is not None else [] if not tensor_transforms: tensor_transforms = [ transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) ] tensor_transforms = transforms.Compose(tensor_transforms) # Initialize datasets val_dataset = train_dataset if val_dataset is None else val_dataset train_dataset = obj_factory(train_dataset, train_dir, pil_transforms=pil_transforms, tensor_transforms=tensor_transforms, visual_augmentations=visual_augmentations, geo_augmentations=geo_augmentations) if val_dir: val_dataset = obj_factory(val_dataset, val_dir, pil_transforms=pil_transforms, tensor_transforms=tensor_transforms) # Initialize data loaders if iterations is None: train_sampler = tutils.data.sampler.WeightedRandomSampler( train_dataset.weights, len(train_dataset), replacement=False) else: train_sampler = tutils.data.sampler.WeightedRandomSampler( train_dataset.weights, iterations, replacement=False) train_loader = tutils.data.DataLoader( train_dataset, batch_size=batch_size, sampler=train_sampler, num_workers=workers, pin_memory=True, drop_last=True, shuffle=False, ) if val_dir: val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size, num_workers=workers) # Create model model = obj_factory(arch) if not pretrained: model.apply(utils.init_weights) model.to(device) # Optimizer and scheduler optimizer = obj_factory(optimizer, model.parameters()) scheduler = obj_factory(scheduler, optimizer) # optionally resume from a checkpoint checkpoint_dir = pretrained_dir model_path = os.path.join(checkpoint_dir, 'model_best.pth') if os.path.isfile(model_path): print("=> loading checkpoint from '{}'".format(checkpoint_dir)) checkpoint = torch.load(model_path) model.load_state_dict(checkpoint['state_dict']) else: print("=> no checkpoint found at '{}'".format(checkpoint_dir)) if pretrained: print("=> using pre-trained model '{}'".format(arch)) else: print("=> randomly initializing model...") # Support multiple GPUs if gpus and len(gpus) > 1: model = nn.DataParallel(model, gpus) # define loss function (criterion) criterion = nn.CrossEntropyLoss().to(device) # define running metrics running_metrics_val = seg_utils.Metrics(train_dataset.n_classes) # if input shape are the same for the dataset then set to True, otherwise False cudnn.benchmark = cudnn_benchmark # training/validation procedure loop for epoch in range(start_epoch, epochs): if not isinstance(scheduler, optim.lr_scheduler.ReduceLROnPlateau): scheduler.step() # train for one epoch train(train_loader, model, criterion, optimizer, epoch, device, logger, log_freq) # evaluate on validation set val_loss, score, report = validate(val_loader, model, criterion, epoch, epochs, device, logger, running_metrics_val) # remember best MeanIoU and save checkpoint iou = score["Mean IoU : \t"] is_best = iou > best_iou best_iou = max(iou, best_iou) utils.save_checkpoint( exp_dir, 'model', { 'epoch': epoch + 1, 'arch': arch, 'state_dict': model.module.state_dict() if gpus and len(gpus) > 1 else model.state_dict(), 'best_iou': best_iou, 'optimizer': optimizer.state_dict(), 'scheduler': scheduler.state_dict(), 'report': report, }, is_best) if isinstance(scheduler, optim.lr_scheduler.ReduceLROnPlateau): scheduler.step(val_loss)
def main(exp_dir='/data/experiments', train_dir=None, val_dir=None, workers=4, iterations=None, epochs=180, start_epoch=0, batch_size=64, resume_dir=None, seed=None, gpus=None, tensorboard=False, train_dataset=None, val_dataset=None, optimizer='optim.SGD(lr=0.1,momentum=0.9,weight_decay=1e-4)', scheduler='lr_scheduler.StepLR(step_size=30,gamma=0.1)', log_freq=20, pil_transforms=None, tensor_transforms=None, arch='resnet18'): best_acc_mean = 0 # Validation if not os.path.isdir(exp_dir): raise RuntimeError('Experiment directory was not found: \'' + exp_dir + '\'') # Seed if seed is not None: random.seed(seed) torch.manual_seed(seed) cudnn.deterministic = True warnings.warn('You have chosen to seed training. ' 'This will turn on the CUDNN deterministic setting, ' 'which can slow down your training considerably! ' 'You may see unexpected behavior when restarting ' 'from checkpoints.') # Check CUDA device availability use_cuda = torch.cuda.is_available() if use_cuda: gpus = list(range(torch.cuda.device_count())) if not gpus else gpus print('=> using GPU devices: {}'.format(', '.join(map(str, gpus)))) else: gpus = None print('=> using CPU device') device = torch.device('cuda:{}'.format(gpus[0])) if gpus else torch.device('cpu') # Initialize loggers logger = SummaryWriter(log_dir=exp_dir) if tensorboard else None # Initialize datasets pil_transforms = [obj_factory(t) for t in pil_transforms] if pil_transforms is not None else [] tensor_transforms = [obj_factory(t) for t in tensor_transforms] if tensor_transforms is not None else [] if not tensor_transforms: tensor_transforms = [transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])] img_transforms = transforms.Compose(pil_transforms + tensor_transforms) val_dataset = train_dataset if val_dataset is None else val_dataset train_dataset = obj_factory(train_dataset, train_dir, pil_transforms=img_transforms) if val_dir: val_dataset = obj_factory(val_dataset, val_dir, pil_transforms=img_transforms) # Initialize data loaders if iterations is None: train_sampler = tutils.data.sampler.WeightedRandomSampler(train_dataset.weights, len(train_dataset)) else: train_sampler = tutils.data.sampler.WeightedRandomSampler(train_dataset.weights, iterations) train_loader = tutils.data.DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler, num_workers=workers, pin_memory=True, drop_last=True, shuffle=False) if val_dir: val_loader = tutils.data.DataLoader(val_dataset, batch_size=batch_size, num_workers=workers, pin_memory=True, drop_last=True, shuffle=False) # Create model model = obj_factory(arch, num_classes=train_dataset.classes).to(device) # Optimizer and scheduler optimizer = obj_factory(optimizer, model.parameters()) scheduler = obj_factory(scheduler, optimizer) # Optionally resume from a checkpoint checkpoint_dir = exp_dir if resume_dir is None else resume_dir model_path = os.path.join(checkpoint_dir, 'model_latest.pth') if os.path.isfile(model_path): print("=> loading checkpoint from '{}'".format(checkpoint_dir)) checkpoint = torch.load(model_path) best_acc_mean = checkpoint['best_acc_mean'] start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) scheduler.load_state_dict(checkpoint['scheduler']) else: print("=> no checkpoint found at '{}'".format(checkpoint_dir)) # Support multiple GPUs if gpus and len(gpus) > 1: model = nn.DataParallel(model, gpus) # define loss function (criterion) criterion = nn.MultiLabelSoftMarginLoss().to(device) running_metrics_val = utils.RunningScore(train_dataset.classes) cudnn.benchmark = True for epoch in range(start_epoch, epochs): if not isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): scheduler.step() # train for one epoch train(train_loader, model, criterion, optimizer, epoch, epochs, device, logger, log_freq) # evaluate on validation set val_loss, acc_mean = validate(val_loader, model, criterion, epoch, epochs, device, logger, running_metrics_val) # remember best prec@1 and save checkpoint is_best = acc_mean > best_acc_mean best_acc_mean = max(acc_mean, best_acc_mean) utils.save_checkpoint(exp_dir, 'model', { 'epoch': epoch + 1, 'arch': arch, 'state_dict': model.module.state_dict() if gpus and len(gpus) > 1 else model.state_dict(), 'best_acc_mean': best_acc_mean, 'optimizer': optimizer.state_dict(), 'scheduler': scheduler.state_dict(), }, is_best) if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): scheduler.step(val_loss) running_metrics_val.reset()
def main(model='res_unet.ResUNet', res=256): from face_tran.utils.obj_factory import obj_factory model = obj_factory(model) img = torch.rand(1, model.in_nc, res, res) pred = model(img) print(pred.shape)
def main(exp_dir='/data/experiments', output_dir='/data/experiments', val_dir=None, workers=4, iterations=None, batch_size=1, seed=None, gpus=None, tensorboard=False, val_dataset=None, pil_transforms=None, tensor_transforms=None, arch='resnet18', cudnn_benchmark=True, use_crf=False): # Validation if not os.path.isdir(exp_dir): raise RuntimeError('Experiment directory was not found: \'' + exp_dir + '\'') # Seed if seed is not None: random.seed(seed) torch.manual_seed(seed) cudnn.deterministic = True warnings.warn('You have chosen to seed training. ' 'This will turn on the CUDNN deterministic setting, ' 'which can slow down your training considerably! ' 'You may see unexpected behavior when restarting ' 'from checkpoints.') # Check CUDA device availability use_cuda = torch.cuda.is_available() if use_cuda: gpus = list(range(torch.cuda.device_count())) if not gpus else gpus print('=> using GPU devices: {}'.format(', '.join(map(str, gpus)))) else: gpus = None print('=> using CPU device') device = torch.device('cuda:{}'.format( gpus[0])) if gpus else torch.device('cpu') # Initialize datasets if pil_transforms is not None: pil_transforms = [obj_factory(t) for t in pil_transforms ] if type(pil_transforms) == list else pil_transforms tensor_transforms = [obj_factory(t) for t in tensor_transforms ] if tensor_transforms is not None else [] if not tensor_transforms: tensor_transforms = [ transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) ] tensor_transforms = transforms.Compose(tensor_transforms) print(val_dir, val_dataset) val_dataset = obj_factory(val_dataset, val_dir, pil_transforms=pil_transforms, tensor_transforms=tensor_transforms) # Initialize data loaders val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size, num_workers=workers) # Create model model = obj_factory(arch) model.to(device) # Load weights checkpoint_dir = exp_dir model_path = os.path.join(checkpoint_dir, 'model_best.pth') if os.path.isfile(model_path): print("=> loading checkpoint from '{}'".format(checkpoint_dir)) checkpoint = torch.load(model_path) model.load_state_dict(checkpoint['state_dict']) # Support multiple GPUs if gpus and len(gpus) > 1: model = nn.DataParallel(model, gpus) # if input shape are the same for the dataset then set to True, otherwise False cudnn.benchmark = cudnn_benchmark # evaluate on validation set validate(val_loader, model, device, val_dir, batch_size)