3 bird eye view plotting

.gitignore

@@ -166,3 +166,5 @@ wandb/
 *.mapping
 *.xml
 src/models/components/utils.py
+src/utils/plotting_dummy.py
+/3dbbox

configs/inference.yaml

@@ -31,17 +31,17 @@ save_det2d: False
 save_result: True
 
 # save result in txt
-save_txt: False
+save_txt: True
 
 # regressor weights
-regressor_weights: ${root}/weights/sample_regressor.xml
+regressor_weights: ${root}/weights/regressor_resnet18.pt
 
 # inference type
-inference_type: openvino # [pytorch, onnx, openvino, tensorrt]
+inference_type: pytorch # [pytorch, onnx, openvino, tensorrt]
 
 # source directory
-source_dir: ${root}/data/demo/images
-# source_dir: ${root}/data/demo/videos/2011_09_26/image_02/data
+# source_dir: ${root}/data/demo/images
+source_dir: ${root}/data/demo/videos/2011_09_26/image_02/data
 
 # device to inference
 device: 'cuda:0'

docs/index.md

@@ -18,7 +18,7 @@
 ## 📼  Demo
 <div align="center">
 
-![demo](./docs/assets/demo.gif)
+![demo](./assets/demo.gif)
 
 </div>
 

inference.py

@@ -12,8 +12,9 @@
 from src.utils import Calib
 from src.utils.ClassAverages import ClassAverages
 from src.utils.Plotting import calc_alpha, plot_3d_box
-from src.utils.Math import calc_location
+from src.utils.Math import calc_location, compute_orientaion, recover_angle, translation_constraints
 from src.utils.Plotting import calc_theta_ray
+from src.utils.Plotting import Plot3DBoxBev
 
 import dotenv
 import hydra
@@ -22,16 +23,183 @@
 import sys
 import pyrootutils
 import src.utils
+from src.utils.utils import KITTIObject
 
-import onnxruntime
-import openvino.runtime as ov
+import time
 
-dotenv.load_dotenv(override=True)
 log = src.utils.get_pylogger(__name__)
+
+try: 
+    import onnxruntime
+    import openvino.runtime as ov
+except ImportError:
+    log.warning("ONNX and OpenVINO not installed")
+
+dotenv.load_dotenv(override=True)
+
 root = pyrootutils.setup_root(__file__, dotenv=True, pythonpath=True)
 
 @hydra.main(version_base="1.2", config_path=root / "configs", config_name="inference.yaml")
 def inference(config: DictConfig):
+    """Inference function"""
+    # ONNX provider
+    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] \
+                if config.get("device") == "cuda" else ['CPUExecutionProvider']
+    # global calibration P2 matrix
+    P2 = Calib.get_P(config.get("calib_file"))
+    # dimension averages #TODO: depricated
+    class_averages = ClassAverages()
+
+    # time
+    avg_time = {
+        "initiate_detector": 0,
+        "initiate_regressor": 0,
+        "detector": 0,
+        "regressor": 0,
+        "plotting": 0,
+    }
+
+    # initialize detector model
+    start_detector = time.time()
+    log.info(f"Instantiating detector <{config.detector._target_}>")
+    detector = hydra.utils.instantiate(config.detector)
+    avg_time["initiate_detector"] = time.time() - start_detector
+
+    # initialize regressor model
+    start_regressor = time.time()
+    if config.get("inference_type") == "pytorch":
+        # pytorch regressor model
+        log.info(f"Instantiating regressor <{config.model._target_}>")
+        regressor: LightningModule = hydra.utils.instantiate(config.model)
+        regressor.load_state_dict(torch.load(config.get("regressor_weights")))
+        regressor.eval().to(config.get("device"))
+    elif config.get("inference_type") == "onnx":
+        # onnx regressor model
+        log.info(f"Instantiating ONNX regressor <{config.get('regressor_weights').split('/')[-1]}>")
+        regressor = onnxruntime.InferenceSession(config.get("regressor_weights"), providers=providers)
+        input_name = regressor.get_inputs()[0].name
+    elif config.get("inference_type") == "openvino":
+        # openvino regressor model
+        log.info(f"Instantiating OpenVINO regressor <{config.get('regressor_weights').split('/')[-1]}>")
+        core = ov.Core()
+        model = core.read_model(config.get("regressor_weights"))
+        regressor = core.compile_model(model, 'CPU') #TODO: change to config.get("device")
+        infer_req = regressor.create_infer_request()
+    avg_time["initiate_regressor"] = time.time() - start_regressor
+
+    # initialize preprocessing transforms
+    log.info(f"Instantiating Preprocessing Transforms")
+    preprocess: List[torch.nn.Module] = []
+    if "augmentation" in config:
+        for _, conf in config.augmentation.items():
+            if "_target_" in conf:
+                preprocess.append(hydra.utils.instantiate(conf))
+    preprocess = transforms.Compose(preprocess)
+
+    # Create output directory
+    os.makedirs(config.get("output_dir"), exist_ok=True)
+
+    # TODO: inference on video
+    # loop thru images
+    imgs_path = sorted(glob(os.path.join(config.get("source_dir"), "*")))
+    for img_path in imgs_path:
+        # Initialize object and plotting modules
+        plot3dbev = Plot3DBoxBev(P2)
+
+        img_name = img_path.split("/")[-1].split(".")[0]
+        img = cv2.imread(img_path, cv2.IMREAD_COLOR)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        # detect object with Detector
+        start_detect = time.time()
+        dets = detector(img).crop(save=config.get("save_det2d"))
+        avg_time["detector"] += time.time() - start_detect
+
+        # dimension averages #TODO: depricated
+        DIMENSION = []
+
+        # loop thru detections
+        for det in dets:
+            # initialize object container
+            obj = KITTIObject()
+            obj.name = det["label"].split(" ")[0].capitalize()
+            obj.truncation = float(0.00)
+            obj.occlusion = int(-1)
+            box = [box.cpu().numpy() for box in det["box"]]
+            obj.xmin, obj.ymin, obj.xmax, obj.ymax = box[0], box[1], box[2], box[3]
+
+            # preprocess img with torch.transforms
+            crop = preprocess(cv2.resize(det["im"], (224, 224)))
+            crop = crop.reshape((1, *crop.shape)).to(config.get("device"))
+
+            start_reg = time.time()
+            # regress 2D bbox with Regressor
+            if config.get("inference_type") == "pytorch":
+                [orient, conf, dim] = regressor(crop)
+                orient = orient.cpu().detach().numpy()[0, :, :]
+                conf = conf.cpu().detach().numpy()[0, :]
+                dim = dim.cpu().detach().numpy()[0, :]
+            elif config.get("inference_type") == "onnx":
+                # TODO: inference with GPU
+                [orient, conf, dim] = regressor.run(None, {input_name: crop.cpu().numpy()})
+                orient = orient[0]
+                conf = conf[0]
+                dim = dim[0]
+            elif config.get("inference_type") == "openvino":
+                infer_req.infer(inputs={0: crop.cpu().numpy()})
+                orient = infer_req.get_output_tensor(0).data[0]
+                conf = infer_req.get_output_tensor(1).data[0]
+                dim = infer_req.get_output_tensor(2).data[0]
+
+            # dimension averages # TODO: depricated
+            try:
+                dim += class_averages.get_item(class_to_labels(det["cls"].cpu().numpy()))
+                DIMENSION.append(dim)
+            except:
+                dim = DIMENSION[-1]
+
+            obj.alpha = recover_angle(orient, conf, 2)
+            obj.h, obj.w, obj.l = dim[0], dim[1], dim[2]
+            obj.rot_global, rot_local = compute_orientaion(P2, obj)
+            obj.tx, obj.ty, obj.tz = translation_constraints(P2, obj, rot_local)
+
+            # output prediction label
+            output_line = obj.member_to_list()
+            output_line.append(1.0)
+            output_line = " ".join([str(i) for i in output_line]) + "\n"
+
+            avg_time["regressor"] += time.time() - start_reg
+
+            # write results
+            if config.get("save_txt"):
+                with open(f"{config.get('output_dir')}/{img_name}.txt", "a") as f:
+                    f.write(output_line)
+
+
+            if config.get("save_result"):
+                start_plot = time.time()
+                plot3dbev.plot(
+                    img=img,
+                    class_object=obj.name.lower(),
+                    bbox=[obj.xmin, obj.ymin, obj.xmax, obj.ymax],
+                    dim=[obj.h, obj.w, obj.l],
+                    loc=[obj.tx, obj.ty, obj.tz],
+                    rot_y=obj.rot_global,
+                )
+                avg_time["plotting"] += time.time() - start_plot
+
+        # save images
+        if config.get("save_result"):
+            # cv2.imwrite(f'{config.get("output_dir")}/{name}.png', img_draw)
+            plot3dbev.save_plot(config.get("output_dir"), img_name)
+
+    # print time
+    for key, value in avg_time.items():
+        if key in ["detector", "regressor", "plotting"]:
+            avg_time[key] = value / len(imgs_path)
+    log.info(f"Average Time: {avg_time}")
+
+@hydra.main(version_base="1.2", config_path=root / "configs", config_name="inference.yaml")
+def inference_old(config: DictConfig):
 
     # ONNX provider
     providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] \
@@ -165,7 +333,6 @@ def inference(config: DictConfig):
                 for i in range(len(RESULTS_TXT)):
                     f.write(f"{RESULTS_TXT[i]}\n")
 
-
 def detector_yolov5(model_path: str, cfg_path: str, classes: int, device: str):
     """YOLOv5 detector model"""
     sys.path.append(str(root / "yolov5"))

scripts/frames_to_video.py

@@ -0,0 +1,37 @@
+"""Generate frames to vid"""
+
+import argparse
+import cv2
+from glob import glob
+import os
+from tqdm import tqdm
+
+def generate(imgs_path, vid_path, fps=30, frame_size=(1242, 375), resize=True):
+    """Generate frames to vid"""
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    vid_writer = cv2.VideoWriter(vid_path, fourcc, fps, frame_size)
+    imgs_glob = sorted(glob(os.path.join(imgs_path, "*.png")))
+    if resize:
+        for img_path in tqdm(imgs_glob):
+            img = cv2.imread(img_path)
+            img = cv2.resize(img, frame_size)
+            vid_writer.write(img)
+    else:
+        for img_path in imgs_glob:
+            img = cv2.imread(img_path, cv2.IMREAD_COLOR)
+            vid_writer.write(img)
+    vid_writer.release()
+    print('[INFO] Video saved to {}'.format(vid_path))
+
+if __name__ == "__main__":
+    # create argparser
+    parser = argparse.ArgumentParser(description="Generate frames to vid")
+    parser.add_argument("--imgs_path", type=str, default="outputs/2022-09-01/22-35-10/inference", help="path to imgs")
+    parser.add_argument("--vid_path", type=str, default="outputs/videos/002.mp4", help="path to vid")
+    parser.add_argument("--fps", type=int, default=24, help="fps")
+    parser.add_argument("--frame_size", type=int, nargs=2, default=(int(1242), int(375)), help="frame size")
+    parser.add_argument("--resize", action="store_true", help="resize")
+    args = parser.parse_args()
+
+    # generate vid
+    generate(args.imgs_path, args.vid_path, args.fps, args.frame_size)

scripts/get_weights.py

@@ -0,0 +1,38 @@
+"""Download pretrained weights from github release"""
+
+from pprint import pprint
+import requests
+import os
+import shutil
+import argparse
+from zipfile import ZipFile
+
+def get_assets(tag):
+    """Get release assets by tag name"""
+    url = 'https://api.github.com/repos/ruhyadi/yolo3d-lightning/releases/tags/' + tag
+    response = requests.get(url)
+    return response.json()['assets']
+
+def download_assets(assets, dir):
+    """Download assets to dir"""
+    for asset in assets:
+        url = asset['browser_download_url']
+        filename = asset['name']
+        print('[INFO] Downloading {}'.format(filename))
+        response = requests.get(url, stream=True)
+        with open(os.path.join(dir, filename), 'wb') as f:
+            shutil.copyfileobj(response.raw, f)
+        del response
+
+        with ZipFile(os.path.join(dir, filename), 'r') as zip_file:
+            zip_file.extractall(dir)
+        os.remove(os.path.join(dir, filename))
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Download pretrained weights')
+    parser.add_argument('--tag', type=str, default='v0.1', help='tag name')
+    parser.add_argument('--dir', type=str, default='./', help='directory to save weights')
+    args = parser.parse_args()
+
+    assets = get_assets(args.tag)
+    download_assets(assets, args.dir)

scripts/post_weights.py

@@ -0,0 +1,71 @@
+"""Upload weights to github release"""
+
+from pprint import pprint
+import requests
+import os
+import dotenv
+import argparse
+from zipfile import ZipFile
+
+dotenv.load_dotenv()
+
+
+def create_release(tag, name, description, target="main"):
+    """Create release"""
+    token = os.environ.get("GITHUB_TOKEN")
+    headers = {
+        "Accept": "application/vnd.github.v3+json",
+        "Authorization": f"token {token}",
+        "Content-Type": "application/zip"
+    }
+    url = "https://api.github.com/repos/ruhyadi/yolo3d-lightning/releases"
+    payload = {
+        "tag_name": tag,
+        "target_commitish": target,
+        "name": name,
+        "body": description,
+        "draft": True,
+        "prerelease": False,
+        "generate_release_notes": True,
+    }
+    print("[INFO] Creating release {}".format(tag))
+    response = requests.post(url, json=payload, headers=headers)
+    print("[INFO] Release created id: {}".format(response.json()["id"]))
+
+    return response.json()
+
+
+def post_assets(assets, release_id):
+    """Post assets to release"""
+    token = os.environ.get("GITHUB_TOKEN")
+    headers = {
+        "Accept": "application/vnd.github.v3+json",
+        "Authorization": f"token {token}",
+        "Content-Type": "application/zip"
+    }
+    for asset in assets:
+        asset_path = os.path.join(os.getcwd(), asset)
+        with ZipFile(f"{asset_path}.zip", "w") as zip_file:
+            zip_file.write(asset)
+        asset_path = f"{asset_path}.zip"
+        filename = asset_path.split("/")[-1]
+        url = (
+            "https://uploads.github.com/repos/ruhyadi/yolo3d-lightning/releases/"
+            + str(release_id)
+            + f"/assets?name={filename}"
+        )
+        print("[INFO] Uploading {}".format(filename))
+        response = requests.post(url, files={"name": open(asset_path, "rb")}, headers=headers)
+        pprint(response.json())
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Upload weights to github release")
+    parser.add_argument("--tag", type=str, default="v0.1", help="tag name")
+    parser.add_argument("--name", type=str, default="Release v0.1", help="release name")
+    parser.add_argument("--description", type=str, default="v0.1", help="release description")
+    parser.add_argument("--assets", type=tuple, default=["weights/detector_yolov5s.pt", "weights/regressor_resnet18.pt"], help="directory to save weights",)
+    args = parser.parse_args()
+
+    release_id = create_release(args.tag, args.name, args.description)["id"]
+    post_assets(args.assets, release_id)