Change target class to Category

opendr-eu · Jun 7, 2021 · 84cb8cf · 84cb8cf
1 parent 59ef9e3
commit 84cb8cf
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/src/opendr/perception/speech_recognition/edgespeechnets/edgespeechnets_learner.py b/src/opendr/perception/speech_recognition/edgespeechnets/edgespeechnets_learner.py
@@ -24,7 +24,7 @@
 
 from opendr.engine.data import Timeseries
 from opendr.engine.learners import Learner
-from opendr.engine.target import SpeechCommand
+from opendr.engine.target import Category
 from opendr.perception.speech_recognition.edgespeechnets.algorithm.audioutils import get_mfcc
 import opendr.perception.speech_recognition.edgespeechnets.algorithm.models as models
 
@@ -218,7 +218,7 @@ def infer(self, batch):
         prediction = output.max(1, keepdim=True)
         batch_predictions = []
         for target, confidence in zip(prediction[1], prediction[0].exp()):
-            batch_predictions.append(SpeechCommand(target.item(), confidence=confidence.item()))
+            batch_predictions.append(Category(target.item(), confidence=confidence.item()))
         return batch_predictions[0] if len(batch_predictions) == 1 else batch_predictions
 
     def save(self, path):

diff --git a/tests/sources/tools/perception/speech_recognition/edgespeechnets/test_edgespeechnets.py b/tests/sources/tools/perception/speech_recognition/edgespeechnets/test_edgespeechnets.py
@@ -23,7 +23,7 @@
 from opendr.perception.speech_recognition.edgespeechnets.edgespeechnets_learner import EdgeSpeechNetsLearner
 from opendr.engine.data import Timeseries
 from opendr.engine.datasets import DatasetIterator
-from opendr.engine.target import SpeechCommand
+from opendr.engine.target import Category
 
 TEST_BATCH_SIZE = 2
 TEST_EPOCHS = 1
@@ -78,12 +78,12 @@ def test_infer_batch(self):
         batch = [Timeseries(np.ones((1, TEST_SIGNAL_LENGTH))) for _ in range(TEST_INFER_LENGTH)]
         results = self.learner.infer(batch)
         self.assertTrue(len(results) == TEST_INFER_LENGTH)
-        self.assertTrue(all([isinstance(x, SpeechCommand) for x in results]))
+        self.assertTrue(all([isinstance(x, Category) for x in results]))
 
     def test_infer_pure_signal(self):
         signal = Timeseries(np.ones((1, TEST_SIGNAL_LENGTH)))
         result = self.learner.infer(signal)
-        self.assertTrue(isinstance(result, SpeechCommand))
+        self.assertTrue(isinstance(result, Category))
 
     def test_reset(self):
         weights_before_reset = list(self.learner.model.parameters())[0].clone()