add icc_profile

app_base.py

@@ -30,6 +30,7 @@ def create_demo() -> gr.Blocks:
         w2 = 1.0
 
         input_image = Image.open(input_image_path)
+        icc_profile = input_image.info.get("icc_profile")
 
         run_model = base_run
         res_image = run_model(
@@ -48,17 +49,17 @@ def create_demo() -> gr.Blocks:
         tmpPrefix = "/tmp/gradio/"
 
         extension = 'png'
-        # if enhanced_image.mode == 'RGBA':
-        #     extension = 'png'
-        # else:
-        #     extension = 'jpg'
+        if enhanced_image.mode == 'RGBA':
+            extension = 'png'
+        else:
+            extension = 'jpg'
 
         targetDir = f"{tmpPrefix}output/"
         if not os.path.exists(targetDir):
             os.makedirs(targetDir)
 
         enhanced_path = f"{targetDir}{uuid.uuid4()}.{extension}"
-        enhanced_image.save(enhanced_path, quality=100)
+        enhanced_image.save(enhanced_path, quality=100, icc_profile=icc_profile)
 
         return enhanced_path