import os
import shutil
from ultralytics import YOLO

# =======================================
# 🟡🔧 User-configurable settings
# =======================================

# 📁 Folder containing images to classify
INPUT_FOLDER = r"C:\Users\ASUS\OneDrive - Binder\Desktop\test.orch"  # ← Modify this path

# 📁 Folder to save classified images
OUTPUT_BASE = os.path.join(INPUT_FOLDER, "classified")  # ← Automatically inside input folder

# 📦 Path to YOLOv8 classification model
MODEL_PATH = r"C:\Users\ASUS\Downloads\best.pt"  # ← Modify this path

# ✅ Allowed classes only
ALLOWED_CLASSES = {"CR1", "CR2", "CR3", "CR4", "CR5", "CR6", "CR7", "b1", "b2", "b3", "b4", "v1", "v2", "v3"}

# =======================================
# 🚀 Load the model
# =======================================
model = YOLO(MODEL_PATH)

# =======================================
# 🧠 Classify a single image
# =======================================
def classify_image_yolo(image_path):
    try:
        results = model(image_path, verbose=False)[0]
        class_id = int(results.probs.top1)
        class_name = model.names.get(class_id, "others")
        return class_name
    except Exception as e:
        print(f"❌ Error classifying image {image_path}: {e}")
        return "others"

# =======================================
# 📂 Classify and move images
# =======================================
processed_files = set()

for filename in os.listdir(INPUT_FOLDER):
    file_path = os.path.join(INPUT_FOLDER, filename)

    if not os.path.isfile(file_path):
        continue

    if filename in processed_files:
        continue

    class_name = classify_image_yolo(file_path)

    # ✅ Check if class is allowed, otherwise assign to "others"
    if class_name not in ALLOWED_CLASSES:
        class_name = "others"

    # Create class folder if not exists
    output_folder = os.path.join(OUTPUT_BASE, class_name)
    os.makedirs(output_folder, exist_ok=True)

    destination_path = os.path.join(output_folder, filename)

    # Avoid duplicate images
    if not os.path.exists(destination_path):
        shutil.move(file_path, destination_path)
        print(f"✅ Moved image {filename} to [{class_name}]")
        processed_files.add(filename)
    else:
        print(f"⚠️ Image {filename} already exists in [{class_name}], skipped")

print("🎉 Finished classifying images using YOLOv8 Classification!")