julius/prototorch
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[FEATURE] Add distribution property to LabeledComponents

Browse Source
This commit is contained in:
Jensun Ravichandran 2021-06-14 21:08:48 +02:00
parent d45e71256c
commit 1f458ac0cc
1 changed files with 8 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
prototorch/core/components.py
View File

@ -116,7 +116,7 @@ class AbstractLabels(torch.nn.Module):
@property
def num_labels(self):
return len(self.labels)
return len(self._labels)
@property
def unique_labels(self):
@ -193,6 +193,13 @@ class LabeledComponents(AbstractComponents):
"""Tensor containing the component labels."""
return self._labels
@property
def distribution(self):
unique, counts = torch.unique(self._labels,
sorted=True,
return_counts=True)
return dict(zip(unique.tolist(), counts.tolist()))
def _register_labels(self, labels):
self.register_buffer("_labels", labels)