multiple skier handling
This commit is contained in:
parent
391620fd5c
commit
8951dd27f0
@ -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))
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user