FROM nvcr.io/nvidia/pytorch:23.03-py3
# FROM nvcr.io/nvidia/pytorch:24.02-py3

ARG HOME_PATH="/home"
WORKDIR ${HOME_PATH}

RUN pip3 install --upgrade pip wheel
RUN pip3 install azure-storage-blob azure-identity

# supervision
RUN git clone https://github.com/roboflow/supervision.git && \
    cd supervision && \
    grep -v "^opencv-python-headless" pyproject.toml > tmp.toml && \
    mv tmp.toml pyproject.toml && \
    pip3 install --no-cache -e .

# ultralytics
ADD https://ultralytics.com/assets/Arial.ttf https://ultralytics.com/assets/Arial.Unicode.ttf /root/.config/Ultralytics/
RUN git clone https://github.com/ultralytics/ultralytics && \
    cd ultralytics && \
    grep -v "opencv-python\|openvino-dev" pyproject.toml > tmp.toml && mv tmp.toml pyproject.toml && \
    pip3 install "opencv-python-headless<4.7" "opencv-contrib-python<4.7" "opencv-contrib-python-headless<4.7" "albumentations<1.4.0" && \
    pip3 install .

# download dataset
ARG CVAT_URL
ARG CVAT_ORG
ARG CVAT_TASKS_YAML
ARG TRAIN_HP_YAML
ARG PYPREPROCESS

COPY . .
# COPY AIEM/trainer /home/trainer
# COPY ${CVAT_TASKS_YAML} ${CVAT_TASKS_YAML}
# COPY ${TRAIN_HP_YAML} ${TRAIN_HP_YAML}

ENV APP_PYPREPROCESS=${PYPREPROCESS}
ENV APP_CVAT_TASKS_YAML=${CVAT_TASKS_YAML}
ENV APP_HOME=${HOME_PATH}
ENV APP_TRAIN_HP_YAML=${TRAIN_HP_YAML}

RUN cd AIEM/trainer && \
    python3 utils/download_cvatdata.py \
        "$CVAT_URL" \
        "$CVAT_ORG"
RUN cd /data && \
    rm -rf *.zip

ENTRYPOINT ["python3", "AIEM/trainer/train_yolov8.py"]