def show_seg_result(points,
gt_seg,
pred_seg,
out_dir,
filename,
palette,
ignore_index=None,
show=True,
snapshot=False):
"""Convert results into format that is directly readable for meshlab.
Args:
points (np.ndarray): Points.
gt_seg (np.ndarray): Ground truth segmentation mask.
pred_seg (np.ndarray): Predicted segmentation mask.
out_dir (str): Path of output directory
filename (str): Filename of the current frame.
palette (np.ndarray): Mapping between class labels and colors.
ignore_index (int, optional): The label index to be ignored, e.g. \
unannotated points. Defaults to None.
show (bool, optional): Visualize the results online. Defaults to False.
snapshot (bool, optional): Whether to save the online results. \
Defaults to False.
"""
# we need 3D coordinates to visualize segmentation mask
if gt_seg is not None or pred_seg is not None:
assert points is not None, \
'3D coordinates are required for segmentation visualization'
# filter out ignored points
if gt_seg is not None and ignore_index is not None:
if points is not None:
points = points[gt_seg != ignore_index]
if pred_seg is not None:
pred_seg = pred_seg[gt_seg != ignore_index]
gt_seg = gt_seg[gt_seg != ignore_index]
if gt_seg is not None:
gt_seg_color = palette[gt_seg]
gt_seg_color = np.concatenate([points[:, :3], gt_seg_color], axis=1)
if pred_seg is not None:
pred_seg_color = palette[pred_seg]
pred_seg_color = np.concatenate([points[:, :3], pred_seg_color],
axis=1)
result_path = osp.join(out_dir, filename)
mmcv.mkdir_or_exist(result_path)
# online visualization of segmentation mask
# we show three masks in a row, scene_points, gt_mask, pred_mask
if show:
from .open3d_vis import Visualizer
mode = 'xyzrgb' if points.shape[1] == 6 else 'xyz'
vis = Visualizer(points, mode=mode)
if gt_seg is not None:
vis.add_seg_mask(gt_seg_color)
if pred_seg is not None:
vis.add_seg_mask(pred_seg_color)
show_path = osp.join(result_path,
f'{filename}_online.png') if snapshot else None
vis.show(show_path)
if points is not None:
_write_obj(points, osp.join(result_path, f'{filename}_points.obj'))
if gt_seg is not None:
_write_obj(gt_seg_color, osp.join(result_path, f'{filename}_gt.obj'))
if pred_seg is not None:
_write_obj(pred_seg_color, osp.join(result_path,
f'{filename}_pred.obj'))
def main():
args = parse_args()
models_root = args.root
models_out = args.out
mmcv.mkdir_or_exist(models_out)
# find all models in the root directory to be gathered
raw_configs = list(mmcv.scandir('./configs', '.py', recursive=True))
# filter configs that is not trained in the experiments dir
used_configs = []
for raw_config in raw_configs:
if osp.exists(osp.join(models_root, raw_config)):
used_configs.append(raw_config)
print(f'Find {len(used_configs)} models to be gathered')
# find final_ckpt and log file for trained each config
# and parse the best performance
model_infos = []
for used_config in used_configs:
exp_dir = osp.join(models_root, used_config)
# get logs
log_json_path = glob.glob(osp.join(exp_dir, '*.log.json'))[0]
log_txt_path = glob.glob(osp.join(exp_dir, '*.log'))[0]
model_performance = get_best_results(log_json_path)
final_epoch = model_performance['epoch']
final_model = 'epoch_{}.pth'.format(final_epoch)
model_path = osp.join(exp_dir, final_model)
# skip if the model is still training
if not osp.exists(model_path):
print(f'Expected {model_path} does not exist!')
continue
if model_performance is None:
print(f'Obtained no performance for model {used_config}')
continue
model_time = osp.split(log_txt_path)[-1].split('.')[0]
model_infos.append(
dict(config=used_config,
results=model_performance,
epochs=final_epoch,
model_time=model_time,
log_json_path=osp.split(log_json_path)[-1]))
# publish model for each checkpoint
publish_model_infos = []
for model in model_infos:
model_publish_dir = osp.join(models_out, model['config'].rstrip('.py'))
mmcv.mkdir_or_exist(model_publish_dir)
model_name = model['config'].split('/')[-1].rstrip(
'.py') + '_' + model['model_time']
publish_model_path = osp.join(model_publish_dir, model_name)
trained_model_path = osp.join(models_root, model['config'],
'epoch_{}.pth'.format(model['epochs']))
# convert model
final_model_path = process_checkpoint(trained_model_path,
publish_model_path)
# copy log
shutil.copy(
osp.join(models_root, model['config'], model['log_json_path']),
osp.join(model_publish_dir, f'{model_name}.log.json'))
shutil.copy(
osp.join(models_root, model['config'],
model['log_json_path'].rstrip('.json')),
osp.join(model_publish_dir, f'{model_name}.log'))
# copy config to guarantee reproducibility
config_path = model['config']
config_path = osp.join(
'configs',
config_path) if 'configs' not in config_path else config_path
target_cconfig_path = osp.split(config_path)[-1]
shutil.copy(config_path,
osp.join(model_publish_dir, target_cconfig_path))
model['model_path'] = final_model_path
publish_model_infos.append(model)
models = dict(models=publish_model_infos)
print(f'Totally gathered {len(publish_model_infos)} models')
mmcv.dump(models, osp.join(models_out, 'model_info.json'))
def generate_ann(root_path, split, image_infos, preserve_vertical, format):
"""Generate cropped annotations and label txt file.
Args:
root_path (str): The root path of the dataset
split (str): The split of dataset. Namely: training or test
image_infos (list[dict]): A list of dicts of the img and
annotation information
preserve_vertical (bool): Whether to preserve vertical texts
format (str): Annotation format, whether be txt or jsonl
"""
print('Cropping images...')
dst_image_root = osp.join(root_path, 'crops', split)
ignore_image_root = osp.join(root_path, 'ignores', split)
if split == 'training':
dst_label_file = osp.join(root_path, f'train_label.{format}')
elif split == 'val':
dst_label_file = osp.join(root_path, f'val_label.{format}')
mmcv.mkdir_or_exist(dst_image_root)
mmcv.mkdir_or_exist(ignore_image_root)
lines = []
for image_info in image_infos:
index = 1
src_img_path = osp.join(root_path, 'imgs', image_info['file_name'])
image = mmcv.imread(src_img_path)
src_img_root = image_info['file_name'].split('.')[0]
for anno in image_info['anno_info']:
word = anno['word']
dst_img = crop_img(image, anno['bbox'], 0, 0)
h, w, _ = dst_img.shape
dst_img_name = f'{src_img_root}_{index}.png'
index += 1
# Skip invalid annotations
if min(dst_img.shape) == 0:
continue
# Skip vertical texts
if not preserve_vertical and h / w > 2 and split == 'training':
dst_img_path = osp.join(ignore_image_root, dst_img_name)
else:
dst_img_path = osp.join(dst_image_root, dst_img_name)
mmcv.imwrite(dst_img, dst_img_path)
if format == 'txt':
lines.append(f'{osp.basename(dst_image_root)}/{dst_img_name} '
f'{word}')
elif format == 'jsonl':
lines.append(
json.dumps(
{
'filename':
f'{osp.basename(dst_image_root)}/{dst_img_name}',
'text': word
},
ensure_ascii=False))
else:
raise NotImplementedError
list_to_file(dst_label_file, lines)
def main(): # noqa: C901
"""Start test."""
args = parse_args()
if args.work_dir is not None:
mmcv.mkdir_or_exist(args.work_dir)
if args.tmpdir is None:
args.tmpdir = osp.join(args.work_dir, 'tmp_dir')
mmcv.mkdir_or_exist(args.tmpdir)
if args.out is None:
args.out = osp.join(args.work_dir, 'result.pkl')
if args.checkpoint is None:
args.checkpoint = osp.join(args.work_dir, 'latest.pth')
fps_file = osp.join(args.work_dir, 'fps.pkl')
mAP_file = osp.join(args.work_dir, 'mAP.pkl')
else:
mAP_file, fps_file = None, None
if args.checkpoint is None:
raise ValueError('Checkpoint file cannot be empty.')
if args.config.endswith(".json"):
load_method = mmcv.load
mmcv.load = json_to_dict
cfg = mmcv.Config.fromfile(args.config)
mmcv.load = load_method
else:
cfg = mmcv.Config.fromfile(args.config)
cfg.model.pretrained = None
cfg.data.test.test_mode = True
if args.dist:
init_dist('pytorch', **cfg.dist_params)
# build the dataloader
dataset = build_dataset(cfg.data.test)
data_loader = build_dataloader(
dataset,
imgs_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=True,
shuffle=False)
# build the model and load checkpoint
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
load_checkpoint(model, args.checkpoint, map_location='cpu')
model.CLASSES = dataset.CLASSES
if args.dist:
model = MMDistributedDataParallel(model.cuda(),
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False)
outputs = multi_gpu_test(model, data_loader, args.tmpdir)
else:
model = MMDataParallel(model, device_ids=[0])
outputs = single_gpu_test(model, data_loader, fps_file)
rank, _ = get_dist_info()
if args.out and rank == 0:
print('\nwriting results to {}'.format(args.out))
mmcv.dump(outputs, args.out)
eval_types = args.eval
if eval_types:
if eval_types == ['lamr']:
print('Starting evaluate {}'.format(' and '.join(eval_types)))
results2frame(dataset, outputs, args.tmpdir)
ecp_eval(args.tmpdir, args.tmpdir, cfg.data_root + '/day/labels/val', 'SP-NAS')
else:
print('Starting evaluate {}'.format(' and '.join(eval_types)))
assert not isinstance(outputs[0], dict)
result_files = results2json(dataset, outputs, args.out)
coco_eval(result_files, eval_types, dataset.coco, dump_file=mAP_file)
# build the model from a config file and a checkpoint file
#model = init_segmentor(config_file, checkpoint_file, device='cuda:0')
import os.path as osp
import numpy as np
from PIL import Image
# convert dataset annotation to semantic segmentation map
data_root = '/data/dentdata'
img_dir = 'images'
ann_dir = 'labels_class'
# define class and plaette for better visualization
classes = ('background', 'damage')
palette = [[255, 0, 255], [0, 255, 0]]
split_dir = 'splits'
mmcv.mkdir_or_exist(osp.join(data_root, split_dir))
filename_list = [osp.splitext(filename)[0] for filename in mmcv.scandir(
osp.join(data_root, ann_dir), suffix='.png')]
with open(osp.join(data_root, split_dir, 'train.txt'), 'w') as f:
# select first 4/5 as train set
train_length = int(len(filename_list)*4/5)
f.writelines(line + '\n' for line in filename_list[:train_length])
with open(osp.join(data_root, split_dir, 'val.txt'), 'w') as f:
# select last 1/5 as train set
f.writelines(line + '\n' for line in filename_list[train_length:])
from mmseg.datasets.builder import DATASETS
from mmseg.datasets.custom import CustomDataset
@DATASETS.register_module()
def main():
args = parse_args()
# NOTE
args.config = '/home/liqiaofei/mmdetection-v1.1.0/configs/detnasnet_detectron2_syncbn/mask_rcnn_dconv_gcb_libra_fpn_1x.py'
cfg = Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
# update configs according to CLI args
if args.work_dir is not None:
cfg.work_dir = args.work_dir
if args.resume_from is not None:
cfg.resume_from = args.resume_from
cfg.gpus = args.gpus
if args.autoscale_lr:
# apply the linear scaling rule (https://arxiv.org/abs/1706.02677)
cfg.optimizer['lr'] = cfg.optimizer['lr'] * cfg.gpus / 8
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# create work_dir
mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
# init the logger before other steps
timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
log_file = osp.join(cfg.work_dir, '{}.log'.format(timestamp))
logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)
# init the meta dict to record some important information such as
# environment info and seed, which will be logged
meta = dict()
# log env info
env_info_dict = collect_env()
env_info = '\n'.join([('{}: {}'.format(k, v))
for k, v in env_info_dict.items()])
dash_line = '-' * 60 + '\n'
logger.info('Environment info:\n' + dash_line + env_info + '\n' +
dash_line)
meta['env_info'] = env_info
# log some basic info
logger.info('Distributed training: {}'.format(distributed))
logger.info('Config:\n{}'.format(cfg.text))
# set random seeds
if args.seed is not None:
logger.info('Set random seed to {}, deterministic: {}'.format(
args.seed, args.deterministic))
set_random_seed(args.seed, deterministic=args.deterministic)
cfg.seed = args.seed
meta['seed'] = args.seed
model = build_detector(cfg.model,
train_cfg=cfg.train_cfg,
test_cfg=cfg.test_cfg)
datasets = [build_dataset(cfg.data.train)]
if len(cfg.workflow) == 2:
val_dataset = copy.deepcopy(cfg.data.val)
val_dataset.pipeline = cfg.data.train.pipeline
datasets.append(build_dataset(val_dataset))
if cfg.checkpoint_config is not None:
# save mmdet version, config file content and class names in
# checkpoints as meta data
cfg.checkpoint_config.meta = dict(mmdet_version=__version__,
config=cfg.text,
CLASSES=datasets[0].CLASSES)
# add an attribute for visualization convenience
model.CLASSES = datasets[0].CLASSES
train_detector(model,
datasets,
cfg,
distributed=distributed,
validate=args.validate,
timestamp=timestamp,
meta=meta)
def main():
args = parse_args()
cfg = Config.fromfile(args.config)
if args.options is not None:
cfg.merge_from_dict(args.options)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
# work_dir is determined in this priority: CLI > segment in file > filename
if args.work_dir is not None:
# update configs according to CLI args if args.work_dir is not None
cfg.work_dir = args.work_dir
elif cfg.get('work_dir', None) is None:
# use config filename as default work_dir if cfg.work_dir is None
cfg.work_dir = osp.join('./work_dirs',
osp.splitext(osp.basename(args.config))[0])
if args.resume_from is not None:
cfg.resume_from = args.resume_from
if args.gpu_ids is not None:
cfg.gpu_ids = args.gpu_ids
else:
cfg.gpu_ids = range(1) if args.gpus is None else range(args.gpus)
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
_, world_size = get_dist_info()
cfg.gpu_ids = range(world_size)
# create work_dir
mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
# init the logger before other steps
timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
log_file = osp.join(cfg.work_dir, f'{timestamp}.log')
logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)
# init the meta dict to record some important information such as
# environment info and seed, which will be logged
meta = dict()
# log env info
env_info_dict = collect_env()
env_info = '\n'.join([(f'{k}: {v}') for k, v in env_info_dict.items()])
dash_line = '-' * 60 + '\n'
logger.info('Environment info:\n' + dash_line + env_info + '\n' +
dash_line)
meta['env_info'] = env_info
# log some basic info
logger.info(f'Distributed training: {distributed}')
logger.info(f'Config:\n{cfg.pretty_text}')
# set random seeds
if args.seed is not None:
logger.info(f'Set random seed to {args.seed}, '
f'deterministic: {args.deterministic}')
set_random_seed(args.seed, deterministic=args.deterministic)
cfg.seed = args.seed
meta['seed'] = args.seed
model = build_classifier(cfg.model)
datasets = [build_dataset(cfg.data.train)]
if len(cfg.workflow) == 2:
val_dataset = copy.deepcopy(cfg.data.val)
val_dataset.pipeline = cfg.data.train.pipeline
datasets.append(build_dataset(val_dataset))
if cfg.checkpoint_config is not None:
# save mmcls version, config file content and class names in
# checkpoints as meta data
cfg.checkpoint_config.meta = dict(mmcls_version=__version__,
config=cfg.pretty_text,
CLASSES=datasets[0].CLASSES)
# add an attribute for visualization convenience
train_model(model,
datasets,
cfg,
distributed=distributed,
validate=(not args.no_validate),
timestamp=timestamp,
meta=meta)
def main(mode='folder'):
"""Test reds dataset.
Args:
mode: There are two modes: 'lmdb', 'folder'.
"""
opt = {}
opt['dist'] = False
opt['phase'] = 'train'
opt['name'] = 'REDS'
opt['type'] = 'REDSDataset'
if mode == 'folder':
opt['dataroot_gt'] = 'datasets/REDS/train_sharp'
opt['dataroot_lq'] = 'datasets/REDS/train_sharp_bicubic'
opt['dataroot_flow'] = None
opt['meta_info_file'] = 'basicsr/data/meta_info/meta_info_REDS_GT.txt'
opt['io_backend'] = dict(type='disk')
elif mode == 'lmdb':
opt['dataroot_gt'] = 'datasets/REDS/train_sharp_with_val.lmdb'
opt['dataroot_lq'] = 'datasets/REDS/train_sharp_bicubic_with_val.lmdb'
opt['dataroot_flow'] = None
opt['meta_info_file'] = 'basicsr/data/meta_info/meta_info_REDS_GT.txt'
opt['io_backend'] = dict(type='lmdb')
opt['val_partition'] = 'REDS4'
opt['num_frame'] = 5
opt['gt_size'] = 256
opt['interval_list'] = [1]
opt['random_reverse'] = True
opt['use_flip'] = True
opt['use_rot'] = True
opt['use_shuffle'] = True
opt['num_worker_per_gpu'] = 1
opt['batch_size_per_gpu'] = 16
opt['scale'] = 4
opt['dataset_enlarge_ratio'] = 1
mmcv.mkdir_or_exist('tmp')
dataset = create_dataset(opt)
data_loader = create_dataloader(dataset,
opt,
num_gpu=0,
dist=opt['dist'],
sampler=None)
nrow = int(math.sqrt(opt['batch_size_per_gpu']))
padding = 2 if opt['phase'] == 'train' else 0
print('start...')
for i, data in enumerate(data_loader):
if i > 5:
break
print(i)
lq = data['lq']
gt = data['gt']
key = data['key']
print(key)
for j in range(opt['num_frame']):
torchvision.utils.save_image(lq[:, j, :, :, :],
f'tmp/lq_{i:03d}_frame{j}.png',
nrow=nrow,
padding=padding,
normalize=False)
torchvision.utils.save_image(gt,
f'tmp/gt_{i:03d}.png',
nrow=nrow,
padding=padding,
normalize=False)
def imshow_det_bboxes(img,
bboxes,
labels,
segms=None,
class_names=None,
score_thr=0,
bbox_color='green',
text_color='green',
mask_color=None,
thickness=2,
font_scale=0.5,
font_size=13,
win_name='',
fig_size=(15, 10),
show=True,
wait_time=0,
out_file=None):
"""Draw bboxes and class labels (with scores) on an image.
Args:
img (str or ndarray): The image to be displayed.
bboxes (ndarray): Bounding boxes (with scores), shaped (n, 4) or
(n, 5).
labels (ndarray): Labels of bboxes.
segms (ndarray or None): Masks, shaped (n,h,w) or None
class_names (list[str]): Names of each classes.
score_thr (float): Minimum score of bboxes to be shown. Default: 0
bbox_color (str or tuple(int) or :obj:`Color`):Color of bbox lines.
The tuple of color should be in BGR order. Default: 'green'
text_color (str or tuple(int) or :obj:`Color`):Color of texts.
The tuple of color should be in BGR order. Default: 'green'
mask_color (None or str or tuple(int) or :obj:`Color`):
Color of masks. The tuple of color should be in BGR order.
Default: None
thickness (int): Thickness of lines. Default: 2
font_scale (float): Font scales of texts. Default: 0.5
font_size (int): Font size of texts. Default: 13
show (bool): Whether to show the image. Default: True
win_name (str): The window name. Default: ''
fig_size (tuple): Figure size of the pyplot figure. Default: (15, 10)
wait_time (float): Value of waitKey param. Default: 0.
out_file (str or None): The filename to write the image. Default: None
Returns:
ndarray: The image with bboxes drawn on it.
"""
warnings.warn('"font_scale" will be deprecated in v2.9.0,'
'Please use "font_size"')
assert bboxes.ndim == 2, \
f' bboxes ndim should be 2, but its ndim is {bboxes.ndim}.'
assert labels.ndim == 1, \
f' labels ndim should be 1, but its ndim is {labels.ndim}.'
assert bboxes.shape[0] == labels.shape[0], \
'bboxes.shape[0] and labels.shape[0] should have the same length.'
assert bboxes.shape[1] == 4 or bboxes.shape[1] == 5,\
f' bboxes.shape[1] should be 4 or 5, but its {bboxes.shape[1]}.'
img = mmcv.imread(img).copy()
if score_thr > 0:
assert bboxes.shape[1] == 5
scores = bboxes[:, -1]
inds = scores > score_thr
bboxes = bboxes[inds, :]
labels = labels[inds]
if segms is not None:
segms = segms[inds, ...]
mask_colors = []
if labels.shape[0] > 0:
if mask_color is None:
# random color
np.random.seed(42)
mask_colors = [
np.random.randint(0, 256, (1, 3), dtype=np.uint8)
for _ in range(max(labels) + 1)
]
else:
# specify color
mask_colors = [
np.array(mmcv.color_val(mask_color)[::-1], dtype=np.uint8)
] * (max(labels) + 1)
bbox_color = color_val_matplotlib(bbox_color)
text_color = color_val_matplotlib(text_color)
img = mmcv.bgr2rgb(img)
img = np.ascontiguousarray(img)
plt.figure(figsize=fig_size)
plt.title(win_name)
plt.axis('off')
ax = plt.gca()
polygons = []
color = []
for i, (bbox, label) in enumerate(zip(bboxes, labels)):
bbox_int = bbox.astype(np.int32)
poly = [[bbox_int[0], bbox_int[1]], [bbox_int[0], bbox_int[3]],
[bbox_int[2], bbox_int[3]], [bbox_int[2], bbox_int[1]]]
np_poly = np.array(poly).reshape((4, 2))
polygons.append(Polygon(np_poly))
color.append(bbox_color)
label_text = class_names[
label] if class_names is not None else f'class {label}'
if len(bbox) > 4:
label_text += f'|{bbox[-1]:.02f}'
ax.text(bbox_int[0],
bbox_int[1],
f'{label_text}',
bbox={
'facecolor': 'black',
'alpha': 0.8,
'pad': 0.7,
'edgecolor': 'none'
},
color=text_color,
fontsize=font_size,
verticalalignment='top',
horizontalalignment='left')
if segms is not None:
color_mask = mask_colors[labels[i]]
mask = segms[i].astype(bool)
img[mask] = img[mask] * 0.5 + color_mask * 0.5
plt.imshow(img)
p = PatchCollection(polygons,
facecolor='none',
edgecolors=color,
linewidths=thickness)
ax.add_collection(p)
if out_file is not None:
dir_name = osp.abspath(osp.dirname(out_file))
mmcv.mkdir_or_exist(dir_name)
plt.savefig(out_file)
if show:
if wait_time == 0:
plt.show()
else:
plt.show(block=False)
plt.pause(wait_time)
plt.close()
return mmcv.rgb2bgr(img)
def main():
args = parse_args()
cfg = Config.fromfile(args.config)
if args.options is not None:
cfg.merge_from_dict(args.options)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
# work_directory is determined in this priority: CLI > segment in file > filename
if args.work_directory is not None:
# update configs according to CLI args if args.work_directory is not None
cfg.work_directory = args.work_directory
elif cfg.get('work_directory', None) is None:
# use config filename as default work_directory if cfg.work_directory is None
cfg.work_directory = osp.join(
'./work_directory/retina',
osp.splitext(osp.basename(args.config))[0])
if args.resume_from is not None:
cfg.resume_from = args.resume_from
if args.gpu_ids is not None:
cfg.gpu_ids = args.gpu_ids
else:
cfg.gpu_ids = range(1) if args.gpus is None else range(args.gpus)
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# create work_directory
mmcv.mkdir_or_exist(osp.abspath(cfg.work_directory))
# dump config
cfg.dump(osp.join(cfg.work_directory, osp.basename(args.config)))
# init the logger before other steps
timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
log_file = osp.join(cfg.work_directory, f'{timestamp}.log')
logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)
# init the meta dict to record some important information such as environment info and seed, which will be logged
meta = dict()
# log env info
env_info_dict = collect_env()
env_info = '\n'.join([f'{k}: {v}' for k, v in env_info_dict.items()])
dash_line = '-' * 60 + '\n'
logger.info('Environment info:\n' + dash_line + env_info + '\n' +
dash_line)
meta['env_info'] = env_info
# log some basic info
logger.info(f'Distributed training: {distributed}')
logger.info(f'Config:\n{cfg.pretty_text}')
# ---------- MI-AOD Training and Test Start Here ---------- #
# set random seeds
if args.seed is not None:
logger.info(
f'Set random seed to {args.seed}, deterministic: {args.deterministic}'
)
set_random_seed(args.seed, deterministic=args.deterministic)
cfg.seed = args.seed
meta['seed'] = args.seed
X_L, X_U, X_all, all_anns = get_X_L_0(cfg)
# # load set and model
# # Please change it to the timestamp directory which you want to load data from.
# last_timestamp = '/20201013_154728'
# # Please change it to the cycle which you want to load data from.
# load_cycle = 0
# X_L = np.load(cfg.work_directory + last_timestamp +'/X_L_' + str(load_cycle) + '.npy')
# X_U = np.load(cfg.work_directory + last_timestamp +'/X_U_' + str(load_cycle) + '.npy')
# cfg.cycles = list(range(load_cycle, 7))
cfg.work_directory = cfg.work_directory + '/' + timestamp
mmcv.mkdir_or_exist(osp.abspath(cfg.work_directory))
np.save(cfg.work_directory + '/X_L_' + '0' + '.npy', X_L)
np.save(cfg.work_directory + '/X_U_' + '0' + '.npy', X_U)
initial_step = cfg.lr_config.step
theta_f_1 = [
'bbox_head.f_1_convs.0.conv.weight', 'bbox_head.f_1_convs.0.conv.bias',
'bbox_head.f_1_convs.1.conv.weight', 'bbox_head.f_1_convs.1.conv.bias',
'bbox_head.f_1_convs.2.conv.weight', 'bbox_head.f_1_convs.2.conv.bias',
'bbox_head.f_1_convs.3.conv.weight', 'bbox_head.f_1_convs.3.conv.bias',
'bbox_head.f_1_retina.weight', 'bbox_head.f_1_retina.bias'
]
theta_f_2 = [
'bbox_head.f_2_convs.0.conv.weight', 'bbox_head.f_2_convs.0.conv.bias',
'bbox_head.f_2_convs.1.conv.weight', 'bbox_head.f_2_convs.1.conv.bias',
'bbox_head.f_2_convs.2.conv.weight', 'bbox_head.f_2_convs.2.conv.bias',
'bbox_head.f_2_convs.3.conv.weight', 'bbox_head.f_2_convs.3.conv.bias',
'bbox_head.f_2_retina.weight', 'bbox_head.f_2_retina.bias'
]
for cycle in cfg.cycles:
# set random seeds
if args.seed is not None:
logger.info(
f'Set random seed to {args.seed}, deterministic: {args.deterministic}'
)
set_random_seed(args.seed, deterministic=args.deterministic)
cfg.seed = args.seed
meta['seed'] = args.seed
# get the config of the labeled dataset
cfg = create_X_L_file(cfg, X_L, all_anns, cycle)
# load model
model = build_detector(cfg.model,
train_cfg=cfg.train_cfg,
test_cfg=cfg.test_cfg)
# # Please change it to the epoch which you want to load model at.
# model_file_name = '/latest.pth'
# model.load_state_dict(torch.load(cfg.work_directory[:16] + last_timestamp + model_file_name)['state_dict'])
# load dataset
datasets = [build_dataset(cfg.data.train)]
if len(cfg.workflow) == 2:
val_dataset = copy.deepcopy(cfg.data.val)
val_dataset.pipeline = cfg.data.train.pipeline
datasets.append(build_dataset(val_dataset))
if cfg.checkpoint_config is not None and cycle == 0:
# save mmdet version, config file content and class names in
# checkpoints as meta data
cfg.checkpoint_config.meta = dict(mmdet_version=__version__ +
get_git_hash()[:7],
config=cfg.pretty_text,
CLASSES=datasets[0].CLASSES)
model.CLASSES = datasets[0].CLASSES
for epoch in range(cfg.epoch):
# Only in the last 3 epoch does the learning rate need to be reduced and the model needs to be evaluated.
if epoch == cfg.epoch - 1:
cfg.lr_config.step = initial_step
cfg.evaluation.interval = cfg.epoch_ratio[0]
else:
cfg.lr_config.step = [1000]
cfg.evaluation.interval = 100
# ---------- Label Set Training ----------
if epoch == 0:
cfg = create_X_L_file(cfg, X_L, all_anns, cycle)
datasets = [build_dataset(cfg.data.train)]
losstype.update_vars(0)
cfg.total_epochs = cfg.epoch_ratio[0]
cfg_bak = cfg.deepcopy()
time.sleep(2)
for name, value in model.named_parameters():
value.requires_grad = True
train_detector(model,
datasets,
cfg,
distributed=distributed,
validate=(not args.no_validate),
timestamp=timestamp,
meta=meta)
cfg = cfg_bak
# ---------- Re-weighting and Minimizing Instance Uncertainty ----------
cfg_u = create_X_U_file(cfg.deepcopy(), X_U, all_anns, cycle)
cfg = create_X_L_file(cfg, X_L, all_anns, cycle)
datasets_u = [build_dataset(cfg_u.data.train)]
datasets = [build_dataset(cfg.data.train)]
losstype.update_vars(1)
cfg_u.total_epochs = cfg_u.epoch_ratio[1]
cfg.total_epochs = cfg.epoch_ratio[1]
cfg_u_bak = cfg_u.deepcopy()
cfg_bak = cfg.deepcopy()
time.sleep(2)
for name, value in model.named_parameters():
if name in theta_f_1:
value.requires_grad = False
elif name in theta_f_2:
value.requires_grad = False
else:
value.requires_grad = True
train_detector(model, [datasets, datasets_u], [cfg, cfg_u],
distributed=distributed,
validate=(not args.no_validate),
timestamp=timestamp,
meta=meta)
cfg_u = cfg_u_bak
cfg = cfg_bak
# ---------- Re-weighting and Maximizing Instance Uncertainty ----------
cfg_u = create_X_U_file(cfg.deepcopy(), X_U, all_anns, cycle)
cfg = create_X_L_file(cfg, X_L, all_anns, cycle)
datasets_u = [build_dataset(cfg_u.data.train)]
datasets = [build_dataset(cfg.data.train)]
losstype.update_vars(2)
cfg_u.total_epochs = cfg_u.epoch_ratio[1]
cfg.total_epochs = cfg.epoch_ratio[1]
cfg_u_bak = cfg_u.deepcopy()
cfg_bak = cfg.deepcopy()
time.sleep(2)
for name, value in model.named_parameters():
if name in theta_f_1:
value.requires_grad = True
elif name in theta_f_2:
value.requires_grad = True
else:
value.requires_grad = False
train_detector(model, [datasets, datasets_u], [cfg, cfg_u],
distributed=distributed,
validate=(not args.no_validate),
timestamp=timestamp,
meta=meta)
cfg_u = cfg_u_bak
cfg = cfg_bak
# ---------- Label Set Training ----------
cfg = create_X_L_file(cfg, X_L, all_anns, cycle)
datasets = [build_dataset(cfg.data.train)]
losstype.update_vars(0)
cfg.total_epochs = cfg.epoch_ratio[0]
cfg_bak = cfg.deepcopy()
for name, value in model.named_parameters():
value.requires_grad = True
time.sleep(2)
train_detector(model,
datasets,
cfg,
distributed=distributed,
validate=args.no_validate,
timestamp=timestamp,
meta=meta)
cfg = cfg_bak
# ---------- Informative Image Selection ----------
if cycle != cfg.cycles[-1]:
# get new labeled data
dataset_al = build_dataset(cfg.data.test)
data_loader = build_dataloader(
dataset_al,
samples_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=False,
shuffle=False)
# set random seeds
if args.seed is not None:
logger.info(
f'Set random seed to {args.seed}, deterministic: {args.deterministic}'
)
set_random_seed(args.seed, deterministic=args.deterministic)
cfg.seed = args.seed
meta['seed'] = args.seed
uncertainty = calculate_uncertainty(cfg,
model,
data_loader,
return_box=False)
# update labeled set
X_L, X_U = update_X_L(uncertainty, X_all, X_L, cfg.X_S_size)
# save set and model
np.save(cfg.work_directory + '/X_L_' + str(cycle + 1) + '.npy',
X_L)
np.save(cfg.work_directory + '/X_U_' + str(cycle + 1) + '.npy',
X_U)
cfg.log_config.interval = 50
# Change the evaluation metric since we use customized dataset.
cfg.evaluation.metric = 'mAP'
# We can set the evaluation interval to reduce the evaluation times
cfg.evaluation.interval = 12
# We can set the checkpoint saving interval to reduce the storage cost
cfg.checkpoint_config.interval = 1
# Set seed thus the results are more reproducible
cfg.seed = 0
set_random_seed(0, deterministic=False)
cfg.gpu_ids = range(1)
# We can initialize the logger for training and have a look
# at the final config used for training
print(f'Config:\n{cfg.pretty_text}')
# Build dataset
datasets = [build_dataset(cfg.data.train)]
# Build the detector
model = build_detector(
cfg.model, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg)
# Add an attribute for visualization convenience
model.CLASSES = datasets[0].CLASSES
# Create work_dir
mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
train_detector(model, datasets, cfg, distributed=False, validate=True)
def build_website_from_mmdet3d_without_pred(cat, level, annotations,
dataset_dir, output_dir,
part_name_list):
dataset_path = Path(dataset_dir)
item_dir = osp.join(output_dir, 'item')
shutil.rmtree(output_dir, True)
mmcv.mkdir_or_exist(item_dir)
num = len(annotations)
num = 10
table = []
item_file_format = """<a href="./item/{page_name}" title="{page_title}"> {page_title} </a>"""
color_map_item_format = """<h2 style="display:block"> <span style = "color:{color_name}"> {class_name} </span> </h2>"""
bar = ProgressBar()
for i in bar(range(num)):
annotation = annotations[i]
gt_box = annotation['annos']['gt_boxes_upright_depth']
num_gt_box = annotation['annos']['gt_num']
box_class = annotation['annos']['class']
point_file = annotation['pts_path']
instance_file = annotation['pts_instance_mask_path']
sem_file = annotation['pts_semantic_mask_path']
#annotation['class']
#annotation['name']
point_file = str(dataset_path / point_file)
sem_file = str(dataset_path / instance_file)
instance_file = str(dataset_path / instance_file)
point = np.fromfile(point_file, np.float32).reshape(-1, 3)
#print(point.shape, point_file)
#exit(0)
#sem = np.fromfile(sem_file).astype(int)
instance = np.fromfile(instance_file, np.long)
pi_path = osp.join(item_dir, f'point_instance_{i}.ply')
bbox_path = osp.join(item_dir, f'gt_box_{i}.ply')
json_bbox_path = osp.join(item_dir, f'gt_box_{i}.json')
gt_box[:, -3:] *= 1.01
colors = write_bbox_color_json(gt_box, box_class, json_bbox_path)
write_bbox_color(gt_box, box_class, bbox_path)
write_ply_color(point, instance, pi_path)
def rgb_to_hex(rgb):
ans = '#%02x%02x%02x' % (rgb[0], rgb[1], rgb[2])
return ans
rpi_path = osp.relpath(pi_path, item_dir)
rbbox_path = osp.relpath(bbox_path, item_dir)
rjbbox_path = osp.relpath(json_bbox_path, item_dir)
next_page_rel_path = osp.relpath(
osp.join(item_dir, f'{(i + 1) % num}.html'), item_dir)
last_page_rel_path = osp.relpath(
osp.join(item_dir, f'{(i - 1 + num) % num}.html'), item_dir)
table.append([
item_file_format.format(page_name=f'{i}.html',
page_title=f'Instance {i}')
] + [0 for i in range(len(part_name_list))])
color_item = ''
for label in colors:
c = rgb_to_hex(colors[label])
l = part_name_list[label]
color_item = color_item + color_map_item_format.format(
color_name=c, class_name=l)
html_from_template(item_template,
out_file=osp.join(item_dir, f'{i}.html'),
file_path=f"'{rpi_path}', '{rjbbox_path}'",
file_name="'pc_instance', 'gt_box'",
file_type="'pcd', 'json_box'",
color_map=color_item,
last_page=last_page_rel_path,
next_page=next_page_rel_path)
df = pd.DataFrame(table, columns=['Obj'] + part_name_list)
main_table = df.to_html(escape=False,
classes="table sortable is-striped is-hoverable",
border=0,
index=False)
html_from_template(table_template,
out_file=osp.join(output_dir, 'main_page.html'),
table_string=main_table,
type=f'{cat}-{level}')
def build_website_from_mmdet3d(cat,
level,
annotations,
dataset_dir,
output_dir,
pred,
part_name_list,
iou_th=(0.5, )):
ret_dict, (aps, precisions, recalls,
aabb_ious) = evaluate(annotations, pred, part_name_list, iou_th)
#print(ret_dict)
#exit(0)
dataset_path = Path(dataset_dir)
item_dir = osp.join(output_dir, 'item')
shutil.rmtree(output_dir, True)
mmcv.mkdir_or_exist(item_dir)
num = len(annotations)
#num = 10
table = []
data_name = []
overall_string = ""
for j, part_name in enumerate(part_name_list):
data_name.append(f'{part_name}_mIOU')
overall_string += f"<th>{np.mean([aabb_ious[i][j] for i in aabb_ious if j in aabb_ious[i]]):.2f}</th>"
for part_name in part_name_list:
for iou in iou_th:
data_name.append(f'{part_name}_AP @ {iou:.2f}')
res = ret_dict[f'{part_name}_AP_{iou:.2f}']
overall_string += f"<th>{res:.2f}</th>"
for part_name in part_name_list:
for iou in iou_th:
data_name.append(f'{part_name}_AR @ {iou:.2f}')
res = ret_dict[f'{part_name}_rec_{iou:.2f}']
overall_string += f"<th>{res:.2f}</th>"
item_file_format = """<a href="./item/{page_name}" title="{page_title}"> {page_title} </a>"""
color_map_item_format = """<h2 style="display:block"> <span style = "color:{color_name}"> {class_name} </span> </h2>"""
bar = ProgressBar()
for i in bar(range(num)):
#if i != 9:
# continue
annotation = annotations[i]
gt_box = annotation['annos']['gt_boxes_upright_depth']
num_gt_box = annotation['annos']['gt_num']
box_class = annotation['annos']['class']
point_file = annotation['pts_path']
instance_file = annotation['pts_instance_mask_path']
sem_file = annotation['pts_semantic_mask_path']
#annotation['class']
#annotation['name']
point_file = str(dataset_path / point_file)
sem_file = str(dataset_path / instance_file)
instance_file = str(dataset_path / instance_file)
point = np.fromfile(point_file, np.float32).reshape(-1, 3)
#print(point.shape, point_file)
#sem = np.fromfile(sem_file).astype(int)
instance = np.fromfile(instance_file, np.long)
pi_path = osp.join(item_dir, f'point_instance_{i}.ply')
bbox_path = osp.join(item_dir, f'gt_box_{i}.ply')
json_bbox_path = osp.join(item_dir, f'gt_box_{i}.json')
colors = write_bbox_color_json(gt_box, box_class, json_bbox_path)
write_bbox_color(gt_box, box_class, bbox_path)
write_ply_color(point, instance, pi_path)
def rgb_to_hex(rgb):
ans = '#%02x%02x%02x' % (rgb[0], rgb[1], rgb[2])
return ans
rpi_path = osp.relpath(pi_path, item_dir)
rbbox_path = osp.relpath(bbox_path, item_dir)
rjbbox_path = osp.relpath(json_bbox_path, item_dir)
next_page_rel_path = osp.relpath(
osp.join(item_dir, f'{(i + 1) % num}.html'), item_dir)
last_page_rel_path = osp.relpath(
osp.join(item_dir, f'{(i - 1 + num) % num}.html'), item_dir)
table_item_i = [
item_file_format.format(page_name=f'{i}.html',
page_title=f'Instance {i}')
]
#for k in range(len(iou_th)):
for j, part_name in enumerate(part_name_list):
table_item_i.append(f'{aabb_ious[i][j]:.2f}')
for k in range(len(iou_th)):
for j, part_name in enumerate(part_name_list):
if j in aps[k][i]:
table_item_i.append(f'{float(aps[k][i][j]):.2f}')
else:
table_item_i.append('N/A')
for k in range(len(iou_th)):
for j, part_name in enumerate(part_name_list):
if j in recalls[k][i]:
table_item_i.append(f'{float(recalls[k][i][j][-1]):.2f}')
else:
table_item_i.append('N/A')
boxes = pred[i]['boxes_3d'].tensor.numpy()[:, :6]
boxes[:, 2] += boxes[:, 5] * 0.5
#print(boxes, gt_box)
#exit(0)
score = pred[i]['scores_3d'].numpy()
label = pred[i]['labels_3d'].numpy()
flag = score > 0.1
# flag = np.logical_and(score > 0.1, label == 0)
boxes = boxes[flag]
label = label[flag]
#print(boxes.shape, label, score[score > 0.1])
table.append(table_item_i)
#print(boxes.shape, label.shape)
pred_bbox_path = osp.join(item_dir, f'pred_box_{i}.ply')
json_pred_bbox_path = osp.join(item_dir, f'pred_box_{i}.json')
write_bbox_color(boxes, label, pred_bbox_path)
colors_ = write_bbox_color_json(boxes, label, json_pred_bbox_path)
colors.update(colors_)
color_item = ''
for label in colors:
c = rgb_to_hex(colors[label])
l = part_name_list[label]
color_item = color_item + color_map_item_format.format(
color_name=c, class_name=l)
pred_rjbbox_path = osp.relpath(json_pred_bbox_path, item_dir)
html_from_template(
item2_template,
out_file=osp.join(item_dir, f'{i}.html'),
file_path=f"'{rpi_path}', '{rjbbox_path}', '{pred_rjbbox_path}'",
file_name="'pc_instance', 'gt_box', 'pred_box'",
file_type="'pcd', 'json_box', 'json_box'",
color_map=color_item,
last_page=last_page_rel_path,
next_page=next_page_rel_path)
df = pd.DataFrame(table, columns=[
'Object',
] + data_name)
main_table = df.to_html(escape=False,
classes="table sortable is-striped is-hoverable",
border=0,
index=False)
html_from_template(table_template,
out_file=osp.join(output_dir, 'main_page.html'),
table_string=main_table,
overal_info=overall_string,
type=f'{cat}-{level}')
def show_multi_modality_result(img,
gt_bboxes,
pred_bboxes,
proj_mat,
out_dir,
filename,
box_mode='lidar',
img_metas=None,
show=True,
gt_bbox_color=(61, 102, 255),
pred_bbox_color=(241, 101, 72)):
"""Convert multi-modality detection results into 2D results.
Project the predicted 3D bbox to 2D image plane and visualize them.
Args:
img (np.ndarray): The numpy array of image in cv2 fashion.
gt_bboxes (:obj:`BaseInstance3DBoxes`): Ground truth boxes.
pred_bboxes (:obj:`BaseInstance3DBoxes`): Predicted boxes.
proj_mat (numpy.array, shape=[4, 4]): The projection matrix
according to the camera intrinsic parameters.
out_dir (str): Path of output directory.
filename (str): Filename of the current frame.
box_mode (str): Coordinate system the boxes are in. Should be one of
'depth', 'lidar' and 'camera'. Defaults to 'lidar'.
img_metas (dict): Used in projecting depth bbox.
show (bool): Visualize the results online. Defaults to False.
gt_bbox_color (str or tuple(int)): Color of bbox lines.
The tuple of color should be in BGR order. Default: (255, 102, 61)
pred_bbox_color (str or tuple(int)): Color of bbox lines.
The tuple of color should be in BGR order. Default: (72, 101, 241)
"""
if box_mode == 'depth':
draw_bbox = draw_depth_bbox3d_on_img
elif box_mode == 'lidar':
draw_bbox = draw_lidar_bbox3d_on_img
elif box_mode == 'camera':
draw_bbox = draw_camera_bbox3d_on_img
else:
raise NotImplementedError(f'unsupported box mode {box_mode}')
result_path = osp.join(out_dir, filename)
mmcv.mkdir_or_exist(result_path)
if show:
show_img = img.copy()
if gt_bboxes is not None:
show_img = draw_bbox(gt_bboxes,
show_img,
proj_mat,
img_metas,
color=gt_bbox_color)
if pred_bboxes is not None:
show_img = draw_bbox(pred_bboxes,
show_img,
proj_mat,
img_metas,
color=pred_bbox_color)
mmcv.imshow(show_img, win_name='project_bbox3d_img', wait_time=0)
if img is not None:
mmcv.imwrite(img, osp.join(result_path, f'{filename}_img.png'))
if gt_bboxes is not None:
gt_img = draw_bbox(gt_bboxes,
img,
proj_mat,
img_metas,
color=gt_bbox_color)
mmcv.imwrite(gt_img, osp.join(result_path, f'{filename}_gt.png'))
if pred_bboxes is not None:
pred_img = draw_bbox(pred_bboxes,
img,
proj_mat,
img_metas,
color=pred_bbox_color)
mmcv.imwrite(pred_img, osp.join(result_path, f'{filename}_pred.png'))
def bbox2result_kitti2d(self,
net_outputs,
class_names,
pklfile_prefix=None,
submission_prefix=None):
"""Convert 2D detection results to kitti format for evaluation and test
submission.
Args:
net_outputs (list[np.ndarray]): List of array storing the \
inferenced bounding boxes and scores.
class_names (list[String]): A list of class names.
pklfile_prefix (str | None): The prefix of pkl file.
submission_prefix (str | None): The prefix of submission file.
Returns:
list[dict]: A list of dictionaries have the kitti format
"""
assert len(net_outputs) == len(self.data_infos), \
'invalid list length of network outputs'
det_annos = []
print('\nConverting prediction to KITTI format')
for i, bboxes_per_sample in enumerate(
mmcv.track_iter_progress(net_outputs)):
annos = []
anno = dict(
name=[],
truncated=[],
occluded=[],
alpha=[],
bbox=[],
dimensions=[],
location=[],
rotation_y=[],
score=[])
sample_idx = self.data_infos[i]['image']['image_idx']
num_example = 0
for label in range(len(bboxes_per_sample)):
bbox = bboxes_per_sample[label]
for i in range(bbox.shape[0]):
anno['name'].append(class_names[int(label)])
anno['truncated'].append(0.0)
anno['occluded'].append(0)
anno['alpha'].append(0.0)
anno['bbox'].append(bbox[i, :4])
# set dimensions (height, width, length) to zero
anno['dimensions'].append(
np.zeros(shape=[3], dtype=np.float32))
# set the 3D translation to (-1000, -1000, -1000)
anno['location'].append(
np.ones(shape=[3], dtype=np.float32) * (-1000.0))
anno['rotation_y'].append(0.0)
anno['score'].append(bbox[i, 4])
num_example += 1
if num_example == 0:
annos.append(
dict(
name=np.array([]),
truncated=np.array([]),
occluded=np.array([]),
alpha=np.array([]),
bbox=np.zeros([0, 4]),
dimensions=np.zeros([0, 3]),
location=np.zeros([0, 3]),
rotation_y=np.array([]),
score=np.array([]),
))
else:
anno = {k: np.stack(v) for k, v in anno.items()}
annos.append(anno)
annos[-1]['sample_idx'] = np.array(
[sample_idx] * num_example, dtype=np.int64)
det_annos += annos
if pklfile_prefix is not None:
# save file in pkl format
pklfile_path = (
pklfile_prefix[:-4] if pklfile_prefix.endswith(
('.pkl', '.pickle')) else pklfile_prefix)
mmcv.dump(det_annos, pklfile_path)
if submission_prefix is not None:
# save file in submission format
mmcv.mkdir_or_exist(submission_prefix)
print(f'Saving KITTI submission to {submission_prefix}')
for i, anno in enumerate(det_annos):
sample_idx = self.data_infos[i]['image']['image_idx']
cur_det_file = f'{submission_prefix}/{sample_idx:06d}.txt'
with open(cur_det_file, 'w') as f:
bbox = anno['bbox']
loc = anno['location']
dims = anno['dimensions'][::-1] # lhw -> hwl
for idx in range(len(bbox)):
print(
'{} -1 -1 {:4f} {:4f} {:4f} {:4f} {:4f} {:4f} '
'{:4f} {:4f} {:4f} {:4f} {:4f} {:4f} {:4f}'.format(
anno['name'][idx],
anno['alpha'][idx],
*bbox[idx], # 4 float
*dims[idx], # 3 float
*loc[idx], # 3 float
anno['rotation_y'][idx],
anno['score'][idx]),
file=f,
)
print('Result is saved to {}'.format(submission_prefix))
return det_annos
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', type=int)
parser.add_argument('--world_size', type=int)
args = parser.parse_args()
print('what is the rank of the current program: ')
print(args.local_rank)
print('world size: ')
print(args.world_size)
# initialize dist
if mp.get_start_method(allow_none=True) is None:
mp.set_start_method('spawn')
torch.cuda.set_device(int(args.local_rank))
dist.init_process_group(backend='nccl', init_method='env://')
# init logger before other steps
logger = logging.getLogger()
if not logger.hasHandlers():
logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s',
level=LOG_LEVEL)
if args.local_rank != 0:
logger.setLevel('ERROR')
logger.info('Starting Distributed training')
# set random seeds
if SEED is not None:
logger.info('Set random seed to {}'.format(SEED))
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)
# build dataset
dental_dataset = DentalDataset(
num_class=NUM_CLASS,
ann_file=ANN_FILE,
img_prefix=IMG_PREFIX,
img_scale=IMG_SCALE,
img_norm_cfg=IMG_NORM_CFG,
multiscale_mode=
'value', # select a scale, rather than random from a range.
flip_ratio=FLIP_RATIO,
with_label=True,
extra_aug=EXTRA_AUG,
test_mode=False,
)
# build model
model = SSDDetector(
# basic
input_size=IMG_SCALE,
num_classes=NUM_CLASS,
in_channels=(512, 1024, 512, 256, 256),
use_dropout=False,
dropout_rate=None,
# anchor generate
anchor_ratios=([1 / 2.0, 1.0,
2.0], [1 / 3.0, 1 / 2.0, 1.0, 2.0,
3.0], [1 / 3.0, 1 / 2.0, 1.0, 2.0,
3.0], [1 / 3.0, 1 / 2.0, 1.0, 2.0,
3.0], [1 / 2.0, 1.0, 2.0]),
anchor_strides=((16, 16), (16, 16), (30, 30), (60, 60), (100, 100)),
basesizes=((12, 12), (16, 16), (24, 24), (30, 30), (36, 36)),
allowed_border=-1,
# regression
target_means=(.0, .0, .0, .0),
target_stds=(0.1, 0.1, 0.2, 0.2),
# box assign
pos_iou_thr=0.5,
neg_iou_thr=0.5,
min_pos_iou=0.,
gt_max_assign_all=False,
# sampling
sampling=False,
# balancing the loss
neg_pos_ratio=3,
# loss
smoothl1_beta=1.,
# inference nms
nms_pre=-1,
score_thr=0.02,
min_size=100.0,
max_scale_ratio=10.0,
nms_cfg=['nms', 0.45, None],
max_per_img=200,
# device
device=None,
)
model.cuda(args.local_rank)
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[args.local_rank],
output_device=args.local_rank,
find_unused_parameters=True)
if hasattr(model, 'module'):
model = model.module
model.device_ids = [
args.local_rank,
]
# build sampler for shuffling, padding, and mixing.
sampler = DistributedNonGroupSampler(
dataset=dental_dataset,
samples_per_gpu=IMGS_PER_GPU,
num_replicas=args.world_size,
rank=args.local_rank,
shuffle=True,
)
# build data loader.
data_loader = DataLoader(
dataset=dental_dataset,
batch_size=IMGS_PER_GPU,
# shuffle should be False when sampler is given.
shuffle=False,
sampler=sampler,
batch_sampler=None,
num_workers=WORKERS_PER_GPU,
collate_fn=partial(collate, samples_per_gpu=IMGS_PER_GPU),
pin_memory=False,
drop_last=False,
timeout=0,
worker_init_fn=None,
)
# optimizer
if TRAIN_HEAD_ONLY:
train_params = []
for name, param in model.named_parameters():
if 'backbone.features.30' in name or 'backbone.features.32' in name or 'backbone.extra' in name or 'bbox_head' in name:
param.requires_grad = True
train_params.append(param)
else:
param.requires_grad = False
else:
train_params = model.parameters()
optimizer = torch.optim.SGD(
params=train_params,
lr=LR,
momentum=MOMENTUM,
weight_decay=WEIGHT_DECAY,
)
# initial folder
if mmcv.is_str(WORK_DIR):
mmcv.mkdir_or_exist(WORK_DIR)
else:
raise TypeError('"work_dir" must be a str or None')
model.train()
# plan B
runner = Runner(
model=model,
batch_processor=batch_processor,
optimizer=optimizer,
work_dir=WORK_DIR,
log_level=logging.INFO,
logger=None,
)
# register hooks: optimization after the forward
optimizer_config = DistOptimizerHook(
grad_clip=grad_clip,
coalesce=True,
bucket_size_mb=-1,
)
# register hooks: along with training
runner.register_training_hooks(lr_config=lr_config,
optimizer_config=optimizer_config,
checkpoint_config=checkpoint_config,
log_config=log_config)
# register hooks: set sampler seed before each epoch
runner.register_hook(DistSamplerSeedHook())
if LOAD_FROM is not None:
runner.load_checkpoint(LOAD_FROM, strict=False)
runner.run(data_loaders=[data_loader],
workflow=workflow,
max_epochs=total_epochs)
def before_run(self, runner):
self.lock_dir = osp.join(runner.work_dir, '.lock_map_hook')
if runner.rank == 0:
if osp.exists(self.lock_dir):
shutil.rmtree(self.lock_dir)
mmcv.mkdir_or_exist(self.lock_dir)
def main():
args = parse_args()
models_root = args.root
models_out = args.out
mmcv.mkdir_or_exist(models_out)
# find all models in the root directory to be gathered
raw_configs = list(mmcv.scandir('./configs', '.py', recursive=True))
# filter configs that is not trained in the experiments dir
used_configs = []
for raw_config in raw_configs:
if osp.exists(osp.join(models_root, raw_config)):
used_configs.append(raw_config)
print(f'Find {len(used_configs)} models to be gathered')
# find final_ckpt and log file for trained each config
# and parse the best performance
model_infos = []
for used_config in used_configs:
exp_dir = osp.join(models_root, used_config)
# check whether the exps is finished
final_epoch = get_final_epoch(used_config)
final_model = 'epoch_{}.pth'.format(final_epoch)
model_path = osp.join(exp_dir, final_model)
# skip if the model is still training
if not osp.exists(model_path):
continue
# get the latest logs
log_json_path = list(sorted(glob.glob(osp.join(exp_dir,
'*.log.json'))))[-1]
log_txt_path = list(sorted(glob.glob(osp.join(exp_dir, '*.log'))))[-1]
cfg = mmcv.Config.fromfile('./configs/' + used_config)
results_lut = cfg.evaluation.metric
if not isinstance(results_lut, list):
results_lut = [results_lut]
# case when using VOC, the evaluation key is only 'mAP'
results_lut = [key + '_mAP' for key in results_lut if 'mAP' not in key]
model_performance = get_final_results(log_json_path, final_epoch,
results_lut)
if model_performance is None:
continue
model_time = osp.split(log_txt_path)[-1].split('.')[0]
model_infos.append(
dict(config=used_config,
results=model_performance,
epochs=final_epoch,
model_time=model_time,
log_json_path=osp.split(log_json_path)[-1]))
# publish model for each checkpoint
publish_model_infos = []
for model in model_infos:
model_publish_dir = osp.join(models_out, model['config'].rstrip('.py'))
mmcv.mkdir_or_exist(model_publish_dir)
model_name = osp.split(model['config'])[-1].split('.')[0]
model_name += '_' + model['model_time']
publish_model_path = osp.join(model_publish_dir, model_name)
trained_model_path = osp.join(models_root, model['config'],
'epoch_{}.pth'.format(model['epochs']))
# convert model
final_model_path = process_checkpoint(trained_model_path,
publish_model_path)
# copy log
shutil.copy(
osp.join(models_root, model['config'], model['log_json_path']),
osp.join(model_publish_dir, f'{model_name}.log.json'))
shutil.copy(
osp.join(models_root, model['config'],
model['log_json_path'].rstrip('.json')),
osp.join(model_publish_dir, f'{model_name}.log'))
# copy config to guarantee reproducibility
config_path = model['config']
config_path = osp.join(
'configs',
config_path) if 'configs' not in config_path else config_path
target_cconfig_path = osp.split(config_path)[-1]
shutil.copy(config_path,
osp.join(model_publish_dir, target_cconfig_path))
model['model_path'] = final_model_path
publish_model_infos.append(model)
models = dict(models=publish_model_infos)
print(f'Totally gathered {len(publish_model_infos)} models')
mmcv.dump(models, osp.join(models_out, 'model_info.json'))
def main(mode='folder'):
"""Test paired image dataset.
Args:
mode: There are three modes: 'lmdb', 'folder', 'meta_info_file'.
"""
opt = {}
opt['dist'] = False
opt['phase'] = 'train'
opt['name'] = 'DIV2K'
opt['type'] = 'PairedImageDataset'
if mode == 'folder':
opt['dataroot_gt'] = 'datasets/DIV2K/DIV2K_train_HR_sub'
opt['dataroot_lq'] = 'datasets/DIV2K/DIV2K_train_LR_bicubic/X4_sub'
opt['filename_tmpl'] = '{}'
opt['io_backend'] = dict(type='disk')
elif mode == 'meta_info_file':
opt['dataroot_gt'] = 'datasets/DIV2K/DIV2K_train_HR_sub'
opt['dataroot_lq'] = 'datasets/DIV2K/DIV2K_train_LR_bicubic/X4_sub'
opt['meta_info_file'] = 'basicsr/data/meta_info/meta_info_DIV2K800sub_GT.txt' # noqa:E501
opt['filename_tmpl'] = '{}'
opt['io_backend'] = dict(type='disk')
elif mode == 'lmdb':
opt['dataroot_gt'] = 'datasets/DIV2K/DIV2K_train_HR_sub.lmdb'
opt['dataroot_lq'] = 'datasets/DIV2K/DIV2K_train_LR_bicubic_X4_sub.lmdb' # noqa:E501
opt['io_backend'] = dict(type='lmdb')
opt['gt_size'] = 128
opt['use_flip'] = True
opt['use_rot'] = True
opt['use_shuffle'] = True
opt['num_worker_per_gpu'] = 2
opt['batch_size_per_gpu'] = 16
opt['scale'] = 4
opt['dataset_enlarge_ratio'] = 1
mmcv.mkdir_or_exist('tmp')
dataset = create_dataset(opt)
data_loader = create_dataloader(dataset,
opt,
num_gpu=0,
dist=opt['dist'],
sampler=None)
nrow = int(math.sqrt(opt['batch_size_per_gpu']))
padding = 2 if opt['phase'] == 'train' else 0
print('start...')
for i, data in enumerate(data_loader):
if i > 5:
break
print(i)
lq = data['lq']
gt = data['gt']
lq_path = data['lq_path']
gt_path = data['gt_path']
print(lq_path, gt_path)
torchvision.utils.save_image(lq,
f'tmp/lq_{i:03d}.png',
nrow=nrow,
padding=padding,
normalize=False)
torchvision.utils.save_image(gt,
f'tmp/gt_{i:03d}.png',
nrow=nrow,
padding=padding,
normalize=False)
def main():
args = parse_args()
dataset_path = args.dataset_path
if args.out_dir is None:
out_dir = osp.join('data', 'CHASE_DB1')
else:
out_dir = args.out_dir
print('Making directories...')
mmcv.mkdir_or_exist(out_dir)
mmcv.mkdir_or_exist(osp.join(out_dir, 'images'))
mmcv.mkdir_or_exist(osp.join(out_dir, 'images', 'training'))
mmcv.mkdir_or_exist(osp.join(out_dir, 'images', 'validation'))
mmcv.mkdir_or_exist(osp.join(out_dir, 'annotations'))
mmcv.mkdir_or_exist(osp.join(out_dir, 'annotations', 'training'))
mmcv.mkdir_or_exist(osp.join(out_dir, 'annotations', 'validation'))
with tempfile.TemporaryDirectory(dir=args.tmp_dir) as tmp_dir:
print('Extracting CHASEDB1.zip...')
zip_file = zipfile.ZipFile(dataset_path)
zip_file.extractall(tmp_dir)
print('Generating training dataset...')
assert len(os.listdir(tmp_dir)) == CHASE_DB1_LEN, \
'len(os.listdir(tmp_dir)) != {}'.format(CHASE_DB1_LEN)
for img_name in sorted(os.listdir(tmp_dir))[:TRAINING_LEN]:
img = mmcv.imread(osp.join(tmp_dir, img_name))
if osp.splitext(img_name)[1] == '.jpg':
mmcv.imwrite(
img,
osp.join(out_dir, 'images', 'training',
osp.splitext(img_name)[0] + '.png'))
else:
# The annotation img should be divided by 128, because some of
# the annotation imgs are not standard. We should set a
# threshold to convert the nonstandard annotation imgs. The
# value divided by 128 is equivalent to '1 if value >= 128
# else 0'
mmcv.imwrite(
img[:, :, 0] // 128,
osp.join(out_dir, 'annotations', 'training',
osp.splitext(img_name)[0] + '.png'))
for img_name in sorted(os.listdir(tmp_dir))[TRAINING_LEN:]:
img = mmcv.imread(osp.join(tmp_dir, img_name))
if osp.splitext(img_name)[1] == '.jpg':
mmcv.imwrite(
img,
osp.join(out_dir, 'images', 'validation',
osp.splitext(img_name)[0] + '.png'))
else:
mmcv.imwrite(
img[:, :, 0] // 128,
osp.join(out_dir, 'annotations', 'validation',
osp.splitext(img_name)[0] + '.png'))
print('Removing the temporary files...')
print('Done!')
def bbox2result_kitti(self,
net_outputs,
class_names,
pklfile_prefix=None,
submission_prefix=None):
"""Convert results to kitti format for evaluation and test submission.
Args:
net_outputs (List[np.ndarray]): list of array storing the
bbox and score
class_nanes (List[String]): A list of class names
pklfile_prefix (str | None): The prefix of pkl file.
submission_prefix (str | None): The prefix of submission file.
Returns:
List[dict]: A list of dict have the kitti 3d format
"""
assert len(net_outputs) == len(self.data_infos), \
'invalid list length of network outputs'
if submission_prefix is not None:
mmcv.mkdir_or_exist(submission_prefix)
det_annos = []
print('\nConverting prediction to KITTI format')
for idx, pred_dicts in enumerate(
mmcv.track_iter_progress(net_outputs)):
annos = []
info = self.data_infos[idx]
sample_idx = info['image']['image_idx']
image_shape = info['image']['image_shape'][:2]
box_dict = self.convert_valid_bboxes(pred_dicts, info)
if len(box_dict['bbox']) > 0:
box_2d_preds = box_dict['bbox']
box_preds = box_dict['box3d_camera']
scores = box_dict['scores']
box_preds_lidar = box_dict['box3d_lidar']
label_preds = box_dict['label_preds']
anno = {
'name': [],
'truncated': [],
'occluded': [],
'alpha': [],
'bbox': [],
'dimensions': [],
'location': [],
'rotation_y': [],
'score': []
}
for box, box_lidar, bbox, score, label in zip(
box_preds, box_preds_lidar, box_2d_preds, scores,
label_preds):
bbox[2:] = np.minimum(bbox[2:], image_shape[::-1])
bbox[:2] = np.maximum(bbox[:2], [0, 0])
anno['name'].append(class_names[int(label)])
anno['truncated'].append(0.0)
anno['occluded'].append(0)
anno['alpha'].append(
-np.arctan2(-box_lidar[1], box_lidar[0]) + box[6])
anno['bbox'].append(bbox)
anno['dimensions'].append(box[3:6])
anno['location'].append(box[:3])
anno['rotation_y'].append(box[6])
anno['score'].append(score)
anno = {k: np.stack(v) for k, v in anno.items()}
annos.append(anno)
if submission_prefix is not None:
curr_file = f'{submission_prefix}/{sample_idx:07d}.txt'
with open(curr_file, 'w') as f:
bbox = anno['bbox']
loc = anno['location']
dims = anno['dimensions'] # lhw -> hwl
for idx in range(len(bbox)):
print('{} -1 -1 {:.4f} {:.4f} {:.4f} {:.4f} '
'{:.4f} {:.4f} {:.4f} '
'{:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f}'.
format(anno['name'][idx], anno['alpha'][idx],
bbox[idx][0], bbox[idx][1],
bbox[idx][2], bbox[idx][3],
dims[idx][1], dims[idx][2],
dims[idx][0], loc[idx][0],
loc[idx][1], loc[idx][2],
anno['rotation_y'][idx],
anno['score'][idx]),
file=f)
else:
annos.append({
'name': np.array([]),
'truncated': np.array([]),
'occluded': np.array([]),
'alpha': np.array([]),
'bbox': np.zeros([0, 4]),
'dimensions': np.zeros([0, 3]),
'location': np.zeros([0, 3]),
'rotation_y': np.array([]),
'score': np.array([]),
})
annos[-1]['sample_idx'] = np.array([sample_idx] *
len(annos[-1]['score']),
dtype=np.int64)
det_annos += annos
if pklfile_prefix is not None:
if not pklfile_prefix.endswith(('.pkl', '.pickle')):
out = f'{pklfile_prefix}.pkl'
mmcv.dump(det_annos, out)
print(f'Result is saved to {out}.')
return det_annos
def main():
args = parse_args()
cfg = Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
# update configs according to CLI args
if args.work_dir is not None:
cfg.work_dir = args.work_dir
if args.resume_from is not None:
cfg.resume_from = args.resume_from
cfg.gpus = args.gpus
if args.autoscale_lr:
# apply the linear scaling rule (https://arxiv.org/abs/1706.02677)
cfg.optimizer['lr'] = cfg.optimizer['lr'] * cfg.gpus / 8
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# create work_dir
mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
# init the logger before other steps
timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
log_file = osp.join(cfg.work_dir, '{}.log'.format(timestamp))
logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)
# log some basic info
logger.info('Distributed training: {}'.format(distributed))
logger.info('MMDetection Version: {}'.format(__version__))
logger.info('Config:\n{}'.format(cfg.text))
# set random seeds
if args.seed is not None:
logger.info('Set random seed to {}, deterministic: {}'.format(
args.seed, args.deterministic))
set_random_seed(args.seed, deterministic=args.deterministic)
model = build_detector(cfg.model,
train_cfg=cfg.train_cfg,
test_cfg=cfg.test_cfg)
datasets = [build_dataset(cfg.data.train)]
if len(cfg.workflow) == 2:
datasets.append(build_dataset(cfg.data.val))
if cfg.checkpoint_config is not None:
# save mmdet version, config file content and class names in
# checkpoints as meta data
cfg.checkpoint_config.meta = dict(mmdet_version=__version__,
config=cfg.text,
CLASSES=datasets[0].CLASSES)
# add an attribute for visualization convenience
model.CLASSES = datasets[0].CLASSES
train_detector(model,
datasets,
cfg,
distributed=distributed,
validate=args.validate,
timestamp=timestamp)
import glob
CENTERNET_PATH = '/datadrive/sangliang/CenterNet/src/lib/'
sys.path.insert(0, CENTERNET_PATH)
from detectors.detector_factory import detector_factory
from opts import opts
thres = 0.3
MODEL_PATH = '/datadrive/sangliang/CenterNet/exp/ctdet_angle/coco_dla_2x_old/model_last.pth'
TASK = 'ctdet_angle' # or 'multi_pose' for human pose estimation
opt = opts().init('{} --load_model {} --dataset retail50k'.format(
TASK, MODEL_PATH).split(' '))
detector = detector_factory[opt.task](opt)
ann_file = '/datadrive/sangliang/CenterNet/data/retail50k/annotations/val.json'
output_folder = '/datadrive/sangliang/CenterNet/data/retail50k/eval'
mmcv.mkdir_or_exist(output_folder)
coco = COCO(ann_file)
cat_ids = coco.getCatIds()
cat2label = {cat_id: i + 1 for i, cat_id in enumerate(cat_ids)}
img_ids = coco.getImgIds()
img_infos = []
predict_bboxes = []
gt_bboxes = []
outfile = codecs.open(
'/datadrive/sangliang/CenterNet/data/retail50k/thres_{}_result.csv'.format(
thres),
'w',
encoding='utf-8')
outfile.write('ImgUrl,prob,x0,y0,x1,y1,x2,y2,x3,y3,label' + '\n')
result_dict = {}
for model_key in config:
model_infos = config[model_key]
if not isinstance(model_infos, list):
model_infos = [model_infos]
for model_info in model_infos:
record_metrics = model_info['metric']
cfg_path = model_info['config'].strip()
cfg = Config.fromfile(cfg_path)
checkpoint = osp.join(args.checkpoint_root,
model_info['checkpoint'].strip())
try:
fps = measure_inferense_speed(cfg, checkpoint, args.max_iter,
args.log_interval,
args.fuse_conv_bn)
print(
f'{cfg_path} fps : {fps:.{args.round_num}f} img / s, '
f'times per image: {1000/fps:.{args.round_num}f} ms / img',
flush=True)
result_dict[cfg_path] = dict(fps=round(fps, args.round_num),
ms_times_pre_image=round(
1000 / fps, args.round_num))
except Exception as e:
print(f'{config} error: {repr(e)}')
result_dict[cfg_path] = 0
if args.out:
mmcv.mkdir_or_exist(args.out)
mmcv.dump(result_dict, osp.join(args.out, 'batch_inference_fps.json'))
def main():
args = parse_args()
cfg = Config.fromfile(args.config)
if args.options is not None:
cfg.merge_from_dict(args.options)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
# work_dir is determined in this priority: CLI > segment in file > filename
if args.work_dir is not None:
# update configs according to CLI args if args.work_dir is not None
cfg.work_dir = args.work_dir
elif cfg.get('work_dir', None) is None:
# use config filename as default work_dir if cfg.work_dir is None
cfg.work_dir = osp.join('./work_dirs',
osp.splitext(osp.basename(args.config))[0])
if args.load_from is not None:
cfg.load_from = args.load_from
if args.resume_from is not None:
cfg.resume_from = args.resume_from
if args.gpu_ids is not None:
cfg.gpu_ids = args.gpu_ids
else:
cfg.gpu_ids = range(1) if args.gpus is None else range(args.gpus)
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# gpu_ids is used to calculate iter when resuming checkpoint
_, world_size = get_dist_info()
cfg.gpu_ids = range(world_size)
# create work_dir
mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
# dump config
cfg.dump(osp.join(cfg.work_dir, osp.basename(args.config)))
# init the logger before other steps
timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
log_file = osp.join(cfg.work_dir, f'{timestamp}.log')
logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)
# init the meta dict to record some important information such as
# environment info and seed, which will be logged
meta = dict()
# log env info
env_info_dict = collect_env()
env_info = '\n'.join([f'{k}: {v}' for k, v in env_info_dict.items()])
dash_line = '-' * 60 + '\n'
logger.info('Environment info:\n' + dash_line + env_info + '\n' +
dash_line)
meta['env_info'] = env_info
# log some basic info
logger.info(f'Distributed training: {distributed}')
logger.info(f'Config:\n{cfg.pretty_text}')
# set random seeds
if args.seed is not None:
logger.info(f'Set random seed to {args.seed}, deterministic: '
f'{args.deterministic}')
set_random_seed(args.seed, deterministic=args.deterministic)
cfg.seed = args.seed
meta['seed'] = args.seed
meta['exp_name'] = osp.basename(args.config)
model = build_segmentor(cfg.model,
train_cfg=cfg.get('train_cfg'),
test_cfg=cfg.get('test_cfg'))
model.init_weights()
# SyncBN is not support for DP
if not distributed:
warnings.warn(
'SyncBN is only supported with DDP. To be compatible with DP, '
'we convert SyncBN to BN. Please use dist_train.sh which can '
'avoid this error.')
model = revert_sync_batchnorm(model)
logger.info(model)
datasets = [build_dataset(cfg.data.train)]
if len(cfg.workflow) == 2:
val_dataset = copy.deepcopy(cfg.data.val)
val_dataset.pipeline = cfg.data.train.pipeline
datasets.append(build_dataset(val_dataset))
if cfg.checkpoint_config is not None:
# save mmseg version, config file content and class names in
# checkpoints as meta data
cfg.checkpoint_config.meta = dict(
mmseg_version=f'{__version__}+{get_git_hash()[:7]}',
config=cfg.pretty_text,
CLASSES=datasets[0].CLASSES,
PALETTE=datasets[0].PALETTE)
# add an attribute for visualization convenience
model.CLASSES = datasets[0].CLASSES
# passing checkpoint meta for saving best checkpoint
meta.update(cfg.checkpoint_config.meta)
train_segmentor(model,
datasets,
cfg,
distributed=distributed,
validate=(not args.no_validate),
timestamp=timestamp,
meta=meta)
def main():
# args = parse_args()
args_config = "C:/_koray/korhun/mmsegmentation/configs/deeplabv3plus/deeplabv3plus_r50-d8_512x1024_80k_cityscapes_koray.py"
args_work_dir = "C:/_koray/korhun/mmsegmentation/data/space/work_dir"
args_launcher = "none"
args_seed = None
args_deterministic = False
args_no_validate = False
cfg = Config.fromfile(args_config)
# if args.options is not None:
# cfg.merge_from_dict(args.options)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
# work_dir is determined in this priority: CLI > segment in file > filename
if args_work_dir is not None:
# update configs according to CLI args if args.work_dir is not None
cfg.work_dir = args_work_dir
elif cfg.get('work_dir', None) is None:
# use config filename as default work_dir if cfg.work_dir is None
cfg.work_dir = osp.join('./work_dirs',
osp.splitext(osp.basename(args_config))[0])
# if args.load_from is not None:
# cfg.load_from = args.load_from
# if args.resume_from is not None:
# cfg.resume_from = args.resume_from
# if args.gpu_ids is not None:
# cfg.gpu_ids = args.gpu_ids
# else:
# cfg.gpu_ids = range(1) if args.gpus is None else range(args.gpus)
# init distributed env first, since logger depends on the dist info.
if args_launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args_launcher, **cfg.dist_params)
# create work_dir
mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
# dump config
cfg.dump(osp.join(cfg.work_dir, osp.basename(args_config)))
# init the logger before other steps
timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
log_file = osp.join(cfg.work_dir, f'{timestamp}.log')
logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)
# init the meta dict to record some important information such as
# environment info and seed, which will be logged
meta = dict()
# log env info
env_info_dict = collect_env()
env_info = '\n'.join([f'{k}: {v}' for k, v in env_info_dict.items()])
dash_line = '-' * 60 + '\n'
logger.info('Environment info:\n' + dash_line + env_info + '\n' +
dash_line)
meta['env_info'] = env_info
# log some basic info
logger.info(f'Distributed training: {distributed}')
logger.info(f'Config:\n{cfg.pretty_text}')
# set random seeds
if args_seed is not None:
logger.info(f'Set random seed to {args_seed}, deterministic: '
f'{args_deterministic}')
set_random_seed(args_seed, deterministic=args_deterministic)
cfg.seed = args_seed
meta['seed'] = args_seed
meta['exp_name'] = osp.basename(args_config)
model = build_segmentor(cfg.model,
train_cfg=cfg.get('train_cfg'),
test_cfg=cfg.get('test_cfg'))
logger.info(model)
datasets = [build_dataset(cfg.data.train)]
if len(cfg.workflow) == 2:
val_dataset = copy.deepcopy(cfg.data.val)
val_dataset.pipeline = cfg.data.train.pipeline
datasets.append(build_dataset(val_dataset))
if cfg.checkpoint_config is not None:
# save mmseg version, config file content and class names in
# checkpoints as meta data
cfg.checkpoint_config.meta = dict(
mmseg_version=f'{__version__}+{get_git_hash()[:7]}',
config=cfg.pretty_text,
CLASSES=datasets[0].CLASSES,
PALETTE=datasets[0].PALETTE)
# add an attribute for visualization convenience
model.CLASSES = datasets[0].CLASSES
train_segmentor(model,
datasets,
cfg,
distributed=distributed,
validate=(not args_no_validate),
timestamp=timestamp,
meta=meta)
def main():
args = parse_args()
image_path = args.image_path
labels_ah = args.labels_ah
labels_vk = args.labels_vk
if args.out_dir is None:
out_dir = osp.join('data', 'STARE')
else:
out_dir = args.out_dir
print('Making directories...')
mmcv.mkdir_or_exist(out_dir)
mmcv.mkdir_or_exist(osp.join(out_dir, 'images'))
mmcv.mkdir_or_exist(osp.join(out_dir, 'images', 'training'))
mmcv.mkdir_or_exist(osp.join(out_dir, 'images', 'validation'))
mmcv.mkdir_or_exist(osp.join(out_dir, 'annotations'))
mmcv.mkdir_or_exist(osp.join(out_dir, 'annotations', 'training'))
mmcv.mkdir_or_exist(osp.join(out_dir, 'annotations', 'validation'))
with tempfile.TemporaryDirectory(dir=args.tmp_dir) as tmp_dir:
mmcv.mkdir_or_exist(osp.join(tmp_dir, 'gz'))
mmcv.mkdir_or_exist(osp.join(tmp_dir, 'files'))
print('Extracting stare-images.tar...')
with tarfile.open(image_path) as f:
f.extractall(osp.join(tmp_dir, 'gz'))
for filename in os.listdir(osp.join(tmp_dir, 'gz')):
un_gz(osp.join(tmp_dir, 'gz', filename),
osp.join(tmp_dir, 'files',
osp.splitext(filename)[0]))
now_dir = osp.join(tmp_dir, 'files')
assert len(os.listdir(now_dir)) == STARE_LEN, \
'len(os.listdir(now_dir)) != {}'.format(STARE_LEN)
for filename in sorted(os.listdir(now_dir))[:TRAINING_LEN]:
img = mmcv.imread(osp.join(now_dir, filename))
mmcv.imwrite(
img,
osp.join(out_dir, 'images', 'training',
osp.splitext(filename)[0] + '.png'))
for filename in sorted(os.listdir(now_dir))[TRAINING_LEN:]:
img = mmcv.imread(osp.join(now_dir, filename))
mmcv.imwrite(
img,
osp.join(out_dir, 'images', 'validation',
osp.splitext(filename)[0] + '.png'))
print('Removing the temporary files...')
with tempfile.TemporaryDirectory(dir=args.tmp_dir) as tmp_dir:
mmcv.mkdir_or_exist(osp.join(tmp_dir, 'gz'))
mmcv.mkdir_or_exist(osp.join(tmp_dir, 'files'))
print('Extracting labels-ah.tar...')
with tarfile.open(labels_ah) as f:
f.extractall(osp.join(tmp_dir, 'gz'))
for filename in os.listdir(osp.join(tmp_dir, 'gz')):
un_gz(osp.join(tmp_dir, 'gz', filename),
osp.join(tmp_dir, 'files',
osp.splitext(filename)[0]))
now_dir = osp.join(tmp_dir, 'files')
assert len(os.listdir(now_dir)) == STARE_LEN, \
'len(os.listdir(now_dir)) != {}'.format(STARE_LEN)
for filename in sorted(os.listdir(now_dir))[:TRAINING_LEN]:
img = mmcv.imread(osp.join(now_dir, filename))
# The annotation img should be divided by 128, because some of
# the annotation imgs are not standard. We should set a threshold
# to convert the nonstandard annotation imgs. The value divided by
# 128 equivalent to '1 if value >= 128 else 0'
mmcv.imwrite(
img[:, :, 0] // 128,
osp.join(out_dir, 'annotations', 'training',
osp.splitext(filename)[0] + '.png'))
for filename in sorted(os.listdir(now_dir))[TRAINING_LEN:]:
img = mmcv.imread(osp.join(now_dir, filename))
mmcv.imwrite(
img[:, :, 0] // 128,
osp.join(out_dir, 'annotations', 'validation',
osp.splitext(filename)[0] + '.png'))
print('Removing the temporary files...')
with tempfile.TemporaryDirectory(dir=args.tmp_dir) as tmp_dir:
mmcv.mkdir_or_exist(osp.join(tmp_dir, 'gz'))
mmcv.mkdir_or_exist(osp.join(tmp_dir, 'files'))
print('Extracting labels-vk.tar...')
with tarfile.open(labels_vk) as f:
f.extractall(osp.join(tmp_dir, 'gz'))
for filename in os.listdir(osp.join(tmp_dir, 'gz')):
un_gz(osp.join(tmp_dir, 'gz', filename),
osp.join(tmp_dir, 'files',
osp.splitext(filename)[0]))
now_dir = osp.join(tmp_dir, 'files')
assert len(os.listdir(now_dir)) == STARE_LEN, \
'len(os.listdir(now_dir)) != {}'.format(STARE_LEN)
for filename in sorted(os.listdir(now_dir))[:TRAINING_LEN]:
img = mmcv.imread(osp.join(now_dir, filename))
mmcv.imwrite(
img[:, :, 0] // 128,
osp.join(out_dir, 'annotations', 'training',
osp.splitext(filename)[0] + '.png'))
for filename in sorted(os.listdir(now_dir))[TRAINING_LEN:]:
img = mmcv.imread(osp.join(now_dir, filename))
mmcv.imwrite(
img[:, :, 0] // 128,
osp.join(out_dir, 'annotations', 'validation',
osp.splitext(filename)[0] + '.png'))
print('Removing the temporary files...')
print('Done!')
def forward_test(self,
masked_img,
mask,
save_image=False,
save_path=None,
iteration=None,
**kwargs):
"""Forward function for testing.
Args:
masked_img (torch.Tensor): Tensor with shape of (n, 3, h, w).
mask (torch.Tensor): Tensor with shape of (n, 1, h, w).
save_image (bool, optional): If True, results will be saved as
image. Defaults to False.
save_path (str, optional): If given a valid str, the reuslts will
be saved in this path. Defaults to None.
iteration (int, optional): Iteration number. Defaults to None.
Returns:
dict: Contain output results and eval metrics (if have).
"""
input_x = torch.cat([masked_img, mask], dim=1)
fake_res = self.generator(input_x)
fake_img = fake_res * mask + masked_img * (1. - mask)
output = dict()
eval_results = {}
if self.eval_with_metrics:
gt_img = kwargs['gt_img']
data_dict = dict(gt_img=gt_img, fake_res=fake_res, mask=mask)
for metric_name in self.test_cfg['metrics']:
if metric_name in ['ssim', 'psnr']:
eval_results[metric_name] = self._eval_metrics[
metric_name](tensor2img(fake_img, min_max=(-1, 1)),
tensor2img(gt_img, min_max=(-1, 1)))
else:
eval_results[metric_name] = self._eval_metrics[
metric_name]()(data_dict).item()
output['eval_results'] = eval_results
else:
output['fake_res'] = fake_res
output['fake_img'] = fake_img
output['meta'] = None if 'meta' not in kwargs else kwargs['meta'][0]
if save_image:
assert save_image and save_path is not None, (
'Save path should been given')
assert output['meta'] is not None, (
'Meta information should be given to save image.')
tmp_filename = output['meta']['gt_img_path']
filestem = Path(tmp_filename).stem
if iteration is not None:
filename = f'{filestem}_{iteration}.png'
else:
filename = f'{filestem}.png'
mmcv.mkdir_or_exist(save_path)
img_list = [kwargs['gt_img']] if 'gt_img' in kwargs else []
img_list.extend(
[masked_img,
mask.expand_as(masked_img), fake_res, fake_img])
img = torch.cat(img_list, dim=3).cpu()
self.save_visualization(img, osp.join(save_path, filename))
output['save_img_path'] = osp.abspath(
osp.join(save_path, filename))
return output
def __init__(self,
model,
batch_processor=None,
optimizer=None,
work_dir=None,
logger=None,
meta=None,
max_iters=None,
max_epochs=None,
amp=False):
if batch_processor is not None:
if not callable(batch_processor):
raise TypeError('batch_processor must be callable, '
f'but got {type(batch_processor)}')
warnings.warn('batch_processor is deprecated, please implement '
'train_step() and val_step() in the model instead.')
# raise an error is `batch_processor` is not None and
# `model.train_step()` exists.
if is_module_wrapper(model):
_model = model.module
else:
_model = model
if hasattr(_model, 'train_step') or hasattr(_model, 'val_step'):
raise RuntimeError(
'batch_processor and model.train_step()/model.val_step() '
'cannot be both available.')
else:
assert hasattr(model, 'train_step')
# check the type of `optimizer`
if isinstance(optimizer, dict):
for name, optim in optimizer.items():
if not isinstance(optim, Optimizer):
raise TypeError(
f'optimizer must be a dict of torch.optim.Optimizers, '
f'but optimizer["{name}"] is a {type(optim)}')
elif not isinstance(optimizer, Optimizer) and optimizer is not None:
raise TypeError(
f'optimizer must be a torch.optim.Optimizer object '
f'or dict or None, but got {type(optimizer)}')
# check the type of `logger`
if not isinstance(logger, logging.Logger):
raise TypeError(f'logger must be a logging.Logger object, '
f'but got {type(logger)}')
# check the type of `meta`
if meta is not None and not isinstance(meta, dict):
raise TypeError(
f'meta must be a dict or None, but got {type(meta)}')
self.model = model
self.batch_processor = batch_processor
self.optimizer = optimizer
self.logger = logger
self.meta = meta
self.amp = amp
# create work_dir
if mmcv.is_str(work_dir):
self.work_dir = osp.abspath(work_dir)
mmcv.mkdir_or_exist(self.work_dir)
elif work_dir is None:
self.work_dir = None
else:
raise TypeError('"work_dir" must be a str or None')
# get model name from the model class
if hasattr(self.model, 'module'):
self._model_name = self.model.module.__class__.__name__
else:
self._model_name = self.model.__class__.__name__
self._rank, self._world_size = get_dist_info()
self.timestamp = get_time_str()
self.mode = None
self._hooks = []
self._epoch = 0
self._iter = 0
self._inner_iter = 0
if max_epochs is not None and max_iters is not None:
raise ValueError(
'Only one of `max_epochs` or `max_iters` can be set.')
self._max_epochs = max_epochs
self._max_iters = max_iters
# TODO: Redesign LogBuffer, it is not flexible and elegant enough
self.log_buffer = LogBuffer()
def main():
args = parse_args()
if args.output_dir is not None:
mkdir_or_exist(args.output_dir)
cfg = build_data_cfg(args.config, args.skip_type, args.cfg_options)
try:
dataset = build_dataset(cfg.data.train,
default_args=dict(filter_empty_gt=False))
except TypeError: # seg dataset doesn't have `filter_empty_gt` key
dataset = build_dataset(cfg.data.train)
data_infos = dataset.data_infos
dataset_type = cfg.dataset_type
# configure visualization mode
vis_type = 'det' # single-modality detection
if dataset_type in ['ScanNetSegDataset', 'S3DISSegDataset']:
vis_type = 'seg' # segmentation
multi_modality = args.multi_modality
if multi_modality:
# check whether dataset really supports multi-modality input
if not is_multi_modality(dataset):
warnings.warn(
f'{dataset_type} with current config does not support multi-'
'modality data loading, only show point clouds here')
multi_modality = False
for idx, data_info in enumerate(track_iter_progress(data_infos)):
if dataset_type in ['KittiDataset', 'WaymoDataset']:
pts_path = data_info['point_cloud']['velodyne_path']
elif dataset_type in [
'ScanNetDataset', 'SUNRGBDDataset', 'ScanNetSegDataset',
'S3DISSegDataset'
]:
pts_path = data_info['pts_path']
elif dataset_type in ['NuScenesDataset', 'LyftDataset']:
pts_path = data_info['lidar_path']
else:
raise NotImplementedError(
f'unsupported dataset type {dataset_type}')
file_name = osp.splitext(osp.basename(pts_path))[0]
if vis_type == 'det':
# show 3D bboxes on 3D point clouds
show_det_data(idx,
dataset,
args.output_dir,
file_name,
show=args.online)
if multi_modality:
# project 3D bboxes to 2D image
show_proj_bbox_img(idx,
dataset,
args.output_dir,
file_name,
show=args.online)
elif vis_type == 'seg':
# show 3D segmentation mask on 3D point clouds
show_seg_data(idx,
dataset,
args.output_dir,
file_name,
show=args.online)
