# YOLOv8 training environment (NVIDIA GPU)
#
# The image contains Python, PyTorch and Ultralytics only. The project
# directory (including data.yaml, train/valid/test and *.pt files) is mounted
# at /workspace when the container is started; this keeps the image small.

FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1 \
    PIP_NO_CACHE_DIR=1 \
    NVIDIA_VISIBLE_DEVICES=all \
    NVIDIA_DRIVER_CAPABILITIES=compute,utility

RUN apt-get update \
    && apt-get install -y --no-install-recommends \
        ca-certificates \
        libglib2.0-0 \
        libgl1 \
        libsm6 \
        libxext6 \
        libxrender1 \
        python3 \
        python3-dev \
        python3-pip \
    && rm -rf /var/lib/apt/lists/*

# PyTorch 2.5.1 is paired with torchvision 0.20.1 and CUDA 12.4.
RUN python3 -m pip install --upgrade pip \
    && python3 -m pip install \
        torch==2.5.1 \
        torchvision==0.20.1 \
        --index-url https://download.pytorch.org/whl/cu124 \
    && python3 -m pip install ultralytics==8.3.0

WORKDIR /workspace

# These files make the image usable for a quick smoke test. The normal run
# command mounts the whole project over /workspace.
COPY train_yolov8.py data.yaml ./

ENTRYPOINT ["python3", "train_yolov8.py"]