Add docker support and fix GPU memory exceed issue

.gitignore

@@ -1,3 +1,4 @@
+splash_results/
 # Data files and directories common in repo root
 datasets/
 logs/

bubble/bubble.py

@@ -46,6 +46,7 @@
 from mrcnn import visualize
 import tensorflow as tf
 import datetime
+from solve_cudnn_error import *
 
 # Root directory of the project
 ROOT_DIR = os.path.abspath("../")
@@ -509,7 +510,7 @@ def splash(model, image_path=None, video_path=None, result_path=RESULTS_DIR):
 
 if __name__ == '__main__':
     import argparse
-
+    solve_cudnn_error()
     # Parse command line arguments
     parser = argparse.ArgumentParser(
         description='Train Mask R-CNN to detect bubbles.')

bubble/solve_cudnn_error.py

@@ -0,0 +1,14 @@
+import tensorflow as tf
+def solve_cudnn_error():
+    gpus = tf.config.experimental.list_physical_devices('GPU')
+    if gpus:
+        try:
+            # Currently, memory growth needs to be the same across GPUs
+            for gpu in gpus:
+                tf.config.experimental.set_memory_growth(gpu, True)
+            logical_gpus = tf.config.experimental.list_logical_devices('GPU')
+            print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
+        except RuntimeError as e:
+            # Memory growth must be set before GPUs have been initialized
+            print(e)
+

docker/Dockerfile

@@ -0,0 +1,23 @@
+FROM tensorflow/tensorflow:1.14.0-gpu-py3
+
+COPY ./requirements.txt /root
+
+RUN pip install --upgrade pip setuptools wheel
+
+RUN cd /root \
+    && pip install -r requirements.txt --verbose \
+    && pip install --no-deps imgaug==0.4.0
+
+RUN rm /etc/apt/sources.list.d/cuda.list \
+    && apt-get update \
+    && apt-get install vim git ffmpeg libsm6 libxext6  -y
+
+RUN cd /root \
+    && git clone https://github.com/ywflow/BubMask.git
+
+RUN cd /root/BubMask \
+    && python setup.py install \
+    && cd .. \
+    && rm -r BubMask
+
+WORKDIR /root/BubMask

docker/README.md

@@ -0,0 +1,27 @@
+# BubMask Docker
+
+## Deploy container
+
+```bash
+docker compose up -d
+```
+
+## Enter container and test sample image
+
+Make sure you have download mask_rcnn_bubble.h5 and place inside `BubMask/weights`
+
+```bash
+docker start bubmask-container
+docker exec -it bubmask-container /bin/bash
+python ./bubble/bubble.py splash --weights=./weights/mask_rcnn_bubble.h5 --image=./sample_image/Expansion_pipe.jpg --results=./splash_results --confidence=0.01
+```
+
+## Build image
+
+```bash
+docker build -t frakw/bubmask:latest .
+```
+
+You can also use pre-build image here:
+
+[https://hub.docker.com/repository/docker/frakw/bubmask/tags](https://hub.docker.com/repository/docker/frakw/bubmask/tags)

docker/docker-compose.yaml

@@ -0,0 +1,24 @@
+services:
+    bubmask:
+        stdin_open: true
+        tty: true
+        devices:
+            - /dev/dri
+        volumes:
+            - ..:/root/BubMask
+        deploy:
+            resources:
+                reservations:
+                    devices:
+                        - driver: nvidia
+                          count: all
+                          capabilities:
+                              - gpu
+        environment:
+            - QT_X11_NO_MITSHM=1
+            - NVIDIA_VISIBLE_DEVICES=all
+            - NVIDIA_DRIVER_CAPABILITIES=all
+        container_name: bubmask-container
+        build: .
+        image: frakw/bubmask:latest
+        command: /bin/bash

docker/requirements.txt

@@ -0,0 +1,16 @@
+opencv-python-headless==4.6.0.66
+opencv-contrib-python-headless==4.6.0.66
+numpy==1.19.5
+scipy==1.5.4
+Pillow==8.4.0
+Cython==3.0.11
+matplotlib==3.3.4
+scikit-image==0.17.2
+Keras==2.0.8
+h5py==2.10.0
+#imgaug==0.4.0
+ipykernel==5.5.6
+ipyparallel==8.2.1
+ipython==7.16.3
+ipython-genutils==0.2.0
+ipywidgets==7.8.5