Some README.md changes + prepare_dataset.py

This commit is contained in:
Hamza Zakaria 2026-04-26 16:02:06 +01:00
parent 099ac2b439
commit 82679c1330
3 changed files with 133 additions and 2 deletions

View File

@ -10,6 +10,12 @@
pip install ultralytics albumentations opencv-python tqdm pyyaml torch
```
or using conda
```bash
conda env create -f enviroment.yml
``
---
## Step-by-step
@ -31,7 +37,7 @@ Put your raw product photos into `dataset/images/train/`.
Use **Roboflow** (recommended) or **LabelImg** to draw bounding boxes and export in YOLO format:
- [Roboflow](https://roboflow.com) — free, web-based, exports directly to YOLO format
- LabelImg: `pip install labelImg && labelImg`
- LabelImg: `pip install labelImg && labelImg` or download it from [here](https://github.com/tzutalin/labelImg/files/2638199/windows_v1.8.1.zip)
Each image needs a matching `.txt` in `dataset/labels/train/`:
@ -127,4 +133,4 @@ With 4 original images + augmentation and 13 classes:
- **mAP50 ≈ 0.650.80** (good for production)
- **Inference speed** ≈ 1530ms/image on CPU, <5ms on GPU
Adding even 23 more labeled images per class typically pushes mAP above 0.85.
Adding even 23 more labeled images per class typically pushes mAP above 0.85.

18
enviroment.yml Normal file
View File

@ -0,0 +1,18 @@
name: yolo-coffee-object-detection
channels:
- conda-forge
dependencies:
- python=3.10
- ultralytics
- albumentations
- opencv
- tqdm
- pyyaml
#- torch
#- numpy
#- opencv-contrib-python
#- scikit-learn # @Todo: Remove it later if we don't need it.
- pip
- pip:
- fastapi
- uvicorn[standard]

107
prepare_dataset.py Normal file
View File

@ -0,0 +1,107 @@
"""
STEP 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(" {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()