fix bug in kwimage.compress

parent 7abd7d4f
Pipeline #171933 passed with stages
in 6 minutes and 25 seconds
......@@ -8,6 +8,7 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
### Fixed
* Fix bug in `Detections.from_coco_annots` when categories are null.
* Fix bug `Detections.compress` when flags are in tensor space.
## Version 0.6.2 - Released 2020-05-01
......
......@@ -1072,10 +1072,20 @@ class Detections(ub.NiceRepr, _DetAlgoMixin, _DetDrawMixin):
if isinstance(flags, np.ndarray):
if flags.dtype.kind == 'b':
flags = flags.astype(np.uint8)
if _TORCH_HAS_BOOL_COMP:
flags = torch.BoolTensor(flags).to(self.device)
if isinstance(flags, torch.Tensor):
if _TORCH_HAS_BOOL_COMP:
if flags.dtype != torch.bool:
flags = flags.bool()
else:
if flags.dtype != torch.uint8:
flags = flags.byte()
if flags.device != flags.device:
flags = flags.to(self.device)
else:
flags = torch.ByteTensor(flags).to(self.device)
if _TORCH_HAS_BOOL_COMP:
flags = torch.BoolTensor(flags).to(self.device)
else:
flags = torch.ByteTensor(flags).to(self.device)
newdata = {k: _generic._safe_compress(v, flags, axis)
for k, v in self.data.items()}
return self.__class__(newdata, self.meta)
......
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