Mirroristas/whisper
2
0
Fork 0
mirror of https://github.com/openai/whisper.git synced 2025-11-26 15:35:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

add colorized text

Browse Source
This commit is contained in:
khaliladib11 2024-05-26 11:28:35 +01:00
parent 34d9a9b6a3
commit 0fd17c99c8
No known key found for this signature in database
1 changed files with 13 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
examples/test_prob.py
View File

@ -48,6 +48,18 @@ def decode_audio(model, audio, language="en", f16=True):
return text_tokens, result.token_probs
def get_colored_text(text_tokens: List[int], token_probs: List[float]):
init(autoreset=False) # Initialize colorama with autoreset=True to reset colors after each print
output_text = ""
for i, (token, prob) in enumerate(zip(text_tokens, token_probs)):
# Interpolate between red and green in the HSV color space
r, g, b = colorsys.hsv_to_rgb(prob * (1/3), 1, 1)
r, g, b = int(r * 255), int(g * 255), int(b * 255)
color_code = f"\033[38;2;{r};{g};{b}m"
colored_token = f"{color_code}{Style.BRIGHT}{str(token)}{Style.RESET_ALL}"
output_text += colored_token
return output_text
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--audio', type=str, help='the path of the audio file')
@ -63,5 +75,4 @@ if __name__ == '__main__':
model = whisper.load_model(model)
audio = load_audio_from_source(audio_source=audio)
text, proba = decode_audio(model=model, audio=audio)
print(text)
print(proba)
print(get_colored_text(text, proba))