ml-talking-face / client_rest.py
astasama's picture
Duplicate from CVPR/ml-talking-face
3cd2bc7
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)