472 lines
18 KiB
Python
472 lines
18 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Phase 1 — Dataset Reorganization for Hierarchical Model Training.
|
|
|
|
Reorganizes flat data/dataset/plant-disease-name/ directories into:
|
|
data/organized/
|
|
train/{species}/{disease}/
|
|
val/{species}/{disease}/
|
|
species_index.json
|
|
class_hierarchy.json
|
|
dataset_stats.json
|
|
|
|
Usage: python3 scripts/organize-dataset.py
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import random
|
|
from collections import Counter, defaultdict
|
|
from pathlib import Path
|
|
|
|
from PIL import Image
|
|
from joblib import Parallel, delayed
|
|
from tqdm import tqdm
|
|
|
|
# ─── Config ───────────────────────────────────────────────────────────────────
|
|
|
|
BASE_DIR = Path(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
DATASET_DIR = BASE_DIR / "data" / "dataset"
|
|
ORGANIZED_DIR = BASE_DIR / "data" / "organized"
|
|
TRAIN_DIR = ORGANIZED_DIR / "train"
|
|
VAL_DIR = ORGANIZED_DIR / "val"
|
|
|
|
RANDOM_SEED = 42
|
|
TRAIN_RATIO = 0.85
|
|
VAL_RATIO = 1.0 - TRAIN_RATIO
|
|
|
|
MAX_DIM = 512
|
|
JPEG_QUALITY = 90
|
|
N_JOBS = 16
|
|
|
|
random.seed(RANDOM_SEED)
|
|
|
|
# Known disease-prefix words — words that start disease names but should NOT
|
|
# be part of a plant name. If a plant part ends with one of these, we know
|
|
# the split point is wrong.
|
|
DISEASE_PREFIX_WORDS = {
|
|
"bacterial", "fungal", "viral", "downy", "powdery",
|
|
"alternaria", "phytophthora", "phoma", "phymatotrichum",
|
|
"pythium", "rhizoctonia", "sclerotinia", "fusarium",
|
|
"verticillium", "cercospora", "septoria", "anthracnose",
|
|
"black", "white", "gray", "brown", "green", "pink", "blue",
|
|
"soft", "hard", "sour", "bitter",
|
|
"southern", "northern", "common", "false", "true",
|
|
"european", "american", "aspen", "bacterial-blight",
|
|
"cercospora-leaf", "septoria-leaf", "alternaria-leaf",
|
|
}
|
|
|
|
# Valid multi-word plant suffixes (these CAN follow a hyphen in plant names)
|
|
VALID_MULTI_WORD_PLANTS = {
|
|
"squash", "bean", "berry", "apple", "fern", "tree", "vine",
|
|
"cactus", "grass", "weed", "mint", "root", "seed", "leaf",
|
|
"flower", "fruit", "bark", "wood", "nut", "pea", "lily",
|
|
"rose", "moss", "palm", "fern", "orchid", "fig", "cress",
|
|
"plant", "sage", "thyme", "leaf-fig", "nest-fern", "tongue",
|
|
"tail", "ear", "eye", "nut-tree", "bean-tree",
|
|
}
|
|
|
|
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tiff", ".tif"}
|
|
|
|
# ─── Load KB Data ─────────────────────────────────────────────────────────────
|
|
|
|
def load_kb():
|
|
with open(BASE_DIR / "src" / "data" / "plants.json") as f:
|
|
plants = json.load(f)
|
|
with open(BASE_DIR / "src" / "data" / "diseases.json") as f:
|
|
diseases = json.load(f)
|
|
return plants, diseases
|
|
|
|
PLANTS, DISEASES = load_kb()
|
|
KB_PLANT_IDS = {p["id"] for p in PLANTS}
|
|
|
|
def get_dataset_dirs():
|
|
"""Get all non-hidden subdirectories in the dataset folder."""
|
|
dirs = sorted([
|
|
d for d in os.listdir(DATASET_DIR)
|
|
if os.path.isdir(DATASET_DIR / d) and not d.startswith(".")
|
|
])
|
|
return dirs
|
|
|
|
def count_images(path):
|
|
"""Count image files in a directory."""
|
|
if not path.exists():
|
|
return 0
|
|
return len([
|
|
f for f in os.listdir(path)
|
|
if os.path.isfile(path / f) and os.path.splitext(f)[1].lower() in IMAGE_EXTS
|
|
])
|
|
|
|
# ─── Phase 1: Parse directory names ────────────────────────────────────────────
|
|
|
|
def build_plant_and_disease_dictionaries(dirs):
|
|
"""
|
|
Build verified plant names and disease suffixes from the dataset.
|
|
Returns (parsed_dict, unmatched_list).
|
|
"""
|
|
# Phase 1: Verify plant names from prefixes that appear with >=3 diseases
|
|
plant_candidates = defaultdict(set)
|
|
for d in dirs:
|
|
parts = d.split("-")
|
|
if len(parts) < 2:
|
|
continue
|
|
for split in range(1, min(len(parts), 6)):
|
|
plant = "-".join(parts[:split])
|
|
disease = "-".join(parts[split:])
|
|
if plant and disease and len(disease) > 2:
|
|
plant_candidates[plant].add(disease)
|
|
|
|
verified_plants = set(KB_PLANT_IDS)
|
|
for plant, diseases in plant_candidates.items():
|
|
if len(diseases) >= 3 and plant not in verified_plants:
|
|
verified_plants.add(plant)
|
|
|
|
print(f" Verified plants: {len(verified_plants)} ({len(verified_plants & KB_PLANT_IDS)} from KB)")
|
|
|
|
# Phase 2: Match dirs by plant prefix (longest plant first)
|
|
sorted_plants = sorted(verified_plants, key=len, reverse=True)
|
|
plant_matched = {}
|
|
not_matched = []
|
|
|
|
for d in dirs:
|
|
matched = False
|
|
for plant in sorted_plants:
|
|
prefix = plant + "-"
|
|
if d.startswith(prefix):
|
|
disease = d[len(prefix):]
|
|
if disease:
|
|
plant_matched[d] = (plant, disease)
|
|
matched = True
|
|
break
|
|
if not matched:
|
|
if d.endswith("-healthy"):
|
|
plant = d[:-len("-healthy")]
|
|
plant_matched[d] = (plant, "healthy")
|
|
else:
|
|
not_matched.append(d)
|
|
|
|
# Collect disease suffixes from Phase 2 matches
|
|
disease_suffixes = set(p[1] for p in plant_matched.values())
|
|
print(f" Plant-matched dirs: {len(plant_matched)}, disease suffixes: {len(disease_suffixes)}")
|
|
|
|
# Phase 3: Match remaining dirs by disease suffix (longest suffix first)
|
|
sorted_disease_suffixes = sorted(disease_suffixes, key=len, reverse=True)
|
|
still_not_matched = []
|
|
|
|
for d in not_matched:
|
|
matched = False
|
|
for suffix in sorted_disease_suffixes:
|
|
if d.endswith("-" + suffix):
|
|
plant_part = d[:-len("-" + suffix)]
|
|
if plant_part and not plant_part.endswith("-"):
|
|
plant_matched[d] = (plant_part, suffix)
|
|
matched = True
|
|
break
|
|
if not matched:
|
|
still_not_matched.append(d)
|
|
|
|
print(f" Phase 3 matched: {len(not_matched) - len(still_not_matched)}")
|
|
print(f" Phase 3 remaining: {len(still_not_matched)}")
|
|
|
|
# Phase 4: Handle trailing-hyphen dirs and healthy parent dir
|
|
final_unmatched = []
|
|
for d in still_not_matched:
|
|
if d.endswith("-"):
|
|
plant = d[:-1]
|
|
if plant:
|
|
plant_matched[d] = (plant, "unlabeled")
|
|
elif d == "healthy":
|
|
healthy_dir = DATASET_DIR / "healthy"
|
|
if healthy_dir.exists():
|
|
plant_subdirs = [
|
|
s for s in os.listdir(healthy_dir)
|
|
if os.path.isdir(healthy_dir / s) and not s.startswith(".")
|
|
]
|
|
for sub_plant in plant_subdirs:
|
|
# Use healthy/{sub_plant} as key so we know where to find the images
|
|
plant_matched[f"healthy/{sub_plant}"] = (sub_plant, "healthy")
|
|
print(f" Healthy dir: {len(plant_subdirs)} per-plant healthy classes")
|
|
else:
|
|
final_unmatched.append(d)
|
|
|
|
print(f" Phase 4 handled {len(still_not_matched) - len(final_unmatched)} edge cases")
|
|
print(f" Final unmatched: {len(final_unmatched)}")
|
|
if final_unmatched:
|
|
print(f" E.g.: {final_unmatched[:10]}")
|
|
|
|
# Phase 5: Post-processing — fix species names that ate disease-prefix words
|
|
fix_count = 0
|
|
for d in list(plant_matched.keys()):
|
|
if d.startswith("healthy/"):
|
|
continue # Skip healthy subdirs — these are correct
|
|
species, disease = plant_matched[d]
|
|
parts = species.split("-")
|
|
if len(parts) >= 2 and parts[-1] in DISEASE_PREFIX_WORDS:
|
|
# Move the last word from species to disease
|
|
new_species = "-".join(parts[:-1])
|
|
new_disease = parts[-1] + "-" + disease
|
|
plant_matched[d] = (new_species, new_disease)
|
|
fix_count += 1
|
|
|
|
print(f" Post-process fixes (species ending with disease-prefix): {fix_count}")
|
|
|
|
return plant_matched, final_unmatched
|
|
|
|
# ─── Image Processing ────────────────────────────────────────────────────────
|
|
|
|
def process_image(args):
|
|
"""Resize and convert a single image to 512px max JPEG q90."""
|
|
src_path, dst_path = args
|
|
try:
|
|
img = Image.open(src_path)
|
|
if img.mode != "RGB":
|
|
img = img.convert("RGB")
|
|
w, h = img.size
|
|
if max(w, h) > MAX_DIM:
|
|
ratio = MAX_DIM / max(w, h)
|
|
img = img.resize((int(w * ratio), int(h * ratio)), Image.LANCZOS)
|
|
os.makedirs(os.path.dirname(dst_path), exist_ok=True)
|
|
img.save(dst_path, "JPEG", quality=JPEG_QUALITY, optimize=True)
|
|
return (src_path, True, None)
|
|
except Exception as e:
|
|
return (src_path, False, str(e))
|
|
|
|
def copy_and_split_class(src_dir, dst_train_dir, dst_val_dir, train_ratio=TRAIN_RATIO):
|
|
"""
|
|
Copy images from src_dir to train/val dirs, splitting at the IMAGE level.
|
|
Returns (train_processed, train_failed, val_processed, val_failed).
|
|
"""
|
|
# Check both possible source paths (regular dir or healthy subdir)
|
|
if not src_dir.exists():
|
|
return (0, 0, 0, 0)
|
|
|
|
src_files = sorted([
|
|
f for f in os.listdir(src_dir)
|
|
if os.path.isfile(src_dir / f) and os.path.splitext(f)[1].lower() in IMAGE_EXTS
|
|
])
|
|
if not src_files:
|
|
return (0, 0, 0, 0)
|
|
|
|
# Split files at IMAGE level
|
|
random.shuffle(src_files)
|
|
split_idx = max(1, int(len(src_files) * train_ratio))
|
|
train_files = src_files[:split_idx]
|
|
val_files = src_files[split_idx:]
|
|
|
|
# Process train images
|
|
train_pairs = [
|
|
(str(src_dir / f), str(dst_train_dir / f"img_{i:04d}.jpg"))
|
|
for i, f in enumerate(train_files)
|
|
]
|
|
val_pairs = [
|
|
(str(src_dir / f), str(dst_val_dir / f"img_{i:04d}.jpg"))
|
|
for i, f in enumerate(val_files)
|
|
]
|
|
|
|
results = Parallel(n_jobs=N_JOBS, prefer="threads")(
|
|
delayed(process_image)(pair) for pair in train_pairs + val_pairs
|
|
)
|
|
|
|
train_ok = sum(1 for i, (_, ok, _) in enumerate(results) if ok and i < len(train_pairs))
|
|
train_fail = sum(1 for i, (_, ok, _) in enumerate(results) if not ok and i < len(train_pairs))
|
|
val_ok = sum(1 for i, (_, ok, _) in enumerate(results) if ok and i >= len(train_pairs))
|
|
val_fail = sum(1 for i, (_, ok, _) in enumerate(results) if not ok and i >= len(train_pairs))
|
|
|
|
return (train_ok, train_fail, val_ok, val_fail)
|
|
|
|
# ─── Build Metadata ──────────────────────────────────────────────────────────
|
|
|
|
def build_metadata(parsed, train_counts, val_counts, unmatched):
|
|
"""Build species_index.json, class_hierarchy.json, dataset_stats.json."""
|
|
species_disease_map = defaultdict(set)
|
|
for species, disease in parsed.values():
|
|
species_disease_map[species].add(disease)
|
|
species_index = {k: sorted(v) for k, v in sorted(species_disease_map.items())}
|
|
|
|
class_hierarchy = {
|
|
"version": "1.0",
|
|
"description": "Hierarchical plant disease classification dataset",
|
|
"num_species": len(species_index),
|
|
"num_classes": len(parsed),
|
|
"species": {species: sorted(diseases) for species, diseases in species_index.items()}
|
|
}
|
|
|
|
# Aggregate counts
|
|
total_train = sum(cnt for sp, di, cnt in train_counts)
|
|
total_val = sum(cnt for sp, di, cnt in val_counts)
|
|
total_all = total_train + total_val
|
|
|
|
all_counts = [cnt for _, _, cnt in (train_counts + val_counts)]
|
|
|
|
species_disease_counts = defaultdict(lambda: defaultdict(int))
|
|
for sp, di, cnt in train_counts + val_counts:
|
|
species_disease_counts[sp][di] += cnt
|
|
|
|
# Also count classes from the parsed dict (unique species/disease combos)
|
|
parsed_classes = set((sp, di) for sp, di in parsed.values())
|
|
|
|
stats = {
|
|
"total_images": total_all,
|
|
"total_species": len(species_index),
|
|
"total_classes": len(parsed_classes),
|
|
"train_images": total_train,
|
|
"val_images": total_val,
|
|
"images_per_class": {
|
|
"min": min(all_counts) if all_counts else 0,
|
|
"max": max(all_counts) if all_counts else 0,
|
|
"mean": round(sum(all_counts) / len(all_counts)) if all_counts else 0,
|
|
"median": sorted(all_counts)[len(all_counts) // 2] if all_counts else 0,
|
|
},
|
|
"train_pct": round(total_train / total_all * 100, 1) if total_all else 0,
|
|
"val_pct": round(total_val / total_all * 100, 1) if total_all else 0,
|
|
"unmatched_dirs": len(unmatched),
|
|
"unmatched_dir_names": unmatched[:100] if unmatched else [],
|
|
"species_disease_counts": {
|
|
species: dict(diseases) for species, diseases in species_disease_counts.items()
|
|
}
|
|
}
|
|
|
|
return species_index, class_hierarchy, stats
|
|
|
|
# ─── Main Pipeline ───────────────────────────────────────────────────────────
|
|
|
|
def main():
|
|
print("=" * 60)
|
|
print("Phase 1 — Dataset Reorganization")
|
|
print("=" * 60)
|
|
print(f"Dataset: {DATASET_DIR}")
|
|
print(f"Output: {ORGANIZED_DIR}")
|
|
print()
|
|
|
|
# Step 1: Scan
|
|
print("─" * 40)
|
|
print("Step 1: Scanning dataset directories...")
|
|
print("─" * 40)
|
|
dirs = get_dataset_dirs()
|
|
print(f" Found {len(dirs)} class directories")
|
|
|
|
# Step 2: Parse directory names into (species, disease) pairs
|
|
print()
|
|
print("─" * 40)
|
|
print("Step 2: Parsing directory names...")
|
|
print("─" * 40)
|
|
parsed, unmatched = build_plant_and_disease_dictionaries(dirs)
|
|
|
|
species_set = set(s for s, _ in parsed.values())
|
|
disease_set = set(d for _, d in parsed.values())
|
|
raw_classes = len(parsed)
|
|
unique_classes = len(set((s, d) for s, d in parsed.values()))
|
|
print(f"\n Parsed: {raw_classes} entries")
|
|
print(f" Unique species: {len(species_set)}")
|
|
print(f" Unique disease labels: {len(disease_set)}")
|
|
print(f" Unique (species, disease) pairs: {unique_classes}")
|
|
|
|
# Step 3: Process images with image-level train/val split
|
|
print()
|
|
print("─" * 40)
|
|
print("Step 3: Processing images (resize + train/val split)...")
|
|
print(f" Max dimension: {MAX_DIM}px, JPEG q{JPEG_QUALITY}")
|
|
print(f" Workers: {N_JOBS}")
|
|
print(f" Split: {TRAIN_RATIO*100:.0f}/{VAL_RATIO*100:.0f} (image-level)")
|
|
print("─" * 40)
|
|
|
|
train_counts = [] # (species, disease, count)
|
|
val_counts = []
|
|
total_skipped = 0
|
|
|
|
# Process regular dirs
|
|
regular_items = [(d, sp, di) for d, (sp, di) in parsed.items()
|
|
if not d.startswith("healthy/") and d in dirs]
|
|
healthy_items = [(d, sp, di) for d, (sp, di) in parsed.items()
|
|
if d.startswith("healthy/")]
|
|
|
|
# Organize healthy items by plant
|
|
healthy_by_plant = {}
|
|
for d, sp, di in healthy_items:
|
|
healthy_by_plant[sp] = d # d is like "healthy/tomato"
|
|
|
|
print(f"\n Processing {len(regular_items)} disease + {len(healthy_items)} healthy classes...")
|
|
|
|
for d, species, disease in tqdm(regular_items, desc=" Disease classes"):
|
|
src_dir = DATASET_DIR / d
|
|
dst_train = TRAIN_DIR / species / disease
|
|
dst_val = VAL_DIR / species / disease
|
|
|
|
# Skip if already done (check a few files)
|
|
if dst_train.exists() and dst_val.exists() and \
|
|
len(os.listdir(dst_train)) + len(os.listdir(dst_val)) >= count_images(src_dir):
|
|
total_skipped += count_images(src_dir)
|
|
continue
|
|
|
|
tr_ok, tr_fail, va_ok, va_fail = copy_and_split_class(src_dir, dst_train, dst_val)
|
|
train_counts.append((species, disease, tr_ok))
|
|
val_counts.append((species, disease, va_ok))
|
|
|
|
# Process healthy subdirs
|
|
for sp, hkey in tqdm(healthy_by_plant.items(), desc=" Healthy classes"):
|
|
src_dir = DATASET_DIR / hkey # e.g. data/dataset/healthy/tomato
|
|
dst_train = TRAIN_DIR / sp / "healthy"
|
|
dst_val = VAL_DIR / sp / "healthy"
|
|
|
|
if dst_train.exists() and dst_val.exists() and \
|
|
len(os.listdir(dst_train)) + len(os.listdir(dst_val)) >= count_images(src_dir):
|
|
total_skipped += count_images(src_dir)
|
|
continue
|
|
|
|
tr_ok, tr_fail, va_ok, va_fail = copy_and_split_class(src_dir, dst_train, dst_val)
|
|
train_counts.append((sp, "healthy", tr_ok))
|
|
val_counts.append((sp, "healthy", va_ok))
|
|
|
|
total_train = sum(c for _, _, c in train_counts)
|
|
total_val = sum(c for _, _, c in val_counts)
|
|
print(f"\n Train images: {total_train:,}")
|
|
print(f" Val images: {total_val:,}")
|
|
print(f" Skipped previously processed: {total_skipped:,}")
|
|
|
|
# Step 4: Build metadata
|
|
print()
|
|
print("─" * 40)
|
|
print("Step 4: Building metadata files...")
|
|
print("─" * 40)
|
|
ORGANIZED_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
species_index, class_hierarchy, stats = build_metadata(
|
|
parsed, train_counts, val_counts, unmatched
|
|
)
|
|
|
|
with open(ORGANIZED_DIR / "species_index.json", "w") as f:
|
|
json.dump(species_index, f, indent=2)
|
|
print(f" ✓ species_index.json ({len(species_index)} species)")
|
|
|
|
with open(ORGANIZED_DIR / "class_hierarchy.json", "w") as f:
|
|
json.dump(class_hierarchy, f, indent=2)
|
|
print(f" ✓ class_hierarchy.json")
|
|
|
|
with open(ORGANIZED_DIR / "dataset_stats.json", "w") as f:
|
|
json.dump(stats, f, indent=2)
|
|
print(f" ✓ dataset_stats.json")
|
|
|
|
# Summary
|
|
print()
|
|
print("=" * 60)
|
|
print("Done!")
|
|
print("=" * 60)
|
|
print(f" Total images: {stats['total_images']:,}")
|
|
print(f" Species: {stats['total_species']}")
|
|
print(f" Classes: {stats['total_classes']}")
|
|
print(f" Train: {stats['train_images']:,} ({stats['train_pct']}%)")
|
|
print(f" Val: {stats['val_images']:,} ({stats['val_pct']}%)")
|
|
print(f" Unmatched dirs: {stats['unmatched_dirs']}")
|
|
print(f" Train dir: {TRAIN_DIR}")
|
|
print(f" Val dir: {VAL_DIR}")
|
|
|
|
if stats['unmatched_dirs'] > 0:
|
|
print(f"\n ⚠ Manual review needed for {stats['unmatched_dirs']} dirs:")
|
|
for u in stats['unmatched_dir_names'][:20]:
|
|
print(f" {u}")
|
|
|
|
return stats
|
|
|
|
if __name__ == "__main__":
|
|
main()
|