Spaces:

SIH
/

building-segmentation

Runtime error

hlydecker commited on Sep 11, 2023

Commit

b942514

•

1 Parent(s): 59bd0c5

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -29,11 +29,12 @@ cfg.MODEL.DEVICE='cpu'
 cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
 cfg.MODEL.WEIGHTS = "model_weights/chatswood_buildings_poc.pth"
 cfg.MODEL.ROI_HEADS.NUM_CLASSES = 8
 def segment_buildings(im):
     outputs = predictor(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')

 cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
 cfg.MODEL.WEIGHTS = "model_weights/chatswood_buildings_poc.pth"
 cfg.MODEL.ROI_HEADS.NUM_CLASSES = 8
+predictor = DefaultPredictor(cfg)
 def segment_buildings(im):
     outputs = predictor(im)
+    v = Visualizer(im[:, :, ::-1], scale=1.2)
     out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
     return Image.fromarray(np.uint8(out.get_image())).convert('RGB')