Bernt
bae705aa97
- Add NFC ePassport roadmap (ICAO 9303, eIDAS) - Add TensorFlow.js edge face detection (BlazeFace) - Add structured audit logger (GDPR-compliant) - Risk scoring support Part of KYC Apple Native UX v1.1.0
60 lines
1.6 KiB
Python
Executable File
60 lines
1.6 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
QUIXZOOM Custom YOLO Training Pipeline
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import yaml
|
|
from pathlib import Path
|
|
|
|
def create_dataset_config(data_dir, output_file='dataset.yaml'):
|
|
config = {
|
|
'path': os.path.abspath(data_dir),
|
|
'train': 'images/train',
|
|
'val': 'images/val',
|
|
'test': 'images/test',
|
|
'names': {
|
|
0: 'street_lamp',
|
|
1: 'traffic_sign',
|
|
2: 'tree',
|
|
3: 'manhole',
|
|
4: 'utility_box',
|
|
5: 'bench',
|
|
},
|
|
'nc': 6,
|
|
}
|
|
|
|
with open(output_file, 'w') as f:
|
|
yaml.dump(config, f, default_flow_style=False)
|
|
|
|
print(f"[TRAIN] Dataset config saved to {output_file}")
|
|
return output_file
|
|
|
|
def prepare_dataset(data_dir):
|
|
dirs = [
|
|
'images/train', 'images/val', 'images/test',
|
|
'labels/train', 'labels/val', 'labels/test',
|
|
]
|
|
|
|
for d in dirs:
|
|
os.makedirs(os.path.join(data_dir, d), exist_ok=True)
|
|
|
|
print(f"[TRAIN] Dataset structure prepared in {data_dir}")
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='QUIXZOOM YOLO Training')
|
|
parser.add_argument('--data', type=str, required=True, help='Dataset directory')
|
|
parser.add_argument('--epochs', type=int, default=100, help='Training epochs')
|
|
|
|
args = parser.parse_args()
|
|
|
|
prepare_dataset(args.data)
|
|
config_file = create_dataset_config(args.data)
|
|
|
|
print(f"[TRAIN] Ready to train with config: {config_file}")
|
|
print(f"[TRAIN] Run: yolo detect train data={config_file} epochs={args.epochs}")
|
|
|
|
if __name__ == '__main__':
|
|
main()
|