Rename orch.py to orchestrator.py

e062f66 verified over 1 year ago

2.5 kB

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!")

	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!")