Map prediction to class label for IoU computation

manole-alexandru · Mar 23, 2023 · 6ba70e5 · 6ba70e5
1 parent 4048228
commit 6ba70e5
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/utils/loss.py b/utils/loss.py
@@ -161,9 +161,7 @@ def __call__(self, preds, targets, seg_masks):  # predictions, targets
                     lcls += self.BCEcls(pcls, t)  # BCE
 
                 # Mask Loss
-                # print('\n-----PRED MASK', pred_mask.shape, '-------\n')
-                # print('\n-----REAL MASK', seg_masks.shape, '-------\n')
-                print('\n----------- PRED VALID: ', torch.all(pred_mask >= 0), '-----------------\n')
+                # print('\n----------- PRED VALID: ', torch.all(pred_mask >= 0), '-----------------\n')
                 print('\n----------- SEG MASK VALID: ', torch.all(seg_masks >= 0), '-----------------\n')
                 seg_loss = nn.functional.binary_cross_entropy_with_logits(pred_mask, seg_masks, reduction='none').mean()
                 lseg += seg_loss

diff --git a/val.py b/val.py
@@ -72,6 +72,8 @@ def save_one_json(predn, jdict, path, class_map):
 def compute_seg_iou(pred, target, n_classes=2):
     ious = []
     pred = torch.sigmoid(pred)
+    pred[pred < 0.5] = 0
+    pred[pred >= 0.5] = 1
     pred = pred.view(-1)
     target = target.view(-1)
     print(target)