195 lines
7.0 KiB
Python
195 lines
7.0 KiB
Python
"""
|
|
Script 1 — Crop Dataset Generator
|
|
===================================
|
|
Runs your fine-tuned YOLOv8n model on a folder of shelf images,
|
|
crops every detection, and dumps everything into a single flat folder.
|
|
|
|
You classify ONCE, then run script 1b_split_dataset.py to create
|
|
the train / val / test split automatically.
|
|
|
|
Output structure
|
|
----------------
|
|
data/
|
|
unknown/
|
|
img001_crop0000.jpg
|
|
img001_crop0001.jpg
|
|
img002_crop0000.jpg
|
|
...
|
|
|
|
After generation
|
|
----------------
|
|
1. Rename / move crops into class subfolders inside data/:
|
|
data/
|
|
cola_can/ img001_crop0000.jpg …
|
|
pepsi_can/ img003_crop0002.jpg …
|
|
…
|
|
|
|
Use Label Studio, a file manager, or any annotation tool.
|
|
The `unknown/` folder holds anything you haven't labelled yet.
|
|
|
|
2. Then run:
|
|
python 1b_split_dataset.py --data_dir data --output_dir crops_dataset
|
|
|
|
Usage
|
|
-----
|
|
python 1_generate_crop_dataset.py \
|
|
--source /path/to/shelf/images \
|
|
--weights /path/to/best.pt \
|
|
--output_dir data \
|
|
--conf 0.25 \
|
|
--iou 0.45 \
|
|
--img_size 640 \
|
|
--padding 10 \
|
|
--min_size 32
|
|
"""
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from ultralytics import YOLO
|
|
from tqdm import tqdm
|
|
|
|
|
|
# ─────────────────────────── helpers ─────────────────────────────────────────
|
|
|
|
def parse_args():
|
|
p = argparse.ArgumentParser(description="Generate flat crop dataset from YOLOv8 detections")
|
|
p.add_argument("--source", required=True , help="Directory with shelf images (jpg/png/…)")
|
|
p.add_argument("--weights", default="detector/best.pt", help="Path to your fine-tuned YOLOv8n .pt file")
|
|
p.add_argument("--output_dir", default="data" , help="Root output dir — crops go into <output_dir>/unknown/")
|
|
p.add_argument("--conf", type=float, default=0.25 , help="Detection confidence threshold")
|
|
p.add_argument("--iou", type=float, default=0.45 , help="NMS IoU threshold")
|
|
p.add_argument("--img_size", type=int, default=640 , help="YOLO inference image size")
|
|
p.add_argument("--padding", type=int, default=10 , help="Extra pixels around each crop (px)")
|
|
p.add_argument("--min_size", type=int, default=32 , help="Discard crops smaller than this (px)")
|
|
p.add_argument("--save_labels", action="store_true",
|
|
help="Save a sidecar .txt with YOLO cls_id / confidence next to each crop")
|
|
return p.parse_args()
|
|
|
|
|
|
def get_image_paths(source: str):
|
|
exts = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"}
|
|
p = Path(source)
|
|
|
|
if p.is_file():
|
|
# Single image passed directly e.g. --source 01.jpg
|
|
if p.suffix.lower() not in exts:
|
|
raise ValueError(f"File is not a supported image type: {p}")
|
|
return [p]
|
|
|
|
if p.is_dir():
|
|
# Folder — recurse for all images
|
|
paths = sorted(f for f in p.rglob("*") if f.suffix.lower() in exts)
|
|
if not paths:
|
|
raise FileNotFoundError(f"No images found in directory: {p}")
|
|
return paths
|
|
|
|
# Glob pattern e.g. --source "frames/*.jpg"
|
|
paths = sorted(f for f in Path(".").glob(source) if f.suffix.lower() in exts)
|
|
if not paths:
|
|
raise FileNotFoundError(f"No images matched: {source}")
|
|
return paths
|
|
|
|
|
|
def crop_detection(image: np.ndarray, box, padding: int, min_size: int):
|
|
"""
|
|
Extract one padded crop from *image* using *box* (xyxy).
|
|
Returns the crop array, or None if it is too small.
|
|
"""
|
|
x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
|
|
h, w = image.shape[:2]
|
|
|
|
x1 = max(0, x1 - padding)
|
|
y1 = max(0, y1 - padding)
|
|
x2 = min(w, x2 + padding)
|
|
y2 = min(h, y2 + padding)
|
|
|
|
if (x2 - x1) < min_size or (y2 - y1) < min_size:
|
|
return None
|
|
|
|
return image[y1:y2, x1:x2]
|
|
|
|
|
|
# ─────────────────────────── main ────────────────────────────────────────────
|
|
|
|
def main():
|
|
args = parse_args()
|
|
out_dir = Path(args.output_dir) / "unknown"
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# ── Load model ────────────────────────────────────────────────────────────
|
|
print(f"[1/2] Loading model: {args.weights}")
|
|
model = YOLO(args.weights)
|
|
|
|
# ── Detect & crop ─────────────────────────────────────────────────────────
|
|
image_paths = get_image_paths(args.source)
|
|
print(f"[2/2] Found {len(image_paths)} images — running detection …\n")
|
|
|
|
total_saved = 0
|
|
total_skip = 0
|
|
|
|
for img_path in tqdm(image_paths, desc="Detecting & cropping"):
|
|
image = cv2.imread(str(img_path))
|
|
if image is None:
|
|
print(f" ⚠ Cannot read {img_path.name} — skipping")
|
|
continue
|
|
|
|
results = model(
|
|
image,
|
|
conf=args.conf,
|
|
iou=args.iou,
|
|
imgsz=args.img_size,
|
|
verbose=False,
|
|
)[0]
|
|
|
|
stem = img_path.stem
|
|
for i, box in enumerate(results.boxes):
|
|
crop = crop_detection(image, box, args.padding, args.min_size)
|
|
|
|
if crop is None:
|
|
total_skip += 1
|
|
continue
|
|
|
|
crop_name = f"{stem}_crop{i:04d}.jpg"
|
|
crop_path = out_dir / crop_name
|
|
cv2.imwrite(str(crop_path), crop)
|
|
total_saved += 1
|
|
|
|
if args.save_labels:
|
|
cls_id = int(box.cls[0])
|
|
conf = float(box.conf[0])
|
|
cls_name = model.names.get(cls_id, str(cls_id))
|
|
crop_path.with_suffix(".txt").write_text(
|
|
f"{cls_id} {cls_name} {conf:.4f}\n"
|
|
)
|
|
|
|
# ── Summary ───────────────────────────────────────────────────────────────
|
|
print("\n" + "=" * 55)
|
|
print(" Crop generation complete!")
|
|
print(f" Saved : {total_saved:,} crops -> {out_dir.resolve()}")
|
|
print(f" Skipped : {total_skip:,} (below --min_size={args.min_size}px)")
|
|
print("=" * 55)
|
|
print("""
|
|
Next steps:
|
|
1. Organise crops into class subfolders inside data/:
|
|
data/
|
|
cola_can/ *.jpg
|
|
pepsi_can/ *.jpg
|
|
...
|
|
Leave unlabelled crops in data/unknown/ — the split
|
|
script ignores that folder automatically.
|
|
|
|
2. Run the splitter:
|
|
python 1b_split_dataset.py \\
|
|
--data_dir data \\
|
|
--output_dir crops_dataset \\
|
|
--val_split 0.15 \\
|
|
--test_split 0.05
|
|
""")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|