utils.py

import os
import sys
import cv2
import json
import random
import time
import requests
import func_timeout
import numpy as np
import gradio as gr


OssUrl = "https://selfit-deploy-1256039085.cos.accelerate.myqcloud.com/"
Regions = "IndiaPakistanBengal"
TOKEN = os.environ['TOKEN']
UKAPIURL = os.environ['UKAPIURL']
POSEToken = os.environ['POSEToken']


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


def get_cloth_examples(hr=0):
    cloth_dir = os.path.join(data_dir, 'ClothImgs')
    examples = []
    files = sorted(os.listdir(cloth_dir))
    hr_clothes = list(range(588, 597))
    # hr_clothes = []
    for f in files:
        if '.jpg' not in f and '.png' not in f:
            continue

        cloth_id = f.split(".")[0]
        if int(cloth_id) in hr_clothes and hr==0:
            continue
        if int(cloth_id) not in hr_clothes and hr==1:
            continue

        cloth_path = os.path.join(cloth_dir, f)
        examples.append(cloth_path)
    examples = examples[::-1]
    return examples

def get_pose_examples():
    pose_dir = os.path.join(data_dir, 'PoseImgs')
    examples = []
    for f in os.listdir(pose_dir):
        if '.jpg' not in f and '.png' not in f:
            continue
        pose_id = f.split(".")[0]
        pose_path = os.path.join(pose_dir, f)
        examples.append(pose_path)
    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 get_tips():
    path1 = OssUrl+'ClothData/Publics/PoseGuide/tip1.jpg'
    path2 = OssUrl+'ClothData/Publics/PoseGuide/tip2.jpg'
    return path1, path2

def upload_pose_img(clientIp, timeId, img):
    fileName = clientIp.replace(".", "")+str(timeId)+".jpg"
    local_path = os.path.join(tmpFolder, fileName)
    img = cv2.imread(img)
    cv2.imwrite(os.path.join(tmpFolder, fileName), img)

    json_data = {
        "token": "c0e69e5d129b11efa10c525400b75156",  
        "input1": fileName,
        "input2": "",
        "protocol": "",
        "cloud": "ali"
    }
    
    session = requests.session()
    ret = requests.post(
        f"{UKAPIURL}/upload",
        headers={'Content-Type': 'application/json'},
        json=json_data
    )
    
    res = ""
    if ret.status_code==200:
        if 'upload1' in ret.json():
            upload_url = ret.json()['upload1']
            headers = {'Content-Type': 'image/jpeg'}
            response = session.put(upload_url, data=open(local_path, 'rb').read(), headers=headers)
            # print(response.status_code)
            if response.status_code == 200:
                res = upload_url
    if os.path.exists(local_path):
        os.remove(local_path)
    return res


def publicClothSwap(image, clothId, is_hr=0):
    json_data = {
        "image": image,
        "task_type": "11", 
        "param1": str(clothId),
        "param2": "",
        "param3": "",
        "param4": str(is_hr),
        "delete_if_complete": "1",
        "force_celery":"0"
    }
    
    headers = {
        'Authorization': f'Bearer {TOKEN}',
        'Content-Type': 'application/json'
    }
    
    ret = requests.post(
        f'{UKAPIURL}/public_advton',
        headers=headers,
        json=json_data
    )
    
    if ret.status_code == 200:
        response = ret.json()
        if 'mid_result' in response and 'id' in response:
            return {'mid_result': response['mid_result'], 'id': response['id'], "msg": response['msg']}
        # print(response)
        return None
    else:
        return None

def getInfRes(taskId):
    headers = {
        'Content-Type': 'application/json'
    }
    json_data = {
        'id': taskId
    }
    ret = requests.post(
        f'{UKAPIURL}/status_advton',
        headers=headers,
        json=json_data
    )
    if ret.status_code == 200:
        response = ret.json()
        if 'status' in response:
            return response
        print(response)
        return None
    else:
        return None

@func_timeout.func_set_timeout(10)
def check_region(ip):
    session = requests.session()
    # ret = requests.get(f"https://api.ip2location.io/?ip={ip}")
    ret = requests.get(f"https://realip.cc/?ip={ip}")
    # print(ret)
    nat = ret.json()['country'].lower()
    if nat in Regions.lower():
        print(nat, 'invalid', ip)
        return False
    else:
        print(nat, 'valid', ip)
    return True

def check_region_warp(ip):
    try:
        return check_region(ip)
    except Exception as e:
        print(e)
        return True

def public_pose_changer(image_url, prompt="Change the pose: two hands on hips.#Change the pose: arms extended to show outfit."):
    """
    发送姿势变换请求到API
    """
    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {TOKEN}'
    }
    
    json_data = {
        'image': image_url,
        'param1': 'pose_change',
        'param2': prompt,
        'param3': '',
        'param4': '',
        'param5': '',
        'is_private': '0',
        'delete_if_complete': '0'
    }
    
    try:
        ret = requests.post(
            f'{UKAPIURL}/public_comfyui',
            headers=headers,
            json=json_data
        )
        
        if ret.status_code == 200:
            response = ret.json()
            if 'id' in response:
                return {'id': response['id'], 'msg': response.get('msg', '')}
            return None
        else:
            print(f"姿势变换请求失败，状态码: {ret.status_code}")
            return None
    except Exception as e:
        print(f"姿势变换请求异常: {e}")
        return None
    
def get_pose_changer_res(task_id):
    """
    获取姿势变换任务的结果
    """
    headers = {
        'Content-Type': 'application/json'
    }
    try:
        ret = requests.post(
            f'{UKAPIURL}/status_comfyui',
            headers=headers,
            json={ 'id': task_id}
        )
        if ret.status_code == 200:
            response = ret.json()
            return response
        else:
            print(f"获取姿势变换结果失败，状态码: {ret.status_code}")
            return None
    except Exception as e:
        print(f"获取姿势变换结果异常: {e}")
        return None

if __name__ == "__main__":
    
    res = public_pose_changer(image_url="https://www.selfitcamera.site/SelfitUpload/2024-10-11/p2.jpg", 
        prompt="Change the pose: two hands on hips.#Change the pose: arms extended to show outfit.")
    print(res)
    for _ in range(100):
        result = get_pose_changer_res(res['id'])
        print(result)
        time.sleep(1)