Add YOLOv8 support and TensorRT for YOLOv8

docker/Dockerfile.jetson_newer

@@ -190,6 +190,35 @@ RUN colcon build \
     rm -rf /var/lib/apt/lists/* && \
     apt-get clean
 
+
+ARG ONNXRUNTIME_VERSION=v1.16.1
+ARG ONNXRUNTIME_FLAGS="--allow_running_as_root"
+ARG CUDA_ARCHITECTURES=53;62;72;87
+
+#ADD https://api.github.com/repos/microsoft/onnxruntime/git/refs/heads/${ONNXRUNTIME_VERSION} /tmp/onnxruntime_version.json
+
+# https://onnxruntime.ai/docs/build/eps.html#nvidia-jetson-tx1tx2nanoxavier
+# https://github.com/microsoft/onnxruntime/commit/b7b8b5b2ce80edb33990c7ae0fedac6ae3c623f4
+RUN pip3 uninstall -y onnxruntime && \
+    git clone https://github.com/microsoft/onnxruntime /tmp/onnxruntime && \
+    cd /tmp/onnxruntime && \
+    git checkout ${ONNXRUNTIME_VERSION} && \
+    git submodule update --init --recursive && \
+    sed -i 's|archive/3.4/eigen-3.4.zip;ee201b07085203ea7bd8eb97cbcb31b07cfa3efb|archive/3.4.0/eigen-3.4.0.zip;ef24286b7ece8737c99fa831b02941843546c081|' cmake/deps.txt || echo "cmake/deps.txt not found" && \
+    ./build.sh --config Release --update --parallel --build --build_wheel \
+        --skip_tests --skip_submodule_sync ${ONNXRUNTIME_FLAGS} \
+        --cmake_extra_defines CMAKE_CXX_FLAGS="-Wno-unused-variable -I/usr/local/cuda/include" CMAKE_CUDA_ARCHITECTURES="${CUDA_ARCHITECTURES}" onnxruntime_BUILD_UNIT_TESTS=OFF \
+        --cuda_home /usr/local/cuda --cudnn_home /usr/lib/$(uname -m)-linux-gnu \
+        --use_tensorrt --tensorrt_home /usr/lib/$(uname -m)-linux-gnu && \
+    cd build/Linux/Release && \
+    make install && \
+    cp dist/onnxruntime_gpu-*.whl /opt && \
+    pip3 install --no-cache-dir --verbose /opt/onnxruntime_gpu-*.whl && \
+    rm -rf /tmp/onnxruntime
+
+# test import and print build info
+RUN python3 -c 'import onnxruntime; print(onnxruntime.__version__);'
+
 WORKDIR /home/${USERNAME}/driverless_ws
 
 # remove driverless repo as deps have now been installed
@@ -204,6 +233,8 @@ COPY ./src/hardware/sensors/camera_files/SN33580719.conf /usr/local/zed/settings
 COPY ./src/machines/jetson_machine/requirements.txt /tmp/requirements.txt
 RUN pip3 install -r /tmp/requirements.txt && \
     rm /tmp/requirements.txt
+# no deps as CV2 is already installed but compiled from source
+RUN pip3 install ultralytics --no-deps
 
 COPY --chown=${USERNAME}:${HOST_GROUP} ./src/ ./src
 
@@ -219,6 +250,8 @@ RUN echo "test -f /home/${USERNAME}/driverless_ws/install/setup.bash && source /
 COPY ./docker/ros_entrypoint.sh /ros_entrypoint.sh
 ENTRYPOINT ["/ros_entrypoint.sh"]
 
+RUN cp -f ./src/machines/jetson_machine/scripts/replace_for_tensor.py /usr/lib/python3.8/dist-packages/tensorrt/__init__.py
+
 RUN rm -rf ./src
 
 RUN chmod -R go=rwx /home/${USERNAME} && \
@@ -227,6 +260,6 @@ RUN chmod -R go=rwx /home/${USERNAME} && \
     chmod -R go=rwx /usr/local/lib/python3.8/
 
 # tell ROS to use Cyclone DDS
-# ENV RMW_IMPLEMENTATION="rmw_cyclonedds_cpp"
-# ENV CYCLONEDDS_URI=file:///home/developer/driverless_ws/src/machines/jetson_machine/cyclonedds.xml
+ENV RMW_IMPLEMENTATION="rmw_cyclonedds_cpp"
+ENV CYCLONEDDS_URI=file:///home/developer/driverless_ws/src/machines/jetson_machine/cyclonedds.xml
 ENV ROS_DOMAIN_ID=69

src/perception/vision_pipeline/vision_pipeline/torch_inference.py

@@ -21,24 +21,6 @@ def torch_init(model_path: str, repo_path: str, conf_thresh: float, iou_thresh:
     return model
 
 
-# initialising function for the YOLOv7 model with confidence threshold
-# def torch_init_v7(model_path: str, repo_path: str, conf_thresh: float, iou_thresh: float) -> torch.nn.Module:
-#     """
-#     Returns a YOLOv7 PyTorch model
-#     """
-#     model = torch.hub.load(
-#         repo_path,
-#         "custom",
-#         model_path,
-#         source="local",
-#         force_reload=True,  # for fixing bad cache
-#     )
-#     model.conf = conf_thresh
-#     model.iou = iou_thresh
-#     model.agnostic = True
-#     return model
-
-
 def infer(colour_frame: np.ndarray, model):
     rgb_frame: np.ndarray = cv2.cvtColor(colour_frame, cv2.COLOR_BGR2RGB)
     results = model(rgb_frame)

src/perception/vision_pipeline/vision_pipeline/yolov8_torch_inference.py

@@ -0,0 +1,60 @@
+import time
+
+import cv2
+import numpy as np
+import torch
+from ultralytics import YOLO
+from ultralytics.engine.results import Results
+
+
+class YOLOv8Wrapper:
+    def __init__(self, model_path: str, conf_thresh: float = 0.7, imgsz: int = 1280, segment: bool = False):
+        """
+        initialising function for the YOLOv8 PyTorch model with confidence threshold
+        """
+        if segment:
+            task = "segment"
+        else:
+            task = "detect"
+
+        self.model = YOLO(model_path, task=task)
+        if model_path.endswith(".pt"):
+            self.model.info(verbose=True, detailed=False)
+        self.model.conf = conf_thresh
+        self.imgsz = imgsz
+        self.segment = segment
+
+    def infer(self, colour_frame: np.ndarray, verbose: bool = False):
+        """
+        function for running inference on a single frame
+        """
+        start = time.time()
+        colour_frame = cv2.cvtColor(colour_frame, cv2.COLOR_RGBA2RGB)
+        frame_result: Results = self.model.predict(colour_frame, verbose=verbose, imgsz=self.imgsz)[0]
+        end = time.time()
+
+        detection_boxes = []
+        detection_masks = []
+        if frame_result.boxes.xyxy.shape[0] == 0:
+            return []
+
+        for i in range(len(frame_result.boxes.xyxy)):
+            # ensure box is int list type
+            class_id = [frame_result.boxes.cls[i].to(device=torch.device("cpu"), dtype=torch.int32).tolist()]
+            box = frame_result.boxes.xyxy[i].to(device=torch.device("cpu"), dtype=torch.int32).tolist()
+
+            detection_boxes.append(class_id + box)
+            # convert mask to numpy array
+            if self.segment:
+                # convert from 0-1 float pixel values to 0-255 int pixel values
+                int_mask = (frame_result.masks.data[0].to(device=torch.device("cpu")).numpy() * 255).astype(np.uint8)
+                # 1080 isnt divisible by 32, so this outputs to 1088. resize to 1080
+                int_mask = cv2.resize(
+                    int_mask, (colour_frame.shape[0], colour_frame.shape[1]), interpolation=cv2.INTER_NEAREST
+                )
+                detection_masks.append(int_mask)
+
+        if self.segment:
+            return detection_boxes, detection_masks
+        else:
+            return detection_boxes