Update app.py

app.py

@@ -8,21 +8,15 @@ from basicsr.archs.rrdbnet_arch import RRDBNet
 from basicsr.utils.download_util import load_file_from_url
 from realesrgan import RealESRGANer
 from realesrgan.archs.srvgg_arch import SRVGGNetCompact
+from torchvision.transforms.functional import rgb_to_grayscale
 
 last_file = None
-img_mode = "RGB"
-
-def image_properties(img):
-    """Returns the dimensions (width and height) and color mode of the input image"""
-    global img_mode
-    if img:
-        img_mode = "RGB"  # Force the image mode to RGB
-        properties = f"Resolution: Width: {img.size[0]}, Height: {img.size[1]}  |  Color Mode: {img_mode}"
-        return properties
+img_mode = "RGBA"
 
 @spaces.GPU(duration=120)
 def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
-    """Real-ESRGAN function to restore (and upscale) images."""
+    """Real-ESRGAN function to restore (and upscale) images.
+    """
     if not img:
         return
 
@@ -56,6 +50,7 @@ def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
     if not os.path.isfile(model_path):
         ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
         for url in file_url:
+            # model_path will be updated
             model_path = load_file_from_url(
                 url=url, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)
 
@@ -89,9 +84,9 @@ def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
             channel_multiplier=2,
             bg_upsampler=upsampler)
 
-    # Convert the input PIL image to cv2 image (forcing RGB mode)
-    cv_img = numpy.array(img.convert("RGB"))
-    img = cv2.cvtColor(cv_img, cv2.COLOR_RGB2BGR)
+    # Convert the input PIL image to cv2 image, so that it can be processed by realesrgan
+    cv_img = numpy.array(img)
+    img = cv2.cvtColor(cv_img, cv2.COLOR_RGBA2BGRA)
 
     # Apply restoration
     try:
@@ -103,8 +98,13 @@ def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
         print('Error', error)
         print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
     else:
-        # Save restored image
-        out_filename = f"output_{rnd_string(8)}.jpg"
+        # Save restored image and return it to the output Image component
+        if img_mode == 'RGBA':  # RGBA images should be saved in png format
+            extension = 'jpg'
+        else:
+            extension = 'jpg'
+
+        out_filename = f"output_{rnd_string(8)}.{extension}"
         cv2.imwrite(out_filename, output)
         global last_file
         last_file = out_filename
@@ -112,14 +112,17 @@ def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
 
 
 def rnd_string(x):
-    """Returns a string of 'x' random characters"""
+    """Returns a string of 'x' random characters
+    """
     characters = "abcdefghijklmnopqrstuvwxyz_0123456789"
     result = "".join((random.choice(characters)) for i in range(x))
     return result
 
 
 def reset():
-    """Resets the Image components of the Gradio interface and deletes the last processed image"""
+    """Resets the Image components of the Gradio interface and deletes
+    the last processed image
+    """
     global last_file
     if last_file:
         print(f"Deleting {last_file} ...")
@@ -128,6 +131,44 @@ def reset():
     return gr.update(value=None), gr.update(value=None)
 
 
+def has_transparency(img):
+    """This function works by first checking to see if a "transparency" property is defined
+    in the image's info -- if so, we return "True". Then, if the image is using indexed colors
+    (such as in GIFs), it gets the index of the transparent color in the palette
+    (img.info.get("transparency", -1)) and checks if it's used anywhere in the canvas
+    (img.getcolors()). If the image is in RGBA mode, then presumably it has transparency in
+    it, but it double-checks by getting the minimum and maximum values of every color channel
+    (img.getextrema()), and checks if the alpha channel's smallest value falls below 255.
+    https://stackoverflow.com/questions/43864101/python-pil-check-if-image-is-transparent
+    """
+    if img.info.get("transparency", None) is not None:
+        return True
+    if img.mode == "P":
+        transparent = img.info.get("transparency", -1)
+        for _, index in img.getcolors():
+            if index == transparent:
+                return True
+    elif img.mode == "RGBA":
+        extrema = img.getextrema()
+        if extrema[3][0] < 255:
+            return True
+    return False
+
+
+def image_properties(img):
+    """Returns the dimensions (width and height) and color mode of the input image and
+    also sets the global img_mode variable to be used by the realesrgan function
+    """
+    global img_mode
+    if img:
+        if has_transparency(img):
+            img_mode = "RGBA"
+        else:
+            img_mode = "RGB"
+        properties = f"Resolution: Width: {img.size[0]}, Height: {img.size[1]}  |  Color Mode: {img_mode}"
+        return properties
+
+
 def main():
     # Gradio Interface
     with gr.Blocks(title="Real-ESRGAN Gradio Demo", theme="dark") as demo:
@@ -156,9 +197,9 @@ def main():
                 
         with gr.Row():
             with gr.Group():
-                input_image = gr.Image(label="Input Image", type="pil", image_mode="RGB")
+                input_image = gr.Image(label="Input Image", type="pil", image_mode="RGBA")
                 input_image_properties = gr.Textbox(label="Image Properties", max_lines=1)
-            output_image = gr.Image(label="Output Image", image_mode="RGB")
+            output_image = gr.Image(label="Output Image", image_mode="RGBA")
         with gr.Row():
             reset_btn = gr.Button("Remove images")
             restore_btn = gr.Button("Upscale")
@@ -169,6 +210,8 @@ def main():
                           inputs=[input_image, model_name, denoise_strength, face_enhance, outscale],
                           outputs=output_image)
         reset_btn.click(fn=reset, inputs=[], outputs=[output_image, input_image])
+        # reset_btn.click(None, inputs=[], outputs=[input_image], _js="() => (null)\n")
+        # Undocumented method to clear a component's value using Javascript
 
         gr.Markdown(
             """Made with love by Ilaria 💖 | Support me on [Ko-Fi](https://ko-fi.com/ilariaowo) | Join [AI Hub](https://discord.gg/aihub)