feat generate rgb混乱问题

This commit is contained in:
zhouchengrong
2024-04-25 17:36:35 +08:00
parent a0c47c3367
commit 5f045f8ac1

View File

@@ -40,6 +40,7 @@ class GenerateImage:
self.grpc_client = grpcclient.InferenceServerClient(url=GI_MODEL_URL)
self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
if request_data.mode == "img2img":
# cv2 读图片是BGR PIL读图片是RGB
self.image = self.get_image(request_data.image_url)
self.prompt = request_data.prompt
else:
@@ -60,12 +61,14 @@ class GenerateImage:
def get_image(self, image_url):
# Get data of an object.
# Read data from response.
# read image use cv2
try:
response = self.minio_client.get_object(image_url.split('/')[0], image_url[image_url.find('/') + 1:])
image_file = BytesIO(response.data)
image_array = np.asarray(bytearray(image_file.read()), dtype=np.uint8)
image_cv2 = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
image = cv2.resize(image_cv2, (1024, 1024))
image_rbg = cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)
image = cv2.resize(image_rbg, (1024, 1024))
except minio.error.S3Error:
image = np.random.randint(0, 256, (1024, 1024, 3), dtype=np.uint8)
return image
@@ -77,11 +80,9 @@ class GenerateImage:
# self.generate_data['data'] = str(error)
self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
else:
# pil图像转成numpy数组
image = result.as_numpy("generated_image")
if self.mode == "img2img":
image_result = np.squeeze(image.astype(np.uint8))
else:
image_result = cv2.cvtColor(np.squeeze(image.astype(np.uint8)), cv2.COLOR_BGR2RGB)
image_result = cv2.cvtColor(np.squeeze(image.astype(np.uint8)), cv2.COLOR_RGB2BGR)
is_smudge = True
if self.category == "sketch":
# 去背景