From 06b7b1a8c5a7e449b9725012fcacc2bb97cb5a89 Mon Sep 17 00:00:00 2001 From: aunyoni88 Date: Mon, 18 Dec 2023 19:34:10 +0600 Subject: [PATCH] Scratch remove API image upscale flag add --- .../scripts/api.py | 29 +++++-------------- 1 file changed, 8 insertions(+), 21 deletions(-) diff --git a/extensions/arifScratchRemoverWebUIExtention/scripts/api.py b/extensions/arifScratchRemoverWebUIExtention/scripts/api.py index b0513d9d..73039945 100644 --- a/extensions/arifScratchRemoverWebUIExtention/scripts/api.py +++ b/extensions/arifScratchRemoverWebUIExtention/scripts/api.py @@ -65,14 +65,14 @@ async def download_model( @app.post('/sdapi/ai/v1/scratch_remove') async def generate_mask_image( input_image: str = Body("", title='scratch remove input image'), - image_name: str = Body("", title='input image name') + upscale: bool = Body(False, title='input image name') ): utc_time = datetime.now(timezone.utc) start_time = time.time() downloadScratchRemoverModelModel() pil_image = api.decode_base64_to_image(input_image) - image_base64_str = remove_scratch_using_mask(pil_image) + image_base64_str = remove_scratch_using_mask(pil_image, upscale) end_time = time.time() server_process_time = end_time - start_time @@ -82,7 +82,7 @@ async def generate_mask_image( "output_image": image_base64_str } - def remove_scratch_using_mask(source_image: Image): + def remove_scratch_using_mask(source_image: Image, upscale : bool): curDir = os.getcwd() fileName = "arif.png" @@ -138,11 +138,11 @@ def remove_scratch_using_mask(source_image: Image): mask_image=main_mask ).images[0] #return base64 image - # opencvImage = cv2.cvtColor(numpy.array(without_scratch_Image_output), cv2.COLOR_RGB2BGR) - # _, encoded_img = cv2.imencode('.jpg', opencvImage) - # img_str = base64.b64encode(encoded_img).decode("utf-8") - # return img_str - + if upscale == False: + opencvImage = cv2.cvtColor(numpy.array(without_scratch_Image_output), cv2.COLOR_RGB2BGR) + _, encoded_img = cv2.imencode('.jpg', opencvImage) + img_str = base64.b64encode(encoded_img).decode("utf-8") + return img_str args = scripts.scripts_postproc.create_args_for_run({ "Upscale": { @@ -164,25 +164,12 @@ def remove_scratch_using_mask(source_image: Image): }, }) - # new_Image = postprocessing.run_postprocessing(0, without_scratch_Image_output, "", "", "", True, *args, save_output=False) - # opencvImage = cv2.cvtColor(numpy.array(new_Image), cv2.COLOR_RGB2BGR) - # _, encoded_img = cv2.imencode('.jpg', opencvImage) - # img_str = base64.b64encode(encoded_img).decode("utf-8") result = postprocessing.run_postprocessing(0, without_scratch_Image_output, "", "", "", True, *args, save_output=False) img_str = api.encode_pil_to_base64(result[0][0]) return img_str - - - - - - - - - def save_file(file: UploadFile, path: str): with open(path, "wb+") as file_object: shutil.copyfileobj(file.file, file_object)