| import ctypes |
| import numpy as np |
| from pyaxdev import _lib, AxDeviceType, AxDevices, check_error |
|
|
|
|
| class SRInit(ctypes.Structure): |
| _fields_ = [ |
| ('dev_type', AxDeviceType), |
| ('devid', ctypes.c_char), |
| ('model_path', ctypes.c_char * 256) |
| ] |
|
|
| class SRImage(ctypes.Structure): |
| _fields_ = [ |
| ('width', ctypes.c_int), |
| ('height', ctypes.c_int), |
| ('pVirAddr', ctypes.POINTER(ctypes.c_ubyte)) |
| ] |
|
|
|
|
| _lib.ax_sr_init.argtypes = [ctypes.POINTER(SRInit), ctypes.POINTER(ctypes.c_void_p)] |
| _lib.ax_sr_init.restype = ctypes.c_int |
|
|
| _lib.ax_sr_deinit.argtypes = [ctypes.c_void_p] |
| _lib.ax_sr_deinit.restype = ctypes.c_int |
|
|
| _lib.ax_sr_run.argtypes = [ctypes.c_void_p, ctypes.POINTER(SRImage), ctypes.POINTER(SRImage)] |
| _lib.ax_sr_run.restype = ctypes.c_int |
|
|
|
|
| class SR: |
| def __init__(self, init_info: dict): |
| self.handle = None |
| self.init_info = SRInit() |
| |
| |
| self.init_info.dev_type = init_info.get('dev_type', AxDeviceType.axcl_device) |
| self.init_info.devid = init_info.get('devid', 0) |
| setattr(self.init_info, 'model_path', init_info['model_path'].encode('utf-8')) |
|
|
| handle = ctypes.c_void_p() |
| check_error(_lib.ax_sr_init(ctypes.byref(self.init_info), ctypes.byref(handle))) |
| self.handle = handle |
|
|
| def __del__(self): |
| if self.handle: |
| _lib.ax_sr_deinit(self.handle) |
|
|
| def __call__(self, image_data: np.ndarray) -> None: |
| image = SRImage() |
| image.width = image_data.shape[1] |
| image.height = image_data.shape[0] |
| image.pVirAddr = ctypes.cast(image_data.ctypes.data, ctypes.POINTER(ctypes.c_ubyte)) |
| |
| np_sr_image = np.zeros((image.height*2, image.width*2, 3), dtype=np.uint8) |
| sr_image = SRImage() |
| sr_image.width = image.width*2 |
| sr_image.height = image.height*2 |
| sr_image.pVirAddr = ctypes.cast(np_sr_image.ctypes.data, ctypes.POINTER(ctypes.c_ubyte)) |
| |
| check_error(_lib.ax_sr_run(self.handle, ctypes.byref(image), ctypes.byref(sr_image))) |
| return np_sr_image |
|
|
