cleanup and documentation

Author: denniswittich

README.md

@@ -7,6 +7,21 @@ This repository is an implementation of Nodes that interact with the Zauberzeug
 
 This node is used to train Yolov5 Models in the Learning Loop. It is based on [this image](https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-23-07.html) running Python 3.10.
 
+## Hyperparameters
+
+We support all native hyperparameters of YOLOv5 (cf. `hyp_det.yaml` / `hyp_cla.yaml` for reference).
+In addition, we support the following hyperparameters:
+
+- `epochs`: The number of epochs to train the model.
+- `detect_nms_conf_thres`: The confidence threshold for the NMS during inference and validation (not relevant for training).
+- `detect_nms_iou_thres`: The IoU threshold for the NMS during inference and validation (not used for training).
+
+Further, we support the following hyperparameters for point detection:
+
+- `reset_points`: Whether to reset the size of the points after data augmentation.
+- `point_sizes_by_id`: A dictionary that maps from point category uuids to the size of the points in the output (fractional size 0-1).
+- `flip_label_pairs`: A list of pairs of point uuids that should be swapped when a horizontal flip is applied during data augmentation.
+
 ## Images
 
 Trainer Docker-Images are published on https://hub.docker.com/r/zauberzeug/yolov5-trainer

trainer/app_code/yolov5/utils/dataloaders.py

@@ -788,10 +788,9 @@ def __getitem__(self, index):
         if reset_points and len(self.point_sizes_by_id) > 0:
             for label in labels:
                 if label[0] in self.point_sizes_by_id:
-                    target_point_size_px = self.point_sizes_by_id[label[0]]
-                    # target_point_size_px = 10
-                    label[3] = target_point_size_px  # / self.img_size
-                    label[4] = target_point_size_px  # / self.img_size
+                    target_point_size_fraction = self.point_sizes_by_id[label[0]]
+                    label[3] = target_point_size_fraction
+                    label[4] = target_point_size_fraction
 
         labels_out = torch.zeros((nl, 6))
         if nl:

trainer/hyp_det.yaml

@@ -34,10 +34,10 @@ copy_paste: 0.0
 
 #own parameters
 epochs: 2000
-reset_points: false
 detect_nms_conf_thres: 0.2
 detect_nms_iou_thres: 0.45
-#
+
 #extra point parameters (not directly used by yolov5 but by the trainer logic)
+reset_points: false
 point_sizes_by_id: "" # e.g "1111-2222-3333-4444:0.03,5555-6666-7777-8888:0.05"
 flip_label_pairs: "" # e.g "1111-2222-3333-4444:5555-6666-7777-8888"