initial implementation framework; model prediction added to node main

docker/perception/camera_segmentation.Dockerfile

@@ -1,11 +1,34 @@
 # ================= Dependencies ===================
-FROM ros:humble AS base
+FROM leungjch/cuda118-tensorrt-base:latest as base
 
-# ADD DEPENDENCIES HERE
+# Segmentation Dependencies
+# Install essential packages and dependencies
+RUN sudo apt-get update && sudo apt-get install -y \
+    git \
+    make build-essential libssl-dev zlib1g-dev \
+    libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \
+    libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev \
+    ffmpeg libsm6 libxext6 cmake libgl1-mesa-glx \
+    && sudo rm -rf /var/lib/apt/lists/*
+
+RUN curl https://pyenv.run | bash
+
+ENV PATH="/root/.pyenv/bin:$PATH"
+ENV PYENV_ROOT="/root/.pyenv"
+
+SHELL ["/bin/bash", "-c"]
+# RUN source "$HOME/.bashrc" && pyenv install 3.8.15 && pyenv global 3.8.15
+SHELL ["/bin/sh", "-c"]
+
+RUN sudo apt-get install -y python3-pip
+
+RUN pip3 install --no-cache-dir --upgrade pip setuptools wheel
+RUN pip3 install ninja
+RUN pip install --no-cache-dir --upgrade requests urllib3
 
 # fix user permissions when deving in container
-COPY docker/fixuid_setup.sh /project/fixuid_setup.sh
-RUN /project/fixuid_setup.sh
+# COPY docker/fixuid_setup.sh /project/fixuid_setup.sh
+# RUN /project/fixuid_setup.sh
 USER docker:docker
 
 ENV DEBIAN_FRONTEND noninteractive
@@ -21,12 +44,17 @@ WORKDIR /home/docker/ament_ws/src
 COPY src/perception/camera_segmentation/ camera_segmentation
 COPY src/wato_msgs/sample_msgs sample_msgs
 
+COPY src/perception/camera_segmentation/requirements.txt camera_segmentation/requirements.txt
+RUN pip3 install --no-cache-dir --upgrade -r camera_segmentation/requirements.txt
+
 WORKDIR /home/docker/ament_ws
-RUN . /opt/ros/$ROS_DISTRO/setup.sh && \
+RUN . /opt/ros/$ROS_DISTRO/setup.bash && \
+    sudo rosdep init && \
     rosdep update && \
-    rosdep install -i --from-path src --rosdistro $ROS_DISTRO -y && \
+    rosdep install -i --from-path src \
+    --rosdistro $ROS_DISTRO -y  && \
     colcon build \
-        --cmake-args -DCMAKE_BUILD_TYPE=Release
+        --cmake-args -DCMAKE_BUILD_TYPE=Release -DCUDA_TOOLKIT_ROOT_DIR=/usr/local/cuda
 
 # Entrypoint will run before any CMD on launch. Sources ~/opt/<ROS_DISTRO>/setup.bash and ~/ament_ws/install/setup.bash
 COPY docker/wato_ros_entrypoint.sh /home/docker/wato_ros_entrypoint.sh

profiles/docker-compose.perception.yaml

@@ -1,16 +1,16 @@
 version: "3.8"
 services:
-  radar_object_detection:
-    build:
-      context: ..
-      dockerfile: docker/perception/radar_object_detection.Dockerfile
-      cache_from: 
-        - "${RADAR_OBJECT_DETECTION_IMAGE:?}:${TAG}"
-        - "${RADAR_OBJECT_DETECTION_IMAGE:?}:main"
-    image: "${RADAR_OBJECT_DETECTION_IMAGE:?}:${TAG}"
-    user: ${FIXUID:?}:${FIXGID:?}
-    volumes:
-      - ../src/perception/radar_object_detection:/home/docker/ament_ws/src/radar_object_detection
+  # radar_object_detection:
+  #   build:
+  #     context: ..
+  #     dockerfile: docker/perception/radar_object_detection.Dockerfile
+  #     cache_from: 
+  #       - "${RADAR_OBJECT_DETECTION_IMAGE:?}:${TAG}"
+  #       - "${RADAR_OBJECT_DETECTION_IMAGE:?}:main"
+  #   image: "${RADAR_OBJECT_DETECTION_IMAGE:?}:${TAG}"
+  #   user: ${FIXUID:?}:${FIXGID:?}
+  #   volumes:
+  #     - ../src/perception/radar_object_detection:/home/docker/ament_ws/src/radar_object_detection
   camera_segmentation:
     build:
       context: ..

src/perception/camera_segmentation/camera_segmentation/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+from . import data  # register all new datasets
+from . import modeling
+
+# config
+from .config import *
+
+# models
+from .oneformer_model import OneFormer

src/perception/camera_segmentation/camera_segmentation/camera_segmentation_node.py

@@ -27,15 +27,122 @@ def image_callback(self):
 
 
 def main(args=None):
-    rclpy.init(args=args)
+    print("Hello World")
 
-    node = CameraSegmentationNode()
+    import detectron2
+    from detectron2.utils.logger import setup_logger
+    setup_logger()
+    setup_logger(name="oneformer")
+
+    # Import libraries
+    import numpy as np
+    import cv2
+    import torch
+    import imutils
+    import os
+    import subprocess
+    from huggingface_hub import hf_hub_download
+
+    # Import detectron2 utilities
+    from detectron2.config import get_cfg
+    from detectron2.projects.deeplab import add_deeplab_config
+    from detectron2.data import MetadataCatalog
+    from .demo.defaults import DefaultPredictor
+    from .demo.visualizer import Visualizer, ColorMode
+
+    from .config import (
+        add_oneformer_config,
+        add_common_config,
+        add_swin_config,
+        add_dinat_config,
+        add_convnext_config,
+    )
+
+    torch.zeros(1).cuda()
+    cpu_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    SWIN_CFG_DICT = {"cityscapes": "/home/docker/ament_ws/src/camera_segmentation/camera_segmentation/configs/cityscapes/oneformer_swin_large_IN21k_384_bs16_90k.yaml",}
+
+    DINAT_CFG_DICT = {"cityscapes": "/home/docker/ament_ws/src/camera_segmentation/camera_segmentation/configs/cityscapes/oneformer_dinat_large_bs16_90k.yaml",}
+
+    def setup_cfg(dataset, model_path, use_swin):
+        # load config from file and command-line arguments
+        cfg = get_cfg()
+        add_deeplab_config(cfg)
+        add_common_config(cfg)
+        add_swin_config(cfg)
+        add_dinat_config(cfg)
+        add_convnext_config(cfg)
+        add_oneformer_config(cfg)
+        if use_swin:
+            cfg_path = SWIN_CFG_DICT[dataset]
+        else:
+            cfg_path = DINAT_CFG_DICT[dataset]
+            cfg.merge_from_file(cfg_path)
+            cfg.MODEL.DEVICE = 'cuda'
+            cfg.MODEL.WEIGHTS = model_path
+            cfg.freeze()
+        return cfg
+
+    def setup_modules(dataset, model_path, use_swin):
+        cfg = setup_cfg(dataset, model_path, use_swin)
+        predictor = DefaultPredictor(cfg)
+        metadata = MetadataCatalog.get(
+            cfg.DATASETS.TEST_PANOPTIC[0] if len(cfg.DATASETS.TEST_PANOPTIC) else "__unused"
+        )
+        if 'cityscapes_fine_sem_seg_val' in cfg.DATASETS.TEST_PANOPTIC[0]:
+            from cityscapesscripts.helpers.labels import labels
+            stuff_colors = [k.color for k in labels if k.trainId != 255]
+            metadata = metadata.set(stuff_colors=stuff_colors)
+
+        return predictor, metadata
+
+    def panoptic_run(img, predictor, metadata):
+        visualizer = Visualizer(img[:, :, ::-1], metadata=metadata, instance_mode=ColorMode.IMAGE)
+        predictions = predictor(img, "panoptic")
+        panoptic_seg, segments_info = predictions["panoptic_seg"]
+        out = visualizer.draw_panoptic_seg_predictions(
+        panoptic_seg.to(cpu_device), segments_info, alpha=0.5
+    )
+        return out
+
+    def instance_run(img, predictor, metadata):
+        visualizer = Visualizer(img[:, :, ::-1], metadata=metadata, instance_mode=ColorMode.IMAGE)
+        predictions = predictor(img, "instance")
+        instances = predictions["instances"].to(cpu_device)
+        out = visualizer.draw_instance_predictions(predictions=instances, alpha=0.5)
+        return out
 
+    def semantic_run(img, predictor, metadata):
+        visualizer = Visualizer(img[:, :, ::-1], metadata=metadata, instance_mode=ColorMode.IMAGE)
+        predictions = predictor(img, "semantic")
+        out = visualizer.draw_sem_seg(
+            predictions["sem_seg"].argmax(dim=0).to(cpu_device), alpha=0.5
+        )
+        return out
+
+    TASK_INFER = {"panoptic": panoptic_run, 
+                "instance": instance_run, 
+                "semantic": semantic_run}
+
+    predictor, metadata = setup_modules("cityscapes", "250_16_dinat_l_oneformer_cityscapes_90k.pth", False)
+
+    img = cv2.imread("/home/docker/ament_ws/src/camera_segmentation/camera_segmentation/samples/test.png")
+    img = imutils.resize(img, width=512)
+    task = "panoptic"
+    out = TASK_INFER[task](img, predictor, metadata).get_image()
+    cv2.imwrite("/home/docker/ament_ws/src/camera_segmentation/camera_segmentation/samples/test_results.png", out)
+
+    print("At the end")
+
+    rclpy.init(args=args)
+    node = CameraSegmentationNode()
     rclpy.spin(node)
 
     # Destroy the node explicitly
     # (optional - otherwise it will be done automatically
     # when the garbage collector destroys the node object)
+
     node.destroy_node()
     rclpy.shutdown()