Skip to content

GitLab

Commit a14e63cc authored by WXinlong's avatar WXinlong
Browse files

fix bug in solov2 head

parent 451933f7
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
mmdet/models/anchor_heads/solov2_head.py
View file @ a14e63cc
...@@ -410,7 +410,7 @@ class SOLOv2Head(nn.Module): ...@@ -410,7 +410,7 @@ class SOLOv2Head(nn.Module):
kernel_preds = kernel_preds.view(I, N, 1, 1) kernel_preds = kernel_preds.view(I, N, 1, 1)
seg_preds = F.conv2d(seg_preds, kernel_preds, stride=1).squeeze(0).sigmoid() seg_preds = F.conv2d(seg_preds, kernel_preds, stride=1).squeeze(0).sigmoid()
# mask. # mask.
seg_masks = seg_preds > 0.5 seg_masks = seg_preds > cfg.mask_thr
sum_masks = seg_masks.sum((1, 2)).float() sum_masks = seg_masks.sum((1, 2)).float()
# filter. # filter.
...@@ -464,5 +464,5 @@ class SOLOv2Head(nn.Module): ...@@ -464,5 +464,5 @@ class SOLOv2Head(nn.Module):
seg_masks = F.interpolate(seg_preds, seg_masks = F.interpolate(seg_preds,
size=ori_shape[:2], size=ori_shape[:2],
mode='bilinear').squeeze(0) mode='bilinear').squeeze(0)
seg_masks = seg_masks > 0.5 seg_masks = seg_masks > cfg.mask_thr
return seg_masks, cate_labels, cate_scores return seg_masks, cate_labels, cate_scores
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment