jvalero commited on
Commit
830a47f
1 Parent(s): 38b08e2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +13 -1
app.py CHANGED
@@ -16,6 +16,18 @@ pipe = pipeline(
16
  device=device,
17
  )
18
 
 
 
 
 
 
 
 
 
 
 
 
 
19
  def get_edm(filepath):
20
  output = pipe(
21
  filepath,
@@ -23,7 +35,7 @@ def get_edm(filepath):
23
  chunk_length_s=30,
24
  batch_size=8,
25
  )
26
- return output[0]["label"]
27
 
28
 
29
  demo = gr.Blocks()
 
16
  device=device,
17
  )
18
 
19
+
20
+ id2label = {
21
+ "LABEL_0" : 'drumbass',
22
+ "LABEL_1" : 'dubtechno',
23
+ "LABEL_2": 'dupstep',
24
+ "LABEL_3":'hardcore_breaks',
25
+ "LABEL_4": 'house',
26
+ "LABEL_5":'psytrance',
27
+ "LABEL_6" : 'techno',
28
+ "LABEL_7":'ukgarage'
29
+ }
30
+
31
  def get_edm(filepath):
32
  output = pipe(
33
  filepath,
 
35
  chunk_length_s=30,
36
  batch_size=8,
37
  )
38
+ return id2label[output[0]["label"]]
39
 
40
 
41
  demo = gr.Blocks()