def __call__(self):
# the 30th layer of features is relu of conv5_3
model = vgg16(pretrained=False)
if self.configer.get('network', 'pretrained') is not None:
Log.info('Loading pretrained model: {}'.format(
self.configer.get('network', 'pretrained')))
model.load_state_dict(
torch.load(self.configer.get('network', 'pretrained')))
features = list(model.features)[:30]
classifier = model.classifier
classifier = list(classifier)
del classifier[6]
if not self.configer.get('network', 'use_drop'):
del classifier[5]
del classifier[2]
classifier = nn.Sequential(*classifier)
# freeze top4 conv
for layer in features[:10]:
for p in layer.parameters():
p.requires_grad = False
return nn.Sequential(*features), classifier
def train(runner):
Log.info('Training start...')
if runner.configer.get('network',
'resume') is not None and runner.configer.get(
'network', 'resume_val'):
runner.val()
if runner.configer.get('solver', 'lr')['metric'] == 'epoch':
while runner.runner_state['epoch'] < runner.configer.get(
'solver', 'max_epoch'):
if runner.configer.get('network.distributed'):
runner.train_loader.sampler.set_epoch(
runner.runner_state['epoch'])
runner.train()
if runner.runner_state['epoch'] == runner.configer.get(
'solver', 'max_epoch'):
runner.val()
break
else:
while runner.runner_state['iters'] < runner.configer.get(
'solver', 'max_iters'):
if runner.configer.get('network.distributed'):
runner.train_loader.sampler.set_epoch(
runner.runner_state['epoch'])
runner.train()
if runner.runner_state['iters'] == runner.configer.get(
'solver', 'max_iters'):
runner.val()
break
Log.info('Training end...')
def train(self):
"""
Train function of every epoch during train phase.
"""
self.det_net.train()
start_time = time.time()
# Adjust the learning rate after every epoch.
self.runner_state['epoch'] += 1
# data_tuple: (inputs, heatmap, maskmap, vecmap)
for i, data_dict in enumerate(self.train_loader):
Trainer.update(self,
warm_list=(0, ),
warm_lr_list=(self.configer.get('solver',
'lr')['base_lr'], ),
solver_dict=self.configer.get('solver'))
self.data_time.update(time.time() - start_time)
# Forward pass.
out_dict = self.det_net(data_dict)
# Compute the loss of the train batch & backward.
loss = out_dict['loss'].mean()
self.train_losses.update(loss.item(),
len(DCHelper.tolist(data_dict['meta'])))
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
# Update the vars of the train phase.
self.batch_time.update(time.time() - start_time)
start_time = time.time()
self.runner_state['iters'] += 1
# Print the log info & reset the states.
if self.runner_state['iters'] % self.configer.get(
'solver', 'display_iter') == 0:
Log.info(
'Train Epoch: {0}\tTrain Iteration: {1}\t'
'Time {batch_time.sum:.3f}s / {2}iters, ({batch_time.avg:.3f})\t'
'Data load {data_time.sum:.3f}s / {2}iters, ({data_time.avg:3f})\n'
'Learning rate = {3}\tLoss = {loss.val:.8f} (ave = {loss.avg:.8f})\n'
.format(self.runner_state['epoch'],
self.runner_state['iters'],
self.configer.get('solver', 'display_iter'),
RunnerHelper.get_lr(self.optimizer),
batch_time=self.batch_time,
data_time=self.data_time,
loss=self.train_losses))
self.batch_time.reset()
self.data_time.reset()
self.train_losses.reset()
if self.configer.get('solver', 'lr')['metric'] == 'iters' \
and self.runner_state['iters'] == self.configer.get('solver', 'max_iters'):
break
# Check to val the current model.
if self.runner_state['iters'] % self.configer.get(
'solver', 'test_interval') == 0:
self.val()
def val(self):
"""
Validation function during the train phase.
"""
self.pose_net.eval()
start_time = time.time()
with torch.no_grad():
for i, data_dict in enumerate(self.val_loader):
# Forward pass.
out = self.pose_net(data_dict)
# Compute the loss of the val batch.
loss_dict = self.pose_loss(out)
self.val_losses.update({key: loss.item() for key, loss in loss_dict.items()}, data_dict['img'].size(0))
# Update the vars of the val phase.
self.batch_time.update(time.time() - start_time)
start_time = time.time()
self.runner_state['val_loss'] = self.val_losses.avg['loss']
RunnerHelper.save_net(self, self.pose_net, val_loss=self.val_losses.avg['loss'])
# Print the log info & reset the states.
Log.info(
'Test Time {batch_time.sum:.3f}s, ({batch_time.avg:.3f})\t'
'Loss {0}\n'.format(self.val_losses.info(), batch_time=self.batch_time))
self.batch_time.reset()
self.val_losses.reset()
self.pose_net.train()
def parse_dir_pose(self, image_dir, json_dir, mask_dir=None):
if image_dir is None or not os.path.exists(image_dir):
Log.error('Image Dir: {} not existed.'.format(image_dir))
return
if json_dir is None or not os.path.exists(json_dir):
Log.error('Json Dir: {} not existed.'.format(json_dir))
return
for image_file in os.listdir(image_dir):
shotname, extension = os.path.splitext(image_file)
Log.info(image_file)
image_canvas = cv2.imread(os.path.join(
image_dir, image_file)) # B, G, R order.
with open(os.path.join(json_dir, '{}.json'.format(shotname)),
'r') as json_stream:
info_tree = json.load(json_stream)
image_canvas = self.draw_points(image_canvas, info_tree)
if self.configer.exists('details', 'limb_seq'):
image_canvas = self.link_points(image_canvas, info_tree)
if mask_dir is not None:
mask_file = os.path.join(mask_dir,
'{}_vis.png'.format(shotname))
mask_canvas = cv2.imread(mask_file)
image_canvas = cv2.addWeighted(image_canvas, 0.6, mask_canvas,
0.4, 0)
cv2.imshow('main', image_canvas)
cv2.waitKey()
def __test_img(self, image_path, save_path):
Log.info('Image Path: {}'.format(image_path))
ori_image = ImageHelper.read_image(
image_path,
tool=self.configer.get('data', 'image_tool'),
mode=self.configer.get('data', 'input_mode'))
ori_width, ori_height = ImageHelper.get_size(ori_image)
ori_img_bgr = ImageHelper.get_cv2_bgr(ori_image,
mode=self.configer.get(
'data', 'input_mode'))
heatmap_avg = np.zeros(
(ori_height, ori_width, self.configer.get('network',
'heatmap_out')))
for i, scale in enumerate(self.configer.get('test', 'scale_search')):
image = self.blob_helper.make_input(ori_image,
input_size=self.configer.get(
'test', 'input_size'),
scale=scale)
with torch.no_grad():
heatmap_out_list = self.pose_net(image)
heatmap_out = heatmap_out_list[-1]
# extract outputs, resize, and remove padding
heatmap = heatmap_out.squeeze(0).cpu().numpy().transpose(
1, 2, 0)
heatmap = cv2.resize(heatmap, (ori_width, ori_height),
interpolation=cv2.INTER_CUBIC)
heatmap_avg = heatmap_avg + heatmap / len(
self.configer.get('test', 'scale_search'))
all_peaks = self.__extract_heatmap_info(heatmap_avg)
image_canvas = self.__draw_key_point(all_peaks, ori_img_bgr)
ImageHelper.save(image_canvas, save_path)
def relabel(self, json_dir, method='ssd'):
submission_file = os.path.join(json_dir, 'person_instances_val2017_{}_results.json'.format(method))
img_id_list = list()
object_list = list()
for json_file in os.listdir(json_dir):
json_path = os.path.join(json_dir, json_file)
shotname, extensions = os.path.splitext(json_file)
shotname = shotname.rstrip().split('_')[-1]
try:
img_id = int(shotname)
except ValueError:
Log.info('Invalid Json file: {}'.format(json_file))
continue
img_id_list.append(img_id)
with open(json_path, 'r') as json_stream:
info_tree = json.load(json_stream)
for object in info_tree['objects']:
object_dict = dict()
object_dict['image_id'] = img_id
object_dict['category_id'] = int(self.configer.get('details', 'coco_cat_seq')[object['label']])
object_dict['score'] = object['score']
object_dict['bbox'] = [object['bbox'][0], object['bbox'][1],
object['bbox'][2] - object['bbox'][0],
object['bbox'][3] - object['bbox'][1]]
object_list.append(object_dict)
with open(submission_file, 'w') as write_stream:
write_stream.write(json.dumps(object_list))
Log.info('Evaluate {} images...'.format(len(img_id_list)))
return submission_file, img_id_list
def load_net(runner, net, model_path=None, map_location='cpu'):
if model_path is not None or runner.configer.get('network', 'resume') is not None:
resume_path = runner.configer.get('network', 'resume')
resume_path = model_path if model_path is not None else resume_path
Log.info('Resuming from {}'.format(resume_path))
resume_dict = torch.load(resume_path, map_location=map_location)
if 'state_dict' in resume_dict:
checkpoint_dict = resume_dict['state_dict']
elif 'model' in resume_dict:
checkpoint_dict = resume_dict['model']
elif isinstance(resume_dict, OrderedDict):
checkpoint_dict = resume_dict
else:
raise RuntimeError(
'No state_dict found in checkpoint file {}'.format(runner.configer.get('network', 'resume')))
# load state_dict
if hasattr(net, 'module'):
RunnerHelper.load_state_dict(net.module, checkpoint_dict,
runner.configer.get('network', 'resume_strict'))
else:
RunnerHelper.load_state_dict(net, checkpoint_dict, runner.configer.get('network', 'resume_strict'))
if runner.configer.get('network', 'resume_continue'):
# runner.configer.resume(resume_dict['config_dict'])
runner.runner_state = resume_dict['runner_state']
net = RunnerHelper._make_parallel(runner, net)
return net
def val(self):
"""
Validation function during the train phase.
"""
self.gan_net.eval()
start_time = time.time()
for j, data_dict in enumerate(self.val_loader):
with torch.no_grad():
# Forward pass.
out_dict = self.gan_net(data_dict)
# Compute the loss of the val batch.
self.val_losses.update(
out_dict['loss_G'].mean().item() +
out_dict['loss_D'].mean().item(),
len(DCHelper.tolist(data_dict['meta'])))
# Update the vars of the val phase.
self.batch_time.update(time.time() - start_time)
start_time = time.time()
RunnerHelper.save_net(self, self.gan_net, val_loss=self.val_losses.avg)
# Print the log info & reset the states.
Log.info('Test Time {batch_time.sum:.3f}s, ({batch_time.avg:.3f})\t'
'Loss {loss.avg:.8f}\n'.format(batch_time=self.batch_time,
loss=self.val_losses))
self.batch_time.reset()
self.val_losses.reset()
self.gan_net.train()
def _make_parallel(runner, net):
if runner.configer.get('network.distributed', default=False):
local_rank = runner.configer.get('local_rank')
torch.cuda.set_device(local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
if runner.configer.get('network.syncbn', default=False):
Log.info('Converting syncbn model...')
net = nn.SyncBatchNorm.convert_sync_batchnorm(net)
net = nn.parallel.DistributedDataParallel(
net.cuda(),
find_unused_parameters=True,
device_ids=[local_rank],
output_device=local_rank)
# if runner.configer.get('network.syncbn', default=False):
# Log.info('Converting syncbn model...')
# from apex.parallel import convert_syncbn_model
# net = convert_syncbn_model(net)
# from apex.parallel import DistributedDataParallel
# net = DistributedDataParallel(net.cuda(), delay_allreduce=True)
return net
net = net.to(
torch.device(
'cpu' if runner.configer.get('gpu') is None else 'cuda'))
from lib.parallel.data_parallel import ParallelModel
return ParallelModel(net,
gather_=runner.configer.get('network', 'gather'))
def init_weights(net, init_type='normal', init_gain=0.02):
"""Initialize network weights.
Parameters:
net (network) -- network to be initialized
init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
init_gain (float) -- scaling factor for normal, xavier and orthogonal.
We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might
work better for some applications. Feel free to try yourself.
"""
def init_func(m): # define the initialization function
classname = m.__class__.__name__
if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
if init_type == 'normal':
init.normal_(m.weight.data, 0.0, init_gain)
elif init_type == 'xavier':
init.xavier_normal_(m.weight.data, gain=init_gain)
elif init_type == 'kaiming':
init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
elif init_type == 'orthogonal':
init.orthogonal_(m.weight.data, gain=init_gain)
else:
raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
if hasattr(m, 'bias') and m.bias is not None:
init.constant_(m.bias.data, 0.0)
elif classname.find(
'BatchNorm2d') != -1: # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
init.normal_(m.weight.data, 1.0, init_gain)
init.constant_(m.bias.data, 0.0)
Log.info('initialize network with {}'.format(init_type))
net.apply(init_func) # apply the initialization function <init_func>
def test(self, test_dir, out_dir):
for _, data_dict in enumerate(
self.test_loader.get_testloader(test_dir=test_dir)):
data_dict['testing'] = True
out_dict = self.det_net(data_dict)
meta_list = DCHelper.tolist(data_dict['meta'])
batch_detections = self.decode(out_dict['loc'], out_dict['conf'],
self.configer, meta_list)
for i in range(len(meta_list)):
ori_img_bgr = ImageHelper.read_image(meta_list[i]['img_path'],
tool='cv2',
mode='BGR')
json_dict = self.__get_info_tree(batch_detections[i])
image_canvas = self.det_parser.draw_bboxes(
ori_img_bgr.copy(),
json_dict,
conf_threshold=self.configer.get('res', 'vis_conf_thre'))
ImageHelper.save(image_canvas,
save_path=os.path.join(
out_dir, 'vis/{}.png'.format(
meta_list[i]['filename'])))
Log.info('Json Path: {}'.format(
os.path.join(
out_dir,
'json/{}.json'.format(meta_list[i]['filename']))))
JsonHelper.save_file(json_dict,
save_path=os.path.join(
out_dir, 'json/{}.json'.format(
meta_list[i]['filename'])))
def save_file(json_dict, save_path):
dir_name = os.path.dirname(save_path)
if not os.path.exists(dir_name):
Log.info('Json Dir: {} not exists.'.format(dir_name))
os.makedirs(dir_name)
with open(save_path, 'w') as write_stream:
write_stream.write(json.dumps(json_dict))
def xml2json(xml_file, json_file):
if not os.path.exists(xml_file):
Log.error('Xml file: {} not exists.'.format(xml_file))
exit(1)
json_dir_name = os.path.dirname(json_file)
if not os.path.exists(json_dir_name):
Log.info('Json Dir: {} not exists.'.format(json_dir_name))
os.makedirs(json_dir_name)
def json2xml(json_file, xml_file):
if not os.path.exists(json_file):
Log.error('Json file: {} not exists.'.format(json_file))
exit(1)
xml_dir_name = os.path.dirname(xml_file)
if not os.path.exists(xml_dir_name):
Log.info('Xml Dir: {} not exists.'.format(xml_dir_name))
os.makedirs(xml_dir_name)
def debug(runner):
Log.info('Debugging start..')
base_dir = os.path.join(runner.configer.get('project_dir'), 'out/vis',
runner.configer.get('task'), runner.configer.get('network', 'model_name'))
if not os.path.exists(base_dir):
os.makedirs(base_dir)
runner.debug(base_dir)
Log.info('Debugging end...')
def train(self):
"""
Train function of every epoch during train phase.
"""
self.seg_net.train()
start_time = time.time()
# Adjust the learning rate after every epoch.
for i, data_dict in enumerate(self.train_loader):
Trainer.update(self, warm_list=(0,), solver_dict=self.configer.get('solver'))
self.data_time.update(time.time() - start_time)
# Forward pass.
data_dict = RunnerHelper.to_device(self, data_dict)
out = self.seg_net(data_dict)
# Compute the loss of the train batch & backward.
loss_dict = self.loss(out)
loss = loss_dict['loss']
self.train_losses.update({key: loss.item() for key, loss in loss_dict.items()}, data_dict['img'].size(0))
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
# Update the vars of the train phase.
self.batch_time.update(time.time() - start_time)
start_time = time.time()
self.runner_state['iters'] += 1
# Print the log info & reset the states.
if self.runner_state['iters'] % self.configer.get('solver', 'display_iter') == 0:
Log.info('Train Epoch: {0}\tTrain Iteration: {1}\t'
'Time {batch_time.sum:.3f}s / {2}iters, ({batch_time.avg:.3f})\t'
'Data load {data_time.sum:.3f}s / {2}iters, ({data_time.avg:3f})\n'
'Learning rate = {4}\tLoss = {3}\n'.format(
self.runner_state['epoch'], self.runner_state['iters'],
self.configer.get('solver', 'display_iter'), self.train_losses.info(),
RunnerHelper.get_lr(self.optimizer), batch_time=self.batch_time, data_time=self.data_time))
self.batch_time.reset()
self.data_time.reset()
self.train_losses.reset()
if self.runner_state['iters'] % self.configer.get('solver.save_iters') == 0 \
and self.configer.get('local_rank') == 0:
RunnerHelper.save_net(self, self.seg_net)
if self.configer.get('solver', 'lr')['metric'] == 'iters' \
and self.runner_state['iters'] == self.configer.get('solver', 'max_iters'):
break
# Check to val the current model.
if self.runner_state['iters'] % self.configer.get('solver', 'test_interval') == 0 \
and not self.configer.get('network.distributed'):
self.val()
self.runner_state['epoch'] += 1
def test(self, test_dir, out_dir):
for _, data_dict in enumerate(
self.test_loader.get_testloader(test_dir=test_dir)):
total_logits = None
if self.configer.get('test', 'mode') == 'ss_test':
total_logits = self.ss_test(data_dict)
elif self.configer.get('test', 'mode') == 'sscrop_test':
total_logits = self.sscrop_test(data_dict,
params_dict=self.configer.get(
'test', 'sscrop_test'))
elif self.configer.get('test', 'mode') == 'ms_test':
total_logits = self.ms_test(data_dict,
params_dict=self.configer.get(
'test', 'ms_test'))
elif self.configer.get('test', 'mode') == 'mscrop_test':
total_logits = self.mscrop_test(data_dict,
params_dict=self.configer.get(
'test', 'mscrop_test'))
else:
Log.error('Invalid test mode:{}'.format(
self.configer.get('test', 'mode')))
exit(1)
meta_list = DCHelper.tolist(data_dict['meta'])
for i in range(len(meta_list)):
label_map = np.argmax(total_logits[i], axis=-1)
label_img = np.array(label_map, dtype=np.uint8)
ori_img_bgr = ImageHelper.read_image(meta_list[i]['img_path'],
tool='cv2',
mode='BGR')
image_canvas = self.seg_parser.colorize(
label_img, image_canvas=ori_img_bgr)
ImageHelper.save(image_canvas,
save_path=os.path.join(
out_dir, 'vis/{}.png'.format(
meta_list[i]['filename'])))
if self.configer.get('data.label_list',
default=None) is not None:
label_img = self.__relabel(label_img)
if self.configer.get('data.reduce_zero_label', default=False):
label_img = label_img + 1
label_img = label_img.astype(np.uint8)
label_img = Image.fromarray(label_img, 'P')
label_path = os.path.join(
out_dir, 'label/{}.png'.format(meta_list[i]['filename']))
Log.info('Label Path: {}'.format(label_path))
ImageHelper.save(label_img, label_path)
def _get_parameters(self):
if self.solver_dict.get('optim.wdall', default=True):
lr_1 = []
lr_2 = []
params_dict = dict(self.cls_net.named_parameters())
for key, value in params_dict.items():
if value.requires_grad:
if 'backbone' in key:
if self.configer.get('solver.lr.bb_lr_scale') == 0.0:
value.requires_grad = False
else:
lr_1.append(value)
else:
lr_2.append(value)
params = [{
'params':
lr_1,
'lr':
self.solver_dict['lr']['base_lr'] *
self.configer.get('solver.lr.bb_lr_scale')
}, {
'params': lr_2,
'lr': self.solver_dict['lr']['base_lr']
}]
else:
no_decay_list = []
decay_list = []
no_decay_name = []
decay_name = []
for m in self.cls_net.modules():
if (hasattr(m, 'groups') and m.groups > 1) or isinstance(m, torch.nn.BatchNorm2d) \
or m.__class__.__name__ == 'GL':
no_decay_list += m.parameters(recurse=False)
for name, p in m.named_parameters(recurse=False):
no_decay_name.append(m.__class__.__name__ + name)
else:
for name, p in m.named_parameters(recurse=False):
if 'bias' in name:
no_decay_list.append(p)
no_decay_name.append(m.__class__.__name__ + name)
else:
decay_list.append(p)
decay_name.append(m.__class__.__name__ + name)
Log.info('no decay list = {}'.format(no_decay_name))
Log.info('decay list = {}'.format(decay_name))
params = [{
'params': no_decay_list,
'weight_decay': 0
}, {
'params': decay_list
}]
return params
def __read_json(self, root_dir, json_path):
item_list = []
for item in JsonHelper.load_file(json_path):
img_path = os.path.join(root_dir, item['image_path'])
if not os.path.exists(img_path) or not ImageHelper.is_img(img_path):
Log.error('Image Path: {} is Invalid.'.format(img_path))
exit(1)
item_list.append((img_path, '.'.join(item['image_path'].split('.')[:-1])))
Log.info('There are {} images..'.format(len(item_list)))
return item_list
def debug(self, vis_dir):
for i, data_dict in enumerate(self.pose_data_loader.get_trainloader()):
inputs = data_dict['img']
maskmap = data_dict['maskmap']
heatmap = data_dict['heatmap']
vecmap = data_dict['vecmap']
for j in range(inputs.size(0)):
count = count + 1
if count > 10:
exit(1)
Log.info(heatmap.size())
image_bgr = self.blob_helper.tensor2bgr(inputs[j])
mask_canvas = maskmap[j].repeat(3, 1,
1).numpy().transpose(1, 2, 0)
mask_canvas = (mask_canvas * 255).astype(np.uint8)
mask_canvas = cv2.resize(
mask_canvas, (0, 0),
fx=self.configer.get('network', 'stride'),
fy=self.configer.get('network', 'stride'),
interpolation=cv2.INTER_CUBIC)
image_bgr = cv2.addWeighted(image_bgr, 0.6, mask_canvas, 0.4,
0)
heatmap_avg = heatmap[j].numpy().transpose(1, 2, 0)
heatmap_avg = cv2.resize(
heatmap_avg, (0, 0),
fx=self.configer.get('network', 'stride'),
fy=self.configer.get('network', 'stride'),
interpolation=cv2.INTER_CUBIC)
paf_avg = vecmap[j].numpy().transpose(1, 2, 0)
paf_avg = cv2.resize(paf_avg, (0, 0),
fx=self.configer.get('network', 'stride'),
fy=self.configer.get('network', 'stride'),
interpolation=cv2.INTER_CUBIC)
self.pose_visualizer.vis_peaks(heatmap_avg, image_bgr)
self.pose_visualizer.vis_paf(paf_avg, image_bgr)
all_peaks = self.__extract_heatmap_info(heatmap_avg)
special_k, connection_all = self.__extract_paf_info(
image_bgr, paf_avg, all_peaks)
subset, candidate = self.__get_subsets(connection_all,
special_k, all_peaks)
json_dict = self.__get_info_tree(image_bgr, subset, candidate)
image_canvas = self.pose_parser.draw_points(
image_bgr, json_dict)
image_canvas = self.pose_parser.link_points(
image_canvas, json_dict)
cv2.imwrite(
os.path.join(vis_dir, '{}_{}_vis.png'.format(i, j)),
image_canvas)
cv2.imshow('main', image_canvas)
cv2.waitKey()
def load_url(url, map_location=None):
model_dir = os.path.join('~', '.TorchCV', 'model')
if not os.path.exists(model_dir):
os.makedirs(model_dir)
filename = url.split('/')[-1]
cached_file = os.path.join(model_dir, filename)
if not os.path.exists(cached_file):
Log.info('Downloading: "{}" to {}\n'.format(url, cached_file))
urlretrieve(url, cached_file)
Log.info('Loading pretrained model:{}'.format(cached_file))
return torch.load(cached_file, map_location=map_location)
def __read_list(self, root_dir, list_path):
item_list = []
with open(list_path, 'r') as f:
for line in f.readlines()[0:]:
filename = line.strip().split()[0]
img_path = os.path.join(root_dir, filename)
if not os.path.exists(img_path) or not ImageHelper.is_img(img_path):
Log.error('Image Path: {} is Invalid.'.format(img_path))
exit(1)
item_list.append((img_path, '.'.join(filename.split('.')[:-1])))
Log.info('There are {} images..'.format(len(item_list)))
return item_list
def val(self):
"""
Validation function during the train phase.
"""
self.det_net.eval()
start_time = time.time()
with torch.no_grad():
for j, data_dict in enumerate(self.val_loader):
# Forward pass.
data_dict = RunnerHelper.to_device(self, data_dict)
out = self.det_net(data_dict)
loss_dict = self.det_loss(out)
# Compute the loss of the train batch & backward.
loss = loss_dict['loss'].mean()
out_dict, _ = RunnerHelper.gather(self, out)
self.val_losses.update(loss.item(),
len(DCHelper.tolist(data_dict['meta'])))
test_indices_and_rois, test_roi_locs, test_roi_scores, test_rois_num = out_dict[
'test_group']
batch_detections = FastRCNNTest.decode(
test_roi_locs, test_roi_scores, test_indices_and_rois,
test_rois_num, self.configer,
DCHelper.tolist(data_dict['meta']))
batch_pred_bboxes = self.__get_object_list(batch_detections)
self.det_running_score.update(batch_pred_bboxes, [
item['ori_bboxes']
for item in DCHelper.tolist(data_dict['meta'])
], [
item['ori_labels']
for item in DCHelper.tolist(data_dict['meta'])
])
# Update the vars of the val phase.
self.batch_time.update(time.time() - start_time)
start_time = time.time()
RunnerHelper.save_net(self,
self.det_net,
iters=self.runner_state['iters'])
# Print the log info & reset the states.
Log.info(
'Test Time {batch_time.sum:.3f}s, ({batch_time.avg:.3f})\t'
'Loss {loss.avg:.8f}\n'.format(batch_time=self.batch_time,
loss=self.val_losses))
Log.info('Val mAP: {}\n'.format(self.det_running_score.get_mAP()))
self.det_running_score.reset()
self.batch_time.reset()
self.val_losses.reset()
self.det_net.train()
def __test_img(self, image_path, json_path, raw_path, vis_path):
Log.info('Image Path: {}'.format(image_path))
ori_image = ImageHelper.read_image(image_path,
tool=self.configer.get('data', 'image_tool'),
mode=self.configer.get('data', 'input_mode'))
ori_width, ori_height = ImageHelper.get_size(ori_image)
ori_img_bgr = ImageHelper.get_cv2_bgr(ori_image, mode=self.configer.get('data', 'input_mode'))
heatmap_avg = np.zeros((ori_height, ori_width, self.configer.get('network', 'heatmap_out')))
paf_avg = np.zeros((ori_height, ori_width, self.configer.get('network', 'paf_out')))
multiplier = [scale * self.configer.get('test', 'input_size')[1] / ori_height
for scale in self.configer.get('test', 'scale_search')]
stride = self.configer.get('network', 'stride')
for i, scale in enumerate(multiplier):
image, border_hw = self._get_blob(ori_image, scale=scale)
with torch.no_grad():
paf_out_list, heatmap_out_list = self.pose_net(image)
paf_out = paf_out_list[-1]
heatmap_out = heatmap_out_list[-1]
# extract outputs, resize, and remove padding
heatmap = heatmap_out.squeeze(0).cpu().numpy().transpose(1, 2, 0)
heatmap = cv2.resize(heatmap, None, fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC)
heatmap = cv2.resize(heatmap[:border_hw[0], :border_hw[1]],
(ori_width, ori_height), interpolation=cv2.INTER_CUBIC)
paf = paf_out.squeeze(0).cpu().numpy().transpose(1, 2, 0)
paf = cv2.resize(paf, None, fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC)
paf = cv2.resize(paf[:border_hw[0], :border_hw[1]],
(ori_width, ori_height), interpolation=cv2.INTER_CUBIC)
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
paf_avg = paf_avg + paf / len(multiplier)
all_peaks = self.__extract_heatmap_info(heatmap_avg)
special_k, connection_all = self.__extract_paf_info(ori_img_bgr, paf_avg, all_peaks)
subset, candidate = self.__get_subsets(connection_all, special_k, all_peaks)
json_dict = self.__get_info_tree(ori_img_bgr, subset, candidate)
image_canvas = self.pose_parser.draw_points(ori_img_bgr.copy(), json_dict)
image_canvas = self.pose_parser.link_points(image_canvas, json_dict)
ImageHelper.save(image_canvas, vis_path)
ImageHelper.save(ori_img_bgr, raw_path)
Log.info('Json Save Path: {}'.format(json_path))
JsonHelper.save_file(json_dict, json_path)
def relabel(self, json_dir):
submission_dir = os.path.join(json_dir, self.configer.get('method'))
if not os.path.exists(submission_dir):
assert os.path.exists(json_dir)
os.makedirs(submission_dir)
img_shotname_list = list()
object_list = list()
for json_file in os.listdir(json_dir):
if 'json' not in json_file:
continue
json_path = os.path.join(json_dir, json_file)
shotname, extensions = os.path.splitext(json_file)
img_shotname_list.append(shotname)
with open(json_path, 'r') as json_stream:
info_tree = json.load(json_stream)
for object in info_tree['objects']:
# 0-indexing
object_list.append([
shotname, object['label'], object['score'],
int(object['bbox'][0]) + 1,
int(object['bbox'][1]) + 1,
int(object['bbox'][2]) + 1,
int(object['bbox'][3]) + 1
])
file_header_list = list()
for i in range(len(self.configer.get('details', 'name_seq'))):
cls = self.configer.get('details', 'name_seq')[i]
Log.info('Writing {:s} VOC results file'.format(cls))
filename = self.get_voc_results_file_template(submission_dir, cls)
file_header = open(filename, 'wt')
file_header_list.append(file_header)
for object in object_list:
file_header_list[object[1]].write(
'{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f}\n'.format(
object[0], object[2], object[3], object[4], object[5],
object[6]))
for file_header in file_header_list:
file_header.close()
Log.info('Evaluate {} images...'.format(len(img_shotname_list)))
return submission_dir
def relabel(self, json_dir, method='mask_rcnn'):
submission_file = os.path.join(
json_dir,
'person_instances_val2017_{}_results.json'.format(method))
img_id_list = list()
object_list = list()
for json_file in os.listdir(json_dir):
json_path = os.path.join(json_dir, json_file)
shotname, extensions = os.path.splitext(json_file)
try:
img_id = int(shotname)
except ValueError:
Log.info('Invalid Json file: {}'.format(json_file))
continue
img_id_list.append(img_id)
with open(json_path, 'r') as json_stream:
info_tree = json.load(json_stream)
for object in info_tree['objects']:
object_dict = dict()
object_dict['image_id'] = img_id
object_dict['category_id'] = int(
self.configer.get('data',
'coco_cat_seq')[object['label']])
object_dict['score'] = object['score']
object_dict['bbox'] = [
object['bbox'][0], object['bbox'][1],
object['bbox'][2] - object['bbox'][0],
object['bbox'][3] - object['bbox'][1]
]
if isinstance(object['segm'], dict):
object_dict['segmentation'] = object['segm']
else:
object_dict['segmentation'] = maskUtils.encode(
np.asfortranarray(
MaskHelper.polys2mask(object['segm'],
info_tree['height'],
info_tree['width'])))
object_list.append(object_dict)
with open(submission_file, 'w') as write_stream:
write_stream.write(json.dumps(object_list))
Log.info('Evaluate {} images...'.format(len(img_id_list)))
return submission_file, img_id_list
def test(runner):
Log.info('Testing start...')
out_dir = os.path.join(runner.configer.get('project_dir'),
'results', runner.configer.get('task'),
runner.configer.get('dataset'),
runner.configer.get('network', 'checkpoints_name'),
runner.configer.get('test', 'out_dir'))
test_dir = runner.configer.get('test', 'test_dir')
if test_dir is None:
Log.error('test_dir not given!!!')
exit(1)
runner.test(test_dir, out_dir)
Log.info('Testing end...')
def test(self, test_dir, out_dir):
if self.configer.exists('test', 'mode') and self.configer.get('test', 'mode') == 'nir2vis':
jsonA_path = os.path.join(test_dir, 'val_label{}A.json'.format(self.configer.get('data', 'tag')))
test_loader_A = self.test_loader.get_testloader(json_path=jsonA_path) if os.path.exists(
jsonA_path) else None
jsonB_path = os.path.join(test_dir, 'val_label{}B.json'.format(self.configer.get('data', 'tag')))
test_loader_B = self.test_loader.get_testloader(json_path=jsonB_path) if os.path.exists(
jsonB_path) else None
elif self.configer.exists('test', 'mode') and self.configer.get('test', 'mode') == 'pix2pix':
imgA_dir = os.path.join(test_dir, 'imageA')
test_loader_A = self.test_loader.get_testloader(test_dir=imgA_dir) if os.path.exists(imgA_dir) else None
imgB_dir = os.path.join(test_dir, 'imageB')
test_loader_B = self.test_loader.get_testloader(test_dir=imgB_dir) if os.path.exists(imgB_dir) else None
else:
imgA_dir = os.path.join(test_dir, 'imageA')
test_loader_A = self.test_loader.get_testloader(test_dir=imgA_dir) if os.path.exists(imgA_dir) else None
imgB_dir = os.path.join(test_dir, 'imageB')
test_loader_B = self.test_loader.get_testloader(test_dir=imgB_dir) if os.path.exists(imgB_dir) else None
if test_loader_A is not None:
for data_dict in test_loader_A:
new_data_dict = dict(imgA=data_dict['img'], testing=True)
with torch.no_grad():
out_dict = self.gan_net(new_data_dict)
meta_list = DCHelper.tolist(data_dict['meta'])
for key, value in out_dict.items():
for i in range(len(value)):
img_bgr = self.blob_helper.tensor2bgr(value[i])
img_path = meta_list[i]['img_path']
Log.info('Image Path: {}'.format(img_path))
ImageHelper.save(img_bgr,
os.path.join(out_dir, '{}_{}.jpg'.format(meta_list[i]['filename'], key)))
if test_loader_B is not None:
for data_dict in test_loader_B:
new_data_dict = dict(imgB=data_dict['img'], testing=True)
with torch.no_grad():
out_dict = self.gan_net(new_data_dict)
meta_list = DCHelper.tolist(data_dict['meta'])
for key, value in out_dict.items():
for i in range(len(value)):
img_bgr = self.blob_helper.tensor2bgr(value[i])
img_path = meta_list[i]['img_path']
Log.info('Image Path: {}'.format(img_path))
ImageHelper.save(img_bgr,
os.path.join(out_dir, '{}_{}.jpg'.format(meta_list[i]['filename'], key)))
def val(self):
"""
Validation function during the train phase.
"""
self.det_net.eval()
start_time = time.time()
with torch.no_grad():
for i, data_dict in enumerate(self.val_loader):
# Forward pass.
out_dict = self.det_net(data_dict)
# Compute the loss of the val batch.
loss = out_dict['loss'].mean()
self.val_losses.update(loss.item(),
len(DCHelper.tolist(data_dict['meta'])))
batch_detections = YOLOv3Test.decode(
out_dict['dets'], self.configer,
DCHelper.tolist(data_dict['meta']))
batch_pred_bboxes = self.__get_object_list(batch_detections)
self.det_running_score.update(batch_pred_bboxes, [
item['ori_bboxes']
for item in DCHelper.tolist(data_dict['meta'])
], [
item['ori_labels']
for item in DCHelper.tolist(data_dict['meta'])
])
# Update the vars of the val phase.
self.batch_time.update(time.time() - start_time)
start_time = time.time()
RunnerHelper.save_net(self,
self.det_net,
iters=self.runner_state['iters'])
# Print the log info & reset the states.
Log.info(
'Test Time {batch_time.sum:.3f}s, ({batch_time.avg:.3f})\t'
'Loss {loss.avg:.8f}\n'.format(batch_time=self.batch_time,
loss=self.val_losses))
Log.info('Val mAP: {}'.format(self.det_running_score.get_mAP()))
self.det_running_score.reset()
self.batch_time.reset()
self.val_losses.reset()
self.det_net.train()
