Merge pull request #166 from williamfzc/fix-tensorflow-26

fix(#164): tensorflow api changed
williamfzc · Oct 9, 2021 · 90987fb · 90987fb
2 parents f522b62 + 243cee5
commit 90987fb
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)