diff --git a/.gitignore b/.gitignore index fe0586a..f1eb742 100644 --- a/.gitignore +++ b/.gitignore @@ -1,134 +1,135 @@ -dataset/ -logs/ -checkpoints/ -results/ -src/.idea/ -# Byte-compiled / optimized / DLL files -__pycache__/ -*.py[cod] -*$py.class - -# C extensions -*.so - -# Distribution / packaging -.Python -build/ -develop-eggs/ -dist/ -downloads/ -eggs/ -.eggs/ -lib/ -lib64/ -parts/ -sdist/ -var/ -wheels/ -pip-wheel-metadata/ -share/python-wheels/ -*.egg-info/ -.installed.cfg -*.egg -MANIFEST - -# PyInstaller -# Usually these files are written by a python script from a template -# before PyInstaller builds the exe, so as to inject date/other infos into it. -*.manifest -*.spec - -# Installer logs -pip-log.txt -pip-delete-this-directory.txt - -# Unit test / coverage reports -htmlcov/ -.tox/ -.nox/ -.coverage -.coverage.* -.cache -nosetests.xml -coverage.xml -*.cover -*.py,cover -.hypothesis/ -.pytest_cache/ - -# Translations -*.mo -*.pot - -# Django stuff: -*.log -local_settings.py -db.sqlite3 -db.sqlite3-journal - -# Flask stuff: -instance/ -.webassets-cache - -# Scrapy stuff: -.scrapy - -# Sphinx documentation -docs/_build/ - -# PyBuilder -target/ - -# Jupyter Notebook -.ipynb_checkpoints - -# IPython -profile_default/ -ipython_config.py - -# pyenv -.python-version - -# pipenv -# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. -# However, in case of collaboration, if having platform-specific dependencies or dependencies -# having no cross-platform support, pipenv may install dependencies that don't work, or not -# install all needed dependencies. -#Pipfile.lock - -# PEP 582; used by e.g. github.com/David-OConnor/pyflow -__pypackages__/ - -# Celery stuff -celerybeat-schedule -celerybeat.pid - -# SageMath parsed files -*.sage.py - -# Environments -.env -.venv -env/ -venv/ -ENV/ -env.bak/ -venv.bak/ - -# Spyder project settings -.spyderproject -.spyproject - -# Rope project settings -.ropeproject - -# mkdocs documentation -/site - -# mypy -.mypy_cache/ -.dmypy.json -dmypy.json - -# Pyre type checker -.pyre/ +dataset/ +logs/ +checkpoints/ +results/ +src/.idea/ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ +.vscode/settings.json diff --git a/src/data_process/kitti_bev_utils.py b/src/data_process/kitti_bev_utils.py index 1585400..47b328c 100644 --- a/src/data_process/kitti_bev_utils.py +++ b/src/data_process/kitti_bev_utils.py @@ -141,7 +141,7 @@ def build_yolo_target(labels): def inverse_yolo_target(targets, bc): labels = [] for t in targets: - c, y, x, w, l, im, re = t + c, y, x, w, l, im, re, conf = t z, h = -1.55, 1.5 if c == 1: h = 1.8 @@ -154,7 +154,7 @@ def inverse_yolo_target(targets, bc): l = l * (bc["maxX"] - bc["minX"]) w -= 0.3 l -= 0.3 - labels.append([c, x, y, z, h, w, l, - np.arctan2(im, re) - 2 * np.pi]) + labels.append([c, x, y, z, h, w, l, - np.arctan2(im, re) - 2 * np.pi, conf]) return np.array(labels) diff --git a/src/data_process/transformation.py b/src/data_process/transformation.py index 58021b4..db1d53f 100644 --- a/src/data_process/transformation.py +++ b/src/data_process/transformation.py @@ -93,15 +93,15 @@ def camera_to_lidar_box(boxes, V2C=None, R0=None, P2=None): def lidar_to_camera_box(boxes, V2C=None, R0=None, P2=None): - # (N, 7) -> (N, 7) x,y,z,h,w,l,r + # (N, 8) -> (N, 8) x,y,z,h,w,l,r,conf ret = [] for box in boxes: - x, y, z, h, w, l, rz = box - (x, y, z), h, w, l, ry = lidar_to_camera( - x, y, z, V2C=V2C, R0=R0, P2=P2), h, w, l, -rz - np.pi / 2 + x, y, z, h, w, l, rz, conf = box + (x, y, z), h, w, l, ry, conf = lidar_to_camera( + x, y, z, V2C=V2C, R0=R0, P2=P2), h, w, l, -rz - np.pi / 2, conf # ry = angle_in_limit(ry) - ret.append([x, y, z, h, w, l, ry]) - return np.array(ret).reshape(-1, 7) + ret.append([x, y, z, h, w, l, ry, conf]) + return np.array(ret).reshape(-1, 8) def center_to_corner_box2d(boxes_center, coordinate='lidar'): diff --git a/src/test.py b/src/test.py index 5f838fe..e5b8fc8 100644 --- a/src/test.py +++ b/src/test.py @@ -24,7 +24,7 @@ from data_process import kitti_data_utils, kitti_bev_utils from data_process.kitti_dataloader import create_test_dataloader from models.model_utils import create_model -from utils.misc import make_folder +from utils.misc import make_folder, dump_predicted_labels from utils.evaluation_utils import post_processing, rescale_boxes, post_processing_v2 from utils.misc import time_synchronized from utils.visualization_utils import show_image_with_boxes, merge_rgb_to_bev, predictions_to_kitti_format @@ -38,7 +38,7 @@ def parse_test_configs(): help='The name of the model architecture') parser.add_argument('--cfgfile', type=str, default='./config/cfg/complex_yolov4.cfg', metavar='PATH', help='The path for cfgfile (only for darknet)') - parser.add_argument('--pretrained_path', type=str, default=None, metavar='PATH', + parser.add_argument('--pretrained_path', type=str, default='None', metavar='PATH', help='the path of the pretrained checkpoint') parser.add_argument('--use_giou_loss', action='store_true', help='If true, use GIoU loss during training. If false, use MSE loss for training') @@ -82,7 +82,13 @@ def parse_test_configs(): if configs.save_test_output: configs.results_dir = os.path.join(configs.working_dir, 'results', configs.saved_fn) + configs.results_dir_images = os.path.join(configs.working_dir, 'results', configs.saved_fn, 'images') + configs.results_dir_labels = os.path.join(configs.working_dir, 'results', configs.saved_fn, 'labels') + make_folder(configs.results_dir) + if configs.output_format == 'image': + make_folder(configs.results_dir_images) + make_folder(configs.results_dir_labels) return configs @@ -95,9 +101,9 @@ def parse_test_configs(): model.print_network() print('\n\n' + '-*=' * 30 + '\n\n') assert os.path.isfile(configs.pretrained_path), "No file at {}".format(configs.pretrained_path) - model.load_state_dict(torch.load(configs.pretrained_path)) configs.device = torch.device('cpu' if configs.no_cuda else 'cuda:{}'.format(configs.gpu_idx)) + model.load_state_dict(torch.load(configs.pretrained_path, map_location=configs.device)) model = model.to(device=configs.device) out_cap = None @@ -144,7 +150,9 @@ def parse_test_configs(): if configs.save_test_output: if configs.output_format == 'image': img_fn = os.path.basename(img_paths[0])[:-4] - cv2.imwrite(os.path.join(configs.results_dir, '{}.jpg'.format(img_fn)), out_img) + cv2.imwrite(os.path.join(configs.results_dir_images, '{}.jpg'.format(img_fn)), out_img) + dump_predicted_labels(os.path.join(configs.results_dir_labels, '{}.txt'.format(img_fn)), objects_pred) + elif configs.output_format == 'video': if out_cap is None: out_cap_h, out_cap_w = out_img.shape[:2] diff --git a/src/utils/misc.py b/src/utils/misc.py index 7ad5225..2502284 100644 --- a/src/utils/misc.py +++ b/src/utils/misc.py @@ -8,6 +8,33 @@ def make_folder(folder_name): # or os.makedirs(folder_name, exist_ok=True) +def dump_predicted_labels(file_path, predictions): + """ Dumps the output labels (KITTI format) on the provided path """ + with open(file_path, 'w') as out_label_file: + for predicted_box_parameters in predictions: + out_label = [] + out_label.append(predicted_box_parameters.type) # BB class + out_label.append(predicted_box_parameters.truncation) # BB truncation + out_label.append(predicted_box_parameters.occlusion) # BB occlusion level + out_label.append(predicted_box_parameters.alpha) # BB observation angle + out_label.append(predicted_box_parameters.box2d[0]) # 2D BB top + out_label.append(predicted_box_parameters.box2d[1]) # 2D BB left + out_label.append(predicted_box_parameters.box2d[2]) # 2D BB bottom + out_label.append(predicted_box_parameters.box2d[3]) # 2D BB right + out_label.append(predicted_box_parameters.h) # 3D BB height + out_label.append(predicted_box_parameters.w) # 3D BB width + out_label.append(predicted_box_parameters.l) # 3D BB length + out_label.append(predicted_box_parameters.t[0]) # 3D BB x + out_label.append(predicted_box_parameters.t[1]) # 3D BB y + out_label.append(predicted_box_parameters.t[2]) # 3D BB z + out_label.append(predicted_box_parameters.ry) # BB rotation around y-axis + out_label.append(predicted_box_parameters.score) # BB prediction score + + for box_param in out_label: + out_label_file.write(str(box_param)+' ') + out_label_file.write('\n') + + class AverageMeter(object): """Computes and stores the average and current value""" diff --git a/src/utils/visualization_utils.py b/src/utils/visualization_utils.py index 09523ee..2cb7ac6 100644 --- a/src/utils/visualization_utils.py +++ b/src/utils/visualization_utils.py @@ -284,8 +284,8 @@ def predictions_to_kitti_format(img_detections, calib, img_shape_2d, img_size, R if detections is None: continue # Rescale boxes to original image - for x, y, w, l, im, re, *_, cls_pred in detections: - predictions.append([cls_pred, x / img_size, y / img_size, w / img_size, l / img_size, im, re]) + for x, y, w, l, im, re, pred_conf, _, cls_pred in detections: + predictions.append([cls_pred, x / img_size, y / img_size, w / img_size, l / img_size, im, re, pred_conf]) predictions = kitti_bev_utils.inverse_yolo_target(np.array(predictions), cnf.boundary) if predictions.shape[0]: @@ -308,6 +308,7 @@ def predictions_to_kitti_format(img_detections, calib, img_shape_2d, img_size, R obj.t = l[1:4] obj.h, obj.w, obj.l = l[4:7] obj.ry = np.arctan2(math.sin(l[7]), math.cos(l[7])) + obj.score = l[8] _, corners_3d = kitti_data_utils.compute_box_3d(obj, calib.P) corners3d.append(corners_3d)