From 28b11aaa44fdeb645412e010d811ee8e2dc4ff03 Mon Sep 17 00:00:00 2001
From: Piv <18462828+Piv200@users.noreply.github.com>
Date: Fri, 10 Sep 2021 22:13:41 +0930
Subject: [PATCH] Use image type for CoreML input, reshape last layer to work
 with metal shader

---
 coreml.py                | 4 ++--
 fast_depth_functional.py | 1 +
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/coreml.py b/coreml.py
index db2fbae..9c091d8 100644
--- a/coreml.py
+++ b/coreml.py
@@ -1,8 +1,8 @@
 import coremltools as ct
 
 
-def convert_coreml(model_path, save_path='../mobilenet_nnconv5.mlmodel'):
-    mlmodel = ct.convert(model_path)
+def convert_coreml(model_path, save_path='mobilenet_nnconv5.mlmodel'):
+    mlmodel = ct.convert(model_path, inputs=[ct.ImageType()])
     mlmodel.save(save_path)
 
 
diff --git a/fast_depth_functional.py b/fast_depth_functional.py
index 672f224..91eb943 100644
--- a/fast_depth_functional.py
+++ b/fast_depth_functional.py
@@ -72,6 +72,7 @@ def mobilenet_nnconv5(weights=None, shape=(224, 224, 3)):
     x = keras.layers.Conv2D(1, 1, padding='same')(x)
     x = keras.layers.BatchNormalization()(x)
     x = keras.layers.ReLU(6.)(x)
+    x = keras.layers.Reshape([shape[0], shape[1]])(x)
     return keras.Model(inputs=input, outputs=x, name="fast_depth")