[Backport] Orca tf example matplotlib failed on macos bugfix (#3236)

* fix matplotlib crash on macos

* fix jenkins error

* mac fix

pyzoo/zoo/examples/orca/learn/tf/image_segmentation/README.md

@@ -24,16 +24,24 @@ You should manually download the dataset from kaggle [carvana-image-masking-chal
 
 ## Run example on local
 ```bash
-python image_segmentation.py --cluster_mode local 
+# linux
+python image_segmentation.py --cluster_mode local
+# macos
+python image_segmentation.py --cluster_mode local --platform mac
 ```
 
 ## Run example on yarn cluster
 ```bash
-python image_segmentation.py --cluster_mode yarn 
+# linux
+python image_segmentation.py --cluster_mode yarn
+# macos
+python image_segmentation.py --cluster_mode yarn --platform mac
 ```
 
 Options
 * `--cluster_mode` The mode for the Spark cluster. local or yarn. Default to be `local`.
 * `--file_path` The path to carvana train.zip, train_mask.zip and train_mask.csv.zip. Default to be `/tmp/carvana/`.
 * `--epochs` The number of epochs to train the model. Default to be 8.
 * `--batch_size` Batch size for training and prediction. Default to be 8.
+* `--platform` The platform you used to run the example. Default to be `linux`. You should pass `mac` if you use macos.
+* `--non_interactive` Flag to not visualize the result.

pyzoo/zoo/examples/orca/learn/tf/image_segmentation/image_segmentation.py

@@ -46,7 +46,11 @@ def load_data(file_path):
  load_data_from_zip(file_path, 'train_masks.csv.zip')
 
 
-def main(cluster_mode, max_epoch, file_path, batch_size):
+def main(cluster_mode, max_epoch, file_path, batch_size, platform, non_interactive):
+ import matplotlib
+ if not non_interactive and platform == "mac":
+ matplotlib.use('qt5agg')
+
  if cluster_mode == "local":
  init_orca_context(cluster_mode="local", cores=4, memory="3g")
  elif cluster_mode == "yarn":
@@ -174,25 +178,27 @@ def bce_dice_loss(y_true, y_pred):
  val_image_label = val_shards.collect()[0]
  val_image = val_image_label["x"]
  val_label = val_image_label["y"]
- # visualize 5 predicted results
- plt.figure(figsize=(10, 20))
- for i in range(5):
- img = val_image[i]
- label = val_label[i]
- predicted_label = pred[i]
-
- plt.subplot(5, 3, 3 * i + 1)
- plt.imshow(img)
- plt.title("Input image")
-
- plt.subplot(5, 3, 3 * i + 2)
- plt.imshow(label[:, :, 0], cmap='gray')
- plt.title("Actual Mask")
- plt.subplot(5, 3, 3 * i + 3)
- plt.imshow(predicted_label, cmap='gray')
- plt.title("Predicted Mask")
- plt.suptitle("Examples of Input Image, Label, and Prediction")
- plt.show()
+ if not non_interactive:
+ # visualize 5 predicted results
+ plt.figure(figsize=(10, 20))
+ for i in range(5):
+ img = val_image[i]
+ label = val_label[i]
+ predicted_label = pred[i]
+
+ plt.subplot(5, 3, 3 * i + 1)
+ plt.imshow(img)
+ plt.title("Input image")
+
+ plt.subplot(5, 3, 3 * i + 2)
+ plt.imshow(label[:, :, 0], cmap='gray')
+ plt.title("Actual Mask")
+ plt.subplot(5, 3, 3 * i + 3)
+ plt.imshow(predicted_label, cmap='gray')
+ plt.title("Predicted Mask")
+ plt.suptitle("Examples of Input Image, Label, and Prediction")
+
+ plt.show()
 
  stop_orca_context()
 
@@ -207,6 +213,11 @@ def bce_dice_loss(y_true, y_pred):
  help="The number of epochs to train the model")
  parser.add_argument('--batch_size', type=int, default=8,
  help="Batch size for training and prediction")
+ parser.add_argument('--platform', type=str, default="linux",
+ help="The platform you used. Only linux and mac are supported.")
+ parser.add_argument('--non_interactive', default=False, action="store_true",
+ help="Flag to not visualize the result.")
 
  args = parser.parse_args()
- main(args.cluster_mode, args.epochs, args.file_path, args.batch_size)
+ main(args.cluster_mode, args.epochs, args.file_path, args.batch_size, args.platform,
+ args.non_interactive)

pyzoo/zoo/examples/run-example-tests-pip.sh

@@ -602,7 +602,7 @@ fi
 # Run the example
 export SPARK_DRIVER_MEMORY=3g
 python ${ANALYTICS_ZOO_ROOT}/pyzoo/zoo/examples/orca/learn/tf/image_segmentation/image_segmentation.py \
- --file_path analytics-zoo-data/data/carvana
+ --file_path analytics-zoo-data/data/carvana --epochs 1 --non_interactive
 exit_status=$?
 if [ $exit_status -ne 0 ];
 then

pyzoo/zoo/examples/run-example-tests.sh

@@ -622,7 +622,7 @@ ${ANALYTICS_ZOO_HOME}/bin/spark-submit-python-with-zoo.sh \
  --driver-memory 3g \
  --executor-memory 3g \
  ${ANALYTICS_ZOO_ROOT}/pyzoo/zoo/examples/orca/learn/tf/image_segmentation/image_segmentation.py \
- --file_path analytics-zoo-data/data/carvana --epochs 1
+ --file_path analytics-zoo-data/data/carvana --epochs 1 --non_interactive
 exit_status=$?
 if [ $exit_status -ne 0 ];
 then