File size: 2,499 Bytes
2db08ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | 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!")
|