ENH: Add model path by default in the docker container

The trained model is download and extracted automatically when building the docker container

Docker/Dockerfile

@@ -16,7 +16,12 @@ RUN	unzip vtk8.2.0.zip
 RUN pip install itk sklearn pandas matplotlib
 
 WORKDIR /app
+
 RUN git clone https://github.com/DCBIA-OrthoLab/fly-by-cnn.git
 
+RUN wget https://github.com/DCBIA-OrthoLab/fly-by-cnn/releases/download/2.1/u_seg_nn_v3.0.zip
+RUN	unzip u_seg_nn_v3.0.zip
+
+
 ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/python3.6/dist-packages/vtkmodules/
 ENV MESA_GL_VERSION_OVERRIDE=3.2

src/py/predict_v3.py

@@ -12,7 +12,7 @@
 
 parser = argparse.ArgumentParser(description='Predict an input with a trained neural network', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
 parser.add_argument('--surf', type=str, help='Input surface mesh to label', required=True)
-parser.add_argument('--model', type=str, help='Model to do segmentation', required=True)
+parser.add_argument('--model', type=str, help='Model to do segmentation', default="/app/u_seg_nn_v3.0")
 parser.add_argument('--out', type=str, help='Output model with labels', default="out.vtk")
 
 args = parser.parse_args()
@@ -42,8 +42,11 @@
 flyby_features.removeActors()
 
 
-model = tf.keras.models.load_model(args.model, custom_objects={'tf': tf})
-model.summary()
+if os.path.exist(args.model):
+	model = tf.keras.models.load_model(args.model, custom_objects={'tf': tf})
+	model.summary()
+else:
+	print("Please set the model directory to a valid path", file=sys.stderr)
 
 print("Predict ...")
 img_predict_np = model.predict(img_np)