AI_anchor_openshot/api/main.py

from fastapi import FastAPI,Cookie, Depends, FastAPI, Query, WebSocket, status, WebSocketDisconnect,File, UploadFile
from os import listdir
from os.path import isfile, isdir, join
import threading
import zhtts
import os 
import urllib
from typing import List
import requests
from pydantic import BaseModel
from bs4 import BeautifulSoup
from PIL import Image,ImageDraw,ImageFont
import pyttsx3
import rpyc
import random
import time
import math
import hashlib
import re
import asyncio
import urllib.request
from fastapi.responses import FileResponse
from websocket import create_connection
from fastapi.middleware.cors import CORSMiddleware
import dataset
from datetime import datetime
from util.swap_face import swap_face
from fastapi.staticfiles import StaticFiles
import shutil
import io
#test

app = FastAPI()
origins = [
    "https://hhh.com.tw"
    "http://172.105.205.52",
    "http://172.105.205.52:8001",
    "http://172.104.93.163",
]

app.add_middleware(
    CORSMiddleware,
    # allow_origins=origins,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.mount("/static", StaticFiles(directory="static"), name="static")
app.mount("/static/img", StaticFiles(directory="static/img"), name="static/img")

tmp_video_dir = '../OpenshotService/tmp_video/'
tmp_avatar_dir = '../../face_swap/tmp_avatar/'  #change source face path here

video_sub_folder = 'ai_anchor_video/'
avatar_sub_folder = 'swap_save/'
tmp_img_sub_folder = 'tmp_img/'
img_upload_folder = '/var/www/html/'+tmp_img_sub_folder
video_dest = '/var/www/html/'+video_sub_folder
avatar_dest = '/var/www/html/'+avatar_sub_folder

class swap_req(BaseModel):
    imgurl: str

class request(BaseModel):
    name: str
    text_content: List[str]
    image_urls: List[str]
    avatar: str
    client_id :str

class request_eng(BaseModel):
    name: str
    text_content: List[str]
    image_urls: List[str]
    sub_titles: List[str]
    avatar: str
    client_id :str


@app.get("/index2")
async def index2():
    return FileResponse('static/index2.html')

@app.get("/index_eng")
async def index2():
    return FileResponse('static/index_eng.html')

@app.get("/gen_avatar")
async def avatar():
    return FileResponse('static/gen_avatar.html')

@app.post("/swapFace")
async def swapFace(req:swap_req):
    '''
    sf = swap_face(req.imgurl)
    result = sf.run()
    #notify_group(result)hi
    '''
    if 'http' not in req.imgurl:
        req.imgurl= 'http://'+req.imgurl
    try:
        im = Image.open(requests.get(req.imgurl, stream=True).raw)
        im= im.convert("RGB")
    except:
        return {'msg':"無法辨別圖片網址"+req.imgurl}
    name_hash = str(time.time()).replace('.','')
    
    x = threading.Thread(target=gen_avatar, args=(name_hash,req.imgurl))
    x.start()
    return {'msg':'人物生成中，請稍候'}

@app.post("/uploadfile/")
async def create_upload_file(file: UploadFile = File(...)):
    img_name = str(time.time()).replace('.','')
    try:
        contents = await file.read()
        image = Image.open(io.BytesIO(contents))
        image= image.convert("RGB")
        image.save(img_upload_folder+img_name+'.jpg')
    except:
        return {'msg':'檔案無法使用'}
    return {"msg": 'www.choozmo.com:8168/'+tmp_img_sub_folder+img_name+'.jpg'}

@app.post("/make_anchor_video_v2")
async def make_anchor_video_v2(req:request):
    for idx in range(len(req.image_urls)):
        if 'http' not in req.image_urls[idx]:
            req.image_urls[idx] = 'http://'+req.image_urls[idx]
    for txt in req.text_content:
        if re.search('[a-zA-Z]', txt) !=None:
            return {'msg':'輸入字串不能包含英文字!'}
    name_hash = str(time.time()).replace('.','')
    for imgu in req.image_urls:
        try:
            if get_url_type(imgu) =='video/mp4':
                r=requests.get(imgu)
            else:
                im = Image.open(requests.get(imgu, stream=True).raw)
                im= im.convert("RGB")
        except:
            return {'msg':"無法辨別圖片網址"+imgu}

    save_history(req,name_hash)
    x = threading.Thread(target=gen_video_queue, args=(name_hash,req.name, req.text_content, req.image_urls,int(req.avatar)))
    x.start()
    return {"msg":"製作影片需要時間，請您耐心等候，成果會傳送至LINE群組中"} 

@app.post("/make_anchor_video_eng")
async def make_anchor_video_eng(req:request_eng):
    for idx in range(len(req.image_urls)):
        if 'http' not in req.image_urls[idx]:
            req.image_urls[idx] = 'http://'+req.image_urls[idx]
    name_hash = str(time.time()).replace('.','')
    for imgu in req.image_urls:
        try:
            if get_url_type(imgu) =='video/mp4':
                r=requests.get(imgu)
            else:
                im = Image.open(requests.get(imgu, stream=True).raw)
                im= im.convert("RGB")
        except:
            return {'msg':"無法辨別圖片網址"+imgu}

    save_history(req,name_hash)
    x = threading.Thread(target=gen_video_eng, args=(name_hash,req.name, req.text_content, req.image_urls,req.sub_titles,int(req.avatar)))
    x.start()
    return {"msg":"製作影片需要時間，請您耐心等候，成果會傳送至LINE群組中"} 

@app.get("/history_input")
async def history_input():
    db = dataset.connect('mysql://choozmo:pAssw0rd@db.ptt.cx:3306/AI_anchor?charset=utf8mb4')
    statement = 'SELECT * FROM history_input ORDER BY timestamp DESC LIMIT 50'
    logs = []
    for row in db.query(statement):
        logs.append({'id':row['id'],'name':row['name'],'text_content':row['text_content'].split(','),'link':row['link'],'image_urls':row['image_urls'].split(',')})
    return logs

def save_history(req,name_hash):
    db = dataset.connect('mysql://choozmo:pAssw0rd@db.ptt.cx:3306/AI_anchor?charset=utf8mb4')
    log_table = db['history_input']
    txt_content_seperate_by_dot = ''
    for txt in req.text_content:
        txt_content_seperate_by_dot += txt+","
    txt_content_seperate_by_dot = txt_content_seperate_by_dot[:-1]
    img_urls_seperate_by_dot = ''
    for iurl in req.image_urls:
        img_urls_seperate_by_dot += iurl+","
    img_urls_seperate_by_dot = img_urls_seperate_by_dot[:-1]
    time_stamp = datetime.fromtimestamp(time.time())
    time_stamp = time_stamp.strftime("%Y-%m-%d %H:%M:%S")
    pk = log_table.insert({'name':req.name,'text_content':txt_content_seperate_by_dot,'image_urls':img_urls_seperate_by_dot,'link':'www.choozmo.com:8168/'+video_sub_folder+name_hash+'.mp4','timestamp':time_stamp})
    
def get_url_type(url):
    req = urllib.request.Request(url, method='HEAD', headers={'User-Agent': 'Mozilla/5.0'})
    r = urllib.request.urlopen(req)
    contentType = r.getheader('Content-Type')
    return contentType
def notify_group(msg):
    glist=['WekCRfnAirSiSxALiD6gcm0B56EejsoK89zFbIaiZQD']
    for gid in glist:
        headers = {
                "Authorization": "Bearer " + gid,
                "Content-Type": "application/x-www-form-urlencoded"
        }
        params = {"message": msg}   
        r = requests.post("https://notify-api.line.me/api/notify",headers=headers, params=params)
def gen_video(name_hash,name,text_content, image_urls,avatar):
    c = rpyc.connect("localhost", 8878)
    c._config['sync_request_timeout'] = None
    remote_svc = c.root
    my_answer = remote_svc.call_video(name_hash,name,text_content, image_urls,avatar) # method call
    shutil.copy(tmp_video_dir+name_hash+'.mp4',video_dest+name_hash+'.mp4')
    os.remove(tmp_video_dir+name_hash+'.mp4')

def gen_video_eng(name_hash,name,text_content, image_urls,sub_titles,avatar):
    c = rpyc.connect("localhost", 8878)
    c._config['sync_request_timeout'] = None
    remote_svc = c.root
    my_answer = remote_svc.call_video_eng(name_hash,name,text_content, image_urls,sub_titles,avatar) # method call
    shutil.copy(tmp_video_dir+name_hash+'.mp4',video_dest+name_hash+'.mp4')
    os.remove(tmp_video_dir+name_hash+'.mp4')

def gen_video_queue(name_hash,name,text_content, image_urls,avatar):
    db = dataset.connect('mysql://choozmo:pAssw0rd@db.ptt.cx:3306/AI_anchor?charset=utf8mb4')
    time_stamp = datetime.fromtimestamp(time.time()).strftime("%Y-%m-%d %H:%M:%S")
    db['video_queue'].insert({'name_hash':name_hash,'name':name,'text_content':text_content,'image_urls':image_urls,'avatar':avatar,'timestamp':time_stamp})
    while True:
        if first(db.query('SELECT * FROM video_queue_status'))['status'] == 1:#only one row in this table, which is the id 1 one
            print('another process running, leave loop')
            break
        if first(db.query('SELECT COUNT(1) FROM video_queue'))['COUNT(1)'] == 0:
            print('all finish, leave loop')
            break
        top1 = first(db.query('SELECT * FROM video_queue'))
        try:
            db.query('UPDATE video_queue_status SET status = 1;')
            c = rpyc.connect("localhost", 8868)
            c._config['sync_request_timeout'] = None
            remote_svc = c.root
            my_answer = remote_svc.call_video(top1['name_hash'],top1['name'],top1['text_content'], top1['image_urls'],top1['avatar']) # method call
            shutil.copy(tmp_video_dir+top1['name_hash']+'.mp4',video_dest+top1['name_hash']+'.mp4')
            os.remove(tmp_video_dir+top1['name_hash']+'.mp4')
        except:
            print('video generation error')
            notify_group('影片錯誤')
        db['video_queue'].delete(id=top1['id'])
        db.query('UPDATE video_queue_status SET status = 0')

    

def gen_avatar(name_hash, imgurl):
    db = dataset.connect('mysql://choozmo:pAssw0rd@db.ptt.cx:3306/AI_anchor?charset=utf8mb4')
    db['avatar_queue'].insert({'name_hash':name_hash,'imgurl':imgurl})
    while True:
        statement = 'SELECT * FROM avatar_service_status'#only one row in this table, which is the id 1 one
        status = -1
        for row in db.query(statement):
            status = row['status']
        if status == 1:
            print('leave process loop')
            break

        statement = 'SELECT * FROM avatar_queue'
        works = []
        for row in db.query(statement):
            works.append({'id':row['id'],'name_hash':row['name_hash'],'imgurl':row['imgurl']})
        if len(works)==0:
            print('leave process loop')
            break
        try:
            statement = 'UPDATE avatar_service_status SET status = 1 WHERE id=1;'
            db.query(statement)
            name_hash = works[0]['name_hash']
            imgurl = works[0]['imgurl']
            c = rpyc.connect("localhost", 8868)
            c._config['sync_request_timeout'] = None
            remote_svc = c.root
            my_answer = remote_svc.call_avatar(name_hash,imgurl) # method call
            shutil.copy(tmp_avatar_dir+name_hash+'.mp4',avatar_dest+name_hash+'.mp4')
            os.remove(tmp_avatar_dir+name_hash+'.mp4')
            
        except:
            print('gen error')
            notify_group('無法辨識人臉')
        db['avatar_queue'].delete(id=works[0]['id'])
        statement = 'UPDATE avatar_service_status SET status = 0 WHERE id=1;'  #only one row in this table, which id 1 one
        db.query(statement)