def __init__(self, model_path, device=None, benchmark: bool = True): ''' Args: module_path: path to pre-trained model (available to download, see README) device: CUDA device to use. str or torch.device instance warning: this is restricted to 'cpu' or 'cuda' only ('cuda:1' won't work due to main package arcitecture) default is choose 'cuda' if available benchmark: to enable cudnn benchmark mode or not ''' self.model_path = model_path if device is None: device = 'cuda' if torch.cuda.is_available() else 'cpu' self.device = device if benchmark: torch.backends.cudnn.enabled = True torch.backends.cudnn.benchmark = True snapshot = torch.load(self.model_path, map_location=self.device) self.param = snapshot['args'] self.transform = transforms.Compose([ transforms.PreCrop(self.param.pre_crop_expand), transforms.TrainScale2WH( (self.param.crop_width, self.param.crop_height)), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), ]) self.net = models.__dict__[self.param.arch](self.param.modelconfig, None) self.net.train(False).to(self.device) weights = models.remove_module_dict(snapshot['state_dict']) self.net.load_state_dict(weights)
def main(): # Init logger if not os.path.isdir(args.save_path): os.makedirs(args.save_path) log = open( os.path.join(args.save_path, 'seed-{}-{}.log'.format(args.manualSeed, time_for_file())), 'w') print_log('save path : {}'.format(args.save_path), log) print_log('------------ Options -------------', log) for k, v in sorted(vars(args).items()): print_log('Parameter : {:20} = {:}'.format(k, v), log) print_log('-------------- End ----------------', log) print_log("Random Seed: {}".format(args.manualSeed), log) print_log("python version : {}".format(sys.version.replace('\n', ' ')), log) print_log("Pillow version : {}".format(PIL.__version__), log) print_log("torch version : {}".format(torch.__version__), log) print_log("cudnn version : {}".format(torch.backends.cudnn.version()), log) # General Data Argumentation mean_fill = tuple([int(x * 255) for x in [0.5, 0.5, 0.5]]) normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) assert args.arg_flip == False, 'The flip is : {}, rotate is {}'.format( args.arg_flip, args.rotate_max) train_transform = [transforms.PreCrop(args.pre_crop_expand)] train_transform += [ transforms.TrainScale2WH((args.crop_width, args.crop_height)) ] train_transform += [ transforms.AugScale(args.scale_prob, args.scale_min, args.scale_max) ] #if args.arg_flip: # train_transform += [transforms.AugHorizontalFlip()] if args.rotate_max: train_transform += [transforms.AugRotate(args.rotate_max)] train_transform += [ transforms.AugCrop(args.crop_width, args.crop_height, args.crop_perturb_max, mean_fill) ] train_transform += [transforms.ToTensor(), normalize] train_transform = transforms.Compose(train_transform) eval_transform = transforms.Compose([ transforms.PreCrop(args.pre_crop_expand), transforms.TrainScale2WH((args.crop_width, args.crop_height)), transforms.ToTensor(), normalize ]) assert ( args.scale_min + args.scale_max ) / 2 == args.scale_eval, 'The scale is not ok : {},{} vs {}'.format( args.scale_min, args.scale_max, args.scale_eval) args.downsample = 8 # By default args.sigma = args.sigma * args.scale_eval train_data = datasets.GeneralDataset(train_transform, args.sigma, args.downsample, args.heatmap_type, args.dataset_name) train_data.load_list(args.train_list, args.num_pts, True) if args.convert68to49: train_data.convert68to49() elif args.convert68to51: train_data.convert68to51() train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True) eval_loaders = [] if args.eval_lists is not None: for eval_list in args.eval_lists: eval_data = datasets.GeneralDataset(eval_transform, args.sigma, args.downsample, args.heatmap_type, args.dataset_name) eval_data.load_list(eval_list, args.num_pts, True) if args.convert68to49: eval_data.convert68to49() elif args.convert68to51: eval_data.convert68to51() eval_loader = torch.utils.data.DataLoader( eval_data, batch_size=args.eval_batch, shuffle=False, num_workers=args.workers, pin_memory=True) eval_loaders.append(eval_loader) if args.convert68to49 or args.convert68to51: assert args.num_pts == 68, 'The format of num-pts is not right : {}'.format( args.num_pts) assert args.convert68to49 + args.convert68to51 == 1, 'Only support one convert' if args.convert68to49: args.num_pts = 49 else: args.num_pts = 51 args.modelconfig = models.ModelConfig(train_data.NUM_PTS + 1, args.cpm_stage, args.pretrain, args.argmax_size) if args.cycle_model_path is None: # define the network itnetwork = models.itn_model(args.modelconfig, args, log) cycledata = datasets.CycleDataset(train_transform, args.dataset_name) cycledata.set_a(args.cycle_a_lists) cycledata.set_b(args.cycle_b_lists) print_log('Cycle-data initialize done : {}'.format(cycledata), log) args.cycle_model_path = procedure.train_cycle_gan( cycledata, itnetwork, args, log) assert osp.isdir( args.cycle_model_path), '{} does not exist or is not dir.'.format( args.cycle_model_path) # start train itn-cpm model itn_cpm = models.__dict__[args.arch](args.modelconfig, args.cycle_model_path) procedure.train_san_epoch(args, itn_cpm, train_loader, eval_loaders, log) log.close()
def evaluate(args): if not args.cpu: assert torch.cuda.is_available(), 'CUDA is not available.' torch.backends.cudnn.enabled = True torch.backends.cudnn.benchmark = True print('The image is {:}'.format(args.image)) print('The model is {:}'.format(args.model)) snapshot = Path(args.model) assert snapshot.exists(), 'The model path {:} does not exist' if args.cpu: snapshot = torch.load(snapshot, map_location='cpu') else: snapshot = torch.load(snapshot) mean_fill = tuple([int(x * 255) for x in [0.5, 0.5, 0.5]]) normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) param = snapshot['args'] eval_transform = transforms.Compose([ transforms.PreCrop(param.pre_crop_expand), transforms.TrainScale2WH((param.crop_width, param.crop_height)), transforms.ToTensor(), normalize ]) net = models.__dict__[param.arch](param.modelconfig, None) if not args.cpu: net = net.cuda() weights = models.remove_module_dict(snapshot['state_dict']) net.load_state_dict(weights) dataset = datasets.GeneralDataset(eval_transform, param.sigma, param.downsample, param.heatmap_type, param.dataset_name) dataset.reset(param.num_pts) print('[{:}] prepare the input data'.format(time_string())) print("Using MT-CNN face detector.") try: face = utils.detect_face_mtcnn(args.image) except utils.mtcnn_detector.BBoxNotFound: print("MT-CNN detector failed! Using default bbox instead.") face = [153.08, 462., 607.78, 1040.42] [image, _, _, _, _, _, cropped_size], meta = dataset.prepare_input(args.image, face) print('[{:}] prepare the input data done'.format(time_string())) print('Net : \n{:}'.format(net)) # network forward with torch.no_grad(): if args.cpu: inputs = image.unsqueeze(0) else: inputs = image.unsqueeze(0).cuda() gan_output = (net.netG_A(inputs) + net.netG_B(inputs)) / 2 gan_output = (gan_output * 0.5 + 0.5).squeeze(0).cpu().permute( 1, 2, 0).numpy() Image.fromarray((gan_output * 255).astype(np.uint8)).save( args.save_path.replace(".jpg", ".gan.jpg")) batch_heatmaps, batch_locs, batch_scos, _ = net(inputs) #print ('input-shape : {:}'.format(inputs.shape)) flops, params = get_model_infos(net, inputs.shape, None) print('\nIN-shape : {:}, FLOPs : {:} MB, Params : {:}.'.format( list(inputs.shape), flops, params)) flops, params = get_model_infos(net, None, inputs) print('\nIN-shape : {:}, FLOPs : {:} MB, Params : {:}.'.format( list(inputs.shape), flops, params)) print('[{:}] the network forward done'.format(time_string())) # obtain the locations on the image in the orignial size cpu = torch.device('cpu') np_batch_locs, np_batch_scos, cropped_size = batch_locs.to( cpu).numpy(), batch_scos.to(cpu).numpy(), cropped_size.numpy() locations, scores = np_batch_locs[0, :-1, :], np.expand_dims( np_batch_scos[0, :-1], -1) scale_h, scale_w = cropped_size[0] * 1. / inputs.size( -2), cropped_size[1] * 1. / inputs.size(-1) locations[:, 0], locations[:, 1] = locations[:, 0] * scale_w + cropped_size[ 2], locations[:, 1] * scale_h + cropped_size[3] prediction = np.concatenate((locations, scores), axis=1).transpose(1, 0) for i in range(param.num_pts): point = prediction[:, i] print( 'The coordinate of {:02d}/{:02d}-th points : ({:.1f}, {:.1f}), score = {:.3f}' .format(i, param.num_pts, float(point[0]), float(point[1]), float(point[2]))) if args.save_path: image = draw_image_by_points(args.image, prediction, 1, (255, 0, 0), False, False) image.save(args.save_path) print('save image with landmarks into {:}'.format(args.save_path)) print('finish san evaluation on a single image : {:}'.format(args.image))
def evaluate(args): #assert torch.cuda.is_available(), 'CUDA is not available.' torch.backends.cudnn.enabled = False torch.backends.cudnn.benchmark = False print('The image is {:}'.format(args.image)) print('The model is {:}'.format(args.model)) snapshot = Path(args.model) assert snapshot.exists(), 'The model path {:} does not exist' print('The face bounding box is {:}'.format(args.face)) assert len(args.face) == 4, 'Invalid face input : {:}'.format(args.face) snapshot = torch.load(snapshot, map_location='cpu') mean_fill = tuple([int(x * 255) for x in [0.5, 0.5, 0.5]]) normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) param = snapshot['args'] eval_transform = transforms.Compose([ transforms.PreCrop(param.pre_crop_expand), transforms.TrainScale2WH((param.crop_width, param.crop_height)), transforms.ToTensor(), normalize ]) net = models.__dict__[param.arch](param.modelconfig, None) #net = net.cuda() weights = models.remove_module_dict(snapshot['state_dict']) net.load_state_dict(weights) dataset = datasets.GeneralDataset(eval_transform, param.sigma, param.downsample, param.heatmap_type, param.dataset_name) dataset.reset(param.num_pts) print('[{:}] prepare the input data'.format(time_string())) [image, _, _, _, _, _, cropped_size], meta = dataset.prepare_input(args.image, args.face) #inputs = image.unsqueeze(0).cuda() inputs = image.unsqueeze(0) print('[{:}] prepare the input data done'.format(time_string())) # network forward with torch.no_grad(): batch_heatmaps, batch_locs, batch_scos, _ = net(inputs) print('[{:}] the network forward done'.format(time_string())) # obtain the locations on the image in the orignial size cpu = torch.device('cpu') np_batch_locs, np_batch_scos, cropped_size = batch_locs.to( cpu).numpy(), batch_scos.to(cpu).numpy(), cropped_size.numpy() locations, scores = np_batch_locs[0, :-1, :], np.expand_dims( np_batch_scos[0, :-1], -1) scale_h, scale_w = cropped_size[0] * 1. / inputs.size( -2), cropped_size[1] * 1. / inputs.size(-1) locations[:, 0], locations[:, 1] = locations[:, 0] * scale_w + cropped_size[ 2], locations[:, 1] * scale_h + cropped_size[3] prediction = np.concatenate((locations, scores), axis=1).transpose(1, 0) for i in range(param.num_pts): point = prediction[:, i] print('{:02d}/{:02d} : ({:.1f}, {:.1f}), score = {:.3f}'.format( i, param.num_pts, float(point[0]), float(point[1]), float(point[2])))
def main(): # Init logger if not os.path.isdir(args.save_path): os.makedirs(args.save_path) log = open( os.path.join( args.save_path, 'cluster_seed_{}_{}.txt'.format(args.manualSeed, time_for_file())), 'w') print_log('save path : {}'.format(args.save_path), log) print_log('------------ Options -------------', log) for k, v in sorted(vars(args).items()): print_log('Parameter : {:20} = {:}'.format(k, v), log) print_log('-------------- End ----------------', log) print_log("Random Seed: {}".format(args.manualSeed), log) print_log("python version : {}".format(sys.version.replace('\n', ' ')), log) print_log("Pillow version : {}".format(PIL.__version__), log) print_log("torch version : {}".format(torch.__version__), log) print_log("cudnn version : {}".format(torch.backends.cudnn.version()), log) # General Data Argumentation mean_fill = tuple([int(x * 255) for x in [0.485, 0.456, 0.406]]) normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform = transforms.Compose([ transforms.PreCrop(args.pre_crop_expand), transforms.TrainScale2WH((args.crop_width, args.crop_height)), transforms.ToTensor(), normalize ]) args.downsample = 8 # By default args.sigma = args.sigma * args.scale_eval data = datasets.GeneralDataset(transform, args.sigma, args.downsample, args.heatmap_type, args.dataset_name) data.load_list(args.train_list, args.num_pts, True) loader = torch.utils.data.DataLoader(data, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) # Load all lists all_lines = {} for file_path in args.train_list: listfile = open(file_path, 'r') listdata = listfile.read().splitlines() listfile.close() for line in listdata: temp = line.split(' ') assert len(temp) == 6 or len( temp) == 7, 'This line has the wrong format : {}'.format(line) image_path = temp[0] all_lines[image_path] = line assert args.n_clusters >= 2, 'The cluster number must be greater than 2' resnet = models.resnet152(True).cuda() all_features = [] for i, (inputs, target, mask, points, image_index, label_sign, ori_size) in enumerate(loader): input_vars = torch.autograd.Variable(inputs.cuda(), volatile=True) features, classifications = resnet(input_vars) features = features.cpu().data.numpy() all_features.append(features) if i % args.print_freq == 0: print_log( '{} {}/{} extract features'.format(time_string(), i, len(loader)), log) all_features = np.concatenate(all_features, axis=0) kmeans_result = KMeans(n_clusters=args.n_clusters, n_jobs=args.workers).fit(all_features) print_log('kmeans [{}] calculate done'.format(args.n_clusters), log) labels = kmeans_result.labels_.copy() cluster_idx = [] for iL in range(args.n_clusters): indexes = np.where(labels == iL)[0] cluster_idx.append(len(indexes)) cluster_idx = np.argsort(cluster_idx) for iL in range(args.n_clusters): ilabel = cluster_idx[iL] indexes = np.where(labels == ilabel) if isinstance(indexes, tuple) or isinstance(indexes, list): indexes = indexes[0] cluster_features = all_features[indexes, :].copy() filtered_index = filter_cluster(indexes.copy(), cluster_features, 0.8) print_log( '{:} [{:2d} / {:2d}] has {:4d} / {:4d} -> {:4d} = {:.2f} images '. format(time_string(), iL, args.n_clusters, indexes.size, len(data), len(filtered_index), indexes.size * 1. / len(data)), log) indexes = filtered_index.copy() save_dir = osp.join( args.save_path, 'cluster-{:02d}-{:02d}'.format(iL, args.n_clusters)) save_path = save_dir + '.lst' #if not osp.isdir(save_path): os.makedirs( save_path ) print_log('save into {}'.format(save_path), log) txtfile = open(save_path, 'w') for idx in indexes: image_path = data.datas[idx] assert image_path in all_lines, 'Not find {}'.format(image_path) txtfile.write('{}\n'.format(all_lines[image_path])) #basename = osp.basename( image_path ) #os.system( 'cp {} {}'.format(image_path, save_dir) ) txtfile.close()
def evaluate(image, model, face, save_path, cpu): org_image = image if not cpu: assert torch.cuda.is_available(), 'CUDA is not available.' torch.backends.cudnn.enabled = True torch.backends.cudnn.benchmark = True print('The image is {:}'.format(image)) print('The model is {:}'.format(model)) snapshot = model assert os.path.exists(snapshot), 'The model path {:} does not exist' print('The face bounding box is {:}'.format(face)) assert len(face) == 4, 'Invalid face input : {:}'.format(face) if cpu: snapshot = torch.load(snapshot, map_location='cpu') else: snapshot = torch.load(snapshot) mean_fill = tuple([int(x * 255) for x in [0.5, 0.5, 0.5]]) normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) param = snapshot['args'] eval_transform = transforms.Compose([ transforms.PreCrop(param.pre_crop_expand), transforms.TrainScale2WH((param.crop_width, param.crop_height)), transforms.ToTensor(), normalize ]) net = models.__dict__[param.arch](param.modelconfig, None) if not cpu: net = net.cuda() weights = models.remove_module_dict(snapshot['state_dict']) net.load_state_dict(weights) dataset = datasets.GeneralDataset(eval_transform, param.sigma, param.downsample, param.heatmap_type, param.dataset_name) dataset.reset(param.num_pts) print('[{:}] prepare the input data'.format(time_string())) [image, _, _, _, _, _, cropped_size], meta = dataset.prepare_input(image, face) print('[{:}] prepare the input data done'.format(time_string())) print('Net : \n{:}'.format(net)) # network forward with torch.no_grad(): if cpu: inputs = image.unsqueeze(0) else: inputs = image.unsqueeze(0).cuda() batch_heatmaps, batch_locs, batch_scos, _ = net(inputs) #print ('input-shape : {:}'.format(inputs.shape)) flops, params = get_model_infos(net, inputs.shape, None) print('\nIN-shape : {:}, FLOPs : {:} MB, Params : {:}.'.format( list(inputs.shape), flops, params)) flops, params = get_model_infos(net, None, inputs) print('\nIN-shape : {:}, FLOPs : {:} MB, Params : {:}.'.format( list(inputs.shape), flops, params)) print('[{:}] the network forward done'.format(time_string())) # obtain the locations on the image in the orignial size cpu = torch.device('cpu') np_batch_locs, np_batch_scos, cropped_size = batch_locs.to( cpu).numpy(), batch_scos.to(cpu).numpy(), cropped_size.numpy() locations, scores = np_batch_locs[0, :-1, :], np.expand_dims( np_batch_scos[0, :-1], -1) scale_h, scale_w = cropped_size[0] * 1. / inputs.size( -2), cropped_size[1] * 1. / inputs.size(-1) locations[:, 0], locations[:, 1] = locations[:, 0] * scale_w + cropped_size[ 2], locations[:, 1] * scale_h + cropped_size[3] prediction = np.concatenate((locations, scores), axis=1).transpose(1, 0) for i in range(param.num_pts): point = prediction[:, i] print( 'The coordinate of {:02d}/{:02d}-th points : ({:.1f}, {:.1f}), score = {:.3f}' .format(i, param.num_pts, float(point[0]), float(point[1]), float(point[2]))) if save_path: image = draw_image_by_points(org_image, prediction, 1, (255, 0, 0), False, False) image.save(save_path) print('save image with landmarks into {:}'.format(save_path)) print('finish san evaluation on a single image : {:}'.format(image)) return locations
def main(): # Init logger if not os.path.isdir(args.save_path): os.makedirs(args.save_path) log = open(os.path.join(args.save_path, 'cluster_seed_{}_{}.txt'.format(args.manualSeed, time_for_file())), 'w') print_log('save path : {}'.format(args.save_path), log) print_log('------------ Options -------------', log) for k, v in sorted(vars(args).items()): print_log('Parameter : {:20} = {:}'.format(k, v), log) print_log('-------------- End ----------------', log) print_log("Random Seed: {}".format(args.manualSeed), log) print_log("python version : {}".format(sys.version.replace('\n', ' ')), log) print_log("Pillow version : {}".format(PIL.__version__), log) print_log("torch version : {}".format(torch.__version__), log) print_log("cudnn version : {}".format(torch.backends.cudnn.version()), log) # finetune resnet-152 to train style-discriminative features resnet = models.resnet152(True, num_classes=4) resnet = torch.nn.DataParallel(resnet) # define loss function (criterion) and optimizer criterion = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(resnet.parameters(), args.learning_rate, momentum=args.momentum, weight_decay=args.decay) cls_train_dir = args.style_train_root cls_eval_dir = args.style_eval_root assert osp.isdir(cls_train_dir), 'Does not know : {}'.format(cls_train_dir) # train data loader vision_normalize = visiontransforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) print_log('Training dir : {}'.format(cls_train_dir), log) print_log('Evaluate dir : {}'.format(cls_eval_dir), log) cls_train_dataset = visiondatasets.ImageFolder( cls_train_dir, visiontransforms.Compose([ visiontransforms.RandomResizedCrop(224), visiontransforms.RandomHorizontalFlip(), visiontransforms.ToTensor(), vision_normalize, ])) cls_train_loader = torch.utils.data.DataLoader( cls_train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True) if cls_eval_dir is not None: assert osp.isdir(cls_eval_dir), 'Does not know : {}'.format(cls_eval_dir) val_loader = torch.utils.data.DataLoader( visiondatasets.ImageFolder(cls_eval_dir, visiontransforms.Compose([ visiontransforms.Resize(256), visiontransforms.CenterCrop(224), visiontransforms.ToTensor(), vision_normalize, ])), batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) else: val_loader = None for epoch in range(args.epochs): learning_rate = adjust_learning_rate(optimizer, epoch, args) print_log('epoch : [{}/{}] lr={}'.format(epoch, args.epochs, learning_rate), log) top1, losses = AverageMeter(), AverageMeter() resnet.train() for i, (inputs, target) in enumerate(cls_train_loader): #target = target.cuda(async=True) # compute output _, output = resnet(inputs) loss = criterion(output, target) # measure accuracy and record loss prec1 = accuracy(output.data, target, topk=[1]) top1.update(prec1.item(), inputs.size(0)) losses.update(loss.item(), inputs.size(0)) # compute gradient and do SGD step optimizer.zero_grad() loss.backward() optimizer.step() if i % args.print_freq == 0 or i+1 == len(cls_train_loader): print_log(' [Train={:03d}] [{:}] [{:3d}/{:3d}] accuracy : {:.1f}, loss : {:.4f}, input:{:}, output:{:}'.format(epoch, time_string(), i, len(cls_train_loader), top1.avg, losses.avg, inputs.size(), output.size()), log) if val_loader is None: continue # evaluate the model resnet.eval() top1, losses = AverageMeter(), AverageMeter() for i, (inputs, target) in enumerate(val_loader): #target = target.cuda(async=True) # compute output with torch.no_grad(): _, output = resnet(inputs) loss = criterion(output, target) # measure accuracy and record loss prec1 = accuracy(output.data, target, topk=[1]) top1.update(prec1.item(), inputs.size(0)) losses.update(loss.item(), inputs.size(0)) if i % args.print_freq_eval == 0 or i+1 == len(val_loader): print_log(' [Evalu={:03d}] [{:}] [{:3d}/{:3d}] accuracy : {:.1f}, loss : {:.4f}, input:{:}, output:{:}'.format(epoch, time_string(), i, len(val_loader), top1.avg, losses.avg, inputs.size(), output.size()), log) # extract features resnet.eval() # General Data Argumentation mean_fill = tuple( [int(x*255) for x in [0.485, 0.456, 0.406] ] ) normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform = transforms.Compose([transforms.PreCrop(args.pre_crop_expand), transforms.TrainScale2WH((args.crop_width, args.crop_height)), transforms.ToTensor(), normalize]) args.downsample = 8 # By default args.sigma = args.sigma * args.scale_eval data = datasets.GeneralDataset(transform, args.sigma, args.downsample, args.heatmap_type, args.dataset_name) data.load_list(args.train_list, args.num_pts, True) loader = torch.utils.data.DataLoader(data, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) # Load all lists all_lines = {} for file_path in args.train_list: listfile = open(file_path, 'r') listdata = listfile.read().splitlines() listfile.close() for line in listdata: temp = line.split(' ') assert len(temp) == 6 or len(temp) == 7, 'This line has the wrong format : {}'.format(line) image_path = temp[0] all_lines[ image_path ] = line assert args.n_clusters >= 2, 'The cluster number must be greater than 2' all_features = [] for i, (inputs, target, mask, points, image_index, label_sign, ori_size) in enumerate(loader): with torch.no_grad(): features, _ = resnet(inputs) features = features.cpu().numpy() all_features.append( features ) if i % args.print_freq == 0: print_log('{} {}/{} extract features'.format(time_string(), i, len(loader)), log) all_features = np.concatenate(all_features, axis=0) kmeans_result = KMeans(n_clusters=args.n_clusters, n_jobs=args.workers).fit( all_features ) print_log('kmeans [{}] calculate done'.format(args.n_clusters), log) labels = kmeans_result.labels_.copy() cluster_idx = [] for iL in range(args.n_clusters): indexes = np.where( labels == iL )[0] cluster_idx.append( len(indexes) ) cluster_idx = np.argsort(cluster_idx) for iL in range(args.n_clusters): ilabel = cluster_idx[iL] indexes = np.where( labels == ilabel ) if isinstance(indexes, tuple) or isinstance(indexes, list): indexes = indexes[0] #cluster_features = all_features[indexes,:].copy() #filtered_index = filter_cluster(indexes.copy(), cluster_features, 0.8) filtered_index = indexes.copy() print_log('{:} [{:2d} / {:2d}] has {:4d} / {:4d} -> {:4d} = {:.2f} images'.format(time_string(), iL, args.n_clusters, indexes.size, len(data), len(filtered_index), indexes.size*1./len(data)), log) indexes = filtered_index.copy() save_dir = osp.join(args.save_path, 'cluster-{:02d}-{:02d}'.format(iL, args.n_clusters)) save_path = save_dir + '.lst' #if not osp.isdir(save_path): os.makedirs( save_path ) print_log('save into {}'.format(save_path), log) txtfile = open( save_path , 'w') for idx in indexes: image_path = data.datas[idx] assert image_path in all_lines, 'Not find {}'.format(image_path) txtfile.write('{}\n'.format(all_lines[image_path])) #basename = osp.basename( image_path ) #os.system( 'cp {} {}'.format(image_path, save_dir) ) txtfile.close()