feat
fix 超分异常
This commit is contained in:
@@ -58,14 +58,6 @@ class SuperResolution:
|
||||
logging.info(f"{self.tasks_id} ===> {status_data}")
|
||||
return status_data
|
||||
|
||||
# @RunTime
|
||||
def infer(self, inputs):
|
||||
return self.triton_client.async_infer(
|
||||
model_name=SR_MODEL_NAME,
|
||||
inputs=inputs,
|
||||
callback=self.callback
|
||||
)
|
||||
|
||||
# @RunTime
|
||||
def sr_result(self):
|
||||
sample = self.read_image()
|
||||
@@ -82,13 +74,16 @@ class SuperResolution:
|
||||
# , binary_data=True
|
||||
)
|
||||
|
||||
ctx = self.infer(inputs)
|
||||
ctx = self.triton_client.async_infer(
|
||||
model_name=SR_MODEL_NAME,
|
||||
inputs=inputs,
|
||||
callback=self.callback
|
||||
)
|
||||
time_out = 60
|
||||
while time_out > 0:
|
||||
generate_data = self.read_tasks_status()
|
||||
if generate_data['status'] in ["REVOKED", "FAILURE"]:
|
||||
ctx.cancel()
|
||||
# noinspection PyTypeChecker
|
||||
self.channel.basic_publish(exchange='', routing_key=SR_RABBITMQ_QUEUES, body=json.dumps(generate_data))
|
||||
logger.info(f" [x] Sent {generate_data}")
|
||||
break
|
||||
@@ -98,16 +93,6 @@ class SuperResolution:
|
||||
time.sleep(1)
|
||||
return self.read_tasks_status()
|
||||
|
||||
# results = self.triton_client.infer(model_name=SR_MODEL_NAME, inputs=inputs)
|
||||
|
||||
# sr_output = torch.from_numpy(results.as_numpy(f"output"))
|
||||
# output = sr_output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
|
||||
# if output.ndim == 3:
|
||||
# output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0)) # CHW-RGB to HCW-BGR
|
||||
# output = (output * 255.0).round().astype(np.uint8)
|
||||
# output_url = self.upload_img_sr(output, generate_uuid())
|
||||
# return output_url
|
||||
|
||||
def upload_img_sr(self, image):
|
||||
try:
|
||||
image_bytes = cv2.imencode('.jpg', image)[1].tobytes()
|
||||
@@ -121,7 +106,6 @@ class SuperResolution:
|
||||
|
||||
def callback(self, result, error):
|
||||
if error:
|
||||
print(error)
|
||||
sr_info_data = json.dumps({'status': 'FAILURE', 'message': f"{error}", 'data': f"{error}"})
|
||||
self.redis_client.set(self.tasks_id, sr_info_data)
|
||||
else:
|
||||
@@ -147,6 +131,6 @@ def infer_cancel(tasks_id):
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
request_data = SuperResolutionModel(sr_image_url="test/512_image/15.png", sr_xn=2, sr_tasks_id="12341556")
|
||||
request_data = SuperResolutionModel(sr_image_url="aida-users/83/print/b77bf4ca-6ca2-44a1-9040-505f359a974c-3-83.png", sr_xn=2, sr_tasks_id="12341556")
|
||||
service = SuperResolution(request_data)
|
||||
result_url = service.sr_result()
|
||||
|
||||
Reference in New Issue
Block a user