Add ROS module for gknet perception and inference (#4)

* add demo script for testing annotation * wip: add new hardcoded values for saving outputs * Update docker and requirements for DCNv2 build * Update fully functional demo_annotate with docker gui support * Add ros module for keypoint messages * Add initial perception module (untested) * Add missing detector definition and update docker for ros modules * Add manual testing using static image publishers * Add tests for container * Update INSTALL with mirrors to all models/datasets * Update dockerfile for standalone inference usage * Add queue-size for approximate time sync to work correctly * Reorgainze docker file * Update default model path to avoid conflict from source directory * Rename dockerfile without suffix * Set queue size to 1 for static image publisher * Set default publish size to 1 and update logging * Optimize caching of dla34 model * Set QT_X11_NO_MITSHM
ivalab · Apr 19, 2023 · 7114b76 · 7114b76
1 parent 66e2a13
commit 7114b76
Show file tree

Hide file tree

Showing 33 changed files with 961 additions and 231 deletions.
diff --git a/.dockerignore b/.dockerignore
@@ -2,3 +2,6 @@
 /exp/
 /datasets/
 /data/
+.cache/
+*.egg-info
+__pycache__/
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -32,4 +32,4 @@ repos:
     rev: v2.2.4
     hooks:
       - id: codespell
-        args: [--ignore-words-list=thw]
+        args: ["--ignore-words-list=thw,fo"]
diff --git a/README.md b/README.md
@@ -72,7 +72,7 @@ Please refer to for [INSTALL.md](readme/INSTALL.md) installation instructions.
 
 The two training datasets are provided here:
 
-- Cornell: [Download link](https://www.dropbox.com/sh/x4t8p2wrqnfevo3/AAC2gLawRtm-986_JWxE0w0Za?dl=0).
+- Cornell: [Download link](https://www.dropbox.com/scl/fo/3tudn7uorjygxd040fcmn/h?dl=0).
   In case the download link expires in the future, you can also use the matlab scripts provided in the `GKNet_ROOT/scripts/data_aug` to generate your own dataset based on the original Cornell dataset.
   You will need to modify the corresponding path for loading the input images and output files.
 - Abridged Jacquard Dataset (AJD): [Download link](https://smartech.gatech.edu/handle/1853/64897).

diff --git a/docker-compose.yml b/docker-compose.yml
@@ -4,16 +4,21 @@ version: "3.8"
 services:
   # base container -- will simply exit once brought up
   # we can run commands via this container for running experiments, etc.
-  base: &base
+  base:
     build:
       context: .
-      dockerfile: docker/Dockerfile.noetic
-    image: ivalab/gknet:latest
-    network_mode: host
-    shm_size: 2gb
-    volumes:
-      - ./:/app/
-      - .cache/torch:/root/.cache/torch
+      dockerfile: docker/Dockerfile
+    <<: &base
+      image: ivalab/gknet:latest
+      network_mode: host
+      shm_size: 2gb
+      volumes:
+        - ./:/catkin_ws/src/app
+        - /tmp/.X11-unix:/tmp/.X11-unix
+        - ~/.Xauthority:/root/.Xauthority
+      environment:
+        DISPLAY: $DISPLAY
+        QT_X11_NO_MITSHM: 1
   gpu:
     <<: *base
     deploy:

diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -0,0 +1,125 @@
+FROM nvidia/cuda:11.7.0-devel-ubuntu20.04
+
+# NOTE: We need cuda devel to build DCNv2, but we don't necessarily need it to
+# run inference. This should probably become a multi-part build if we care about
+# deployment size.
+
+# take the core ros image on top of the cuda focal image
+# https://github.com/osrf/docker_images/blob/master/ros/noetic/ubuntu/focal/ros-core/Dockerfile
+
+# noninteractive to avoid tzdata prompt
+ENV DEBIAN_FRONTEND noninteractive
+ENV ROS_DISTRO noetic
+RUN echo "deb http://packages.ros.org/ros/ubuntu focal main" > /etc/apt/sources.list.d/ros1-latest.list
+RUN apt-key adv --keyserver hkp://keyserver.ubuntu.com:80 --recv-keys C1CF6E31E6BADE8868B172B4F42ED6FBAB17C654
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+    ros-noetic-ros-core=1.5.0-1* \
+    python3-catkin-tools \
+    python3-rosdep \
+    && \
+    apt clean
+
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+    apt-utils \
+    build-essential \
+    git \
+    python3-dev \
+    python3-pip \
+    python-is-python3 \
+    ninja-build \
+    libgl1-mesa-glx \
+    libsm-dev \
+    wget \
+    && \
+    apt clean
+
+RUN pip install --upgrade pip wheel cython
+
+RUN apt update && \
+    apt install -y --no-install-recommends \
+    # install ros packages
+    # rosdep install --from-paths src -r -y --simulate
+    ros-noetic-ros-pytest \
+    ros-noetic-cv-bridge \
+    # additional package for viewing images
+    ros-noetic-image-view \
+    && \
+    apt clean
+
+# install cocoapi
+WORKDIR /opt
+RUN git clone https://github.com/cocodataset/cocoapi.git && \
+    cd cocoapi &&\
+    git checkout 8c9bcc3 && \
+    cd PythonAPI && \
+    make && \
+    python setup.py install
+
+# install nms
+WORKDIR /app
+ADD vendor/ vendor/
+RUN cd vendor/nms && \
+    python setup.py build install
+
+# install requirements (most importantly pytorch) and DCNv2
+ADD requirements.txt /app/requirements.txt
+RUN pip install -r requirements.txt
+
+# this is fairly slow, and the ABI is dependent on both cuda and pytorch. We
+# need to make sure we build and run with the same version of pytorch.
+WORKDIR /opt
+# fork of https://github.com/lbin/DCNv2.git with ability to compile with cuda
+# without an nvidia gpu attached
+RUN git clone https://github.com/acmiyaguchi/DCNv2.git && \
+    cd DCNv2 && \
+    git checkout pytorch_1.11 && \
+    # https://github.com/facebookresearch/pytorch3d/issues/318
+    FORCE_CUDA=1 \
+    # https://github.com/pytorch/extension-cpp/issues/71
+    TORCH_CUDA_ARCH_LIST="3.5;5.0;6.0;6.1;7.0;7.5;8.0;8.6+PTX" \
+    ./make.sh
+
+# download the dla34 models
+# NOTE: if these links are broken for some reason, this can safely be removed
+# and instead have the models mounted from the host. You will also probably want
+# to download the full set of pretrained models to mount onto the container
+# during development.
+RUN mkdir -p /opt/models && \
+    cd /opt/models && \
+    wget https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/models/model_dla34_cornell.pth && \
+    wget https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/models/model_dla34_ajd.pth
+
+# add the source code to the image
+# we also link /app to /catkin_ws/src/app as a shortcut to the source
+WORKDIR /catkin_ws/src/app
+RUN ln -s /catkin_ws/src/app /app
+
+# let's only copy in the package, in case we don't change much of it
+
+ADD gknet/ ./gknet/
+ADD scripts/ ./scripts/
+ADD setup.py ./setup.py
+RUN pip install -e .
+
+# cache the model weights for default inference
+RUN python scripts/load_detector.py dbmctdet_cornell \
+    --load_model /opt/models/model_dla34_cornell.pth
+
+# Now we can add the rest of the app
+ADD ./ .
+
+WORKDIR /catkin_ws
+RUN catkin config --init --extend /opt/ros/noetic && \
+    rosdep init && \
+    rosdep update && \
+    rosdep install --from-paths src --ignore-src -r -y && \
+    catkin build
+
+WORKDIR /catkin_ws/src/app
+
+# setup entrypoint
+COPY ./docker/bin/ros_entrypoint.sh /
+ENTRYPOINT ["/ros_entrypoint.sh",  "stdbuf", "-o", "L"]
+CMD ["bash"]
diff --git a/docker/Dockerfile.noetic b/docker/Dockerfile.noetic
diff --git a/docker/bin/ros_entrypoint.sh b/docker/bin/ros_entrypoint.sh
@@ -2,5 +2,5 @@
 set -e
 
 # setup ros environment
-source "/opt/ros/$ROS_DISTRO/setup.bash" --
+source "/catkin_ws/devel/setup.bash"
 exec "$@"
diff --git a/docs/INSTALL.md b/docs/INSTALL.md
@@ -5,6 +5,7 @@
 Docker is a self-contained environment that can be used to run the code without having to install the dependencies.
 We have built a docker image with Ubuntu 20.04, ROS Noetic, and PyTorch 1.13.1.
 Make sure you have [nvidia-docker](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html) installed.
+The newer version of PyTorch means that the host machine must have a CUDA 11.7 compatible driver.
 At the root of the repository, run the following command to build the container:
 
 ```bash
@@ -86,3 +87,48 @@ Reference the [Dockerfile](../docker/Dockerfile.noetic) for build instructions a
    cd $GKNet_ROOT/vendor/nms
    make
    ```
+
+## Downloading models
+
+In addition to the model weight and dataset links referenced in the main README, we have hosted the following on a public mirror.
+The purpose is to provide a stable link for a subset of models used during model inference.
+In general, prefix the filename with `https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/` to download the file.
+Please use discretion when downloading the larger files.
+
+- Pretrained_models (zipped, see https://www.dropbox.com/sh/eicrmhhay2wi8fy/AAAGrToUcdp0tO-F732Xhsxwa?dl=0)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/Pretrained_models.zip (5.1 GB, sha1 4ac3a5ded51c46b6abdd552d2b709d84663e4cfd)
+- Pretrained_models
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/models/ctdet_coco_dla_2x.pth (77.2 MB, sha1 f77895ebefe56bd4a6abd173f919d17e6ef60c79)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/models/model_alexnet_ajd.pth (26.5 MB, sha1 5de7fdac3325f5012bc98a4e5516c4d463e8c1ea)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/models/model_alexnet_cornell.pth (26.5 MB, sha1 4aedef106634c2bb7242ef5aa2659c8deb69ccfd)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/models/model_dla34_ajd.pth (80 MB, sha1 cb874a470d3f54801bdd53468411da3ee2d2bbf7)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/models/model_dla34_cornell.pth (235.5 MB, sha1 0c19ab75c54a685514500b10336c84d9b81a5c56)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/models/model_hg104_ajd.pth (2258.8 MB, sha1 1b3ac9a363a9095608c8f952c6f7ae82ccdc3a9a)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/models/model_hg104_cornell.pth (753.3 MB, sha1 4d9235d2f85a61e902cf643db275d8a327f9806a)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/models/model_hg52_ajd.pth (376.2 MB, sha1 d9b9478a1f60dc3cc19c0dc576c9da293d5e191c)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/models/model_hg52_cornell.pth (376.2 MB, sha1 6aeca75cc4276c05933451ea8a2500853ebcc3b8)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/models/model_resnet18_ajd.pth (63.7 MB, sha1 f9bb4a0dff92197b19ad9d387c209abf7164bf2d)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/models/model_resnet18_cornell.pth (63.7 MB, sha1 14a6be95db8b032bc3aac4f14720b752d10c5b88)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/models/model_resnet50_ajd.pth (135 MB, sha1 3220c95b4dfbf8f0b49370517fea8f071f19b838)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/models/model_resnet50_cornell.pth (135 MB, sha1 58ef7b9d0a990faa2be53e1293ec4a5c0f876cca)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/models/model_vgg16_ajd.pth (77.3 MB, sha1 ccc57a04f20b304a4661c03c5bcabbb849118984)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/models/model_vgg16_cornell.pth (77.3 MB, sha1 f0d3d5275b46c923c6b5bb48171577a51f746ac5)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/models/model_dla34_cornell.pth (80 MB, sha1 efa86a33444c59e3bde21e4314104d133cda310c)
+- Cornell_256_augmented (see https://www.dropbox.com/scl/fo/3tudn7uorjygxd040fcmn/h?dl=0)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/Cornell_256_augmented.tar.gz (5.8 GB, sha1 4385f47eeea8cc020f79f5c2248108606ea8b689)
+- Abridged_Jacquard (see https://smartech.gatech.edu/handle/1853/64897)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/Abridged_Jacquard/instances_grasps_test2018.json (0.2 GB, sha1 e9bf723f3f914872084d05ba93e27e21e431b968)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/Abridged_Jacquard/instances_grasps_test2018_edge_denseanno_filter.json (0.3 GB, sha1 cff53c4d152ea7a7eb1fca7b9169e52a10b6b9ef)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/Abridged_Jacquard/instances_grasps_test2018_filter.json (0.2 GB, sha1 98ebbf969f37367d7e48b4437db09b9888df99c1)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/Abridged_Jacquard/instances_grasps_train2018_edge_denseanno_filter.json (1.3 GB, sha1 86698fbfb3f5f34c4344a3684bd43809652f53f4)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/Abridged_Jacquard/instances_grasps_train2018_filter.json (1 GB, sha1 47c5e116bd0172faa8e21a092fa55f60ce000ecc)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/Abridged_Jacquard/README.doc (0 GB, sha1 52a60290b299945527b3fe05666eb1c269347468)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/Abridged_Jacquard/README.md (0 GB, sha1 42bb95169a14bdc4f3142b5197faae05507dbe7a)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/Abridged_Jacquard/test_annotations_0_5.tar.gz (0.1 GB, sha1 bb532c43b54ed0e81151cf73996eaaecfc764e23)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/Abridged_Jacquard/test_annotations_6_11.tar.gz (0.1 GB, sha1 3fb969ccad48ec1d3ae51ec64cbf3d7b2957ac4a)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/Abridged_Jacquard/train_annotations_0_5.tar.gz (0.6 GB, sha1 813402cbaf4e769514be37854a8e4863fdfd3d1e)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/Abridged_Jacquard/train_annotations_6_11.tar.gz (0.5 GB, sha1 f8ac6f5afedd2dc688b737bbc9f997cac393f68d)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/Abridged_Jacquard/train_grasps_test2018_0_5.tar.gz (0.7 GB, sha1 c17bd238fd900e2129190c8109d1c271f5ca7dbd)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/Abridged_Jacquard/train_grasps_test2018_6_11.tar.gz (0.7 GB, sha1 5d100cba83fa4ca35be9fd22d5f06c5530cbf55e)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/Abridged_Jacquard/train_grasps_train2018_0_5.tar.gz (2.9 GB, sha1 a7af278749ecfc84dfa1e3f63bf7303aa110b92c)
+  - https://f004.backblazeb2.com/file/acm-ivalab/GraspKpNet/Abridged_Jacquard/train_grasps_train2018_6_11.tar.gz (2.7 GB, sha1 595fa7429ef613ac3caa7550dbc5052455f03be7)
diff --git a/gknet/detectors/base_detector.py b/gknet/detectors/base_detector.py
@@ -73,13 +73,13 @@ def post_process(self, dets, meta, scale=1):
     def merge_outputs(self, detections):
         raise NotImplementedError
 
-    def debug(self, debugger, images, dets, output, scale=1):
+    def debug(self, debugger, images, dets, output, scale=1, save_dir=None):
         raise NotImplementedError
 
     def show_results(self, debugger, image, results):
         raise NotImplementedError
 
-    def run(self, image_or_path_or_tensor, meta=None):
+    def run(self, image_or_path_or_tensor, meta=None, save_dir=None):
         load_time, pre_time, net_time, dec_time, post_time = 0, 0, 0, 0, 0
         merge_time, tot_time = 0, 0
         debugger = Debugger(
@@ -126,7 +126,7 @@ def run(self, image_or_path_or_tensor, meta=None):
             dec_time += decode_time - forward_time
 
             if self.opt.debug >= 2:
-                self.debug(debugger, images, dets, output, scale)
+                self.debug(debugger, images, dets, output, scale, save_dir)
 
             if self.opt.task.split("_")[0] == "dbmctdet":
                 dets = self.post_process(dets, meta, scale)

diff --git a/gknet/detectors/dbmctdet.py b/gknet/detectors/dbmctdet.py
@@ -11,9 +11,6 @@
 
 
 class DbMCtdetDetector(BaseDetector):
-    def __init__(self, opt):
-        super(DbMCtdetDetector, self).__init__(opt)
-
     def process(self, images, return_time=False):
         with torch.no_grad():
             output = self.model(images)[-1]
@@ -92,28 +89,36 @@ def merge_outputs(self, detections):
                 results[j] = results[j][keep_inds]
         return results
 
-    def debug(self, debugger, images, dets, output, scale=1):
+    def debug(self, debugger, images, dets, output, scale=1, save_dir=None):
         detection = dets.detach().cpu().numpy().copy()
         detection[:, :, :4] *= self.opt.down_ratio
         for i in range(1):
             img = images[i].detach().cpu().numpy().transpose(1, 2, 0)
             img = ((img * self.std + self.mean) * 255).astype(np.uint8)
-            pred = debugger.gen_colormap(output["hm"][i].detach().cpu().numpy())
-            debugger.add_blend_img(img, pred, "pred_hm_{:.1f}".format(scale))
+            for key in ["lm", "rm", "ct"]:
+                pred = debugger.gen_colormap(output[key][i].detach().cpu().numpy())
+                debugger.add_blend_img(
+                    img, pred, "pred_{}_{:.1f}".format(key, scale)
+                )
             debugger.add_img(img, img_id="out_pred_{:.1f}".format(scale))
             for k in range(len(dets[i])):
-                if detection[i, k, 4] > self.opt.center_thresh:
+                if detection[i, k, 4] > self.opt.center_threshold:
                     debugger.add_coco_bbox(
                         detection[i, k, :4],
                         detection[i, k, -1],
                         detection[i, k, 4],
                         img_id="out_pred_{:.1f}".format(scale),
                     )
+        if save_dir:
+            debugger.save_all_imgs(save_dir)
 
     def show_results(self, debugger, image, results):
         debugger.add_img(image, img_id="ctdet")
         for j in range(1, self.num_classes + 1):
             for bbox in results[j]:
                 if bbox[4] > self.opt.vis_thresh:
                     debugger.add_coco_bbox(bbox[:4], j - 1, bbox[4], img_id="ctdet")
-        debugger.show_all_imgs(pause=self.pause)
+        try:
+            debugger.show_all_imgs(pause=self.pause)
+        except:
+            print("visual debugger does not work in headless mode")