From 4d1f2cef39688871d2caafce42a09316381a27ae Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Mon, 5 Jul 2021 23:05:25 +0200 Subject: Refactor with attr, working on cnn+transformer network --- text_recognizer/data/base_dataset.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) (limited to 'text_recognizer/data/base_dataset.py') diff --git a/text_recognizer/data/base_dataset.py b/text_recognizer/data/base_dataset.py index 8d644d4..4318dfb 100644 --- a/text_recognizer/data/base_dataset.py +++ b/text_recognizer/data/base_dataset.py @@ -1,11 +1,13 @@ """Base PyTorch Dataset class.""" from typing import Any, Callable, Dict, Sequence, Tuple, Union +import attr import torch from torch import Tensor from torch.utils.data import Dataset +@attr.s class BaseDataset(Dataset): """ Base Dataset class that processes data and targets through optional transfroms. @@ -18,19 +20,17 @@ class BaseDataset(Dataset): target transforms. """ - def __init__( - self, - data: Union[Sequence, Tensor], - targets: Union[Sequence, Tensor], - transform: Callable = None, - target_transform: Callable = None, - ) -> None: - if len(data) != len(targets): + data: Union[Sequence, Tensor] = attr.ib() + targets: Union[Sequence, Tensor] = attr.ib() + transform: Callable = attr.ib() + target_transform: Callable = attr.ib() + + def __attrs_pre_init__(self) -> None: + super().__init__() + + def __attrs_post_init__(self) -> None: + if len(self.data) != len(self.targets): raise ValueError("Data and targets must be of equal length.") - self.data = data - self.targets = targets - self.transform = transform - self.target_transform = target_transform def __len__(self) -> int: """Return the length of the dataset.""" -- cgit v1.2.3-70-g09d2