75 lines
2.7 KiB
Python
75 lines
2.7 KiB
Python
import requests
|
|
import json
|
|
import base64
|
|
import argparse
|
|
|
|
VIDEO_WIDTH = 1080
|
|
VIDEO_HEIGHT = 1920
|
|
SPEAKER_ID = 0
|
|
|
|
class RestAPIApplication:
|
|
def __init__(self, ip, port):
|
|
|
|
if port < 0:
|
|
self.post_request_addr = f"http://{ip}/register/"
|
|
self.post_headers = {"Content-Type": "application/json"}
|
|
self.generate_addr = (lambda id_: f'http://{ip}/generate/{id_}')
|
|
else:
|
|
self.post_request_addr = f"http://{ip}:{port}/register/"
|
|
self.post_headers = {"Content-Type": "application/json"}
|
|
self.generate_addr = (lambda id_: f'http://{ip}:{port}/generate/{id_}')
|
|
|
|
@staticmethod
|
|
def _get_json_request(text, lang, duration_rate, action, background_data=None, is_video_background=False):
|
|
request_form = dict()
|
|
|
|
request_form['text'] = text
|
|
request_form['speaker'] = SPEAKER_ID
|
|
request_form['width'] = VIDEO_WIDTH
|
|
request_form['height'] = VIDEO_HEIGHT
|
|
|
|
request_form['action'] = action
|
|
|
|
if background_data is not None:
|
|
background_base64 = base64.b64encode(background_data).decode("UTF-8")
|
|
else:
|
|
background_base64 = ""
|
|
|
|
request_form['background'] = background_base64
|
|
request_form['durationRate'] = duration_rate
|
|
request_form['isVideoBackground'] = is_video_background
|
|
request_form['lang'] = lang
|
|
|
|
request_as_json = json.dumps(request_form)
|
|
return request_as_json
|
|
|
|
@staticmethod
|
|
def _get_video_id(results):
|
|
return json.loads(bytes.decode(results.content))['id']
|
|
|
|
def get_video(self, text, lang, duration_rate, action, background_data=None, is_video_background=False):
|
|
request_json = self._get_json_request(text, lang, duration_rate, action, background_data, is_video_background)
|
|
|
|
# POST request with jsonified request
|
|
results = requests.post(self.post_request_addr, headers=self.post_headers, data=request_json)
|
|
|
|
# GET video with the given id
|
|
video_id = self._get_video_id(results)
|
|
video_results = requests.get(self.generate_addr(video_id))
|
|
|
|
return video_results.content
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(
|
|
description='REST API interface for talking face generation submitted to CVPR2022')
|
|
parser.add_argument('-i', '--ip', dest='rest_ip', type=str, default="127.0.0.1", help="IP for REST API")
|
|
parser.add_argument('-p', '--port', dest='rest_port', type=int, default=8080, help="Port for REST API")
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
if __name__ == '__main__':
|
|
args = parse_args()
|
|
rest_api_application = RestAPIApplication(args.rest_ip, args.rest_port)
|