Spaces:
Running
Running
File size: 3,235 Bytes
6a4ae4a 5cee08a 57d7f02 5cee08a 6a4ae4a 5cee08a |
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 75 76 77 78 79 |
<!doctype html>
<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> |