Spaces:
Runtime error
Runtime error
Added slider component
Browse files
app.py
CHANGED
@@ -32,17 +32,21 @@ cfg.MODEL.WEIGHTS = "model_weights/model_final.pth"
|
|
32 |
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 8
|
33 |
predictor = DefaultPredictor(cfg)
|
34 |
|
35 |
-
def segment_buildings(im):
|
36 |
-
|
37 |
im = np.array(im)
|
38 |
outputs = predictor(im)
|
|
|
|
|
|
|
|
|
|
|
|
|
39 |
v = Visualizer(im[:, :, ::-1],
|
40 |
scale=0.5,
|
41 |
instance_mode=ColorMode.SEGMENTATION
|
42 |
)
|
43 |
-
|
44 |
-
|
45 |
-
|
46 |
return Image.fromarray(out.get_image()[:, :, ::-1])
|
47 |
|
48 |
# gradio components
|
|
|
32 |
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 8
|
33 |
predictor = DefaultPredictor(cfg)
|
34 |
|
35 |
+
def segment_buildings(im, confidence_threshold):
|
|
|
36 |
im = np.array(im)
|
37 |
outputs = predictor(im)
|
38 |
+
|
39 |
+
instances = outputs["instances"].to("cpu")
|
40 |
+
scores = instances.scores
|
41 |
+
selected_indices = scores > confidence_threshold
|
42 |
+
selected_instances = instances[selected_indices]
|
43 |
+
|
44 |
v = Visualizer(im[:, :, ::-1],
|
45 |
scale=0.5,
|
46 |
instance_mode=ColorMode.SEGMENTATION
|
47 |
)
|
48 |
+
out = v.draw_instance_predictions(selected_instances)
|
49 |
+
|
|
|
50 |
return Image.fromarray(out.get_image()[:, :, ::-1])
|
51 |
|
52 |
# gradio components
|