Update app.py

app.py

@@ -6,14 +6,15 @@ def process_image(img):
     imageio.imwrite("output_image.png", img["composite"])
 
     alpha_channel = img["layers"][0][:, :, 3]
-    mask = np.where(alpha_channel == 0, 0, 255).astype(np.uint8)
+    mask = np.where(alpha_channel == 0, 0, 255)
+    # mask = np.where(alpha_channel == 0, 0, 255).astype(np.uint8)
 
     # return img["background"], mask
     # Save the mask as a PNG file
-    output_path = "mask_output.png"
-    imageio.imwrite(output_path, mask)  # Save the mask as a PNG
+    # output_path = "mask_output.png"
+    # imageio.imwrite(output_path, mask)  # Save the mask as a PNG
 
-    return output_path
+    return mask
 
 # with gr.Blocks() as demo:
 #   with gr.Row():
@@ -29,7 +30,7 @@ demo = gr.Interface(
     inputs=gr.ImageMask(sources=["upload"], layers=False, transforms=[], format="png", label="base image", show_label=True),
     outputs=[
         # gr.Image(label="Background Image"),
-        gr.Image(label="Mask Image"),
+        gr.Image(label="Mask Image", format="png"),
     ],
 )