Spaces:
Running
Running
| <html> | |
| <head> | |
| <meta charset="utf-8" /> | |
| <meta name="viewport" content="width=device-width" /> | |
| <title>My static Space</title> | |
| <link rel="stylesheet" href="style.css" /> | |
| </head> | |
| <body> | |
| <div class="card"> | |
| <h1>Компьютерное зрение :</h1> | |
| <pre><code> | |
| 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) | |
| </code></pre> | |
| <p>You can modify this app directly by editing <i>index.html</i> in the Files and versions tab.</p> | |
| <p> | |
| Also don't forget to check the | |
| <a href="https://huggingface.co/docs/hub/spaces" target="_blank">Spaces documentation</a>. | |
| </p> | |
| </div> | |
| </body> | |
| </html> |