Fix bugs with predict.

nasa · May 4, 2020 · a68c195 · a68c195
1 parent c3bde5b
commit a68c195
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/delta/imagery/sources/tiff.py b/delta/imagery/sources/tiff.py
@@ -429,7 +429,8 @@ def close(self):
         """
         Finish writing.
         """
-        self._tiff_w.close()
+        if self._tiff_w is not None:
+            self._tiff_w.close()
 
     def abort(self):
         self.close()

diff --git a/delta/ml/predict.py b/delta/ml/predict.py
@@ -183,7 +183,8 @@ def _initialize(self, shape, label, image):
         if self._prob_image:
             self._prob_image.initialize((shape[0], shape[1], self._num_classes), np.float32, image.metadata())
         if self._error_image:
-            self._prob_image.initialize((shape[0], shape[1], self._num_classes), np.float32, image.metadata())
+            self._error_image.initialize((shape[0], shape[1], self._error_colors.shape[1]),
+                                         self._error_colors.dtype, image.metadata())
 
     def _complete(self):
         if self._output_image:
@@ -214,7 +215,7 @@ def _process_block(self, pred_image, x, y, labels):
                 self._output_image.write(pred_image, x, y)
 
         if labels is not None:
-            self._error_image.write(labels != pred_image, x, y)
+            self._error_image.write(self._error_colors[(labels != pred_image).astype(int)], x, y)
             cm = tf.math.confusion_matrix(np.ndarray.flatten(labels),
                                           np.ndarray.flatten(pred_image),
                                           self._num_classes)