Merge pull request #453 from josephvanpeltkw/dev/tensorrt_hands

Dev/tensorrt hands
PTG-Kitware · Oct 17, 2024 · fec1b91 · fec1b91
2 parents 298ce03 + 6d61f1c
commit fec1b91
Show file tree

Hide file tree

Showing 5 changed files with 127 additions and 41 deletions.
diff --git a/ansible/roles/provision-files/vars/main.yml b/ansible/roles/provision-files/vars/main.yml
@@ -61,10 +61,10 @@ girder_file_downloads:
       sha512: 7ded5cb71ef0efe2444a19c7452073c63a34d65e3038b24236b304836abd47e8911e3d3023d2ba7e5a8d8c1e337e6f4b8e3080709827eb2b0daaedb0b6248561
       dest: "{{ stage_dirs.pose_estimation }}/pose_model.pth"
     # Hand detector
-    - file_id: 6605ca2e8b763ca20ae99f77
-      sha512: 42b56f34c0c443ad00d71e66334d6e852a7811011b22a60da9b5e2721565c21d1b5daad17c8a708908c9658ded81a663729090526feb5286bd0a591cf4fb5d8f
-      dest: "{{ stage_dirs.object_detector }}/hands_model.pt"
-    
+    - file_id: 6710037cba16447d41e1663a
+      sha512: cf6a20a966dcaee46354545ef5d2086f2dcf31112eb90282e5dd6a24824a584d5e549274112f6f028911050a48bcaba6c944e5d3d870edcafe2fcd34e61d95f5
+      dest: "{{ stage_dirs.object_detector }}/hands_model.engine"
+
     # ---- M2 ----
     # Object detector
     - file_id: 6617e7ce2357cf6b55ca8dd6
@@ -98,7 +98,7 @@ girder_file_downloads:
     - file_id: 660ebb1baa5c8de3874c43c3
       sha512: fae5d696d6b1a03daa52a8094b8be8f6cc675cd3487270dbfee212551813869af22eedbcc1dc89bd62bbebdfa411819a8584083e198c3d6ab0b35f266a524054
       dest: "{{ stage_dirs.task_monitor }}/m3_test_activity_preds.mscoco.json"
-    
+
     # ---- M5 ----
     # Object detector
     - file_id: 6605e0808b763ca20ae99f92
@@ -115,7 +115,7 @@ girder_file_downloads:
     - file_id: 660ebb1caa5c8de3874c43c6
       sha512: 8e2f77310455460b9e0769d78469d752fca2313a31ed498ddd4eeb4b6b501cbee6f069e4f3c9921f6db2e77f07a2e2ad47714afb80617d08df775bcdcbd15168
       dest: "{{ stage_dirs.task_monitor }}/m5_test_activity_preds.mscoco.json"
-    
+
     # ---- R18 ----
     # Object detector
     - file_id: 6605e0958b763ca20ae99f95
@@ -132,7 +132,7 @@ girder_file_downloads:
     - file_id: 66464bf9687336214e7cdeae
       sha512: bc7884c258cb8704372dd69a3e7d999eaf70d716d271982077c7216ef27ab01219ef1e488f6f9669e11a7e8aa6ffb9d8e07d74edc47294f90cc9e6a723f5a822
       dest: "{{ stage_dirs.task_monitor }}/global_step_predictor_act_avgs_R18.npy"
-    
+
 # List of git repositories to check out at a specific ref and then archive.
 # Destination files will be written as GZipped TAR files, so please suffix
 # appropriately.

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -37,7 +37,7 @@ timm = "^0.5.4"
 transitions = ">=0.8.11"
 trimesh = {extras = ["easy"], version = "^3.10.2"}
 moviepy = "==1.0.3"
-ultralytics = "==8.1.27"
+ultralytics = "==8.3.9"
 rootutils = "==1.0.7"
 torchmetrics = "==0.11.4"
 rich = "==13.7.1"
@@ -48,7 +48,10 @@ mmcv = "==1.3.8"
 torch = {url = "https://download.pytorch.org/whl/cu111/torch-1.9.1%2Bcu111-cp38-cp38-linux_x86_64.whl"}
 torchvision = {url = "https://download.pytorch.org/whl/cu111/torchvision-0.10.1%2Bcu111-cp38-cp38-linux_x86_64.whl"}
 detectron2 = {url = "https://dl.fbaipublicfiles.com/detectron2/wheels/cu111/torch1.9/detectron2-0.6%2Bcu111-cp38-cp38-linux_x86_64.whl"}
-
+tensorrt = "==10.1.0"
+tensorrt-cu12 = "==10.1.0"
+tensorrt-cu12-bindings = "==10.1.0"
+tensorrt-cu12-libs = "==10.1.0"
 
 ## For UHO Activity Classifier
 ## - see arisia: /data/dawei.du/ptg-activity-recognition/activity_hydra/requirements.txt

diff --git a/ros/angel_system_nodes/angel_system_nodes/object_detection/object_and_hand_detection.py b/ros/angel_system_nodes/angel_system_nodes/object_detection/object_and_hand_detection.py
@@ -136,9 +136,9 @@ def __init__(self):
 
         ##########################################
         # Hand model
-        self.hand_model = YOLOv8(self._hand_model_chpt_fp)
+        self.hand_model = YOLOv8(self._hand_model_chpt_fp, task="detect")
         log.info(
-            f"Loaded hand model with classes:\n"
+            "Loaded hand model with classes:\n"
             + "\n".join(f'\t- "{n}"' for n in self.hand_model.names)
         )
 

diff --git a/tmux/demos/medical/Kitware-R18.yml b/tmux/demos/medical/Kitware-R18.yml
@@ -105,7 +105,7 @@ windows:
       -p det_topic:=ObjectDetections2d
       -p object_net_checkpoint:=${MODEL_DIR}/object_detector/r18_det.pt
       -p inference_img_size:=768
-      -p hand_net_checkpoint:=${MODEL_DIR}/object_detector/hands_model.pt
+      -p hand_net_checkpoint:=${MODEL_DIR}/object_detector/hands_model.engine
       -p cuda_device_id:=0
 
   - activity_classifier: ros2 run angel_system_nodes activity_classifier_tcn --ros-args