FROM nvidia/cuda:11.1.1-cudnn8-devel-ubuntu18.04

RUN apt-get update && apt-get install -y \
	git wget sudo ca-certificates ninja-build

# Install miniconda - copy from existing Docker image 
COPY --from=continuumio/miniconda3:latest /opt/conda /opt/conda
ENV PATH=/opt/conda/bin:$PATH
RUN /opt/conda/bin/conda clean -a -y \
    && ln -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh \
    && echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc 

# create conda environment 'maskfreevis', make 'maskfreevis' active
RUN . ~/.bashrc && conda create --name maskfreevis python=3.8 -y \
    && echo "conda activate maskfreevis" >> ~/.bashrc

RUN . ~/.bashrc && conda activate maskfreevis && conda install pytorch==1.10 torchvision==0.11.1 -c pytorch -c nvidia 
RUN . ~/.bashrc && conda activate maskfreevis && pip install numpy==1.20.3
RUN . ~/.bashrc && conda activate maskfreevis && pip install opencv-python-headless
RUN . ~/.bashrc && conda activate maskfreevis && pip install jupyter

# Set environment variables for detectron2; 'from torch.utils.cpp_extension import CUDA_HOME' sets CUDA_HOME as '/usr/local/cuda'
ENV FORCE_CUDA="1"		
# This will by default build detectron2 for all common cuda architectures and take a lot more time,
# because inside `docker build`, there is no way to tell which architecture will be used.
ARG TORCH_CUDA_ARCH_LIST="Kepler;Kepler+Tesla;Maxwell;Maxwell+Tegra;Pascal;Volta;Turing"
ENV TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}"
ENV FVCORE_CACHE="/tmp"

# compile CUDA kernel for MSDeformAttn
RUN git clone https://github.com/SysCV/MaskFreeVIS.git
RUN cd ./MaskFreeVIS/mask2former/modeling/pixel_decoder/ops && . ~/.bashrc && conda activate maskfreevis && sh make.sh
# validate MultiScaleDeformableAttention is built
RUN . ~/.bashrc && conda activate maskfreevis && python -c "import torch; import MultiScaleDeformableAttention as MSDA"

# install detectron2
RUN . ~/.bashrc && conda activate maskfreevis && pip install --user --no-cache-dir tensorboard cmake   # cmake from apt-get is too old
RUN . ~/.bashrc && conda activate maskfreevis && pip install --user 'git+https://github.com/facebookresearch/fvcore'
RUN . ~/.bashrc && conda activate maskfreevis && git clone --depth 1 https://github.com/facebookresearch/detectron2 detectron2_repo && pip install --user -e detectron2_repo

# install requirements.txt from MaskFreeVIS
RUN cd /MaskFreeVIS && . ~/.bashrc && conda activate maskfreevis && pip install --user -r requirements.txt

WORKDIR /MaskFreeVIS
EXPOSE 8888
ENTRYPOINT . ~/.bashrc && conda activate maskfreevis && jupyter notebook --ip=0.0.0.0 --no-browser --allow-root

# build:
# 	docker build --progress=plain --rm --force-rm -t maskfreevis .
# usage (bash):    docker run --rm -it --entrypoint="/bin/bash" maskfreevis
# usage (bash, GPU, mapped drive):    docker run --gpus all -v /c/tmp/dataout:/dataout --rm -it --entrypoint="/bin/bash" maskfreevis
# usage (Jupyter, GPU, mapped drive): docker run --gpus all -p 8888:8888 -v /c/tmp/dataout:/dataout --rm -it maskfreevis

# detectron2 demo
# cd /detectron2_repo
# wget http://images.cocodataset.org/val2017/000000439715.jpg -O input.jpg
# python3 demo/demo.py  \
#    --config-file configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml \
#    --input input.jpg --output outputs/ \
#    --opts MODEL.WEIGHTS detectron2://COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl

# MaskFreeVIS demo_video with video input: before running, modify line 162 in demo_video/demo.py as
#    predictions, visualized_output = demo.run_on_video(vid_frames, args.confidence_threshold) # add a missing parameter
# python3 demo_video/demo.py --config-file configs/youtubevis_2019/swin/video_maskformer2_swin_large_IN21k_384_bs16_8ep.yaml --video-input /dataout/videoin/basketball.mp4 --output /dataout/videoout/ --opts MODEL.WEIGHTS /dataout/MaskFreeVis/mfvis_models/model_final_swinl_0560.pth

# MaskFreeVIS demo_video with images input: before running, modify line 140 in demo_video/demo.py by replacing 'fps=5' with 'duration = 200'
# python3 demo_video/demo.py --config-file configs/youtubevis_2019/swin/video_maskformer2_swin_large_IN21k_384_bs16_8ep.yaml --save-frames True --input /dataout/MaskFreeVis/frames/valid/JPEGImages/ --output /dataout/videoout/ --opts MODEL.WEIGHTS /dataout/MaskFreeVis/mfvis_models/model_final_swinl_0560.pth