main.py 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175
  1. import datetime
  2. from concurrent.futures import ThreadPoolExecutor
  3. from io import BytesIO
  4. from typing import Optional, List
  5. from urllib.parse import quote
  6. import hashlib
  7. import pandas as pd
  8. import uvicorn
  9. import yaml
  10. from fastapi import FastAPI
  11. from fastapi.middleware.cors import CORSMiddleware
  12. from fastapi.responses import StreamingResponse
  13. from pydantic import BaseModel, Field
  14. from asr_client import send_asr_request, send_task_request
  15. from common_func import get_db_engine
  16. from config.url import toutiao_static_video_url
  17. from database import insert, update, query, Task
  18. with open('/data/pythonProject/video_to_word/config/config.yaml', mode='r', encoding='utf-8') as f:
  19. config = yaml.load(f.read(), Loader=yaml.FullLoader)
  20. source_name_map = config['source_name_map']
  21. ai_word_engine = get_db_engine(config['ai_word_dev_db'])
  22. threadPool = ThreadPoolExecutor(max_workers=4)
  23. app = FastAPI()
  24. origins = [
  25. "http://192.168.1.34",
  26. "http://192.168.1.34:8000",
  27. "http://192.168.1.105",
  28. "http://192.168.1.105:3000",
  29. "http://111.206.86.186",
  30. "http://111.206.86.186:3000",
  31. "http://adsp.tjyourong.com.cn",
  32. "http://adsp.tjyourong.com.cn:3000",
  33. "http://adsp.c-top.com.cn",
  34. "http://adsp.c-top.com.cn:3000"
  35. ]
  36. app.add_middleware(
  37. CORSMiddleware,
  38. allow_origins=origins,
  39. allow_credentials=True,
  40. allow_methods=["*"],
  41. allow_headers=["*"],
  42. )
  43. class QueryItem():
  44. signature: Optional[str] = None
  45. url: Optional[str] = None
  46. @app.get('/')
  47. def index():
  48. return {'message': '你已经正确创建 FastApi 服务!'}
  49. @app.post('/asr/task/submit')
  50. def task_submit(signature: str, url: str):
  51. json = send_asr_request(url)
  52. task = Task(signature=signature, task_id=json.Data.TaskId, task_result=json.to_json_string(), task_status=1)
  53. insert(task)
  54. return {'code': 0, 'taskId': json.Data.TaskId}
  55. @app.post('/asr/task/result')
  56. def task_submit(task_id: int):
  57. json = send_task_request(task_id)
  58. task = query(None, None, task_id)[0]
  59. task.task_status = json.Data.Status
  60. task.task_result = json.to_json_string()
  61. try:
  62. if json.Data.Status == 2:
  63. task.word_text = json.Data.ResultDetail[0].FinalSentence
  64. task.word_split = json.Data.ResultDetail[0].SliceSentence
  65. task.word_text_md5 = hashlib.md5(task.word_text.encode('utf-8')).hexdigest()
  66. except:
  67. # 提取原始文本内容和分词内容发生异常,把 task_status 置为 -1
  68. task.task_status = -1
  69. update(task)
  70. return {'code': 0, 'status': json.Data.StatusStr}
  71. @app.post('/asr/task/list')
  72. def task_submit(task_status: int):
  73. task = query(None, task_status, None)
  74. return {'code': 0, 'data': task}
  75. class QueryWordItem(BaseModel):
  76. query_word: str = Field(..., description="查询词", min_length=1)
  77. stat_date: str = Field(..., description="日期", min_length=10, max_length=10)
  78. source: int = Field(..., description="来源,")
  79. @app.post('/export_excel/')
  80. def export_excel(item: List[QueryWordItem]):
  81. video_df = pd.DataFrame()
  82. if len(item) == 1:
  83. # 单个条目,直接导出
  84. pass
  85. else:
  86. # 1 从数据库获取视频数据
  87. # 多个条目,如果同一个素材有多个查询词,则合并打上这多个查询词
  88. for obj in item:
  89. query_word = obj.query_word
  90. stat_date = obj.stat_date
  91. source = obj.source
  92. sql = f"select signature, video_url, query_word, stat_date, {source} source from {source_name_map[source]['table']} " \
  93. f"where query_word = '{query_word}' " \
  94. f"and stat_date = '{stat_date}'"
  95. df = pd.read_sql(sql, ai_word_engine)
  96. video_df = video_df.append(df)
  97. # 按 'signature' + 'query_word' + 'stat_date' 进行去重
  98. video_df.drop_duplicates(['signature', 'query_word', 'stat_date', 'source'], keep='last', inplace=True)
  99. g = video_df.groupby('signature')
  100. query_word_lst_df = g.apply(lambda x: x['query_word'].unique())
  101. query_word_lst_df.name = 'query_word_lst'
  102. url_df = g.apply(lambda x: x['video_url'].values[0])
  103. url_df.name = 'video_url'
  104. source_df = g.apply(lambda x: x['source'].values[0])
  105. source_df.name = 'source'
  106. video_query_word_df = pd.concat([query_word_lst_df, url_df, source_df], axis=1)
  107. video_query_word_df.reset_index(inplace=True, drop=False)
  108. video_query_word_df['video_url'] = video_query_word_df.apply(
  109. lambda row: toutiao_static_video_url + row['signature'] if row.get('source') == 2 else row['video_url'], axis=1)
  110. # 2 根据第一步的视频数据获取脚本
  111. if not video_query_word_df.empty:
  112. sql = f"select signature, word_text from tb_asr_result where signature in " \
  113. f"{tuple(video_query_word_df.signature.values) if len(video_query_word_df.signature.values) > 1 else tuple(list(video_query_word_df.signature.values) * 2)} " \
  114. f"and task_status = 2"
  115. script_df = pd.read_sql(sql, ai_word_engine)
  116. out_df = video_query_word_df.merge(script_df, on='signature', how='inner')
  117. else:
  118. pass
  119. # 3 返回流数据
  120. if not out_df.empty:
  121. bio = BytesIO()
  122. writer = pd.ExcelWriter(bio, engine='xlsxwriter')
  123. out_df[['signature', 'query_word_lst', 'word_text', 'video_url']].to_excel(writer, index=False, encoding='utf8mb4')
  124. writer.save()
  125. bio.seek(0)
  126. # 组装header
  127. now_date = datetime.date.today().strftime('%Y-%m-%d')
  128. headers = {"content-type": "application/vnd.ms-excel",
  129. "content-disposition": f"attachment;filename={quote('优质素材脚本_')}{now_date}.xlsx"
  130. }
  131. return StreamingResponse(bio, media_type='xlsx', headers=headers)
  132. return None
  133. if __name__ == '__main__':
  134. # 1 读取配置文件
  135. # test_items = [{'query_word': '红包', 'stat_date': '2021-10-28', 'source': 2},
  136. # {'query_word': '红包', 'stat_date': '2021-10-28', 'source': 3},
  137. # {'query_word': '赚钱', 'stat_date': '2021-10-28', 'source': 2},
  138. # {'query_word': '赚钱', 'stat_date': '2021-10-28', 'source': 3}]
  139. # export_excel(test_items)
  140. uvicorn.run(app='main:app', host="0.0.0.0", port=31013, reload=True, debug=True)
  141. # gunicorn main:app -w 4 -k uvicorn.workers.UvicornWorker #线上启动命令