add omnivore

Signed-off-by: gexy5 <xinyu.ge@zilliz.com>
2 years ago · 55ad855f26
1 changed files with 2 additions and 2 deletions
--- a/omnivore.py
+++ b/omnivore.py
@ -99,9 +99,9 @@ class Omnivore(NNOperator):
                    )
        inputs = data.to(self.device)[None, ...]

-        outs = self.model(inputs)
+        outs = self.model(inputs,input_type="video")
        post_act = torch.nn.Softmax(dim=1)
-        preds = post_act(outs,input_type="video")
+        preds = post_act(outs)
        pred_scores, pred_classes = preds.topk(k=self.topk)
        labels = [self.classmap[int(i)] for i in pred_classes[0]]
        scores = [round(float(x), 5) for x in pred_scores[0]]