fix(#164): tensorflow api changed

williamfzc · Sep 24, 2021 · 243cee5 · 243cee5
1 parent f522b62
commit 243cee5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/stagesepx/classifier/keras.py b/stagesepx/classifier/keras.py
@@ -204,7 +204,7 @@ def predict_with_object(self, frame: np.ndarray) -> str:
         frame = cv2.resize(frame, dsize=self.data_size)
         frame = np.expand_dims(frame, axis=[0, -1])
 
-        return str(self._model.predict_classes(frame)[0])
+        return str(np.argmax(self._model.predict(frame), axis=1)[0])
 
     def _classify_frame(self, frame: VideoFrame, *_, **__) -> str:
         return self.predict_with_object(frame.data)