omermazig commited on
Commit
596e2ad
1 Parent(s): ed24b85

Get labels in original order (Got them in the wrong order before which effected the results)

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -50,7 +50,8 @@ inference_transform = Compose(
50
  ]
51
  )
52
 
53
- labels = list(trained_model.config.label2id.keys())
 
54
 
55
 
56
  def parse_video_to_clips(video_file):
 
50
  ]
51
  )
52
 
53
+ num_labels = trained_model.config.num_labels
54
+ labels = [trained_model.config.id2label[i] for i in range(num_labels)]
55
 
56
 
57
  def parse_video_to_clips(video_file):