Update app.py
Browse files
app.py
CHANGED
@@ -123,7 +123,7 @@ def predict(input_image):
|
|
123 |
segment_id = segment["id"]
|
124 |
label_name = model.config.id2label[label_id]
|
125 |
|
126 |
-
if score > 0.
|
127 |
# Create a binary mask for this segment
|
128 |
mask = (segmentation == segment_id)
|
129 |
|
|
|
123 |
segment_id = segment["id"]
|
124 |
label_name = model.config.id2label[label_id]
|
125 |
|
126 |
+
if score > 0.5 and label_name in TARGET_CLASSES:
|
127 |
# Create a binary mask for this segment
|
128 |
mask = (segmentation == segment_id)
|
129 |
|