server.py 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113
  1. from typing import Union
  2. from fairseq import tasks
  3. from fairseq.data.dictionary import Dictionary
  4. from fastapi import FastAPI, File, UploadFile
  5. from pydantic import BaseModel
  6. import logging
  7. import argparse
  8. import torch
  9. import torch.nn as nn
  10. import pickle
  11. import soundfile as sf
  12. import torch.nn.functional as F
  13. import yaml
  14. import os, sys
  15. import numpy as np
  16. # from decoder_exps.decode_common import W2V2Decoder, Wav2VecCtc
  17. # from wav2vecEncoder import Wav2VecCtc as CustomWav2VecCtc
  18. # from fairseq.models.wav2vec.wav2vec2_asr import Wav2VecCtc
  19. from inference import inference, inference_file, inference_online
  20. # from inference import inference_file
  21. logging.basicConfig()
  22. logging.root.setLevel(logging.INFO)
  23. logging.basicConfig(level=logging.INFO)
  24. logger = logging.getLogger(__name__)
  25. class Recording(BaseModel):
  26. filename: str
  27. content_type: str
  28. content: list
  29. class AudioClip(BaseModel):
  30. device: str
  31. time: str
  32. recording: Recording
  33. # YAML_FILE = "config/base_org.yaml"
  34. # args = dict()
  35. # os.path.abspath(os.path.dirname(__file__))
  36. # with open(YAML_FILE, 'r') as f:
  37. # args.update(yaml.safe_load(f))
  38. ###############################################################################
  39. ## FastAPI
  40. ###############################################################################
  41. app = FastAPI()
  42. @app.get("/")
  43. def root():
  44. return {"message" : "Hello World!!!"}
  45. @app.post("/test")
  46. def post_test(audioClip: AudioClip):
  47. output = ""
  48. print("input audio? = ", type(audioClip.recording.content))
  49. audio = np.array(audioClip.recording.content).squeeze()
  50. print("in test func, audio = ", type(audio), audio.shape)
  51. feats = get_feature(audio)
  52. print("in test section, feats = ", type(feats), feats.shape)
  53. output = inference(feats)
  54. return {"output" : output}
  55. @app.post("/inference")
  56. def post_inference(audioClip: AudioClip):
  57. output = inference(audioClip.recording.content)
  58. return {"output" : output}
  59. @app.post("/online")
  60. def post_inference(audioClip: AudioClip):
  61. output = inference_online(audioClip.recording.content)
  62. return {"output" : output}
  63. @app.post("/inference_file")
  64. # def inference_file(file: UploadFile = File(...)):
  65. def post_inference_file():
  66. # data = pickle.load(file)
  67. '''
  68. control formatting
  69. if data.format != 'wav':
  70. do_formatting()
  71. '''
  72. ## run model
  73. print('in Inference Start')
  74. output = ''
  75. output = inference_file()
  76. return {"output" : output}
  77. if __name__ == '__main__':
  78. print('this is main')
  79. print(inference(args["wav_path"]))