Update Dockerfile

Dockerfile

@@ -1,29 +1,24 @@
-FROM python:3.9
-
-WORKDIR /code
-
-RUN apt-get update && apt-get install ffmpeg libsm6 libxext6  -y
-
-COPY ./requirements.txt /code/requirements.txt
-RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
-
-RUN pip install -qq -U openmim
-RUN mim install -U mmengine 'mmcv>=2.0.0rc1'
-RUN pip install -qq openmim 'mmdet>=3.0.0rc0' 'mmocr>=1.0.0rc0'
-
-# Set up a new user named "user" with user ID 1000
-RUN useradd -m -u 1000 user
-# Switch to the "user" user
-USER user
-# Set home to the user's home directory
-ENV HOME=/home/user \
-    PATH=/home/user/.local/bin:$PATH
-
-# Set the working directory to the user's home directory
-WORKDIR $HOME/app
-
-# Copy the current directory contents into the container at $HOME/app setting the owner to the user
-COPY --chown=user . $HOME/app
-RUN ls
-
-CMD ["python", "main.py"]
+ARG PYTORCH="1.9.0"
+ARG CUDA="11.1"
+ARG CUDNN="8"
+
+FROM pytorch/pytorch:${PYTORCH}-cuda${CUDA}-cudnn${CUDNN}-devel
+
+# To fix GPG key error when running apt-get update
+RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub
+RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/7fa2af80.pub
+
+RUN apt-get update && apt-get install -y git ninja-build libglib2.0-0 libsm6 libxrender-dev libxext6 libgl1-mesa-glx \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN pip install -U openmim
+RUN mim install mmengine
+RUN mim install 'mmcv>=2.0.0rc1'
+RUN pip install 'mmdet>=3.0.0rc0'
+RUN git clone https://github.com/open-mmlab/mmocr.git /mmocr
+WORKDIR /mmocr
+# ENV FORCE_CUDA="1"
+RUN pip install -r requirements.txt
+RUN pip install --no-cache-dir -e .
+RUN pip install -r requirements/albu.txt