Dockerfile for building pytorch, pyplot & cv2

Dockerfile

##############################################################################
# Dockerfile for building an image with non-root privilege as default
# In addition, it's capable of displaying images created by
# `cv2.imshow()` and `plt.show()`

# In order to use the full functionality, one should 
#
# 1. Run the command `xhost +local:docker`
#
# 2. build an image by including
# `-e DISPLAY=$DISPLAY`, `-e XAUTHORITY=/tmp/.docker.xauth`
# and `-v /tmp/.X11-unix:/tmp/.X11.unix` in the `docker build` command
##############################################################################

# `TAG`: tag of the PyTorch base image
# A list of all available tags are shown in this website:
# https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch/tags
ARG TAG
FROM nvcr.io/nvidia/pytorch:$TAG

# `TZ`: specifies the timezoned. Necessary for installing `python3-tk` non-interactively
ARG TZ=Europe/Berlin
# `UID`: desired uid of the login user
ARG UID=1001
# `USERNAME`: desired uid of the login user
ARG USERNAME
# `PSW`: desired password for the login user
ARG PSW

RUN apt-get update \
    && apt-get upgrade -y \
    # `python3-tk` is necessary for `plt.show()`
    # `libgtk2.0-dev`, `pkg-config` and `libgl1-mesa-glx` are necessary for `cv2.imshow()`
    && DEBIAN_FRONTEND=noninteractive TZ=$TZ \
      apt-get install sudo python3-tk libgtk2.0-dev pkg-config libgl1-mesa-glx -y \
    # Clean up
    && apt-get autoremove -y \
    && apt-get clean -y \
    && rm -rf /var/lib/apt/list* \
    # Create a non-root user and grant him / her the sudo permission
    && useradd -u $UID -d /home/$USERNAME -m -k /etc/skel $USERNAME \
    && echo "$USERNAME:$PSW" | chpasswd \
    && adduser $USERNAME sudo

# Change to non-root privilege
USER $USERNAME

# Update `pip` and re-install `opencv-python` with non-root permission
# Re-installing `opencv-python` is necessary for `cv2.imshow()`
RUN python -m pip install --user -U --no-cache-dir pip opencv-python