xwl0809 commited on
Commit
24797c0
1 Parent(s): 5525934

Added slider component

Browse files
Files changed (1) hide show
  1. app.py +9 -5
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
- print(len(outputs["instances"])," buildings detected.")
44
- out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
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