# YOLOv8 training image (NVIDIA GPU)
#
# Image contains Python / PyTorch / Ultralytics + train script.
# Dataset (*.pt weights, /data/facepp, runs/) are mounted at runtime.

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-pip \
    && rm -rf /var/lib/apt/lists/*

# PyTorch 2.5.1 + torchvision 0.20.1 (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

COPY train_yolov8.py data.yaml ./

ENTRYPOINT ["python3", "train_yolov8.py"]
CMD ["--model", "yolov8n.pt", "--device", "0"]