import os, sys
import waifu2x_vulkan as waifu2x
import time

if __name__ == "__main__":
    os.chdir(os.path.dirname(__file__))
    
    # 设置models所在目录
    # set models dir
    waifu2x.setDefaultPath("..")

    # 初始化ncnn
    # init ncnn
    waifu2x.init()

    # 获得Gpu列表
    # get gpu list
    print(waifu2x.getGpuInfo())

    # 选择Gpu,设置线程数
    # select gpu, set thread num
    waifu2x.initSet(0, 1)

    # 开启打印
    # open debug log
    waifu2x.setDebug(True)

    f = open("0.jpg", "rb")
    data = f.read()
    f.close()
    backId = 1
    count = 0
    
    # 设置长宽缩放2.5倍
    # start convert, by setting scale
    if waifu2x.add(data, waifu2x.MODEL_CUNET_NOISE3, backId, scale=2.5) > 0:
        count += 1
    backId = 2

    # 固定长宽
    # start convert, by setting width and high
    if waifu2x.add(data, waifu2x.MODEL_CUNET_NOISE3, backId, format="png", width=900, high=800) > 0:
        count += 1
    
    saveName = {
        1 : "1.jpg",
        2 : "2.png"
    }

    while count > 0:
        time.sleep(1)

        # 阻塞获取，也可放入到到线程中
        # block
        info = waifu2x.load(0)
        if not info:
            continue 
        count -= 1
        newData, status, backId, tick = info
        f = open(saveName.get(backId), "wb+")
        f.write(newData)
        f.close
    
    # free ncnn, close thread
    waifu2x.stop()
