108 lines
2.9 KiB
Python
108 lines
2.9 KiB
Python
"""
|
|
STEP 1 ─ 1_prepare_dataset.py
|
|
===========================
|
|
Run this FIRST. It creates the required folder structure and
|
|
a dataset.yaml that YOLOv8 expects.
|
|
|
|
Usage:
|
|
python prepare_dataset.py
|
|
|
|
After running. your folder tree will look like this:
|
|
dataset/
|
|
images/
|
|
train/ <- put your labeled images here.
|
|
val/ <- validation split (auto-filled by augment.py)
|
|
|
|
labels/
|
|
train/ <- put your .txt label files here.
|
|
val/ <- auto-generated.
|
|
dataset.yaml <- auto-generated.
|
|
|
|
Label format (one .txt per image, same name):
|
|
<class_id> <cx> <cy> <w> <h> (all values 0.0-1.0, normalised)
|
|
|
|
Recommended free labeling tools:
|
|
• Roboflow https://roboflow.com (easiest, web-based)
|
|
• LabelImg pip install labelImg (desktop, saves YOLO format directly)
|
|
"""
|
|
import os
|
|
import yaml
|
|
|
|
# ── Edit this list to match your product classes ────────
|
|
CLASS_NAMES = [
|
|
"1001",
|
|
"Aroma",
|
|
"Aroma Gold",
|
|
"Bonal",
|
|
"Facto",
|
|
"Famico",
|
|
"Famico Exclusive",
|
|
"Fegalo",
|
|
"Gosto",
|
|
"Molino",
|
|
"Oscar",
|
|
"Primo",
|
|
"Siglo"
|
|
]
|
|
|
|
DATASET_ROOT = "dataset"
|
|
|
|
DIRS = [
|
|
f"{DATASET_ROOT}/images/train",
|
|
f"{DATASET_ROOT}/images/val",
|
|
f"{DATASET_ROOT}/labels/train",
|
|
f"{DATASET_ROOT}/labels/val",
|
|
]
|
|
|
|
def Structure_Create():
|
|
for d in DIRS:
|
|
os.makedirs(d, exist_ok=True)
|
|
print("[1/3] Folder structure created:")
|
|
|
|
for d in DIRS:
|
|
print(f" {d}/")
|
|
|
|
def YAML_Write():
|
|
cfg = {
|
|
"path" : os.path.abspath(DATASET_ROOT),
|
|
"train": "images/train",
|
|
"val" : "images/val",
|
|
"nc" : len(CLASS_NAMES),
|
|
"names": CLASS_NAMES,
|
|
}
|
|
|
|
yaml_path = "dataset.yml"
|
|
with open(yaml_path, "w") as f:
|
|
yaml.dump(cfg, f, default_flow_style=False, sort_keys=False)
|
|
|
|
print("[2/3] dataset.yaml written -> {yaml_path}")
|
|
return yaml_path
|
|
|
|
def Print_NextSetups():
|
|
print("""
|
|
[3/3] NEXT STEPS
|
|
─────────────────────────────────────────────────────────────
|
|
1. Put your RAW images (one per product view) into:
|
|
dataset/images/train/
|
|
|
|
2. Label them using Roboflow or LabelImg (YOLO format).
|
|
Each image → a matching .txt in dataset/labels/train/
|
|
Example label file (famico_veloute_01.txt):
|
|
0 0.512 0.483 0.320 0.541
|
|
|
|
3. Run the augmentation script:
|
|
python 2_augment.py
|
|
|
|
4. Train:
|
|
python 3_train.py
|
|
|
|
5. Detect on new scenes:
|
|
python 4_detect.py --source scene.jpg
|
|
─────────────────────────────────────────────────────────────
|
|
""")
|
|
|
|
if __name__ == "__main__":
|
|
Structure_Create()
|
|
YAML_Write()
|
|
Print_NextSetups()
|