utils.py

import os
import sys
import cv2
import json
import random
import time
import pickle
import requests
import func_timeout
import numpy as np
import gradio as gr
from collections import OrderedDict


apiUrl = os.environ['apiUrl']
uploadToken = os.environ['uploadToken']
openId = os.environ['openId']
apiKey = os.environ['apiKey']
Regions = os.environ['Regions']
LimitTask = int(os.environ['LimitTask'])


proj_dir = os.path.dirname(os.path.abspath(__file__))
data_dir = os.path.join(proj_dir, 'Datas')
tmpFolder = os.path.join(proj_dir, 'tmp')
os.makedirs(tmpFolder, exist_ok=True)



def load_pkl(path):
    with open(path, 'rb') as f:
        return pickle.load(f)

def save_pkl(data, path, reweite=False):
    if not os.path.exists(path) or reweite: # 不存在或者强制重写
        with open(path,'wb') as file:
          pickle.dump(data, file, protocol=4)
        return data
    else:
        load_data = load_pkl(path)
        for k in data:
            load_data[k] = data[k]
        save_pkl(load_data, path, reweite=True)
        return load_data

class UserRecorder(object):

    def __init__(self, ):
        super(UserRecorder, self).__init__()
        record_dir = os.path.join(data_dir, 'UserRecord')
        self.ip_dir = os.path.join(record_dir, 'Ips')
        self.token_dir = os.path.join(record_dir, 'Tokens')
        os.makedirs(self.ip_dir, exist_ok=True)
        os.makedirs(self.token_dir, exist_ok=True)

    def save_record(self, taskRes, ip=None, token=None):
        if ip is None and token is None: return
        if ip is not None:
            record_path = os.path.join(self.ip_dir, f'{ip}.pkl')
        else:
            record_path = os.path.join(self.token_dir, f'{token}.pkl')
        taskId = taskRes['id']
        status = taskRes['status']
        if 'output' in taskRes:
            input1 = taskRes['output']['job_results']['input1']
            output1 = taskRes['output']['job_results']['output1']
        else:
            input1, output1 = None, None
        data = OrderedDict()
        data[taskId] = {'input1':input1, 'output1':output1, 'status':status, }
        save_data = save_pkl(data, record_path, reweite=False)
        return save_data

    def get_record(self, ip=None, token=None):
        if ip is None and token is None: return
        if ip is not None:
            identity = ip
            record_path = os.path.join(self.ip_dir, f'{ip}.pkl')
        else:
            identity = token
            record_path = os.path.join(self.token_dir, f'{token}.pkl')
        if os.path.exists(record_path):
            record_data = load_pkl(record_path)
        else:
            record_data = {}
        total_n = len(record_data)
        success_n, fail_n, process_n = 0, 0, 0
        shows = [None]*6
        show_i = 0
        print(record_data)
        for key in reversed(record_data):
            status = record_data[key]['status']
            if status in ['FAILED', 'CANCELLED', 'TIMED_OUT', ]:
                fail_n += 1

            elif status in ['COMPLETED', ]:
                success_n += 1
                if record_data[key]['input1'] is not None:
                    input1 = record_data[key]['input1']
                    output1 = record_data[key]['output1']
                    shows[show_i*2] = f"<img src=\"{input1}\" >"
                    shows[show_i*2+1] = f"<img src=\"{output1}\" >"
                    show_i += 1
            elif status in ['IN_QUEUE', 'IN_PROGRESS', 'IN_QUEUE', ]:
                process_n += 1

        msg = f"Dear {identity}, You have {total_n} tasks, {success_n} successed, {fail_n} failed, {process_n} processing,  "

        return shows, total_n, msg

        
def get_temps_examples(taskType):
    temp_dir = os.path.join(data_dir, f'task{taskType}/temps')
    examples = []
    files = [f for f in sorted(os.listdir(temp_dir)) if '.' in f]
    for f in files:
        temp_name = f.split(".")[0]
        if len(temp_name)==0: continue
        temp_path = os.path.join(temp_dir, f)
        examples.append([temp_path])
    return examples

def get_user_examples(taskType):
    user_dir = os.path.join(data_dir, f'task{taskType}/srcs')
    examples = []
    files = [f for f in sorted(os.listdir(user_dir)) if '.' in f]
    for f in files:
        user_id = f.split(".")[0]
        if len(user_id)==0: continue
        user_path = os.path.join(user_dir, f)
        examples.append([user_path])
    return examples

def get_showcase_examples(taskType):
    if taskType=="3":
        examples=[
            ["task3/temps/flow-water.jpg", "task3/srcs/src01.jpg", "task3/showcases/src01_flower-water.jpg"],
            ["task3/temps/mountain-water.jpg", "task3/srcs/src01.jpg", "task3/showcases/src01_mountain-water.jpg"],
            ["task3/temps/rock-on-water.jpg", "task3/srcs/src01.jpg", "task3/showcases/src01_rock-on-water.jpg"],
        ]
    for i in range(len(examples)):
        for j in range(len(examples[i])):
            examples[i][j] = os.path.join(data_dir, examples[i][j])
            assert os.path.exists(examples[i][j]), examples[i][j]
    return examples

def get_result_example(cloth_id, pose_id):
    result_dir = os.path.join(data_dir, 'ResultImgs')
    res_path = os.path.join(result_dir, f"{cloth_id}_{pose_id}.jpg")
    return res_path

def upload_user_img(clientIp, img):
    timeId = int(  str(time.time()).replace(".", "")  )+random.randint(1000, 9999)
    fileName = clientIp.replace(".", "")+str(timeId)+".jpg"
    local_path = os.path.join(tmpFolder, fileName)
    cv2.imwrite(os.path.join(tmpFolder, fileName), img[:,:,::-1])
    params = {'token':uploadToken, 'input1':fileName, 'input2':''}
    session = requests.session()
    ret = requests.post(f"{apiUrl}/upload", data=json.dumps(params))
    res = ""
    if ret.status_code==200:
        if 'upload1' in ret.json():
            upload_url = ret.json()['upload1']
            with open(local_path, 'rb') as file:
                response = requests.put(upload_url, data=file)
                if response.status_code == 200:
                    res = upload_url
    else:
        print(ret.json(), ret.status_code, 'call upload failed')
    if os.path.exists(local_path):
        os.remove(local_path)
    return res


def publicSelfitTask(image, temp_image, caption_text):
    temp_name = os.path.basename(temp_image).split('.')[0]
    params = {'openId':openId, 'apiKey':apiKey, 'image':image, 'mask':"",
        "image_type":"2", "task_type":"3", 'param1':temp_name, 
        'param2':str(caption_text), 'param3':"1", 'param4':"", 'param5':""}
    session = requests.session()
    ret = requests.post(f"{apiUrl}/public", data=json.dumps(params))
    print(ret)
    if ret.status_code==200:
        if 'id' in ret.json():
            print(ret.json())
            return ret.json()['id']
    else:
        print(ret.json(), ret.status_code, 'call public failed')

def getTaskRes(taskId):
    params = {'id':taskId}
    session = requests.session()
    ret = requests.post(f"{apiUrl}/status", data=json.dumps(params))
    print(ret)
    if ret.status_code==200:
        if 'status' in ret.json():
            print(ret.json())
            return ret.json()
    else:
        print(ret.json(), ret.status_code, 'call status failed')
        return None

@func_timeout.func_set_timeout(10)
def check_region(ip):
    session = requests.session()
    ret = requests.get(f"https://webapi-pc.meitu.com/common/ip_location?ip={ip}")
    for k in ret.json()['data']:
        nat = ret.json()['data'][k]['nation']
        if nat in Regions:
            print(nat, 'invalid')
            return False
        else:
            print(nat, 'valid')
    return True
def check_region_warp(ip):
    try:
        return check_region(ip)
    except Exception as e:
        print(e)
        return True