from pysr import SR
from pyaxdev import enum_devices, sys_init, sys_deinit, AxDeviceType
import cv2
import argparse

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, default='edsr_x2_small_1.axmodel')
    parser.add_argument('--image', type=str)
    args = parser.parse_args()

    devices_info = enum_devices()
    dev_type = None
    devid = None
    print("可用设备:", devices_info)
    if devices_info['host']['available']:
        print("host device available")
        sys_init(AxDeviceType.host_device, -1)
        dev_type = AxDeviceType.host_device
        devid = -1
    elif devices_info['devices']['count'] > 0:
        print("axcl device available, use device-0")
        sys_init(AxDeviceType.axcl_device, 0)
        dev_type = AxDeviceType.axcl_device
        devid = 0
    else:
        raise Exception("No available device")
    
    sr = SR({
            'dev_type': dev_type,
            'devid': devid,
            'model_path': args.model
        })
    image = cv2.imread(args.image)
    sr_image = sr(image)
    cv2.imwrite('sr_image.jpg', sr_image)
    
    if devices_info['host']['available']:
        sys_deinit(AxDeviceType.host_device, -1)
    elif devices_info['devices']['count'] > 0:
        sys_deinit(AxDeviceType.axcl_device, 0)