Skip to main content

Image Models Example

Image Models Example

Файлы: examples/image_models/train.py и predict.py. Демонстрируют YoloImageModel на подвыборке COCO (4 изображения, examples/image_models/data/).

train.py — обучение всех pretrained-моделей в папке

model_files = list(pretrained_dir.glob("*.pt"))   # examples/image_models/pretrained_models/*.pt

for model_path in model_files:
    base_model = YOLO(str(model_path))
    wrapper = YoloImageModel(base_model)
    wrapper.train(data=str(dataset_yaml), epochs=5, imgsz=640)
    wrapper.save(str(trained_dir / f"{model_path.stem}_trained.pt"))

В репозитории уже лежат три готовые pretrained-модели (yolo11n.pt, yolov5nu.pt, yolov8n.pt в pretrained_models/) — скрипт дообучает каждую на локальном мини-датасете (5 эпох) и сохраняет результат в trained_models/. Там же уже лежат готовые *_trained.pt — то есть скрипт был прогнан раньше, и его можно перезапускать повторно (перезапишет результат).

predict.py — инференс всеми обученными моделями

trained_models = list(trained_dir.glob("*.pt"))
for model_path in trained_models:
    model = YOLO(str(model_path))
    wrapper = YoloImageModel(model)
    results = wrapper.predict(str(image_path), save=False)
    results[0].plot(save=True, filename=str(out_img))

Берёт одно фиксированное изображение (data/images/000000000009.jpg), прогоняет через каждую модель из trained_models/ и сохраняет визуализацию в predictions/ ({model}_prediction.jpg) — готовые примеры результатов уже есть в репозитории.

⚠️ Датасет содержит абсолютные пути

examples/image_models/data/dataset.yaml содержит:

train: /home/niashin/mephi/mlumina/examples/image_models/data/images
val: /home/niashin/mephi/mlumina/examples/image_models/data/images

Это абсолютные пути конкретной машины разработчика, а не относительные/портативные. Перед запуском train.py на любой другой машине эти два пути нужно поправить на актуальные — иначе ultralytics.YOLO.train не найдёт изображения (или найдёт чужие/несуществующие).

Как запустить

cd examples/image_models
python train.py     # дообучение (после правки dataset.yaml)
python predict.py    # инференс уже обученными моделями

Оба скрипта используют YoloImageModel в однопроцессном режиме (train/predict, не train_ray/predict_ray) — Ray не требуется и не инициализируется.

Back to top