Update app.py

app.py

@@ -16,8 +16,12 @@ def process_image(input_image, torchscript_jit, output_type):
     if output_type == "Mask only":
         # Process the image and get only the mask
         output = remover.process(input_image, type='map')
-        # Convert the mask to a black and white image
-        mask = Image.fromarray((output * 255).astype(np.uint8))
+        if isinstance(output, Image.Image):
+            # If output is already a PIL Image, convert to grayscale
+            mask = output.convert('L')
+        else:
+            # If output is a numpy array, convert to PIL Image
+            mask = Image.fromarray((output * 255).astype(np.uint8), mode='L')
         return mask
     else:
         # Process the image and return the RGBA result