FROM nvidia/cuda:12.1.0-runtime-ubuntu20.04
LABEL maintainer="hoge@fuga"

# Set timezone to avoid tzdata prompt
ENV TZ=Asia/Tokyo
RUN apt-get update && \
    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends tzdata

# Install base tools
RUN apt-get update && apt-get install -y --no-install-recommends \
    bash curl fish git nano sudo tmux software-properties-common

# Install OpenCV dependencies
RUN apt-get install -y --no-install-recommends libopencv-dev

# Install Python 3.10 via deadsnakes PPA
RUN add-apt-repository ppa:deadsnakes/ppa && \
    apt-get update && \
    apt-get install -y --no-install-recommends \
    python3.10 python3.10-dev python3.10-distutils

# Set Python 3.10 as default
RUN ln -sf /usr/bin/python3.10 /usr/bin/python3 && \
    ln -sf /usr/bin/python3.10 /usr/bin/python

# Install pip for Python 3.10
RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.10 && \
    pip install --upgrade pip

# Add user
ARG UID
ARG USER
ARG PASSWORD
RUN groupadd -g 1000 ${USER}_group && \
    useradd -m --uid=${UID} --gid=${USER}_group --groups=sudo ${USER} && \
    echo ${USER}:${PASSWORD} | chpasswd && \
    echo 'root:root' | chpasswd

ENV PATH=${PATH}:/home/${USER}/.local/bin

# Setup working directory
ENV WORK_DIR=/workspace
RUN mkdir ${WORK_DIR} && \
    chown ${USER}:${USER}_group ${WORK_DIR}
WORKDIR ${WORK_DIR}

# Switch to user
USER ${USER}

# Install PyTorch with CUDA 12.1
RUN pip install torch==2.4.1 torchvision==0.19.1 torchaudio==2.4.1 --index-url https://download.pytorch.org/whl/cu121

# Install Python libraries
RUN pip install --no-cache-dir \
    numpy \
    matplotlib \
    pandas \
    seaborn \
    scipy \
    natsort \
    scikit-learn \
    tensorboard \
    opencv-python \
    h5py \
    tqdm \
    pyyaml \
    open_clip_torch \
    scikit-image \
    timm \
    transformers \
    xformers \
    pot