Update app.py

app.py

@@ -33,22 +33,11 @@ predictor = DefaultPredictor(cfg)
 
 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]
-
+    im = cv2.imread(im)
     outputs = predictor(im)
-
-    # We can use `Visualizer` to draw the predictions on the image.
-    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
-    out_im = Image.fromarray(out_im)
-    return out_im
+    v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
+    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
+    return Image.fromarray(np.uint8(out.get_image())).convert('RGB')
 
 # gradio components 
 """