def visual(clist, cdir, num_pts):
if not cdir.exists(): os.makedirs(str(cdir))
shape = 256
transform = transforms.Compose(
[transforms.PreCrop(0.2),
transforms.TrainScale2WH((shape, shape))])
data = datasets.GeneralDataset(transform, 2, 1, 'gaussian', 'test')
data.load_list(clist, num_pts, True)
for i, tempx in enumerate(data):
image = tempx[0]
heats = models.variable2np(tempx[1]).transpose(1, 2, 0)
xheat = generate_color_from_heatmaps(heats, index=-1)
xheat = PIL.Image.fromarray(np.uint8(xheat * 255))
cimage = overlap_two_pil_image(image, xheat)
basename = osp.basename(data.datas[i]).split('.')[0]
basename = str(cdir) + '/' + basename + '-{:}.jpg'
image.save(basename.format('ori'))
xheat.save(basename.format('heat'))
cimage.save(basename.format('over'))
if i % PRINT_GAP == 0:
print('--->>> process the {:4d}/{:4d}-th image'.format(
i, len(data)))
def calculate_mean(list_file, num_pts, save_path):
#style = 'Original'
#save_dir = 'cache/{}'.format(style)
save_dir = osp.dirname(save_path)
print('crop face images into {} <-> {}'.format(save_dir, save_path))
if not osp.isdir(save_dir): os.makedirs(save_dir)
transform = transforms.Compose(
[transforms.PreCrop(0.2),
transforms.TrainScale2WH((256, 256))])
data = datasets.GeneralDataset(transform, 1, 8, 'gaussian', 'test')
data.load_list(list_file, num_pts, True)
ok_faces, ok_basenames, ok_points = [], [], []
for i, tempx in enumerate(data):
image, mask, points = tempx[0], tempx[2], tempx[3]
#points = points[[0, 8, 16, 36, 39, 42, 45, 33, 48, 54, 27, 57],:]
basename = osp.basename(data.datas[i])
if torch.sum(mask) == num_pts + 1:
ok_faces.append(image)
ok_basenames.append(basename)
ok_points.append(points.numpy())
print('extract done {:} -> {:}'.format(len(data), len(ok_faces)))
mean_landmark = np.array(ok_points).mean(axis=0)
all_faces = []
save_dir = save_dir + '-all'
if not osp.isdir(save_dir): os.makedirs(save_dir)
for face, point, basename in zip(ok_faces, ok_points, ok_basenames):
aligned_face = face_align(face, point, mean_landmark)
aligned_face.save(osp.join(save_dir, basename))
all_faces.append(np.array(aligned_face))
all_faces = np.array(all_faces).mean(axis=0)
mean_face = Image.fromarray(np.uint8(all_faces))
mean_face.save(save_path)
def evaluate(args):
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'
print ('The face bounding box is {:}'.format(args.face))
assert len(args.face) == 4, 'Invalid face input : {:}'.format(args.face)
snapshot = torch.load(snapshot)
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])
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()
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 crop_style(list_file, num_pts, save_dir):
#style = 'Original'
#save_dir = 'cache/{}'.format(style)
print ('crop face images into {}'.format(save_dir))
if not osp.isdir(save_dir): os.makedirs(save_dir)
transform = transforms.Compose([transforms.PreCrop(0.2), transforms.TrainScale2WH((256, 256))])
data = datasets.GeneralDataset(transform, 1, 8, 'gaussian', 'test')
data.load_list(list_file, num_pts, True)
#loader = torch.utils.data.DataLoader(data, batch_size=1, shuffle=False, num_workers=12, pin_memory=True)
for i, tempx in enumerate(data):
image = tempx[0]
#points = tempx[3]
basename = osp.basename(data.datas[i])
save_name = osp.join(save_dir, basename)
image.save(save_name)
if i % PRINT_GAP == 0:
print ('--->>> process the {:4d}/{:4d}-th image'.format(i, len(data)))
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 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()
