Update app.py

app.py

@@ -36,10 +36,14 @@ def segment_buildings(im):
     # Convert PIL Image to NumPy array
     im = np.array(im)
 
+    # Ensure that the image has shape H x W x C (height x width x channels)
+    if im.shape[-1] == 4:  # If it has 4 channels (RGBA), remove the alpha channel
+        im = im[:, :, :3]
+
     outputs = predictor(im)
 
     # We can use `Visualizer` to draw the predictions on the image.
-    v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
+    v = Visualizer(im, MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
     out_im = v.draw_instance_predictions(outputs["instances"].to("cpu")).get_image()
 
     # Convert the output image back to PIL Image