Browse Source
modify
Signed-off-by: gexy5 <xinyu.ge@zilliz.com>
main
gexy5
3 years ago
2 changed files with
4 additions and
4 deletions
-
kinetics_400.txt
-
uniformer.py
|
|
@ -50,14 +50,14 @@ class Uniformer(NNOperator): |
|
|
if 'k400' in model_name: |
|
|
if 'k400' in model_name: |
|
|
self.dataset_name = 'kinetics_400' |
|
|
self.dataset_name = 'kinetics_400' |
|
|
if classmap is None: |
|
|
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(): |
|
|
for line in f.readlines(): |
|
|
line = line.split('\t') |
|
|
line = line.split('\t') |
|
|
k = str(line[0]) |
|
|
k = str(line[0]) |
|
|
v = int(line[1]) |
|
|
v = int(line[1]) |
|
|
classmap[v] = str(k).replace('"', '') |
|
|
|
|
|
|
|
|
self.classmap[v] = str(k).replace('"', '') |
|
|
else: |
|
|
else: |
|
|
self.classmap = classmap |
|
|
self.classmap = classmap |
|
|
self.device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
|
self.device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
|
|
