import requests, json
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import base64
import io

def request_post(url, data, timeout=600, headers = None):
    if headers is None:
        headers = {
            # 'content-type': 'application/json'
            # 'Connection': 'keep-alive',
            'Accept': '*/*', # 接受任何类型的返回数据
            'Content-Type': 'application/json;charset=UTF-8', # 发送数据为json
            # 'Content-Length': '156',
            # 'Accept-Encoding': 'gzip, deflate',
            # 'Accept-Language': 'zh-CN,zh;q=0.9',
            # 'User-Agent': 'SamClub/5.0.45 (iPhone; iOS 15.4; Scale/3.00)',
            # 'device-name': 'iPhone14,3',
            # 'device-os-version': '15.4',
            # 'device-type': 'ios',
            # 'auth-token': authtoken,
            # 'app-version': '5.0.45.1'
        }
    try:
        response = requests.post(url=url, headers=headers, data=json.dumps(data), timeout=timeout)
        response_data = response.json()
        return response_data
    except Exception as e:
        print(f'request_post[Error]:' + str(e))
        print(f'url: {url}')
        print(f'data: {data}')
        print(f'response: {response}')
        return None

url = "http://127.0.0.1:7860/imgCLeaner"

def imgFile_to_base64(image_file):
    with open(image_file, "rb") as f:
        im_bytes = f.read()  
    im_b64_encode = base64.b64encode(im_bytes)
    im_b64 = im_b64_encode.decode("utf8")
    return im_b64

def base64_to_bytes(im_b64):
    im_b64_encode = im_b64.encode("utf-8")
    im_bytes = base64.b64decode(im_b64_encode)
    return im_bytes

def base64_to_PILImage(im_b64):
    im_bytes = base64_to_bytes(im_b64)
    pil_img = Image.open(io.BytesIO(im_bytes))
    return pil_img

def cleaner_img(image_file, remove_texts, mask_extend=20, disp_debug=True):
    data = {'remove_texts': remove_texts,
        'mask_extend': mask_extend,
        'img': imgFile_to_base64(image_file),
        }
    ret = request_post(url, data, timeout=600, headers = None)  
    if ret['code'] == 0:
        if disp_debug:
            for img in ret['result']['imgs']:
                pilImage = base64_to_PILImage(img)
                plt.imshow(pilImage)
                plt.show()
                plt.clf()
                plt.close('all')    
        img_len = len(ret['result']['imgs'])
        pilImage = base64_to_PILImage(ret['result']['imgs'][img_len-1])  
    else:
        pilImage = None
    return pilImage, ret

image_file = 'dog.png'
remove_texts = "小狗 . 椅子"

mask_extend = 20
pil_image, ret = cleaner_img(image_file, remove_texts, mask_extend, disp_debug=False) 

plt.imshow(pil_image)
plt.show()
plt.clf()
plt.close()