jlvdoorn commited on
Commit
bffc737
1 Parent(s): 2348aa3

added app.py

Browse files
Files changed (1) hide show
  1. app.py +50 -0
app.py ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import pipeline
2
+ import gradio as gr
3
+ import os
4
+
5
+ whisper = pipeline(model='jlvdoorn/whisper-large-v2-atco2-asr-atcosim', use_auth_token=os.environ['HUGGINGFACE_TOKEN'])
6
+ bert_atco_ner = pipeline(model='Jzuluaga/bert-base-ner-atc-en-atco2-1h')
7
+
8
+ def transcribe(audio_mic, audio_file):
9
+ if audio_file is not None:
10
+ return whisper(audio_file)['text']
11
+ if audio_mic is not None:
12
+ return whisper(audio_mic)['text']
13
+ else:
14
+ return 'There was no audio to transcribe...'
15
+
16
+ def extractCallSignCommand(transcription):
17
+ if type(transcription) is str:
18
+ result = bert_atco_ner(transcription)
19
+ callsigns = []
20
+ commands = []
21
+ values = []
22
+ for item in result:
23
+ if 'callsign' in item['entity']:
24
+ callsigns.append(item['word'])
25
+ if 'command' in item['entity']:
26
+ commands.append(item['word'])
27
+ if 'value' in item['entity']:
28
+ values.append(item['word'])
29
+
30
+ return 'Callsigns: ' + ', '.join(callsigns) + '\nCommands: ' + ', '.join(commands) + '\nValues: ' + ', '.join(values)
31
+ else:
32
+ return 'There was no transcription to extract a callsign or command from...'
33
+
34
+ def transcribeAndExtract(audio_mic, audio_file, transcribe_only):
35
+ transcription = transcribe(audio_mic, audio_file)
36
+ if not transcribe_only:
37
+ callSignCommandValues = extractCallSignCommand(transcription)
38
+ else:
39
+ callSignCommandValues = ''
40
+ return transcription, callSignCommandValues
41
+
42
+ iface = gr.Interface(
43
+ fn=transcribeAndExtract,
44
+ inputs=[gr.Audio(source='microphone', type='filepath'), gr.Audio(source='upload', type='filepath'), gr.Checkbox(label='Transcribe only', default=False)],
45
+ outputs=[gr.Text(label='Transcription'), gr.Text(label='Callsigns, commands and values')],
46
+ title='Whisper Large v2 - ATCO2-ASR-ATCOSIM',
47
+ description='Whisper Large v2 model fine-tuned on the ATCO2-ASR and ATCOSIM datasets.'
48
+ )
49
+
50
+ iface.launch()