Browse Source

[ENH] add abstract data interface

pull/3/head
Gao Enhao 2 years ago
parent
commit
6a63f5e5f3
19 changed files with 1427 additions and 196 deletions
  1. +26
    -14
      abl/bridge/base_bridge.py
  2. +82
    -69
      abl/bridge/simple_bridge.py
  3. +2
    -1
      abl/dataset/__init__.py
  4. +2
    -1
      abl/dataset/bridge_dataset.py
  5. +2
    -1
      abl/dataset/classification_dataset.py
  6. +56
    -0
      abl/dataset/prediction_dataset.py
  7. +2
    -1
      abl/dataset/regression_dataset.py
  8. +1
    -1
      abl/evaluation/__init__.py
  9. +2
    -2
      abl/evaluation/base_metric.py
  10. +8
    -11
      abl/evaluation/semantics_metric.py
  11. +2
    -1
      abl/evaluation/symbol_metric.py
  12. +36
    -63
      abl/learning/abl_model.py
  13. +56
    -24
      abl/learning/basic_nn.py
  14. +2
    -0
      abl/structures/__init__.py
  15. +629
    -0
      abl/structures/base_data_element.py
  16. +321
    -0
      abl/structures/list_data.py
  17. +2
    -1
      abl/utils/__init__.py
  18. +112
    -0
      abl/utils/cache.py
  19. +84
    -6
      abl/utils/utils.py

+ 26
- 14
abl/bridge/base_bridge.py View File

@@ -1,52 +1,64 @@
from abc import ABCMeta, abstractmethod
from typing import Any, List, Tuple
from typing import Any, List, Optional, Tuple, Union

from ..learning import ABLModel
from ..reasoning import ReasonerBase
from ..structures import ListData

DataSet = Tuple[List[List[Any]], Optional[List[List[Any]]], List[List[Any]]]

class BaseBridge(metaclass=ABCMeta):

class BaseBridge(metaclass=ABCMeta):
def __init__(self, model: ABLModel, abducer: ReasonerBase) -> None:
if not isinstance(model, ABLModel):
raise TypeError("Expected an ABLModel")
raise TypeError(
"Expected an instance of ABLModel, but received type: {}".format(
type(model)
)
)
if not isinstance(abducer, ReasonerBase):
raise TypeError("Expected an ReasonerBase")
raise TypeError(
"Expected an instance of ReasonerBase, but received type: {}".format(
type(abducer)
)
)

self.model = model
self.abducer = abducer

@abstractmethod
def predict(self, X: List[List[Any]]) -> Tuple[List[List[Any]], List[List[Any]]]:
def predict(
self, data_samples: ListData
) -> Tuple[List[List[Any]], List[List[Any]]]:
"""Placeholder for predict labels from input."""
pass

@abstractmethod
def abduce_pseudo_label(self, pseudo_label: List[List[Any]]) -> List[List[Any]]:
def abduce_pseudo_label(self, data_samples: ListData) -> List[List[Any]]:
"""Placeholder for abduce pseudo labels."""
pass

@abstractmethod
def idx_to_pseudo_label(self, idx: List[List[Any]]) -> List[List[Any]]:
def idx_to_pseudo_label(self, data_samples: ListData) -> List[List[Any]]:
"""Placeholder for map label space to symbol space."""
pass

@abstractmethod
def pseudo_label_to_idx(self, pseudo_label: List[List[Any]]) -> List[List[Any]]:
def pseudo_label_to_idx(self, data_samples: ListData) -> List[List[Any]]:
"""Placeholder for map symbol space to label space."""
pass
@abstractmethod
def train(self, train_data):
def train(self, train_data: Union[ListData, DataSet]):
"""Placeholder for train loop of ABductive Learning."""
pass

@abstractmethod
def test(self, test_data):
def valid(self, valid_data: Union[ListData, DataSet]) -> None:
"""Placeholder for model test."""
pass

@abstractmethod
def valid(self, valid_data):
def test(self, test_data: Union[ListData, DataSet]) -> None:
"""Placeholder for model validation."""
pass

+ 82
- 69
abl/bridge/simple_bridge.py View File

@@ -1,13 +1,14 @@
from ..learning import ABLModel
from ..reasoning import ReasonerBase
from ..evaluation import BaseMetric
from .base_bridge import BaseBridge
from typing import List, Union, Any, Tuple, Dict, Optional
import os.path as osp
from typing import Any, Dict, List, Optional, Tuple, Union

from numpy import ndarray

from torch.utils.data import DataLoader
from ..dataset import BridgeDataset
from ..utils.logger import print_log
from ..evaluation import BaseMetric
from ..learning import ABLModel
from ..reasoning import ReasonerBase
from ..structures import ListData
from ..utils import print_log
from .base_bridge import BaseBridge, DataSet


class SimpleBridge(BaseBridge):
@@ -20,85 +21,99 @@ class SimpleBridge(BaseBridge):
super().__init__(model, abducer)
self.metric_list = metric_list

def predict(self, X) -> Tuple[List[List[Any]], ndarray]:
pred_res = self.model.predict(X)
pred_idx, pred_prob = pred_res["label"], pred_res["prob"]
return pred_idx, pred_prob
# TODO: add abducer.mapping to the property of SimpleBridge

def predict(self, data_samples: ListData) -> Tuple[List[ndarray], List[ndarray]]:
self.model.predict(data_samples)
return data_samples["pred_idx"], data_samples.get("pred_prob", None)

def abduce_pseudo_label(
self,
pred_prob: ndarray,
pred_pseudo_label: List[List[Any]],
Y: List[Any],
data_samples: ListData,
max_revision: int = -1,
require_more_revision: int = 0,
) -> List[List[Any]]:
return self.abducer.batch_abduce(pred_prob, pred_pseudo_label, Y, max_revision, require_more_revision)
self.abducer.batch_abduce(data_samples, max_revision, require_more_revision)
return data_samples["abduced_pseudo_label"]

def idx_to_pseudo_label(
self, idx: List[List[Any]], mapping: Dict = None
self, data_samples: ListData, mapping: Optional[Dict] = None
) -> List[List[Any]]:
if mapping is None:
mapping = self.abducer.mapping
return [[mapping[_idx] for _idx in sub_list] for sub_list in idx]
pred_idx = data_samples.pred_idx
data_samples.pred_pseudo_label = [
[mapping[_idx] for _idx in sub_list] for sub_list in pred_idx
]
return data_samples["pred_pseudo_label"]

def pseudo_label_to_idx(
self, pseudo_label: List[List[Any]], mapping: Dict = None
self, data_samples: ListData, mapping: Optional[Dict] = None
) -> List[List[Any]]:
if mapping is None:
mapping = self.abducer.remapping
return [
[mapping[_pseudo_label] for _pseudo_label in sub_list]
for sub_list in pseudo_label
abduced_idx = [
[mapping[_abduced_pseudo_label] for _abduced_pseudo_label in sub_list]
for sub_list in data_samples.abduced_pseudo_label
]
data_samples.abduced_idx = abduced_idx
return data_samples["abduced_idx"]

def data_preprocess(self, X: List[Any], gt_pseudo_label: List[Any], Y: List[Any]) -> ListData:
data_samples = ListData()

data_samples.X = X
data_samples.gt_pseudo_label = gt_pseudo_label
data_samples.Y = Y

return data_samples

def train(
self,
train_data: Tuple[List[List[Any]], Optional[List[List[Any]]], List[List[Any]]],
epochs: int = 50,
batch_size: Union[int, float] = -1,
train_data: Union[ListData, DataSet],
loops: int = 50,
segment_size: Union[int, float] = -1,
eval_interval: int = 1,
save_interval: Optional[int] = None,
save_dir: Optional[str] = None,
):
dataset = BridgeDataset(*train_data)
data_loader = DataLoader(
dataset,
batch_size=batch_size,
collate_fn=lambda data_list: [list(data) for data in zip(*data_list)],
)

for epoch in range(epochs):
for seg_idx, (X, Z, Y) in enumerate(data_loader):
pred_idx, pred_prob = self.predict(X)
pred_pseudo_label = self.idx_to_pseudo_label(pred_idx)
abduced_pseudo_label = self.abduce_pseudo_label(
pred_prob, pred_pseudo_label, Y
)
abduced_label = self.pseudo_label_to_idx(abduced_pseudo_label)
loss = self.model.train(X, abduced_label)
if isinstance(train_data, ListData):
data_samples = train_data
else:
data_samples = self.data_preprocess(*train_data)

for loop in range(loops):
for seg_idx in range((len(data_samples) - 1) // segment_size + 1):
sub_data_samples = data_samples[
seg_idx * segment_size : (seg_idx + 1) * segment_size
]
self.predict(sub_data_samples)
self.idx_to_pseudo_label(sub_data_samples)
self.abduce_pseudo_label(sub_data_samples)
self.pseudo_label_to_idx(sub_data_samples)
loss = self.model.train(sub_data_samples)

print_log(
f"Epoch(train) [{epoch + 1}] [{(seg_idx + 1):3}/{len(data_loader)}] model loss is {loss:.5f}",
f"loop(train) [{loop + 1}/{loops}] segment(train) [{(seg_idx + 1)}/{(len(data_samples) - 1) // segment_size + 1}] model loss is {loss:.5f}",
logger="current",
)

if (epoch + 1) % eval_interval == 0 or epoch == epochs - 1:
print_log(f"Evaluation start: Epoch(val) [{epoch}]", logger="current")
if (loop + 1) % eval_interval == 0 or loop == loops - 1:
print_log(f"Evaluation start: loop(val) [{loop + 1}]", logger="current")
self.valid(train_data)

def _valid(self, data_loader):
for X, Z, Y in data_loader:
pred_idx, pred_prob = self.predict(X)
pred_pseudo_label = self.idx_to_pseudo_label(pred_idx)
data_samples = dict(
pred_idx=pred_idx,
pred_prob=pred_prob,
pred_pseudo_label=pred_pseudo_label,
gt_pseudo_label=Z,
Y=Y,
logic_forward=self.abducer.kb.logic_forward,
)
if save_interval is not None and ((loop + 1) % save_interval == 0 or loop == loops - 1):
print_log(f"Saving model: loop(save) [{loop + 1}]", logger="current")
self.model.save(save_path=osp.join(save_dir, f"model_checkpoint_loop_{loop + 1}.pth"))

def _valid(self, data_samples: ListData, batch_size: int = 128) -> None:
for seg_idx in range((len(data_samples) - 1) // batch_size + 1):
sub_data_samples = data_samples[seg_idx * batch_size : (seg_idx + 1) * batch_size]
self.predict(sub_data_samples)
self.idx_to_pseudo_label(sub_data_samples)

for metric in self.metric_list:
metric.process(data_samples)
metric.process(sub_data_samples)

res = dict()
for metric in self.metric_list:
@@ -108,14 +123,12 @@ class SimpleBridge(BaseBridge):
msg += k + f": {v:.3f} "
print_log(msg, logger="current")

def valid(self, valid_data, batch_size=1000):
dataset = BridgeDataset(*valid_data)
data_loader = DataLoader(
dataset,
batch_size=batch_size,
collate_fn=lambda data_list: [list(data) for data in zip(*data_list)],
)
self._valid(data_loader)

def test(self, test_data, batch_size=1000):
self.valid(test_data, batch_size)
def valid(self, valid_data: Union[ListData, DataSet], batch_size: int = 128) -> None:
if not isinstance(valid_data, ListData):
data_samples = self.data_preprocess(*valid_data)
else:
data_samples = valid_data
self._valid(data_samples, batch_size=batch_size)

def test(self, test_data: Union[ListData, DataSet], batch_size: int = 128) -> None:
self.valid(test_data, batch_size=batch_size)

+ 2
- 1
abl/dataset/__init__.py View File

@@ -1,3 +1,4 @@
from .bridge_dataset import BridgeDataset
from .classification_dataset import ClassificationDataset
from .regression_dataset import RegressionDataset
from .prediction_dataset import PredictionDataset
from .regression_dataset import RegressionDataset

+ 2
- 1
abl/dataset/bridge_dataset.py View File

@@ -1,5 +1,6 @@
from typing import Any, List, Tuple

from torch.utils.data import Dataset
from typing import List, Any, Tuple


class BridgeDataset(Dataset):


+ 2
- 1
abl/dataset/classification_dataset.py View File

@@ -1,6 +1,7 @@
from typing import Any, Callable, List, Tuple

import torch
from torch.utils.data import Dataset
from typing import List, Any, Tuple, Callable


class ClassificationDataset(Dataset):


+ 56
- 0
abl/dataset/prediction_dataset.py View File

@@ -0,0 +1,56 @@
from typing import Any, Callable, List, Tuple

import torch
from torch.utils.data import Dataset


class PredictionDataset(Dataset):
def __init__(self, X: List[Any], transform: Callable[..., Any] = None):
"""
Initialize the dataset used for classification task.

Parameters
----------
X : List[Any]
The input data.
transform : Callable[..., Any], optional
A function/transform that takes in an object and returns a transformed version. Defaults to None.
"""
if not isinstance(X, list):
raise ValueError("X should be of type list.")

self.X = X
self.transform = transform

def __len__(self) -> int:
"""
Return the length of the dataset.

Returns
-------
int
The length of the dataset.
"""
return len(self.X)

def __getitem__(self, index: int) -> Tuple[Any, torch.Tensor]:
"""
Get the item at the given index.

Parameters
----------
index : int
The index of the item to get.

Returns
-------
Tuple[Any, torch.Tensor]
A tuple containing the object and its label.
"""
if index >= len(self):
raise ValueError("index range error")

x = self.X[index]
if self.transform is not None:
x = self.transform(x)
return x

+ 2
- 1
abl/dataset/regression_dataset.py View File

@@ -1,6 +1,7 @@
from typing import Any, List, Tuple

import torch
from torch.utils.data import Dataset
from typing import List, Any, Tuple


class RegressionDataset(Dataset):


+ 1
- 1
abl/evaluation/__init__.py View File

@@ -1,3 +1,3 @@
from .base_metric import BaseMetric
from .symbol_metric import SymbolMetric
from .semantics_metric import SemanticsMetric
from .symbol_metric import SymbolMetric

+ 2
- 2
abl/evaluation/base_metric.py View File

@@ -1,8 +1,8 @@
import logging
from abc import ABCMeta, abstractmethod
from typing import Any, List, Optional, Sequence
from ..utils import print_log

import logging
from ..utils import print_log


class BaseMetric(metaclass=ABCMeta):


+ 8
- 11
abl/evaluation/semantics_metric.py View File

@@ -1,25 +1,22 @@
from typing import Optional, Sequence

from ..reasoning import BaseKB
from .base_metric import BaseMetric

class ABLMetric():
pass

class SemanticsMetric(BaseMetric):
def __init__(self, prefix: Optional[str] = None) -> None:
def __init__(self, kb: BaseKB = None, prefix: Optional[str] = None) -> None:
super().__init__(prefix)
self.kb = kb

def process(self, data_samples: Sequence[dict]) -> None:
pred_pseudo_label = data_samples["pred_pseudo_label"]
gt_Y = data_samples["Y"]
logic_forward = data_samples["logic_forward"]

for pred_z, y in zip(pred_pseudo_label, gt_Y):
if logic_forward(pred_z) == y:
for data_sample in data_samples:
if self.kb.check_equal(data_sample, data_sample["Y"][0]):
self.results.append(1)
else:
self.results.append(0)

def compute_metrics(self, results: list) -> dict:
metrics = dict()
metrics["semantics_accuracy"] = sum(results) / len(results)
return metrics
return metrics

+ 2
- 1
abl/evaluation/symbol_metric.py View File

@@ -1,4 +1,5 @@
from typing import Optional, Sequence, Callable
from typing import Optional, Sequence

from .base_metric import BaseMetric




+ 36
- 63
abl/learning/abl_model.py View File

@@ -10,8 +10,10 @@
#
# ================================================================#
import pickle
from utils import flatten, reform_idx
from typing import List, Any, Optional
from typing import Any, Dict

from ..structures import ListData
from ..utils import reform_idx


class ABLModel:
@@ -30,7 +32,7 @@ class ABLModel:

Methods
-------
predict(X: List[List[Any]], mapping: Optional[dict] = None) -> dict
predict(X: List[List[Any]], mapping: Optional[Dict] = None) -> Dict
Predict the labels and probabilities for the given data.
valid(X: List[List[Any]], Y: List[Any]) -> float
Calculate the accuracy score for the given data.
@@ -42,20 +44,13 @@ class ABLModel:
Load the model from a file.
"""

def __init__(self, base_model) -> None:
self.classifier_list = []
self.classifier_list.append(base_model)
def __init__(self, base_model: Any) -> None:
if not (hasattr(base_model, "fit") and hasattr(base_model, "predict")):
raise NotImplementedError("The base_model should implement fit and predict methods.")

if not (
hasattr(base_model, "fit")
and hasattr(base_model, "predict")
and hasattr(base_model, "score")
):
raise NotImplementedError(
"base_model should have fit, predict and score methods."
)
self.base_model = base_model

def predict(self, X: List[List[Any]], mapping: Optional[dict] = None) -> dict:
def predict(self, data_samples: ListData) -> Dict:
"""
Predict the labels and probabilities for the given data.

@@ -63,53 +58,30 @@ class ABLModel:
----------
X : List[List[Any]]
The data to predict on.
mapping : Optional[dict], optional
A mapping dictionary to map labels to their original values, by default None.

Returns
-------
dict
A dictionary containing the predicted labels and probabilities.
"""
model = self.classifier_list[0]
data_X = flatten(X)
model = self.base_model
data_X = data_samples.flatten("X")
if hasattr(model, "predict_proba"):
prob = model.predict_proba(X=data_X)
label = prob.argmax(axis=1)
prob = reform_idx(prob, X)
prob = reform_idx(prob, data_samples["X"])
else:
prob = None
label = model.predict(X=data_X)
label = reform_idx(label, data_samples["X"])

if mapping is not None:
label = [mapping[y] for y in label]

label = reform_idx(label, X)
data_samples.pred_idx = label
if prob is not None:
data_samples.pred_prob = prob

return {"label": label, "prob": prob}

def valid(self, X: List[List[Any]], Y: List[Any]) -> float:
"""
Calculate the accuracy for the given data.

Parameters
----------
X : List[List[Any]]
The data to calculate the accuracy on.
Y : List[Any]
The true labels for the given data.

Returns
-------
float
The accuracy score for the given data.
"""
data_X = flatten(X)
data_Y = flatten(Y)
score = self.classifier_list[0].score(X=data_X, y=data_Y)
return score

def train(self, X: List[List[Any]], Y: List[Any]) -> float:
def train(self, data_samples: ListData) -> float:
"""
Train the model on the given data.

@@ -125,29 +97,30 @@ class ABLModel:
float
The loss value of the trained model.
"""
data_X = flatten(X)
data_Y = flatten(Y)
return self.classifier_list[0].fit(X=data_X, y=data_Y)
data_X = data_samples.flatten("X")
data_y = data_samples.flatten("abduced_idx")
return self.base_model.fit(X=data_X, y=data_y)

def _model_operation(self, operation: str, *args, **kwargs):
model = self.classifier_list[0]
model = self.base_model
if hasattr(model, operation):
method = getattr(model, operation)
method(*args, **kwargs)
else:
try:
if not f"{operation}_path" in kwargs.keys():
raise ValueError(f"'{operation}_path' should not be None")
if operation == "save":
with open(kwargs["save_path"], 'wb') as file:
pickle.dump(model, file, protocol=pickle.HIGHEST_PROTOCOL)
elif operation == "load":
with open(kwargs["load_path"], 'rb') as file:
self.classifier_list[0] = pickle.load(file)
except:
raise NotImplementedError(
f"{type(model).__name__} object doesn't have the {operation} method"
)
if not f"{operation}_path" in kwargs.keys():
raise ValueError(f"'{operation}_path' should not be None")
else:
try:
if operation == "save":
with open(kwargs["save_path"], "wb") as file:
pickle.dump(model, file, protocol=pickle.HIGHEST_PROTOCOL)
elif operation == "load":
with open(kwargs["load_path"], "rb") as file:
self.base_model = pickle.load(file)
except:
raise NotImplementedError(
f"{type(model).__name__} object doesn't have the {operation} method and the default pickle-based {operation} method failed."
)

def save(self, *args, **kwargs) -> None:
"""


+ 56
- 24
abl/learning/basic_nn.py View File

@@ -10,14 +10,16 @@
#
# ================================================================#

import torch
import os
import logging
from typing import Any, Callable, List, Optional, T, Tuple

import numpy
import torch
from torch.utils.data import DataLoader
from ..utils.logger import print_log
from ..dataset import ClassificationDataset

import os
from typing import List, Any, T, Optional, Callable, Tuple
from ..dataset import ClassificationDataset, PredictionDataset
from ..utils.logger import print_log


class BasicNN:
@@ -99,9 +101,7 @@ class BasicNN:
loss_value = self.train_epoch(data_loader)
if self.save_interval is not None and (epoch + 1) % self.save_interval == 0:
if self.save_dir is None:
raise ValueError(
"save_dir should not be None if save_interval is not None."
)
raise ValueError("save_dir should not be None if save_interval is not None.")
self.save(epoch + 1)
if self.stop_loss is not None and loss_value < self.stop_loss:
break
@@ -191,7 +191,7 @@ class BasicNN:

with torch.no_grad():
results = []
for data, _ in data_loader:
for data in data_loader:
data = data.to(device)
out = model(data)
results.append(out)
@@ -199,7 +199,10 @@ class BasicNN:
return torch.cat(results, axis=0)

def predict(
self, data_loader: DataLoader = None, X: List[Any] = None
self,
data_loader: DataLoader = None,
X: List[Any] = None,
test_transform: Callable[..., Any] = None,
) -> numpy.ndarray:
"""
Predict the class of the input data.
@@ -218,11 +221,28 @@ class BasicNN:
"""

if data_loader is None:
data_loader = self._data_loader(X)
if test_transform is None:
print_log(
"Transform used in the training phase will be used in prediction.",
"current",
level=logging.WARNING,
)
dataset = PredictionDataset(X, self.transform)
else:
dataset = PredictionDataset(X, test_transform)
data_loader = DataLoader(
dataset,
batch_size=self.batch_size,
num_workers=int(self.num_workers),
collate_fn=self.collate_fn,
)
return self._predict(data_loader).argmax(axis=1).cpu().numpy()

def predict_proba(
self, data_loader: DataLoader = None, X: List[Any] = None
self,
data_loader: DataLoader = None,
X: List[Any] = None,
test_transform: Callable[..., Any] = None,
) -> numpy.ndarray:
"""
Predict the probability of each class for the input data.
@@ -241,7 +261,21 @@ class BasicNN:
"""

if data_loader is None:
data_loader = self._data_loader(X)
if test_transform is None:
print_log(
"Transform used in the training phase will be used in prediction.",
"current",
level=logging.WARNING,
)
dataset = PredictionDataset(X, self.transform)
else:
dataset = PredictionDataset(X, test_transform)
data_loader = DataLoader(
dataset,
batch_size=self.batch_size,
num_workers=int(self.num_workers),
collate_fn=self.collate_fn,
)
return self._predict(data_loader).softmax(axis=1).cpu().numpy()

def _score(self, data_loader) -> Tuple[float, float]:
@@ -313,15 +347,14 @@ class BasicNN:
if data_loader is None:
data_loader = self._data_loader(X, y)
mean_loss, accuracy = self._score(data_loader)
print_log(
f"mean loss: {mean_loss:.3f}, accuray: {accuracy:.3f}", logger="current"
)
print_log(f"mean loss: {mean_loss:.3f}, accuray: {accuracy:.3f}", logger="current")
return accuracy

def _data_loader(
self,
X: List[Any],
y: List[int] = None,
shuffle: bool = True,
) -> DataLoader:
"""
Generate a DataLoader for user-provided input and target data.
@@ -350,7 +383,7 @@ class BasicNN:
data_loader = DataLoader(
dataset,
batch_size=self.batch_size,
shuffle=True,
shuffle=shuffle,
num_workers=int(self.num_workers),
collate_fn=self.collate_fn,
)
@@ -368,14 +401,13 @@ class BasicNN:
The path to save the model, by default None.
"""
if self.save_dir is None and save_path is None:
raise ValueError(
"'save_dir' and 'save_path' should not be None simultaneously."
)
raise ValueError("'save_dir' and 'save_path' should not be None simultaneously.")

if save_path is None:
save_path = os.path.join(
self.save_dir, f"model_checkpoint_epoch_{epoch_id}.pth"
)
if save_path is not None:
if not os.path.exists(os.path.dirname(save_path)):
os.makedirs(os.path.dirname(save_path))
else:
save_path = os.path.join(self.save_dir, f"model_checkpoint_epoch_{epoch_id}.pth")
if not os.path.exists(self.save_dir):
os.makedirs(self.save_dir)



+ 2
- 0
abl/structures/__init__.py View File

@@ -0,0 +1,2 @@
from .base_data_element import BaseDataElement
from .list_data import ListData

+ 629
- 0
abl/structures/base_data_element.py View File

@@ -0,0 +1,629 @@
# Copyright (c) OpenMMLab. All rights reserved.
import copy
from typing import Any, Iterator, Optional, Tuple, Type, Union

import numpy as np
import torch


class BaseDataElement:
"""A base data interface that supports Tensor-like and dict-like
operations.

A typical data elements refer to predicted results or ground truth labels
on a task, such as predicted bboxes, instance masks, semantic
segmentation masks, etc. Because groundtruth labels and predicted results
often have similar properties (for example, the predicted bboxes and the
groundtruth bboxes), MMEngine uses the same abstract data interface to
encapsulate predicted results and groundtruth labels, and it is recommended
to use different name conventions to distinguish them, such as using
``gt_instances`` and ``pred_instances`` to distinguish between labels and
predicted results. Additionally, we distinguish data elements at instance
level, pixel level, and label level. Each of these types has its own
characteristics. Therefore, MMEngine defines the base class
``BaseDataElement``, and implement ``InstanceData``, ``PixelData``, and
``LabelData`` inheriting from ``BaseDataElement`` to represent different
types of ground truth labels or predictions.

Another common data element is sample data. A sample data consists of input
data (such as an image) and its annotations and predictions. In general,
an image can have multiple types of annotations and/or predictions at the
same time (for example, both pixel-level semantic segmentation annotations
and instance-level detection bboxes annotations). All labels and
predictions of a training sample are often passed between Dataset, Model,
Visualizer, and Evaluator components. In order to simplify the interface
between components, we can treat them as a large data element and
encapsulate them. Such data elements are generally called XXDataSample in
the OpenMMLab. Therefore, Similar to `nn.Module`, the `BaseDataElement`
allows `BaseDataElement` as its attribute. Such a class generally
encapsulates all the data of a sample in the algorithm library, and its
attributes generally are various types of data elements. For example,
MMDetection is assigned by the BaseDataElement to encapsulate all the data
elements of the sample labeling and prediction of a sample in the
algorithm library.

The attributes in ``BaseDataElement`` are divided into two parts,
the ``metainfo`` and the ``data`` respectively.

- ``metainfo``: Usually contains the
information about the image such as filename,
image_shape, pad_shape, etc. The attributes can be accessed or
modified by dict-like or object-like operations, such as
``.`` (for data access and modification), ``in``, ``del``,
``pop(str)``, ``get(str)``, ``metainfo_keys()``,
``metainfo_values()``, ``metainfo_items()``, ``set_metainfo()`` (for
set or change key-value pairs in metainfo).

- ``data``: Annotations or model predictions are
stored. The attributes can be accessed or modified by
dict-like or object-like operations, such as
``.``, ``in``, ``del``, ``pop(str)``, ``get(str)``, ``keys()``,
``values()``, ``items()``. Users can also apply tensor-like
methods to all :obj:`torch.Tensor` in the ``data_fields``,
such as ``.cuda()``, ``.cpu()``, ``.numpy()``, ``.to()``,
``to_tensor()``, ``.detach()``.

Args:
metainfo (dict, optional): A dict contains the meta information
of single image, such as ``dict(img_shape=(512, 512, 3),
scale_factor=(1, 1, 1, 1))``. Defaults to None.
kwargs (dict, optional): A dict contains annotations of single image or
model predictions. Defaults to None.

Examples:
>>> import torch
>>> from mmengine.structures import BaseDataElement
>>> gt_instances = BaseDataElement()
>>> bboxes = torch.rand((5, 4))
>>> scores = torch.rand((5,))
>>> img_id = 0
>>> img_shape = (800, 1333)
>>> gt_instances = BaseDataElement(
... metainfo=dict(img_id=img_id, img_shape=img_shape),
... bboxes=bboxes, scores=scores)
>>> gt_instances = BaseDataElement(
... metainfo=dict(img_id=img_id, img_shape=(640, 640)))

>>> # new
>>> gt_instances1 = gt_instances.new(
... metainfo=dict(img_id=1, img_shape=(640, 640)),
... bboxes=torch.rand((5, 4)),
... scores=torch.rand((5,)))
>>> gt_instances2 = gt_instances1.new()

>>> # add and process property
>>> gt_instances = BaseDataElement()
>>> gt_instances.set_metainfo(dict(img_id=9, img_shape=(100, 100)))
>>> assert 'img_shape' in gt_instances.metainfo_keys()
>>> assert 'img_shape' in gt_instances
>>> assert 'img_shape' not in gt_instances.keys()
>>> assert 'img_shape' in gt_instances.all_keys()
>>> print(gt_instances.img_shape)
(100, 100)
>>> gt_instances.scores = torch.rand((5,))
>>> assert 'scores' in gt_instances.keys()
>>> assert 'scores' in gt_instances
>>> assert 'scores' in gt_instances.all_keys()
>>> assert 'scores' not in gt_instances.metainfo_keys()
>>> print(gt_instances.scores)
tensor([0.5230, 0.7885, 0.2426, 0.3911, 0.4876])
>>> gt_instances.bboxes = torch.rand((5, 4))
>>> assert 'bboxes' in gt_instances.keys()
>>> assert 'bboxes' in gt_instances
>>> assert 'bboxes' in gt_instances.all_keys()
>>> assert 'bboxes' not in gt_instances.metainfo_keys()
>>> print(gt_instances.bboxes)
tensor([[0.0900, 0.0424, 0.1755, 0.4469],
[0.8648, 0.0592, 0.3484, 0.0913],
[0.5808, 0.1909, 0.6165, 0.7088],
[0.5490, 0.4209, 0.9416, 0.2374],
[0.3652, 0.1218, 0.8805, 0.7523]])

>>> # delete and change property
>>> gt_instances = BaseDataElement(
... metainfo=dict(img_id=0, img_shape=(640, 640)),
... bboxes=torch.rand((6, 4)), scores=torch.rand((6,)))
>>> gt_instances.set_metainfo(dict(img_shape=(1280, 1280)))
>>> gt_instances.img_shape # (1280, 1280)
>>> gt_instances.bboxes = gt_instances.bboxes * 2
>>> gt_instances.get('img_shape', None) # (1280, 1280)
>>> gt_instances.get('bboxes', None) # 6x4 tensor
>>> del gt_instances.img_shape
>>> del gt_instances.bboxes
>>> assert 'img_shape' not in gt_instances
>>> assert 'bboxes' not in gt_instances
>>> gt_instances.pop('img_shape', None) # None
>>> gt_instances.pop('bboxes', None) # None

>>> # Tensor-like
>>> cuda_instances = gt_instances.cuda()
>>> cuda_instances = gt_instances.to('cuda:0')
>>> cpu_instances = cuda_instances.cpu()
>>> cpu_instances = cuda_instances.to('cpu')
>>> fp16_instances = cuda_instances.to(
... device=None, dtype=torch.float16, non_blocking=False,
... copy=False, memory_format=torch.preserve_format)
>>> cpu_instances = cuda_instances.detach()
>>> np_instances = cpu_instances.numpy()

>>> # print
>>> metainfo = dict(img_shape=(800, 1196, 3))
>>> gt_instances = BaseDataElement(
... metainfo=metainfo, det_labels=torch.LongTensor([0, 1, 2, 3]))
>>> sample = BaseDataElement(metainfo=metainfo,
... gt_instances=gt_instances)
>>> print(sample)
<BaseDataElement(
META INFORMATION
img_shape: (800, 1196, 3)
DATA FIELDS
gt_instances: <BaseDataElement(
META INFORMATION
img_shape: (800, 1196, 3)
DATA FIELDS
det_labels: tensor([0, 1, 2, 3])
) at 0x7f0ec5eadc70>
) at 0x7f0fea49e130>

>>> # inheritance
>>> class DetDataSample(BaseDataElement):
... @property
... def proposals(self):
... return self._proposals
... @proposals.setter
... def proposals(self, value):
... self.set_field(value, '_proposals', dtype=BaseDataElement)
... @proposals.deleter
... def proposals(self):
... del self._proposals
... @property
... def gt_instances(self):
... return self._gt_instances
... @gt_instances.setter
... def gt_instances(self, value):
... self.set_field(value, '_gt_instances',
... dtype=BaseDataElement)
... @gt_instances.deleter
... def gt_instances(self):
... del self._gt_instances
... @property
... def pred_instances(self):
... return self._pred_instances
... @pred_instances.setter
... def pred_instances(self, value):
... self.set_field(value, '_pred_instances',
... dtype=BaseDataElement)
... @pred_instances.deleter
... def pred_instances(self):
... del self._pred_instances
>>> det_sample = DetDataSample()
>>> proposals = BaseDataElement(bboxes=torch.rand((5, 4)))
>>> det_sample.proposals = proposals
>>> assert 'proposals' in det_sample
>>> assert det_sample.proposals == proposals
>>> del det_sample.proposals
>>> assert 'proposals' not in det_sample
>>> with self.assertRaises(AssertionError):
... det_sample.proposals = torch.rand((5, 4))
"""

def __init__(self, *, metainfo: Optional[dict] = None, **kwargs) -> None:
self._metainfo_fields: set = set()
self._data_fields: set = set()

if metainfo is not None:
self.set_metainfo(metainfo=metainfo)
if kwargs:
self.set_data(kwargs)

def set_metainfo(self, metainfo: dict) -> None:
"""Set or change key-value pairs in ``metainfo_field`` by parameter
``metainfo``.

Args:
metainfo (dict): A dict contains the meta information
of image, such as ``img_shape``, ``scale_factor``, etc.
"""
assert isinstance(
metainfo, dict
), f"metainfo should be a ``dict`` but got {type(metainfo)}"
meta = copy.deepcopy(metainfo)
for k, v in meta.items():
self.set_field(name=k, value=v, field_type="metainfo", dtype=None)

def set_data(self, data: dict) -> None:
"""Set or change key-value pairs in ``data_field`` by parameter
``data``.

Args:
data (dict): A dict contains annotations of image or
model predictions.
"""
assert isinstance(data, dict), f"data should be a `dict` but got {data}"
for k, v in data.items():
# Use `setattr()` rather than `self.set_field` to allow `set_data`
# to set property method.
setattr(self, k, v)

def update(self, instance: "BaseDataElement") -> None:
"""The update() method updates the BaseDataElement with the elements
from another BaseDataElement object.

Args:
instance (BaseDataElement): Another BaseDataElement object for
update the current object.
"""
assert isinstance(
instance, BaseDataElement
), f"instance should be a `BaseDataElement` but got {type(instance)}"
self.set_metainfo(dict(instance.metainfo_items()))
self.set_data(dict(instance.items()))

def new(self, *, metainfo: Optional[dict] = None, **kwargs) -> "BaseDataElement":
"""Return a new data element with same type. If ``metainfo`` and
``data`` are None, the new data element will have same metainfo and
data. If metainfo or data is not None, the new result will overwrite it
with the input value.

Args:
metainfo (dict, optional): A dict contains the meta information
of image, such as ``img_shape``, ``scale_factor``, etc.
Defaults to None.
kwargs (dict): A dict contains annotations of image or
model predictions.

Returns:
BaseDataElement: A new data element with same type.
"""
new_data = self.__class__()

if metainfo is not None:
new_data.set_metainfo(metainfo)
else:
new_data.set_metainfo(dict(self.metainfo_items()))
if kwargs:
new_data.set_data(kwargs)
else:
new_data.set_data(dict(self.items()))
return new_data

def clone(self):
"""Deep copy the current data element.

Returns:
BaseDataElement: The copy of current data element.
"""
clone_data = self.__class__()
clone_data.set_metainfo(dict(self.metainfo_items()))
clone_data.set_data(dict(self.items()))
return clone_data

def keys(self) -> list:
"""
Returns:
list: Contains all keys in data_fields.
"""
# We assume that the name of the attribute related to property is
# '_' + the name of the property. We use this rule to filter out
# private keys.
# TODO: Use a more robust way to solve this problem
private_keys = {
"_" + key
for key in self._data_fields
if isinstance(getattr(type(self), key, None), property)
}
return list(self._data_fields - private_keys)

def metainfo_keys(self) -> list:
"""
Returns:
list: Contains all keys in metainfo_fields.
"""
return list(self._metainfo_fields)

def values(self) -> list:
"""
Returns:
list: Contains all values in data.
"""
return [getattr(self, k) for k in self.keys()]

def metainfo_values(self) -> list:
"""
Returns:
list: Contains all values in metainfo.
"""
return [getattr(self, k) for k in self.metainfo_keys()]

def all_keys(self) -> list:
"""
Returns:
list: Contains all keys in metainfo and data.
"""
return self.metainfo_keys() + self.keys()

def all_values(self) -> list:
"""
Returns:
list: Contains all values in metainfo and data.
"""
return self.metainfo_values() + self.values()

def all_items(self) -> Iterator[Tuple[str, Any]]:
"""
Returns:
iterator: An iterator object whose element is (key, value) tuple
pairs for ``metainfo`` and ``data``.
"""
for k in self.all_keys():
yield (k, getattr(self, k))

def items(self) -> Iterator[Tuple[str, Any]]:
"""
Returns:
iterator: An iterator object whose element is (key, value) tuple
pairs for ``data``.
"""
for k in self.keys():
yield (k, getattr(self, k))

def metainfo_items(self) -> Iterator[Tuple[str, Any]]:
"""
Returns:
iterator: An iterator object whose element is (key, value) tuple
pairs for ``metainfo``.
"""
for k in self.metainfo_keys():
yield (k, getattr(self, k))

@property
def metainfo(self) -> dict:
"""dict: A dict contains metainfo of current data element."""
return dict(self.metainfo_items())

def __setattr__(self, name: str, value: Any):
"""setattr is only used to set data."""
if name in ("_metainfo_fields", "_data_fields"):
if not hasattr(self, name):
super().__setattr__(name, value)
else:
raise AttributeError(
f"{name} has been used as a "
"private attribute, which is immutable."
)
else:
self.set_field(name=name, value=value, field_type="data", dtype=None)

def __delattr__(self, item: str):
"""Delete the item in dataelement.

Args:
item (str): The key to delete.
"""
if item in ("_metainfo_fields", "_data_fields"):
raise AttributeError(
f"{item} has been used as a " "private attribute, which is immutable."
)
super().__delattr__(item)
if item in self._metainfo_fields:
self._metainfo_fields.remove(item)
elif item in self._data_fields:
self._data_fields.remove(item)

# dict-like methods
__delitem__ = __delattr__

def get(self, key, default=None) -> Any:
"""Get property in data and metainfo as the same as python."""
# Use `getattr()` rather than `self.__dict__.get()` to allow getting
# properties.
return getattr(self, key, default)

def pop(self, *args) -> Any:
"""Pop property in data and metainfo as the same as python."""
assert len(args) < 3, "``pop`` get more than 2 arguments"
name = args[0]
if name in self._metainfo_fields:
self._metainfo_fields.remove(args[0])
return self.__dict__.pop(*args)

elif name in self._data_fields:
self._data_fields.remove(args[0])
return self.__dict__.pop(*args)

# with default value
elif len(args) == 2:
return args[1]
else:
# don't just use 'self.__dict__.pop(*args)' for only popping key in
# metainfo or data
raise KeyError(f"{args[0]} is not contained in metainfo or data")

def __contains__(self, item: str) -> bool:
"""Whether the item is in dataelement.

Args:
item (str): The key to inquire.
"""
return item in self._data_fields or item in self._metainfo_fields

def set_field(
self,
value: Any,
name: str,
dtype: Optional[Union[Type, Tuple[Type, ...]]] = None,
field_type: str = "data",
) -> None:
"""Special method for set union field, used as property.setter
functions."""
assert field_type in ["metainfo", "data"]
if dtype is not None:
assert isinstance(
value, dtype
), f"{value} should be a {dtype} but got {type(value)}"

if field_type == "metainfo":
if name in self._data_fields:
raise AttributeError(
f"Cannot set {name} to be a field of metainfo "
f"because {name} is already a data field"
)
self._metainfo_fields.add(name)
else:
if name in self._metainfo_fields:
raise AttributeError(
f"Cannot set {name} to be a field of data "
f"because {name} is already a metainfo field"
)
self._data_fields.add(name)
super().__setattr__(name, value)

# Tensor-like methods
def to(self, *args, **kwargs) -> "BaseDataElement":
"""Apply same name function to all tensors in data_fields."""
new_data = self.new()
for k, v in self.items():
if hasattr(v, "to"):
v = v.to(*args, **kwargs)
data = {k: v}
new_data.set_data(data)
return new_data

# Tensor-like methods
def cpu(self) -> "BaseDataElement":
"""Convert all tensors to CPU in data."""
new_data = self.new()
for k, v in self.items():
if isinstance(v, (torch.Tensor, BaseDataElement)):
v = v.cpu()
data = {k: v}
new_data.set_data(data)
return new_data

# Tensor-like methods
def cuda(self) -> "BaseDataElement":
"""Convert all tensors to GPU in data."""
new_data = self.new()
for k, v in self.items():
if isinstance(v, (torch.Tensor, BaseDataElement)):
v = v.cuda()
data = {k: v}
new_data.set_data(data)
return new_data

# Tensor-like methods
def npu(self) -> "BaseDataElement":
"""Convert all tensors to NPU in data."""
new_data = self.new()
for k, v in self.items():
if isinstance(v, (torch.Tensor, BaseDataElement)):
v = v.npu()
data = {k: v}
new_data.set_data(data)
return new_data

def mlu(self) -> "BaseDataElement":
"""Convert all tensors to MLU in data."""
new_data = self.new()
for k, v in self.items():
if isinstance(v, (torch.Tensor, BaseDataElement)):
v = v.mlu()
data = {k: v}
new_data.set_data(data)
return new_data

# Tensor-like methods
def detach(self) -> "BaseDataElement":
"""Detach all tensors in data."""
new_data = self.new()
for k, v in self.items():
if isinstance(v, (torch.Tensor, BaseDataElement)):
v = v.detach()
data = {k: v}
new_data.set_data(data)
return new_data

# Tensor-like methods
def numpy(self) -> "BaseDataElement":
"""Convert all tensors to np.ndarray in data."""
new_data = self.new()
for k, v in self.items():
if isinstance(v, (torch.Tensor, BaseDataElement)):
v = v.detach().cpu().numpy()
data = {k: v}
new_data.set_data(data)
return new_data

def to_tensor(self) -> "BaseDataElement":
"""Convert all np.ndarray to tensor in data."""
new_data = self.new()
for k, v in self.items():
data = {}
if isinstance(v, np.ndarray):
v = torch.from_numpy(v)
data[k] = v
elif isinstance(v, BaseDataElement):
v = v.to_tensor()
data[k] = v
new_data.set_data(data)
return new_data

def to_dict(self) -> dict:
"""Convert BaseDataElement to dict."""
return {
k: v.to_dict() if isinstance(v, BaseDataElement) else v
for k, v in self.all_items()
}

def __repr__(self) -> str:
"""Represent the object."""

def _addindent(s_: str, num_spaces: int) -> str:
"""This func is modified from `pytorch` https://github.com/pytorch/
pytorch/blob/b17b2b1cc7b017c3daaeff8cc7ec0f514d42ec37/torch/nn/modu
les/module.py#L29.

Args:
s_ (str): The string to add spaces.
num_spaces (int): The num of space to add.

Returns:
str: The string after add indent.
"""
s = s_.split("\n")
# don't do anything for single-line stuff
if len(s) == 1:
return s_
first = s.pop(0)
s = [(num_spaces * " ") + line for line in s]
s = "\n".join(s) # type: ignore
s = first + "\n" + s # type: ignore
return s # type: ignore

def dump(obj: Any) -> str:
"""Represent the object.

Args:
obj (Any): The obj to represent.

Returns:
str: The represented str.
"""
_repr = ""
if isinstance(obj, dict):
for k, v in obj.items():
_repr += f"\n{k}: {_addindent(dump(v), 4)}"
elif isinstance(obj, BaseDataElement):
_repr += "\n\n META INFORMATION"
metainfo_items = dict(obj.metainfo_items())
_repr += _addindent(dump(metainfo_items), 4)
_repr += "\n\n DATA FIELDS"
items = dict(obj.items())
_repr += _addindent(dump(items), 4)
classname = obj.__class__.__name__
_repr = f"<{classname}({_repr}\n) at {hex(id(obj))}>"
else:
_repr += repr(obj)
return _repr

return dump(self)

+ 321
- 0
abl/structures/list_data.py View File

@@ -0,0 +1,321 @@
# Copyright (c) OpenMMLab. All rights reserved.
import itertools
from collections.abc import Sized
from typing import Any, List, Union

import numpy as np
import torch

from ..utils import flatten as flatten_list
from ..utils import to_hashable
from .base_data_element import BaseDataElement

BoolTypeTensor = Union[torch.BoolTensor, torch.cuda.BoolTensor]
LongTypeTensor = Union[torch.LongTensor, torch.cuda.LongTensor]

IndexType = Union[str, slice, int, list, LongTypeTensor, BoolTypeTensor, np.ndarray]


# Modified from
# https://github.com/open-mmlab/mmdetection/blob/master/mmdet/core/data_structures/instance_data.py # noqa
class ListData(BaseDataElement):
"""Data structure for instance-level annotations or predictions.

Subclass of :class:`BaseDataElement`. All value in `data_fields`
should have the same length. This design refer to
https://github.com/facebookresearch/detectron2/blob/master/detectron2/structures/instances.py # noqa E501
ListData also support extra functions: ``index``, ``slice`` and ``cat`` for data field. The type of value
in data field can be base data structure such as `torch.Tensor`, `numpy.ndarray`, `list`, `str`, `tuple`,
and can be customized data structure that has ``__len__``, ``__getitem__`` and ``cat`` attributes.

Examples:
>>> # custom data structure
>>> class TmpObject:
... def __init__(self, tmp) -> None:
... assert isinstance(tmp, list)
... self.tmp = tmp
... def __len__(self):
... return len(self.tmp)
... def __getitem__(self, item):
... if isinstance(item, int):
... if item >= len(self) or item < -len(self): # type:ignore
... raise IndexError(f'Index {item} out of range!')
... else:
... # keep the dimension
... item = slice(item, None, len(self))
... return TmpObject(self.tmp[item])
... @staticmethod
... def cat(tmp_objs):
... assert all(isinstance(results, TmpObject) for results in tmp_objs)
... if len(tmp_objs) == 1:
... return tmp_objs[0]
... tmp_list = [tmp_obj.tmp for tmp_obj in tmp_objs]
... tmp_list = list(itertools.chain(*tmp_list))
... new_data = TmpObject(tmp_list)
... return new_data
... def __repr__(self):
... return str(self.tmp)
>>> from mmengine.structures import ListData
>>> import numpy as np
>>> import torch
>>> img_meta = dict(img_shape=(800, 1196, 3), pad_shape=(800, 1216, 3))
>>> instance_data = ListData(metainfo=img_meta)
>>> 'img_shape' in instance_data
True
>>> instance_data.det_labels = torch.LongTensor([2, 3])
>>> instance_data["det_scores"] = torch.Tensor([0.8, 0.7])
>>> instance_data.bboxes = torch.rand((2, 4))
>>> instance_data.polygons = TmpObject([[1, 2, 3, 4], [5, 6, 7, 8]])
>>> len(instance_data)
2
>>> print(instance_data)
<ListData(
META INFORMATION
img_shape: (800, 1196, 3)
pad_shape: (800, 1216, 3)
DATA FIELDS
det_labels: tensor([2, 3])
det_scores: tensor([0.8000, 0.7000])
bboxes: tensor([[0.4997, 0.7707, 0.0595, 0.4188],
[0.8101, 0.3105, 0.5123, 0.6263]])
polygons: [[1, 2, 3, 4], [5, 6, 7, 8]]
) at 0x7fb492de6280>
>>> sorted_results = instance_data[instance_data.det_scores.sort().indices]
>>> sorted_results.det_scores
tensor([0.7000, 0.8000])
>>> print(instance_data[instance_data.det_scores > 0.75])
<ListData(
META INFORMATION
img_shape: (800, 1196, 3)
pad_shape: (800, 1216, 3)
DATA FIELDS
det_labels: tensor([2])
det_scores: tensor([0.8000])
bboxes: tensor([[0.4997, 0.7707, 0.0595, 0.4188]])
polygons: [[1, 2, 3, 4]]
) at 0x7f64ecf0ec40>
>>> print(instance_data[instance_data.det_scores > 1])
<ListData(
META INFORMATION
img_shape: (800, 1196, 3)
pad_shape: (800, 1216, 3)
DATA FIELDS
det_labels: tensor([], dtype=torch.int64)
det_scores: tensor([])
bboxes: tensor([], size=(0, 4))
polygons: []
) at 0x7f660a6a7f70>
>>> print(instance_data.cat([instance_data, instance_data]))
<ListData(
META INFORMATION
img_shape: (800, 1196, 3)
pad_shape: (800, 1216, 3)
DATA FIELDS
det_labels: tensor([2, 3, 2, 3])
det_scores: tensor([0.8000, 0.7000, 0.8000, 0.7000])
bboxes: tensor([[0.4997, 0.7707, 0.0595, 0.4188],
[0.8101, 0.3105, 0.5123, 0.6263],
[0.4997, 0.7707, 0.0595, 0.4188],
[0.8101, 0.3105, 0.5123, 0.6263]])
polygons: [[1, 2, 3, 4], [5, 6, 7, 8], [1, 2, 3, 4], [5, 6, 7, 8]]
) at 0x7f203542feb0>
"""

def __setattr__(self, name: str, value: list):
"""setattr is only used to set data.

The value must have the attribute of `__len__` and have the same length
of `ListData`.
"""
if name in ("_metainfo_fields", "_data_fields"):
if not hasattr(self, name):
super().__setattr__(name, value)
else:
raise AttributeError(
f"{name} has been used as a "
"private attribute, which is immutable."
)

else:
assert isinstance(value, list), "value must be of type `list`"

if len(self) > 0:
assert len(value) == len(self), (
"The length of "
f"values {len(value)} is "
"not consistent with "
"the length of this "
":obj:`ListData` "
f"{len(self)}"
)
super().__setattr__(name, value)

__setitem__ = __setattr__

def __getitem__(self, item: IndexType) -> "ListData":
"""
Args:
item (str, int, list, :obj:`slice`, :obj:`numpy.ndarray`,
:obj:`torch.LongTensor`, :obj:`torch.BoolTensor`):
Get the corresponding values according to item.

Returns:
:obj:`ListData`: Corresponding values.
"""
assert isinstance(item, IndexType.__args__)
if isinstance(item, list):
item = np.array(item)
if isinstance(item, np.ndarray):
# The default int type of numpy is platform dependent, int32 for
# windows and int64 for linux. `torch.Tensor` requires the index
# should be int64, therefore we simply convert it to int64 here.
# More details in https://github.com/numpy/numpy/issues/9464
item = item.astype(np.int64) if item.dtype == np.int32 else item
item = torch.from_numpy(item)

if isinstance(item, str):
return getattr(self, item)

if isinstance(item, int):
if item >= len(self) or item < -len(self): # type:ignore
raise IndexError(f"Index {item} out of range!")
else:
# keep the dimension
item = slice(item, None, len(self))

new_data = self.__class__(metainfo=self.metainfo)
if isinstance(item, torch.Tensor):
assert item.dim() == 1, (
"Only support to get the" " values along the first dimension."
)
if isinstance(item, BoolTypeTensor.__args__):
assert len(item) == len(self), (
"The shape of the "
"input(BoolTensor) "
f"{len(item)} "
"does not match the shape "
"of the indexed tensor "
"in results_field "
f"{len(self)} at "
"first dimension."
)

for k, v in self.items():
if isinstance(v, torch.Tensor):
new_data[k] = v[item]
elif isinstance(v, np.ndarray):
new_data[k] = v[item.cpu().numpy()]
elif isinstance(v, (str, list, tuple)) or (
hasattr(v, "__getitem__") and hasattr(v, "cat")
):
# convert to indexes from BoolTensor
if isinstance(item, BoolTypeTensor.__args__):
indexes = torch.nonzero(item).view(-1).cpu().numpy().tolist()
else:
indexes = item.cpu().numpy().tolist()
slice_list = []
if indexes:
for index in indexes:
slice_list.append(slice(index, None, len(v)))
else:
slice_list.append(slice(None, 0, None))
r_list = [v[s] for s in slice_list]
if isinstance(v, (str, list, tuple)):
new_value = r_list[0]
for r in r_list[1:]:
new_value = new_value + r
else:
new_value = v.cat(r_list)
new_data[k] = new_value
else:
raise ValueError(
f"The type of `{k}` is `{type(v)}`, which has no "
"attribute of `cat`, so it does not "
"support slice with `bool`"
)

else:
# item is a slice
for k, v in self.items():
new_data[k] = v[item]
return new_data # type:ignore

@staticmethod
def cat(instances_list: List["ListData"]) -> "ListData":
"""Concat the instances of all :obj:`ListData` in the list.

Note: To ensure that cat returns as expected, make sure that
all elements in the list must have exactly the same keys.

Args:
instances_list (list[:obj:`ListData`]): A list
of :obj:`ListData`.

Returns:
:obj:`ListData`
"""
assert all(isinstance(results, ListData) for results in instances_list)
assert len(instances_list) > 0
if len(instances_list) == 1:
return instances_list[0]

# metainfo and data_fields must be exactly the
# same for each element to avoid exceptions.
field_keys_list = [instances.all_keys() for instances in instances_list]
assert len({len(field_keys) for field_keys in field_keys_list}) == 1 and len(
set(itertools.chain(*field_keys_list))
) == len(field_keys_list[0]), (
"There are different keys in "
"`instances_list`, which may "
"cause the cat operation "
"to fail. Please make sure all "
"elements in `instances_list` "
"have the exact same key."
)

new_data = instances_list[0].__class__(metainfo=instances_list[0].metainfo)
for k in instances_list[0].keys():
values = [results[k] for results in instances_list]
v0 = values[0]
if isinstance(v0, torch.Tensor):
new_values = torch.cat(values, dim=0)
elif isinstance(v0, np.ndarray):
new_values = np.concatenate(values, axis=0)
elif isinstance(v0, (str, list, tuple)):
new_values = v0[:]
for v in values[1:]:
new_values += v
elif hasattr(v0, "cat"):
new_values = v0.cat(values)
else:
raise ValueError(
f"The type of `{k}` is `{type(v0)}` which has no "
"attribute of `cat`"
)
new_data[k] = new_values
return new_data # type:ignore

def flatten(self, item: IndexType) -> List:
"""Flatten self[item].

Returns:
list: Flattened data fields.
"""
return flatten_list(self[item])
def elements_num(self, item: IndexType) -> int:
"""int: The number of elements in self[item]."""
return len(self.flatten(item))
def to_tuple(self, item: IndexType) -> tuple:
"""tuple: The data fields in self[item] converted to tuple."""
return to_hashable(self[item])
def __len__(self) -> int:
"""int: The length of ListData."""
if len(self._data_fields) > 0:
one_element = next(iter(self._data_fields))
return len(getattr(self, one_element))
# return len(self.values()[0])
else:
return 0

+ 2
- 1
abl/utils/__init__.py View File

@@ -1,2 +1,3 @@
from .cache import Cache
from .logger import ABLLogger, print_log
from .utils import *
from .utils import *

+ 112
- 0
abl/utils/cache.py View File

@@ -0,0 +1,112 @@
import pickle
from os import PathLike
from pathlib import Path
from typing import Callable, Generic, Hashable, TypeVar, Union

from .logger import print_log

K = TypeVar("K")
T = TypeVar("T")
PREV, NEXT, KEY, RESULT = 0, 1, 2, 3 # names for the link fields


class Cache(Generic[K, T]):
def __init__(
self,
func: Callable[[K], T],
cache: bool,
cache_file: Union[None, str, PathLike],
key_func: Callable[[K], Hashable] = lambda x: x,
max_size: int = 4096,
):
"""Create cache

:param func: Function this cache evaluates
:param cache: If true, do in memory caching.
:param cache_root: If not None, cache to files at the provided path.
:param key_func: Convert the key into a hashable object if needed
"""
self.func = func
self.key_func = key_func
self.cache = cache
if cache is True or cache_file is not None:
print_log("Caching is activated", logger="current")
self._init_cache(cache_file, max_size)
self.first = self.get_from_dict
else:
self.first = self.func

def __getitem__(self, item: K, *args) -> T:
return self.first(item, *args)

def invalidate(self):
"""Invalidate entire cache."""
self.cache_dict.clear()
if self.cache_file:
for p in self.cache_root.iterdir():
p.unlink()

def _init_cache(self, cache_file, max_size):
self.cache = True
self.cache_dict = dict()

self.hits, self.misses, self.maxsize = 0, 0, max_size
self.full = False
self.root = [] # root of the circular doubly linked list
self.root[:] = [self.root, self.root, None, None]

if cache_file is not None:
with open(cache_file, "rb") as f:
cache_dict_from_file = pickle.load(f)
self.maxsize += len(cache_dict_from_file)
print_log(
f"Max size of the cache has been enlarged to {self.maxsize}.", logger="current"
)
for cache_key, result in cache_dict_from_file.items():
last = self.root[PREV]
link = [last, self.root, cache_key, result]
last[NEXT] = self.root[PREV] = self.cache_dict[cache_key] = link

def get(self, item: K, *args) -> T:
return self.first(item, *args)

def get_from_dict(self, item: K, *args) -> T:
"""Implements dict based cache."""
cache_key = (self.key_func(item), *args)
link = self.cache_dict.get(cache_key)
if link is not None:
# Move the link to the front of the circular queue
link_prev, link_next, _key, result = link
link_prev[NEXT] = link_next
link_next[PREV] = link_prev
last = self.root[PREV]
last[NEXT] = self.root[PREV] = link
link[PREV] = last
link[NEXT] = self.root
self.hits += 1
return result
self.misses += 1

result = self.func(item, *args)

if self.full:
# Use the old root to store the new key and result.
oldroot = self.root
oldroot[KEY] = cache_key
oldroot[RESULT] = result
# Empty the oldest link and make it the new root.
self.root = oldroot[NEXT]
oldkey = self.root[KEY]
oldresult = self.root[RESULT]
self.root[KEY] = self.root[RESULT] = None
# Now update the cache dictionary.
del self.cache_dict[oldkey]
self.cache_dict[cache_key] = oldroot
else:
# Put result in a new link at the front of the queue.
last = self.root[PREV]
link = [last, self.root, cache_key, result]
last[NEXT] = self.root[PREV] = self.cache_dict[cache_key] = link
if isinstance(self.maxsize, int):
self.full = len(self.cache_dict) >= self.maxsize
return result

+ 84
- 6
abl/utils/utils.py View File

@@ -1,6 +1,7 @@
import numpy as np
from itertools import chain

import numpy as np


def flatten(nested_list):
"""
@@ -15,6 +16,11 @@ def flatten(nested_list):
-------
list
A flattened version of the input list.

Raises
------
TypeError
If the input object is not a list.
"""
if not isinstance(nested_list, list):
raise TypeError("Input must be of type list.")
@@ -41,6 +47,9 @@ def reform_idx(flattened_list, structured_list):
list
A reformed list that mimics the structure of structured_list.
"""
# if not isinstance(flattened_list, list):
# raise TypeError("Input must be of type list.")

if not isinstance(structured_list[0], (list, tuple)):
return flattened_list

@@ -80,7 +89,7 @@ def hamming_dist(pred_pseudo_label, candidates):
return np.sum(pred_pseudo_label != candidates, axis=1)


def confidence_dist(pred_prob, candidates_idx):
def confidence_dist(pred_prob, candidates):
"""
Compute the confidence distance between prediction probabilities and candidates.

@@ -89,7 +98,7 @@ def confidence_dist(pred_prob, candidates_idx):
pred_prob : list of numpy.ndarray
Prediction probability distributions, each element is an ndarray
representing the probability distribution of a particular prediction.
candidates_idx : list of list of int
candidates : list of list of int
Index of candidate labels, each element is a list of indexes being considered
as a candidate correction.

@@ -99,8 +108,8 @@ def confidence_dist(pred_prob, candidates_idx):
Confidence distances computed for each candidate.
"""
pred_prob = np.clip(pred_prob, 1e-9, 1)
_, cols = np.indices((len(candidates_idx), len(candidates_idx[0])))
return 1 - np.prod(pred_prob[cols, candidates_idx], axis=1)
_, cols = np.indices((len(candidates), len(candidates[0])))
return 1 - np.prod(pred_prob[cols, candidates], axis=1)


def block_sample(X, Z, Y, sample_num, seg_idx):
@@ -135,6 +144,34 @@ def block_sample(X, Z, Y, sample_num, seg_idx):
return (data[start_idx:end_idx] for data in (X, Z, Y))


def check_equal(a, b, max_err=0):
"""
Check whether two numbers a and b are equal within a maximum allowable error.

Parameters
----------
a, b : int or float
The numbers to compare.
max_err : int or float, optional
The maximum allowable absolute difference between a and b for them to be considered equal.
Default is 0, meaning the numbers must be exactly equal.

Returns
-------
bool
True if a and b are equal within the allowable error, False otherwise.

Raises
------
TypeError
If a or b are not of type int or float.
"""
if not (isinstance(a, (int, float)) and isinstance(b, (int, float))):
raise TypeError("Input values must be int or float.")

return abs(a - b) <= max_err


def to_hashable(x):
"""
Convert a nested list to a nested tuple so it is hashable.
@@ -154,7 +191,8 @@ def to_hashable(x):
return tuple(to_hashable(item) for item in x)
return x

def restore_from_hashable(x):

def hashable_to_list(x):
"""
Convert a nested tuple back to a nested list.

@@ -174,6 +212,45 @@ def restore_from_hashable(x):
return x


def calculate_revision_num(parameter, total_length):
"""
Convert a float parameter to an integer, based on a total length.

Parameters
----------
parameter : int or float
The parameter to convert. If float, it should be between 0 and 1.
If int, it should be non-negative. If -1, it will be replaced with total_length.
total_length : int
The total length to calculate the parameter from if it's a fraction.

Returns
-------
int
The calculated parameter.

Raises
------
TypeError
If parameter is not an int or a float.
ValueError
If parameter is a float not in [0, 1] or an int below 0.
"""
if not isinstance(parameter, (int, float)):
raise TypeError("Parameter must be of type int or float.")

if parameter == -1:
return total_length
elif isinstance(parameter, float):
if not (0 <= parameter <= 1):
raise ValueError("If parameter is a float, it must be between 0 and 1.")
return round(total_length * parameter)
else:
if parameter < 0:
raise ValueError("If parameter is an int, it must be non-negative.")
return parameter


if __name__ == "__main__":
A = np.array(
[
@@ -227,4 +304,5 @@ if __name__ == "__main__":
)
B = [[0, 9, 3], [0, 11, 4]]

print(ori_confidence_dist(A, B))
print(confidence_dist(A, B))

Loading…
Cancel
Save