Add mask output to Gradio app; now returns both overlay and mask images

app.py

@@ -59,7 +59,8 @@ def infer_gradio(model, image):
     color_mask.putalpha(alpha)
     image_rgba = image.convert("RGBA")
     overlay_img = Image.alpha_composite(image_rgba, color_mask).convert("RGB")
-    return overlay_img
+    # Return both overlay and mask
+    return overlay_img, mask_img
 
 def lane_detection(image, model_type):
     model = get_model(model_type)
@@ -68,7 +69,10 @@ def lane_detection(image, model_type):
 demo = gr.Interface(
     fn=lane_detection,
     inputs=[gr.Image(type="pil"), gr.Radio(["unet", "unet_depthwise", "unet_depthwise_small", "unet_depthwise_nano"], label="Model Type")],
-    outputs=gr.Image(type="pil", label="Lane Detection Result"),
+    outputs=[
+        gr.Image(type="pil", label="Lane Detection Result (Overlay)"),
+        gr.Image(type="pil", label="Mask Output")
+    ],
     title="Lane Detection UNet",
     description="Upload a road image and select a model to see lane detection results."
 )