From f7783be52e46190b55ca7cc0cae06c806e6d1738 Mon Sep 17 00:00:00 2001 From: gexy5 Date: Mon, 20 Jun 2022 15:37:37 +0800 Subject: [PATCH] modify Signed-off-by: gexy5 --- kinetics_400_categroies.txt => kinetics_400.txt | 0 uniformer.py | 8 ++++---- 2 files changed, 4 insertions(+), 4 deletions(-) rename kinetics_400_categroies.txt => kinetics_400.txt (100%) diff --git a/kinetics_400_categroies.txt b/kinetics_400.txt similarity index 100% rename from kinetics_400_categroies.txt rename to kinetics_400.txt diff --git a/uniformer.py b/uniformer.py index 02e0b4a..28d6a02 100644 --- a/uniformer.py +++ b/uniformer.py @@ -50,14 +50,14 @@ class Uniformer(NNOperator): if 'k400' in model_name: self.dataset_name = 'kinetics_400' if classmap is None: - class_file = os.path.join(str(Path(__file__).parent), self.dataset_name+'.json') - classmap = {} - with open("kinetics_400_categroies.txt", "r") as f: + class_file = os.path.join(str(Path(__file__).parent), self.dataset_name+'.txt') + self.classmap = {} + with open(class_file, "r") as f: for line in f.readlines(): line = line.split('\t') k = str(line[0]) v = int(line[1]) - classmap[v] = str(k).replace('"', '') + self.classmap[v] = str(k).replace('"', '') else: self.classmap = classmap self.device = 'cuda' if torch.cuda.is_available() else 'cpu'