Update app.py

app.py

@@ -43,15 +43,17 @@ def predict(img_name, model=None, device=None, reduce_labels=True):
     my_model.load_checkpoints(checkpoints)
     with torch.no_grad():
         output = my_model(pre_img_data)
+        print(output.shape)
     overlay = visualize_instances_map(pre_img_data,output)
+    print(pre_img_data.shape,overlay.shape)
 #cv2.imwrite('prediction.png', cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))
-    return overlay
+    return pre_img_data,overlay
 gr.Interface(
     predict,
     inputs=[gr.components.Image(label="Upload Input Image", type="filepath"),
             gr.components.Textbox(label='Model Name', value='sribd_med', max_lines=1)],
     outputs=[gr.Image(label="Processed Image"),
-             gr.Image(label="Label Image"),
+             gr.Image(label="Pred Image"),
              ],
     title="Cell Segmentation Results",
 ).launch()