multiple skier handling

This commit is contained in:
schneuwl 2026-01-11 19:44:48 +01:00
parent 391620fd5c
commit 8951dd27f0

View File

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