From dea8c716c5a8bbb3469e80718ebe3bf9810c2c9a Mon Sep 17 00:00:00 2001
From: Philipp Donn <30521025+phinik@users.noreply.github.com>
Date: Tue, 11 Apr 2023 21:16:01 +0200
Subject: [PATCH 1/8] nms across multiple robot classes

---
 yoeo/detect.py      | 40 +++++++++++++++++++++++++++++++++-------
 yoeo/test.py        | 28 ++++++++++++++++++++++++----
 yoeo/train.py       | 11 ++++++++++-
 yoeo/utils/utils.py | 12 ++++++++++--
 4 files changed, 77 insertions(+), 14 deletions(-)

diff --git a/yoeo/detect.py b/yoeo/detect.py
index 76163fc..d362913 100755
--- a/yoeo/detect.py
+++ b/yoeo/detect.py
@@ -26,7 +26,8 @@
 
 
 def detect_directory(model_path, weights_path, img_path, classes, output_path,
-                     batch_size=8, img_size=416, n_cpu=8, conf_thres=0.5, nms_thres=0.5):
+                     batch_size=8, img_size=416, n_cpu=8, conf_thres=0.5, nms_thres=0.5,
+                     multi_robot=False, first_robot_id=0):
     """Detects objects on all images in specified directory and saves output images with drawn detections.
 
     :param model_path: Path to model definition file (.cfg)
@@ -49,6 +50,10 @@ def detect_directory(model_path, weights_path, img_path, classes, output_path,
     :type conf_thres: float, optional
     :param nms_thres: IOU threshold for non-maximum suppression, defaults to 0.5
     :type nms_thres: float, optional
+    :param multi_robot: set to 'True' if multiple robot classes exist and nms shall be performed across all classes.
+    :type multi_robot: bool, optional
+    :param first_robot_id: first class ID of robot classes. Only effective if multi_robot=True.
+    :type first_robot_id: int, optional
     """
     dataloader = _create_data_loader(img_path, batch_size, img_size, n_cpu)
     model = load_model(model_path, weights_path)
@@ -58,14 +63,17 @@ def detect_directory(model_path, weights_path, img_path, classes, output_path,
         dataloader,
         output_path,
         conf_thres,
-        nms_thres)
+        nms_thres,
+        multi_robot,
+        first_robot_id
+    )
     _draw_and_save_output_images(
         img_detections, segmentations, imgs, img_size, output_path, classes)
 
     print(f"---- Detections were saved to: '{output_path}' ----")
 
 
-def detect_image(model, image, img_size=416, conf_thres=0.5, nms_thres=0.5):
+def detect_image(model, image, img_size=416, conf_thres=0.5, nms_thres=0.5, multi_robot=False, first_robot_id=0):
     """Inferences one image with model.
 
     :param model: Model for inference
@@ -78,6 +86,10 @@ def detect_image(model, image, img_size=416, conf_thres=0.5, nms_thres=0.5):
     :type conf_thres: float, optional
     :param nms_thres: IOU threshold for non-maximum suppression, defaults to 0.5
     :type nms_thres: float, optional
+    :param multi_robot: set to 'True' if multiple robot classes exist and nms shall be performed across all classes.
+    :type multi_robot: bool, optional
+    :param first_robot_id: first class ID of robot classes. Only effective if multi_robot=True.
+    :type first_robot_id: int, optional
     :return: Detections on image with each detection in the format: [x1, y1, x2, y2, confidence, class], Segmentation as 2d numpy array with the coresponding class id in each cell
     :rtype: nd.array, nd.array
     """
@@ -97,13 +109,13 @@ def detect_image(model, image, img_size=416, conf_thres=0.5, nms_thres=0.5):
     # Get detections
     with torch.no_grad():
         detections, segmentations = model(input_img)
-        detections = non_max_suppression(detections, conf_thres, nms_thres)
+        detections = non_max_suppression(detections, conf_thres, nms_thres, multi_robot, first_robot_id)
         detections = rescale_boxes(detections[0], img_size, image.shape[0:2])
         segmentations = rescale_segmentation(segmentations, image.shape[0:2])
     return detections.numpy(), segmentations.cpu().detach().numpy()
 
 
-def detect(model, dataloader, output_path, conf_thres, nms_thres):
+def detect(model, dataloader, output_path, conf_thres, nms_thres, multi_robot=False, first_robot_id=0):
     """Inferences images with model.
 
     :param model: Model for inference
@@ -116,6 +128,10 @@ def detect(model, dataloader, output_path, conf_thres, nms_thres):
     :type conf_thres: float, optional
     :param nms_thres: IOU threshold for non-maximum suppression, defaults to 0.5
     :type nms_thres: float, optional
+    :param multi_robot: set to 'True' if multiple robot classes exist and nms shall be performed across all classes.
+    :type multi_robot: bool, optional
+    :param first_robot_id: first class ID of robot classes. Only effective if multi_robot=True.
+    :type first_robot_id: int, optional
     :return: List of detections. The coordinates are given for the padded image that is provided by the dataloader.
         Use `utils.rescale_boxes` to transform them into the desired input image coordinate system before its transformed by the dataloader),
         List of input image paths
@@ -139,7 +155,7 @@ def detect(model, dataloader, output_path, conf_thres, nms_thres):
         # Get detections
         with torch.no_grad():
             detections, segmentations = model(input_imgs)
-            detections = non_max_suppression(detections, conf_thres, nms_thres)
+            detections = non_max_suppression(detections, conf_thres, nms_thres, multi_robot, first_robot_id)
 
         # Store image and detections
         img_detections.extend(detections)
@@ -300,12 +316,19 @@ def run():
     parser.add_argument("--n_cpu", type=int, default=8, help="Number of cpu threads to use during batch generation")
     parser.add_argument("--conf_thres", type=float, default=0.5, help="Object confidence threshold")
     parser.add_argument("--nms_thres", type=float, default=0.4, help="IOU threshold for non-maximum suppression")
+    parser.add_argument("--multiple_robot_classes", type=bool, default=False, help="Set to True if multiple robot classes exist and nms shall be performed across all robot classes")
     args = parser.parse_args()
     print(f"Command line arguments: {args}")
 
     # Extract class names from file
     classes = load_classes(args.classes)['detection']  # List of class names
 
+    first_robot_class_id = -1
+    for idx, c in enumerate(classes):
+        if "robot" in c:
+            first_robot_class_id = idx
+            break
+
     detect_directory(
         args.model,
         args.weights,
@@ -316,7 +339,10 @@ def run():
         img_size=args.img_size,
         n_cpu=args.n_cpu,
         conf_thres=args.conf_thres,
-        nms_thres=args.nms_thres)
+        nms_thres=args.nms_thres,
+        multi_robot=args.multiple_robot_classes,
+        first_robot_id=first_robot_class_id
+    )
 
 
 if __name__ == '__main__':
diff --git a/yoeo/test.py b/yoeo/test.py
index 4106ec5..a9cedfb 100755
--- a/yoeo/test.py
+++ b/yoeo/test.py
@@ -21,7 +21,8 @@
 
 
 def evaluate_model_file(model_path, weights_path, img_path, class_names, batch_size=8, img_size=416,
-                        n_cpu=8, iou_thres=0.5, conf_thres=0.5, nms_thres=0.5, verbose=True):
+                        n_cpu=8, iou_thres=0.5, conf_thres=0.5, nms_thres=0.5, verbose=True,
+                        multi_robot=False, first_robot_id=0):
     """Evaluate model on validation dataset.
 
     :param model_path: Path to model definition file (.cfg)
@@ -46,6 +47,10 @@ def evaluate_model_file(model_path, weights_path, img_path, class_names, batch_s
     :type nms_thres: float, optional
     :param verbose: If True, prints stats of model, defaults to True
     :type verbose: bool, optional
+    :param multi_robot: set to 'True' if multiple robot classes exist and nms shall be performed across all classes.
+    :type multi_robot: bool, optional
+    :param first_robot_id: first class ID of robot classes. Only effective if multi_robot=True.
+    :type first_robot_id: int, optional
     :return: Returns precision, recall, AP, f1, ap_class
     """
     dataloader = _create_validation_data_loader(
@@ -90,7 +95,7 @@ def print_eval_stats(metrics_output, seg_class_ious, class_names, verbose):
     print(f"----Average IoU {mean_seg_class_ious:.5f} ----")
 
 
-def _evaluate(model, dataloader, class_names, img_size, iou_thres, conf_thres, nms_thres, verbose):
+def _evaluate(model, dataloader, class_names, img_size, iou_thres, conf_thres, nms_thres, verbose, multi_robot=False, first_robot_id=0):
     """Evaluate model on validation dataset.
 
     :param model: Model to evaluate
@@ -109,6 +114,10 @@ def _evaluate(model, dataloader, class_names, img_size, iou_thres, conf_thres, n
     :type nms_thres: float
     :param verbose: If True, prints stats of model
     :type verbose: bool
+    :param multi_robot: set to 'True' if multiple robot classes exist and nms shall be performed across all classes.
+    :type multi_robot: bool, optional
+    :param first_robot_id: first class ID of robot classes. Only effective if multi_robot=True.
+    :type first_robot_id: int, optional
     :return: Returns precision, recall, AP, f1, ap_class
     """
     model.eval()  # Set model to evaluation mode
@@ -133,7 +142,7 @@ def _evaluate(model, dataloader, class_names, img_size, iou_thres, conf_thres, n
             t1 = time.time()
             yolo_outputs, segmentation_outputs = model(imgs)
             times.append(time.time() - t1)
-            yolo_outputs = non_max_suppression(yolo_outputs, conf_thres=conf_thres, iou_thres=nms_thres)
+            yolo_outputs = non_max_suppression(yolo_outputs, conf_thres=conf_thres, iou_thres=nms_thres, multi_robot=multi_robot, first_robot_id=first_robot_id)
 
         sample_metrics += get_batch_statistics(yolo_outputs, bb_targets, iou_threshold=iou_thres)
 
@@ -207,6 +216,8 @@ def run():
     parser.add_argument("--iou_thres", type=float, default=0.5, help="IOU threshold required to qualify as detected")
     parser.add_argument("--conf_thres", type=float, default=0.01, help="Object confidence threshold")
     parser.add_argument("--nms_thres", type=float, default=0.4, help="IOU threshold for non-maximum suppression")
+    parser.add_argument("--multiple_robot_classes", type=bool, default=False, help="Set to True if multiple robot classes exist and nms shall be performed across all robot classes")
+
     args = parser.parse_args()
     print(f"Command line arguments: {args}")
 
@@ -216,6 +227,12 @@ def run():
     valid_path = data_config["valid"]
     class_names = load_classes(data_config["names"])  # Detection and segmentation class names
 
+    first_robot_class_id = -1
+    for idx, c in enumerate(class_names["detection"]):
+        if "robot" in c:
+            first_robot_class_id = idx
+            break
+
     evaluate_model_file(
         args.model,
         args.weights,
@@ -227,7 +244,10 @@ def run():
         iou_thres=args.iou_thres,
         conf_thres=args.conf_thres,
         nms_thres=args.nms_thres,
-        verbose=True)
+        verbose=True,
+        multi_robot=args.multiple_robot_classes,
+        first_robot_id=first_robot_class_id
+    )
 
 
 if __name__ == "__main__":
diff --git a/yoeo/train.py b/yoeo/train.py
index d8afdbf..a6f3bae 100755
--- a/yoeo/train.py
+++ b/yoeo/train.py
@@ -78,6 +78,7 @@ def run():
     parser.add_argument("--nms_thres", type=float, default=0.5, help="Evaluation: IOU threshold for non-maximum suppression")
     parser.add_argument("--logdir", type=str, default="logs", help="Directory for training log files (e.g. for TensorBoard)")
     parser.add_argument("--seed", type=int, default=-1, help="Makes results reproducable. Set -1 to disable.")
+    parser.add_argument("--multiple_robot_classes", type=bool, default=False, help="Set to True if multiple robot classes exist and nms shall be performed across all robot classes")
     args = parser.parse_args()
     print(f"Command line arguments: {args}")
 
@@ -95,6 +96,12 @@ def run():
     train_path = data_config["train"]
     valid_path = data_config["valid"]
     class_names = load_classes(data_config["names"])
+    first_robot_class_id = -1
+    for idx, c in enumerate(class_names["detection"]):
+        if "robot" in c:
+            first_robot_class_id = idx
+            break
+
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     # ############
@@ -249,7 +256,9 @@ def run():
                 iou_thres=args.iou_thres,
                 conf_thres=args.conf_thres,
                 nms_thres=args.nms_thres,
-                verbose=args.verbose
+                verbose=args.verbose,
+                multi_robot=args.multiple_robot_classes,
+                first_robot_id=first_robot_class_id
             )
 
             if metrics_output is not None:
diff --git a/yoeo/utils/utils.py b/yoeo/utils/utils.py
index c384d6b..55b908e 100644
--- a/yoeo/utils/utils.py
+++ b/yoeo/utils/utils.py
@@ -418,7 +418,7 @@ def box_area(box):
     return inter / (area1[:, None] + area2 - inter)
 
 
-def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None):
+def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, multi_robot=False, first_robot_id=0):
     """Performs Non-Maximum Suppression (NMS) on inference results
     Returns:
          detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
@@ -473,7 +473,15 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
             x = x[x[:, 4].argsort(descending=True)[:max_nms]]
 
         # Batched NMS
-        c = x[:, 5:6] * max_wh  # classes
+        if not multi_robot:
+            c = x[:, 5:6] * max_wh  # classes
+        else:
+            # If multiple robot classes are present, all robot classes are treated as one class in order to perform
+            # nms across all classes and not per class. For this, all robot classes get the same offset.
+            c = torch.clone(x[:, 5:6])
+            c[c > first_robot_id] = first_robot_id
+            c *= max_wh
+
         # boxes (offset by class), scores
         boxes, scores = x[:, :4] + c, x[:, 4]
         i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS

From 1f0a28179c35d4578c4b7b08ee00a41a1967e6ea Mon Sep 17 00:00:00 2001
From: Philipp Donn <30521025+phinik@users.noreply.github.com>
Date: Tue, 11 Apr 2023 21:32:15 +0200
Subject: [PATCH 2/8] fix argparse parameter

---
 yoeo/detect.py | 2 +-
 yoeo/test.py   | 2 +-
 yoeo/train.py  | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/yoeo/detect.py b/yoeo/detect.py
index d362913..df8b3fa 100755
--- a/yoeo/detect.py
+++ b/yoeo/detect.py
@@ -316,7 +316,7 @@ def run():
     parser.add_argument("--n_cpu", type=int, default=8, help="Number of cpu threads to use during batch generation")
     parser.add_argument("--conf_thres", type=float, default=0.5, help="Object confidence threshold")
     parser.add_argument("--nms_thres", type=float, default=0.4, help="IOU threshold for non-maximum suppression")
-    parser.add_argument("--multiple_robot_classes", type=bool, default=False, help="Set to True if multiple robot classes exist and nms shall be performed across all robot classes")
+    parser.add_argument("--multiple_robot_classes", action="store_true", help="Set to True if multiple robot classes exist and nms shall be performed across all robot classes")
     args = parser.parse_args()
     print(f"Command line arguments: {args}")
 
diff --git a/yoeo/test.py b/yoeo/test.py
index a9cedfb..9c8e511 100755
--- a/yoeo/test.py
+++ b/yoeo/test.py
@@ -216,7 +216,7 @@ def run():
     parser.add_argument("--iou_thres", type=float, default=0.5, help="IOU threshold required to qualify as detected")
     parser.add_argument("--conf_thres", type=float, default=0.01, help="Object confidence threshold")
     parser.add_argument("--nms_thres", type=float, default=0.4, help="IOU threshold for non-maximum suppression")
-    parser.add_argument("--multiple_robot_classes", type=bool, default=False, help="Set to True if multiple robot classes exist and nms shall be performed across all robot classes")
+    parser.add_argument("--multiple_robot_classes", action="store_true", help="Set to True if multiple robot classes exist and nms shall be performed across all robot classes")
 
     args = parser.parse_args()
     print(f"Command line arguments: {args}")
diff --git a/yoeo/train.py b/yoeo/train.py
index a6f3bae..3335aeb 100755
--- a/yoeo/train.py
+++ b/yoeo/train.py
@@ -78,7 +78,7 @@ def run():
     parser.add_argument("--nms_thres", type=float, default=0.5, help="Evaluation: IOU threshold for non-maximum suppression")
     parser.add_argument("--logdir", type=str, default="logs", help="Directory for training log files (e.g. for TensorBoard)")
     parser.add_argument("--seed", type=int, default=-1, help="Makes results reproducable. Set -1 to disable.")
-    parser.add_argument("--multiple_robot_classes", type=bool, default=False, help="Set to True if multiple robot classes exist and nms shall be performed across all robot classes")
+    parser.add_argument("--multiple_robot_classes", action="store_true", help="Set to True if multiple robot classes exist and nms shall be performed across all robot classes")
     args = parser.parse_args()
     print(f"Command line arguments: {args}")
 

From b353c16e9389411cf8c3b3f3cf6aa3de66f6f215 Mon Sep 17 00:00:00 2001
From: Philipp Donn <30521025+phinik@users.noreply.github.com>
Date: Tue, 11 Apr 2023 21:41:14 +0200
Subject: [PATCH 3/8] fix missing arguments

---
 yoeo/test.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/yoeo/test.py b/yoeo/test.py
index 9c8e511..3b52000 100755
--- a/yoeo/test.py
+++ b/yoeo/test.py
@@ -64,7 +64,9 @@ def evaluate_model_file(model_path, weights_path, img_path, class_names, batch_s
         iou_thres,
         conf_thres,
         nms_thres,
-        verbose)
+        verbose,
+        multi_robot=multi_robot,
+        first_robot_id=first_robot_id)
     return metrics_output, seg_class_ious
 
 

From 06e4dad5e18ceaaf5b916c6b340db39ac8cf2ba8 Mon Sep 17 00:00:00 2001
From: Philipp Donn <30521025+phinik@users.noreply.github.com>
Date: Tue, 11 Apr 2023 21:44:29 +0200
Subject: [PATCH 4/8] change argument description

---
 yoeo/detect.py | 2 +-
 yoeo/test.py   | 2 +-
 yoeo/train.py  | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/yoeo/detect.py b/yoeo/detect.py
index df8b3fa..1f355cf 100755
--- a/yoeo/detect.py
+++ b/yoeo/detect.py
@@ -316,7 +316,7 @@ def run():
     parser.add_argument("--n_cpu", type=int, default=8, help="Number of cpu threads to use during batch generation")
     parser.add_argument("--conf_thres", type=float, default=0.5, help="Object confidence threshold")
     parser.add_argument("--nms_thres", type=float, default=0.4, help="IOU threshold for non-maximum suppression")
-    parser.add_argument("--multiple_robot_classes", action="store_true", help="Set to True if multiple robot classes exist and nms shall be performed across all robot classes")
+    parser.add_argument("--multiple_robot_classes", action="store_true", help="If multiple robot classes exist and nms shall be performed across all robot classes")
     args = parser.parse_args()
     print(f"Command line arguments: {args}")
 
diff --git a/yoeo/test.py b/yoeo/test.py
index 3b52000..a3aa1d1 100755
--- a/yoeo/test.py
+++ b/yoeo/test.py
@@ -218,7 +218,7 @@ def run():
     parser.add_argument("--iou_thres", type=float, default=0.5, help="IOU threshold required to qualify as detected")
     parser.add_argument("--conf_thres", type=float, default=0.01, help="Object confidence threshold")
     parser.add_argument("--nms_thres", type=float, default=0.4, help="IOU threshold for non-maximum suppression")
-    parser.add_argument("--multiple_robot_classes", action="store_true", help="Set to True if multiple robot classes exist and nms shall be performed across all robot classes")
+    parser.add_argument("--multiple_robot_classes", action="store_true", help="If multiple robot classes exist and nms shall be performed across all robot classes")
 
     args = parser.parse_args()
     print(f"Command line arguments: {args}")
diff --git a/yoeo/train.py b/yoeo/train.py
index 3335aeb..2bdaa32 100755
--- a/yoeo/train.py
+++ b/yoeo/train.py
@@ -78,7 +78,7 @@ def run():
     parser.add_argument("--nms_thres", type=float, default=0.5, help="Evaluation: IOU threshold for non-maximum suppression")
     parser.add_argument("--logdir", type=str, default="logs", help="Directory for training log files (e.g. for TensorBoard)")
     parser.add_argument("--seed", type=int, default=-1, help="Makes results reproducable. Set -1 to disable.")
-    parser.add_argument("--multiple_robot_classes", action="store_true", help="Set to True if multiple robot classes exist and nms shall be performed across all robot classes")
+    parser.add_argument("--multiple_robot_classes", action="store_true", help="If multiple robot classes exist and nms shall be performed across all robot classes")
     args = parser.parse_args()
     print(f"Command line arguments: {args}")
 

From 73818123b27120c2b6a26c29e6cbb54b02bfb629 Mon Sep 17 00:00:00 2001
From: Philipp Donn <30521025+phinik@users.noreply.github.com>
Date: Thu, 13 Apr 2023 20:04:43 +0200
Subject: [PATCH 5/8] multi robot nms based on list of class ids

---
 yoeo/detect.py      | 48 ++++++++++++++++--------------------
 yoeo/test.py        | 60 ++++++++++++++++++++++++---------------------
 yoeo/train.py       | 18 ++++++++------
 yoeo/utils/utils.py |  7 +++---
 4 files changed, 67 insertions(+), 66 deletions(-)

diff --git a/yoeo/detect.py b/yoeo/detect.py
index 1f355cf..a60ce5c 100755
--- a/yoeo/detect.py
+++ b/yoeo/detect.py
@@ -27,7 +27,7 @@
 
 def detect_directory(model_path, weights_path, img_path, classes, output_path,
                      batch_size=8, img_size=416, n_cpu=8, conf_thres=0.5, nms_thres=0.5,
-                     multi_robot=False, first_robot_id=0):
+                     robot_class_ids: Optional[List[int]] = None):
     """Detects objects on all images in specified directory and saves output images with drawn detections.
 
     :param model_path: Path to model definition file (.cfg)
@@ -50,10 +50,8 @@ def detect_directory(model_path, weights_path, img_path, classes, output_path,
     :type conf_thres: float, optional
     :param nms_thres: IOU threshold for non-maximum suppression, defaults to 0.5
     :type nms_thres: float, optional
-    :param multi_robot: set to 'True' if multiple robot classes exist and nms shall be performed across all classes.
-    :type multi_robot: bool, optional
-    :param first_robot_id: first class ID of robot classes. Only effective if multi_robot=True.
-    :type first_robot_id: int, optional
+    :param robot_class_ids: List of class IDs of robot classes if multiple robot classes exist.
+    :type robot_class_ids: List[int], optional
     """
     dataloader = _create_data_loader(img_path, batch_size, img_size, n_cpu)
     model = load_model(model_path, weights_path)
@@ -64,8 +62,7 @@ def detect_directory(model_path, weights_path, img_path, classes, output_path,
         output_path,
         conf_thres,
         nms_thres,
-        multi_robot,
-        first_robot_id
+        robot_class_ids=robot_class_ids
     )
     _draw_and_save_output_images(
         img_detections, segmentations, imgs, img_size, output_path, classes)
@@ -73,7 +70,7 @@ def detect_directory(model_path, weights_path, img_path, classes, output_path,
     print(f"---- Detections were saved to: '{output_path}' ----")
 
 
-def detect_image(model, image, img_size=416, conf_thres=0.5, nms_thres=0.5, multi_robot=False, first_robot_id=0):
+def detect_image(model, image, img_size=416, conf_thres=0.5, nms_thres=0.5, robot_class_ids: Optional[List[int]] = None):
     """Inferences one image with model.
 
     :param model: Model for inference
@@ -86,10 +83,8 @@ def detect_image(model, image, img_size=416, conf_thres=0.5, nms_thres=0.5, mult
     :type conf_thres: float, optional
     :param nms_thres: IOU threshold for non-maximum suppression, defaults to 0.5
     :type nms_thres: float, optional
-    :param multi_robot: set to 'True' if multiple robot classes exist and nms shall be performed across all classes.
-    :type multi_robot: bool, optional
-    :param first_robot_id: first class ID of robot classes. Only effective if multi_robot=True.
-    :type first_robot_id: int, optional
+    :param robot_class_ids: List of class IDs of robot classes if multiple robot classes exist.
+    :type robot_class_ids: List[int], optional
     :return: Detections on image with each detection in the format: [x1, y1, x2, y2, confidence, class], Segmentation as 2d numpy array with the coresponding class id in each cell
     :rtype: nd.array, nd.array
     """
@@ -109,13 +104,13 @@ def detect_image(model, image, img_size=416, conf_thres=0.5, nms_thres=0.5, mult
     # Get detections
     with torch.no_grad():
         detections, segmentations = model(input_img)
-        detections = non_max_suppression(detections, conf_thres, nms_thres, multi_robot, first_robot_id)
+        detections = non_max_suppression(detections, conf_thres, nms_thres, robot_class_ids=robot_class_ids)
         detections = rescale_boxes(detections[0], img_size, image.shape[0:2])
         segmentations = rescale_segmentation(segmentations, image.shape[0:2])
     return detections.numpy(), segmentations.cpu().detach().numpy()
 
 
-def detect(model, dataloader, output_path, conf_thres, nms_thres, multi_robot=False, first_robot_id=0):
+def detect(model, dataloader, output_path, conf_thres, nms_thres, robot_class_ids: Optional[List[int]] = None):
     """Inferences images with model.
 
     :param model: Model for inference
@@ -128,10 +123,8 @@ def detect(model, dataloader, output_path, conf_thres, nms_thres, multi_robot=Fa
     :type conf_thres: float, optional
     :param nms_thres: IOU threshold for non-maximum suppression, defaults to 0.5
     :type nms_thres: float, optional
-    :param multi_robot: set to 'True' if multiple robot classes exist and nms shall be performed across all classes.
-    :type multi_robot: bool, optional
-    :param first_robot_id: first class ID of robot classes. Only effective if multi_robot=True.
-    :type first_robot_id: int, optional
+    :param robot_class_ids: List of class IDs of robot classes if multiple robot classes exist.
+    :type robot_class_ids: List[int], optional
     :return: List of detections. The coordinates are given for the padded image that is provided by the dataloader.
         Use `utils.rescale_boxes` to transform them into the desired input image coordinate system before its transformed by the dataloader),
         List of input image paths
@@ -155,7 +148,7 @@ def detect(model, dataloader, output_path, conf_thres, nms_thres, multi_robot=Fa
         # Get detections
         with torch.no_grad():
             detections, segmentations = model(input_imgs)
-            detections = non_max_suppression(detections, conf_thres, nms_thres, multi_robot, first_robot_id)
+            detections = non_max_suppression(detections, conf_thres, nms_thres, robot_class_ids=robot_class_ids)
 
         # Store image and detections
         img_detections.extend(detections)
@@ -316,18 +309,20 @@ def run():
     parser.add_argument("--n_cpu", type=int, default=8, help="Number of cpu threads to use during batch generation")
     parser.add_argument("--conf_thres", type=float, default=0.5, help="Object confidence threshold")
     parser.add_argument("--nms_thres", type=float, default=0.4, help="IOU threshold for non-maximum suppression")
-    parser.add_argument("--multiple_robot_classes", action="store_true", help="If multiple robot classes exist and nms shall be performed across all robot classes")
+    parser.add_argument("--multiple_robot_classes", action="store_true",
+                        help="If multiple robot classes exist and nms shall be performed across all robot classes")
     args = parser.parse_args()
     print(f"Command line arguments: {args}")
 
     # Extract class names from file
     classes = load_classes(args.classes)['detection']  # List of class names
 
-    first_robot_class_id = -1
-    for idx, c in enumerate(classes):
-        if "robot" in c:
-            first_robot_class_id = idx
-            break
+    robot_class_ids = None
+    if args.multiple_robot_classes:
+        robot_class_ids = []
+        for idx, c in enumerate(classes):
+            if "robot" in c:
+                robot_class_ids.append(idx)
 
     detect_directory(
         args.model,
@@ -340,8 +335,7 @@ def run():
         n_cpu=args.n_cpu,
         conf_thres=args.conf_thres,
         nms_thres=args.nms_thres,
-        multi_robot=args.multiple_robot_classes,
-        first_robot_id=first_robot_class_id
+        robot_class_ids=robot_class_ids
     )
 
 
diff --git a/yoeo/test.py b/yoeo/test.py
index a3aa1d1..c31a5d7 100755
--- a/yoeo/test.py
+++ b/yoeo/test.py
@@ -14,7 +14,8 @@
 from torch.autograd import Variable
 
 from yoeo.models import load_model
-from yoeo.utils.utils import load_classes, ap_per_class, get_batch_statistics, non_max_suppression, to_cpu, xywh2xyxy, print_environment_info, seg_iou
+from yoeo.utils.utils import load_classes, ap_per_class, get_batch_statistics, non_max_suppression, to_cpu, xywh2xyxy, \
+    print_environment_info, seg_iou
 from yoeo.utils.datasets import ListDataset
 from yoeo.utils.transforms import DEFAULT_TRANSFORMS
 from yoeo.utils.parse_config import parse_data_config
@@ -22,7 +23,7 @@
 
 def evaluate_model_file(model_path, weights_path, img_path, class_names, batch_size=8, img_size=416,
                         n_cpu=8, iou_thres=0.5, conf_thres=0.5, nms_thres=0.5, verbose=True,
-                        multi_robot=False, first_robot_id=0):
+                        robot_class_ids: Optional[List[int]] = None):
     """Evaluate model on validation dataset.
 
     :param model_path: Path to model definition file (.cfg)
@@ -47,10 +48,8 @@ def evaluate_model_file(model_path, weights_path, img_path, class_names, batch_s
     :type nms_thres: float, optional
     :param verbose: If True, prints stats of model, defaults to True
     :type verbose: bool, optional
-    :param multi_robot: set to 'True' if multiple robot classes exist and nms shall be performed across all classes.
-    :type multi_robot: bool, optional
-    :param first_robot_id: first class ID of robot classes. Only effective if multi_robot=True.
-    :type first_robot_id: int, optional
+    :param robot_class_ids: List of class IDs of robot classes if multiple robot classes exist.
+    :type robot_class_ids: List[int], optional
     :return: Returns precision, recall, AP, f1, ap_class
     """
     dataloader = _create_validation_data_loader(
@@ -65,8 +64,7 @@ def evaluate_model_file(model_path, weights_path, img_path, class_names, batch_s
         conf_thres,
         nms_thres,
         verbose,
-        multi_robot=multi_robot,
-        first_robot_id=first_robot_id)
+        robot_class_ids=robot_class_ids)
     return metrics_output, seg_class_ious
 
 
@@ -84,7 +82,6 @@ def print_eval_stats(metrics_output, seg_class_ious, class_names, verbose):
     else:
         print("---- mAP not measured (no detections found by model) ----")
 
-
     # Print segmentation statistics
     if verbose:
         # Print IoU per segmentation class
@@ -97,7 +94,8 @@ def print_eval_stats(metrics_output, seg_class_ious, class_names, verbose):
     print(f"----Average IoU {mean_seg_class_ious:.5f} ----")
 
 
-def _evaluate(model, dataloader, class_names, img_size, iou_thres, conf_thres, nms_thres, verbose, multi_robot=False, first_robot_id=0):
+def _evaluate(model, dataloader, class_names, img_size, iou_thres, conf_thres, nms_thres, verbose,
+              robot_class_ids: Optional[List[int]] = None):
     """Evaluate model on validation dataset.
 
     :param model: Model to evaluate
@@ -116,10 +114,8 @@ def _evaluate(model, dataloader, class_names, img_size, iou_thres, conf_thres, n
     :type nms_thres: float
     :param verbose: If True, prints stats of model
     :type verbose: bool
-    :param multi_robot: set to 'True' if multiple robot classes exist and nms shall be performed across all classes.
-    :type multi_robot: bool, optional
-    :param first_robot_id: first class ID of robot classes. Only effective if multi_robot=True.
-    :type first_robot_id: int, optional
+    :param robot_class_ids: List of class IDs of robot classes if multiple robot classes exist.
+    :type robot_class_ids: List[int], optional
     :return: Returns precision, recall, AP, f1, ap_class
     """
     model.eval()  # Set model to evaluation mode
@@ -130,7 +126,7 @@ def _evaluate(model, dataloader, class_names, img_size, iou_thres, conf_thres, n
     sample_metrics = []  # List of tuples (TP, confs, pred)
     seg_ious = []
     import time
-    times=[]
+    times = []
     for _, imgs, bb_targets, mask_targets in tqdm.tqdm(dataloader, desc="Validating"):
         # Extract labels
         labels += bb_targets[:, 1].tolist()
@@ -144,7 +140,12 @@ def _evaluate(model, dataloader, class_names, img_size, iou_thres, conf_thres, n
             t1 = time.time()
             yolo_outputs, segmentation_outputs = model(imgs)
             times.append(time.time() - t1)
-            yolo_outputs = non_max_suppression(yolo_outputs, conf_thres=conf_thres, iou_thres=nms_thres, multi_robot=multi_robot, first_robot_id=first_robot_id)
+            yolo_outputs = non_max_suppression(
+                yolo_outputs,
+                conf_thres=conf_thres,
+                iou_thres=nms_thres,
+                robot_class_ids=robot_class_ids
+            )
 
         sample_metrics += get_batch_statistics(yolo_outputs, bb_targets, iou_threshold=iou_thres)
 
@@ -154,7 +155,7 @@ def _evaluate(model, dataloader, class_names, img_size, iou_thres, conf_thres, n
         print("---- No detections over whole validation set ----")
         return None
 
-    print(f"Times: Mean {1/np.array(times).mean()}fps | Std: {np.array(times).std()} ms")
+    print(f"Times: Mean {1 / np.array(times).mean()}fps | Std: {np.array(times).std()} ms")
 
     # Concatenate sample statistics
     true_positives, pred_scores, pred_labels = [
@@ -170,7 +171,7 @@ def seg_iou_mean_without_nan(seg_iou: List[float]) -> np.ndarray:
         :return: Segmentation IOUs without NaN
         """
         seg_iou = np.asarray(seg_iou)
-        return seg_iou[~np.isnan(seg_iou)].mean() 
+        return seg_iou[~np.isnan(seg_iou)].mean()
 
     seg_class_ious = [seg_iou_mean_without_nan(class_ious) for class_ious in list(zip(*seg_ious))]
 
@@ -208,8 +209,10 @@ def _create_validation_data_loader(img_path, batch_size, img_size, n_cpu):
 def run():
     print_environment_info()
     parser = argparse.ArgumentParser(description="Evaluate validation data.")
-    parser.add_argument("-m", "--model", type=str, default="config/yoeo.cfg", help="Path to model definition file (.cfg)")
-    parser.add_argument("-w", "--weights", type=str, default="weights/yoeo.pth", help="Path to weights or checkpoint file (.weights or .pth)")
+    parser.add_argument("-m", "--model", type=str, default="config/yoeo.cfg",
+                        help="Path to model definition file (.cfg)")
+    parser.add_argument("-w", "--weights", type=str, default="weights/yoeo.pth",
+                        help="Path to weights or checkpoint file (.weights or .pth)")
     parser.add_argument("-d", "--data", type=str, default="config/torso.data", help="Path to data config file (.data)")
     parser.add_argument("-b", "--batch_size", type=int, default=8, help="Size of each image batch")
     parser.add_argument("-v", "--verbose", action='store_true', help="Makes the validation more verbose")
@@ -218,7 +221,8 @@ def run():
     parser.add_argument("--iou_thres", type=float, default=0.5, help="IOU threshold required to qualify as detected")
     parser.add_argument("--conf_thres", type=float, default=0.01, help="Object confidence threshold")
     parser.add_argument("--nms_thres", type=float, default=0.4, help="IOU threshold for non-maximum suppression")
-    parser.add_argument("--multiple_robot_classes", action="store_true", help="If multiple robot classes exist and nms shall be performed across all robot classes")
+    parser.add_argument("--multiple_robot_classes", action="store_true",
+                        help="If multiple robot classes exist and nms shall be performed across all robot classes")
 
     args = parser.parse_args()
     print(f"Command line arguments: {args}")
@@ -229,11 +233,12 @@ def run():
     valid_path = data_config["valid"]
     class_names = load_classes(data_config["names"])  # Detection and segmentation class names
 
-    first_robot_class_id = -1
-    for idx, c in enumerate(class_names["detection"]):
-        if "robot" in c:
-            first_robot_class_id = idx
-            break
+    robot_class_ids = None
+    if args.multiple_robot_classes:
+        robot_class_ids = []
+        for idx, c in enumerate(class_names["detection"]):
+            if "robot" in c:
+                robot_class_ids.append(idx)
 
     evaluate_model_file(
         args.model,
@@ -247,8 +252,7 @@ def run():
         conf_thres=args.conf_thres,
         nms_thres=args.nms_thres,
         verbose=True,
-        multi_robot=args.multiple_robot_classes,
-        first_robot_id=first_robot_class_id
+        robot_class_ids=robot_class_ids
     )
 
 
diff --git a/yoeo/train.py b/yoeo/train.py
index 2bdaa32..a29fd92 100755
--- a/yoeo/train.py
+++ b/yoeo/train.py
@@ -78,7 +78,8 @@ def run():
     parser.add_argument("--nms_thres", type=float, default=0.5, help="Evaluation: IOU threshold for non-maximum suppression")
     parser.add_argument("--logdir", type=str, default="logs", help="Directory for training log files (e.g. for TensorBoard)")
     parser.add_argument("--seed", type=int, default=-1, help="Makes results reproducable. Set -1 to disable.")
-    parser.add_argument("--multiple_robot_classes", action="store_true", help="If multiple robot classes exist and nms shall be performed across all robot classes")
+    parser.add_argument("--multiple_robot_classes", action="store_true",
+                        help="If multiple robot classes exist and nms shall be performed across all robot classes")
     args = parser.parse_args()
     print(f"Command line arguments: {args}")
 
@@ -96,11 +97,13 @@ def run():
     train_path = data_config["train"]
     valid_path = data_config["valid"]
     class_names = load_classes(data_config["names"])
-    first_robot_class_id = -1
-    for idx, c in enumerate(class_names["detection"]):
-        if "robot" in c:
-            first_robot_class_id = idx
-            break
+
+    robot_class_ids = None
+    if args.multiple_robot_classes:
+        robot_class_ids = []
+        for idx, c in enumerate(class_names["detection"]):
+            if "robot" in c:
+                robot_class_ids.append(idx)
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -257,8 +260,7 @@ def run():
                 conf_thres=args.conf_thres,
                 nms_thres=args.nms_thres,
                 verbose=args.verbose,
-                multi_robot=args.multiple_robot_classes,
-                first_robot_id=first_robot_class_id
+                robot_class_ids=robot_class_ids
             )
 
             if metrics_output is not None:
diff --git a/yoeo/utils/utils.py b/yoeo/utils/utils.py
index 55b908e..12d1ac5 100644
--- a/yoeo/utils/utils.py
+++ b/yoeo/utils/utils.py
@@ -418,7 +418,8 @@ def box_area(box):
     return inter / (area1[:, None] + area2 - inter)
 
 
-def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, multi_robot=False, first_robot_id=0):
+def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None,
+                        robot_class_ids: Optional[List[int]] = None):
     """Performs Non-Maximum Suppression (NMS) on inference results
     Returns:
          detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
@@ -473,13 +474,13 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
             x = x[x[:, 4].argsort(descending=True)[:max_nms]]
 
         # Batched NMS
-        if not multi_robot:
+        if not robot_class_ids:
             c = x[:, 5:6] * max_wh  # classes
         else:
             # If multiple robot classes are present, all robot classes are treated as one class in order to perform
             # nms across all classes and not per class. For this, all robot classes get the same offset.
             c = torch.clone(x[:, 5:6])
-            c[c > first_robot_id] = first_robot_id
+            c[torch.isin(c, robot_class_ids)] = robot_class_ids[0]
             c *= max_wh
 
         # boxes (offset by class), scores

From 501f1de9e6e5f8c1378020e16e712e2952ad5d27 Mon Sep 17 00:00:00 2001
From: Philipp Donn <30521025+phinik@users.noreply.github.com>
Date: Thu, 13 Apr 2023 20:17:32 +0200
Subject: [PATCH 6/8] added missing imports for type hints

---
 yoeo/detect.py      | 5 +++--
 yoeo/test.py        | 4 ++--
 yoeo/train.py       | 4 +++-
 yoeo/utils/utils.py | 2 +-
 4 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/yoeo/detect.py b/yoeo/detect.py
index a60ce5c..7f99ac6 100755
--- a/yoeo/detect.py
+++ b/yoeo/detect.py
@@ -1,7 +1,6 @@
 #! /usr/bin/env python3
 
-from __future__ import division
-
+from __future__ import division, annotations
 import os
 import argparse
 import tqdm
@@ -13,6 +12,8 @@
 from torch.utils.data import DataLoader
 from torch.autograd import Variable
 
+from typing import Optional, List
+
 from imgaug.augmentables.segmaps import SegmentationMapsOnImage
 
 from yoeo.models import load_model
diff --git a/yoeo/test.py b/yoeo/test.py
index c31a5d7..431ceb2 100755
--- a/yoeo/test.py
+++ b/yoeo/test.py
@@ -1,7 +1,7 @@
 #! /usr/bin/env python3
 
-from __future__ import division
-from typing import List
+from __future__ import division, annotaions
+from typing import List, Optional
 
 import argparse
 import tqdm
diff --git a/yoeo/train.py b/yoeo/train.py
index a29fd92..ef328cd 100755
--- a/yoeo/train.py
+++ b/yoeo/train.py
@@ -1,6 +1,6 @@
 #! /usr/bin/env python3
 
-from __future__ import division
+from __future__ import division, annotations
 
 import os
 import argparse
@@ -13,6 +13,8 @@
 import torch.optim as optim
 from torch.autograd import Variable
 
+from typing import List, Optional
+
 from yoeo.models import load_model
 from yoeo.utils.logger import Logger
 from yoeo.utils.utils import to_cpu, load_classes, print_environment_info, provide_determinism, worker_seed_set
diff --git a/yoeo/utils/utils.py b/yoeo/utils/utils.py
index 12d1ac5..710408e 100644
--- a/yoeo/utils/utils.py
+++ b/yoeo/utils/utils.py
@@ -11,7 +11,7 @@
 import numpy as np
 import subprocess
 import random
-from typing import List
+from typing import List, Optional
 import yaml
 
 

From 9f0589d4d5d7c5fc45c4cea543f36f93efa48569 Mon Sep 17 00:00:00 2001
From: Philipp Donn <30521025+phinik@users.noreply.github.com>
Date: Thu, 13 Apr 2023 20:21:40 +0200
Subject: [PATCH 7/8] fix typo

---
 yoeo/test.py        | 2 +-
 yoeo/utils/utils.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/yoeo/test.py b/yoeo/test.py
index 431ceb2..d266572 100755
--- a/yoeo/test.py
+++ b/yoeo/test.py
@@ -1,6 +1,6 @@
 #! /usr/bin/env python3
 
-from __future__ import division, annotaions
+from __future__ import division, annotations
 from typing import List, Optional
 
 import argparse
diff --git a/yoeo/utils/utils.py b/yoeo/utils/utils.py
index 710408e..afa0790 100644
--- a/yoeo/utils/utils.py
+++ b/yoeo/utils/utils.py
@@ -1,4 +1,4 @@
-from __future__ import division
+from __future__ import division, annotations
 
 from typing import Tuple
 

From eb6bd857df945465641662cf7edd07dbe1763623 Mon Sep 17 00:00:00 2001
From: Philipp Donn <30521025+phinik@users.noreply.github.com>
Date: Thu, 13 Apr 2023 20:50:26 +0200
Subject: [PATCH 8/8] fix torch.isin(..) call

---
 yoeo/utils/utils.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/yoeo/utils/utils.py b/yoeo/utils/utils.py
index afa0790..a11bbe8 100644
--- a/yoeo/utils/utils.py
+++ b/yoeo/utils/utils.py
@@ -437,6 +437,8 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
 
     t = time.time()
     output = [torch.zeros((0, 6), device="cpu")] * prediction.shape[0]
+    if robot_class_ids:
+        robot_class_ids = torch.tensor(robot_class_ids, device=prediction.device, dtype=prediction.dtype)
 
     for xi, x in enumerate(prediction):  # image index, image inference
         # Apply constraints
@@ -474,7 +476,7 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
             x = x[x[:, 4].argsort(descending=True)[:max_nms]]
 
         # Batched NMS
-        if not robot_class_ids:
+        if robot_class_ids is None:
             c = x[:, 5:6] * max_wh  # classes
         else:
             # If multiple robot classes are present, all robot classes are treated as one class in order to perform