from run import process
import time
import subprocess
import os
import argparse
import cv2
import sys
from PIL import Image
import torch
import gradio as gr


TESTdevice = "cpu"

index = 1


"""
main.py

 How to run:
 python main.py

"""


def mainTest(inputpath, outpath):
    watermark = deep_nude_process(inputpath)
    watermark1 = cv2.cvtColor(watermark, cv2.COLOR_BGRA2RGBA)
    #cv2.imwrite(outpath, watermark1)
    return watermark1
    #


def deep_nude_process(inputpath):
    dress = cv2.imread(inputpath)
    h = dress.shape[0]
    w = dress.shape[1]
    dress = cv2.resize(dress, (512, 512), interpolation=cv2.INTER_CUBIC)
    watermark = process(dress)
    watermark = cv2.resize(watermark, (w, h), interpolation=cv2.INTER_CUBIC)
    return watermark


def inference(img):
    global index
    bgra = cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA)
    inputpath = "input_" + str(index) + ".jpg"
    cv2.imwrite(inputpath, bgra)

    outputpath = "out_" + str(index) + ".jpg"
    index += 1
    print(time.strftime("开始!!!!!!!!! %Y-%m-%d %H:%M:%S", time.localtime()))
    output = mainTest(inputpath, outputpath)
    print(time.strftime("结束!!!!!!!!! %Y-%m-%d %H:%M:%S", time.localtime()))
    return output


title = "AI脱衣"
description = "传入人物照片,类似最下方测试图的那种,将制作脱衣图,一张图可能等30秒,别传私人照片.\n有队列系统,根据先来先做的逻辑,一次只做一张,\n图片必须至少能看出是人体轮廓"

examples = [
    ['input.png', '测试图'],
    ['input.jpg', '测试图'],
]


web = gr.Interface(inference,
                   inputs="image",
                   outputs="image",
                   title=title,
                   description=description,
                   examples=examples,
                   )

if __name__ == '__main__':
    web.launch(
        enable_queue=True
    )