"git@developer.sourcefind.cn:jerrrrry/infinicore.git" did not exist on "975559ee1b1d570db4dca48ab857e70e6b4cb40b"
Commit 12515f18 authored by mibaumgartner's avatar mibaumgartner
Browse files

bug fix: correct CE num classes

parent 579f2091
...@@ -338,7 +338,7 @@ class CEClassifier(BaseClassifier): ...@@ -338,7 +338,7 @@ class CEClassifier(BaseClassifier):
num_convs=num_convs, num_convs=num_convs,
add_norm=add_norm, add_norm=add_norm,
internal_channels=internal_channels, internal_channels=internal_channels,
num_classes=num_classes, num_classes=num_classes + 1, # add one channel for background
anchors_per_pos=anchors_per_pos, anchors_per_pos=anchors_per_pos,
num_levels=num_levels, num_levels=num_levels,
**kwargs, **kwargs,
...@@ -361,7 +361,7 @@ class CEClassifier(BaseClassifier): ...@@ -361,7 +361,7 @@ class CEClassifier(BaseClassifier):
Returns: Returns:
Tensor: probabilities Tensor: probabilities
""" """
return self.logits_convert_fn(box_logits)[:, 1:] return self.logits_convert_fn(box_logits)[:, 1:] # remove background predictions
class FocalClassifier(BaseClassifier): class FocalClassifier(BaseClassifier):
......
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