import torch
import pandas as pd
import os
import shutil
import random
from ultralytics import YOLO
# Настройки
DATA_DIR = "dataset/"
TRAIN_DIR = os.path.join(DATA_DIR, "train")
VAL_DIR = os.path.join(DATA_DIR, "val")
BATCH_SIZE = 16
EPOCHS = 10
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MODEL_PATH = "yolov8n-cls.pt" # Используем предобученную YOLOv8 для классификации
VAL_SPLIT = 0.2 # 20% данных отдаем на валидацию
# Функция для разбиения данных на train/val
if not os.path.exists(VAL_DIR):
os.makedirs(VAL_DIR)
for class_name in os.listdir(TRAIN_DIR):
class_path = os.path.join(TRAIN_DIR, class_name)
val_class_path = os.path.join(VAL_DIR, class_name)
if os.path.isdir(class_path):
os.makedirs(val_class_path, exist_ok=True)
images = os.listdir(class_path)
random.shuffle(images)
val_size = int(len(images) * VAL_SPLIT)
for img in images[:val_size]:
shutil.move(os.path.join(class_path, img), os.path.join(val_class_path, img))
# Загрузка модели YOLO (при появлении YOLOv11 заменить путь на соответствующую версию)
model = YOLO(MODEL_PATH)
# Обучение модели на кастомном датасете с валидационной выборкой
model.train(data=TRAIN_DIR, epochs=EPOCHS, imgsz=224, batch=BATCH_SIZE, device=DEVICE, val=VAL_DIR)
# Функция предсказания на тестовых данных
def predict():
test_images = sorted(os.listdir(os.path.join(DATA_DIR, "test"))) # Загружаем список изображений из тестовой папки
predictions = []
for img_name in test_images:
img_path = os.path.join(DATA_DIR, "test", img_name)
results = model(img_path) # Делаем предсказание YOLO
predicted_label = results[0].probs.top1 # Получаем индекс предсказанного класса
predictions.append((img_name, model.names[predicted_label])) # Преобразуем индекс в имя класса
return predictions
# Сохранение предсказаний в CSV
def save_predictions(predictions, filename="submission.csv"):
df = pd.DataFrame(predictions, columns=["name", "label"])
df.to_csv(filename, index=False)
print(f"Predictions saved to {filename}")
# Запуск
if __name__ == "__main__":
preds = predict()
save_predictions(preds)
You can modify this app directly by editing index.html in the Files and versions tab.
Also don't forget to check the Spaces documentation.