diff --git a/createDataset.py b/createDataset.py
index 3354964..812d03b 100644
--- a/createDataset.py
+++ b/createDataset.py
@@ -53,7 +53,6 @@ def __download_from_minio(s3_path, local_path):
 def createYOLOdataset():
     __setup_directories()
 
-    # read annotations
     with open(JSON_PATH, 'r', encoding='utf-8') as f:
         data = json.load(f)
 
@@ -65,7 +64,6 @@ def createYOLOdataset():
     for i, entry in enumerate(tqdm(data, desc="Importing Images", unit="img")):
         split = 'train' if i < split_idx else 'val'
 
-        # get image name
         image_s3_path = entry['data']['image']
         filename = os.path.basename(image_s3_path)
         base_name = os.path.splitext(filename)[0]
@@ -79,59 +77,84 @@ def createYOLOdataset():
             tqdm.write(f"Error treating {filename}: {e}")
             continue
 
-        yolo_lines = []
-
-        # check if annotations, otherwise skip
         if not entry.get('annotations'):
             continue
 
         results = entry['annotations'][0].get('result', [])
 
         # dummy vars
-        kp_map = {}  # ID -> {label, x, y}
-        visibility_map = {}  # ID -> v_status (1 oder 2)
-        bboxes = []  # Liste aller gefundenen BBoxes
+        temp_bboxes = []
+        all_keypoints = []
+        visibility_map = {}
 
         for res in results:
             res_id = res['id']
             res_type = res['type']
             val = res.get('value', {})
 
-            if res_type == 'keypointlabels':
-                kp_map[res_id] = {
+            if res_type == 'rectanglelabels':
+                # save bboxes
+                w = val['width'] / 100.0
+                h = val['height'] / 100.0
+                x = val['x'] / 100.0
+                y = val['y'] / 100.0
+                temp_bboxes.append({
+                    'x_center': x + (w / 2.0),
+                    'y_center': y + (h / 2.0),
+                    'width': w,
+                    'height': h,
+                    'x_min': x,
+                    'y_min': y,
+                    'x_max': x + w,
+                    'y_max': y + h,
+                    'assigned_kps': {}  # store corresponding kp
+                })
+
+            elif res_type == 'keypointlabels':
+                all_keypoints.append({
+                    'id': res_id,
                     'label': val['keypointlabels'][0],
                     'x': val['x'] / 100.0,
                     'y': val['y'] / 100.0
-                }
+                })
 
             elif res_type == 'choices':
-                if "1" in val.get('choices', []):
-                    visibility_map[res_id] = 1
+                # Label Studio links via if
+                parent_id = res.get('parent_id')
+                if parent_id and "hidden" in val.get('choices', []):
+                    visibility_map[parent_id] = 1
 
-            elif res_type == 'rectanglelabels':
-                # BBox normalisieren
-                bw = val['width'] / 100.0
-                bh = val['height'] / 100.0
-                bx = (val['x'] / 100.0) + (bw / 2.0)
-                by = (val['y'] / 100.0) + (bh / 2.0)
-                bboxes.append(f"{bx:.6f} {by:.6f} {bw:.6f} {bh:.6f}")
+        # assign kp to bboxes
+        for kp in all_keypoints:
+            for bbox in temp_bboxes:
+                # check inside box
+                if (bbox['x_min'] <= kp['x'] <= bbox['x_max'] and
+                        bbox['y_min'] <= kp['y'] <= bbox['y_max']):
+                    bbox['assigned_kps'][kp['label']] = kp
+                    break  # Punkt gehört zu dieser Box, nächster Punkt
 
-        # create yolo data
-        for bbox_coords in bboxes:
-            line = f"0 {bbox_coords}"
+        # create yolo lines
+        yolo_lines = []
+        for bbox in temp_bboxes:
+            # Class (0) + BBox
+            line_parts = [
+                "0",
+                f"{bbox['x_center']:.6f}", f"{bbox['y_center']:.6f}",
+                f"{bbox['width']:.6f}", f"{bbox['height']:.6f}"
+            ]
 
+            # kp in right order
             for kp_name in KP_ORDER:
-                target_id = next((id for id, d in kp_map.items() if d['label'] == kp_name), None)
-
-                if target_id:
-                    coords = kp_map[target_id]
-                    # visibility, 0 missing, 1 invisible, 2 visible
-                    v = visibility_map.get(target_id, 2)
-                    line += f" {coords['x']:.6f} {coords['y']:.6f} {v}"
+                if kp_name in bbox['assigned_kps']:
+                    kp = bbox['assigned_kps'][kp_name]
+                    # check visibility
+                    v = visibility_map.get(kp['id'], 2)
+                    line_parts.extend([f"{kp['x']:.6f}", f"{kp['y']:.6f}", str(v)])
                 else:
-                    line += " 0.000000 0.000000 0"
+                    # 0 (missing)
+                    line_parts.extend(["0.000000", "0.000000", "0"])
 
-            yolo_lines.append(line)
+            yolo_lines.append(" ".join(line_parts))
 
         with open(label_local_path, 'w', encoding='utf-8') as f:
             f.write('\n'.join(yolo_lines))