tensorneko 0.3.4__tar.gz → 0.3.5__tar.gz

{tensorneko-0.3.4/src/tensorneko.egg-info → tensorneko-0.3.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tensorneko
-Version: 0.3.4
+Version: 0.3.5
 Summary: Tensor Neural Engine Kompanion. An util library based on PyTorch and PyTorch Lightning.
 Home-page: https://github.com/ControlNet/tensorneko
 Author: ControlNet
@@ -38,7 +38,7 @@ License-File: LICENSE
 <div align="center">    
     <a href="https://www.python.org/"><img src="https://img.shields.io/pypi/pyversions/tensorneko?style=flat-square"></a>
     <a href="https://pytorch.org/"><img src="https://img.shields.io/badge/PyTorch-%3E%3D1.9.0-EE4C2C?style=flat-square&logo=pytorch"></a>
-    <a href="https://www.pytorchlightning.ai/"><img src="https://img.shields.io/badge/Lightning-2.0.*-792EE5?style=flat-square&logo=lightning"></a>
+    <a href="https://www.pytorchlightning.ai/"><img src="https://img.shields.io/badge/Lightning-2.0.*/2.1.*-792EE5?style=flat-square&logo=lightning"></a>
 </div>
 
 <div align="center">

{tensorneko-0.3.4 → tensorneko-0.3.5}/README.md

@@ -12,7 +12,7 @@
 <div align="center">    
     <a href="https://www.python.org/"><img src="https://img.shields.io/pypi/pyversions/tensorneko?style=flat-square"></a>
     <a href="https://pytorch.org/"><img src="https://img.shields.io/badge/PyTorch-%3E%3D1.9.0-EE4C2C?style=flat-square&logo=pytorch"></a>
-    <a href="https://www.pytorchlightning.ai/"><img src="https://img.shields.io/badge/Lightning-2.0.*-792EE5?style=flat-square&logo=lightning"></a>
+    <a href="https://www.pytorchlightning.ai/"><img src="https://img.shields.io/badge/Lightning-2.0.*/2.1.*-792EE5?style=flat-square&logo=lightning"></a>
 </div>
 
 <div align="center">

{tensorneko-0.3.4 → tensorneko-0.3.5}/src/tensorneko/arch/__init__.py

@@ -2,10 +2,12 @@ from .gan import GAN
 from .vqvae import VQVAE
 from .wgan import WGAN
 from .auto_encoder import AutoEncoder
+from .binary_classifier import BinaryClassifier
 
 __all__ = [
     "GAN",
     "VQVAE",
     "WGAN",
-    "AutoEncoder"
+    "AutoEncoder",
+    "BinaryClassifier",
 ]

tensorneko-0.3.5/src/tensorneko/arch/binary_classifier.py

@@ -0,0 +1,54 @@
+from abc import ABC
+from typing import Optional, Union, Sequence, Dict
+
+from torch import Tensor
+from torch.nn import BCEWithLogitsLoss
+from torch.optim import Adam
+from torchmetrics import Accuracy, F1Score, AUROC
+
+from ..neko_model import NekoModel
+
+
+class BinaryClassifier(NekoModel, ABC):
+
+    def __init__(self, model=None, learning_rate: float = 1e-4, distributed: bool = False):
+        super().__init__()
+        self.save_hyperparameters()
+        self.model = model
+        self.learning_rate = learning_rate
+        self.distributed = distributed
+        self.loss_fn = BCEWithLogitsLoss()
+        self.acc_fn = Accuracy(task="binary")
+        self.f1_fn = F1Score(task="binary")
+        self.auc_fn = AUROC(task="binary")
+
+    @classmethod
+    def from_module(cls, model, learning_rate: float = 1e-4, distributed=False):
+        return cls(model, learning_rate, distributed)
+
+    def forward(self, x):
+        return self.model(x)
+
+    def step(self, batch: Optional[Union[Tensor, Sequence[Tensor]]]) -> Dict[str, Tensor]:
+        x, y = batch
+        y_hat = self(x).squeeze(1)
+        loss = self.loss_fn(y_hat, y)
+        prob = y_hat.sigmoid()
+        acc = self.acc_fn(prob, y)
+        f1 = self.f1_fn(prob, y)
+        auc = self.auc_fn(prob, y)
+        return {"loss": loss, "acc": acc, "f1": f1, "auc": auc}
+
+    def training_step(self, batch: Optional[Union[Tensor, Sequence[Tensor]]] = None, batch_idx: Optional[int] = None,
+        optimizer_idx: Optional[int] = None, hiddens: Optional[Tensor] = None
+    ) -> Dict[str, Tensor]:
+        return self.step(batch)
+
+    def validation_step(self, batch: Optional[Union[Tensor, Sequence[Tensor]]] = None, batch_idx: Optional[int] = None,
+        dataloader_idx: Optional[int] = None
+    ) -> Dict[str, Tensor]:
+        return self.step(batch)
+
+    def configure_optimizers(self):
+        optimizer = Adam(self.parameters(), lr=self.learning_rate)
+        return [optimizer]

tensorneko-0.3.5/src/tensorneko/dataset/__init__.py

@@ -0,0 +1,11 @@
+from .round_robin_dataset import RoundRobinDataset
+from .nested_dataset import NestedDataset
+from .list_dataset import ListDataset
+from . import sampler
+
+__all__ = [
+    "RoundRobinDataset",
+    "NestedDataset",
+    "ListDataset",
+    "sampler"
+]

tensorneko-0.3.5/src/tensorneko/dataset/list_dataset.py

@@ -0,0 +1,19 @@
+from typing import List
+
+from torch.utils.data.dataset import Dataset, T_co
+
+
+class ListDataset(Dataset[T_co]):
+    """
+    A dataset wrapping a list of data.
+    """
+
+    def __init__(self, data: List[T_co]):
+        super().__init__()
+        self.data = data
+
+    def __getitem__(self, index: int) -> T_co:
+        return self.data[index]
+
+    def __len__(self):
+        return len(self.data)

{tensorneko-0.3.4 → tensorneko-0.3.5}/src/tensorneko/dataset/round_robin_dataset.py

@@ -4,7 +4,7 @@ from typing import List, Optional
 from torch.utils.data import Dataset
 from torch.utils.data.dataset import T_co
 
-from tensorneko.util import circular_pad
+from ..util import circular_pad
 
 
 class RoundRobinDataset(Dataset[T_co]):

tensorneko-0.3.5/src/tensorneko/dataset/sampler/__init__.py

@@ -0,0 +1,5 @@
+from .sequential_iter_sampler import SequentialIterSampler
+
+__all__ = [
+    "SequentialIterSampler"
+]

tensorneko-0.3.5/src/tensorneko/dataset/sampler/sequential_iter_sampler.py

@@ -0,0 +1,26 @@
+from typing import Sized
+
+from torch.utils.data.sampler import Sampler, T_co
+
+
+class SequentialIterSampler(Sampler[T_co]):
+    """
+    Use to split the large scale data into small subsets for each epochs
+    For example, if the dataset size is 1M, and the num_samples = 1000, then each epoch will only use 1000 samples, and
+    the next epoch will use the next 1000 samples.
+    """
+
+    def __init__(self, data_source: Sized, num_samples: int):
+        super().__init__(data_source)
+        self.data_source = data_source
+        self.num_samples = num_samples
+        self.total_size = len(data_source)
+        self.current_position = 0
+
+    def __iter__(self):
+        yield from map(lambda x: x % self.total_size,
+            range(self.current_position, self.current_position + self.num_samples))
+        self.current_position = (self.current_position + self.num_samples) % self.total_size
+
+    def __len__(self):
+        return self.num_samples

{tensorneko-0.3.4 → tensorneko-0.3.5}/src/tensorneko/evaluation/__init__.py

@@ -1,6 +1,7 @@
 from .iou import iou_1d, iou_2d
 from .psnr import psnr_video, psnr_image
 from .ssim import ssim_video, ssim_image
+from .secs import secs
 from .fid import FID
 
 __all__ = [
@@ -10,5 +11,6 @@ __all__ = [
     "psnr_image",
     "ssim_video",
     "ssim_image",
+    "secs",
     "FID",
 ]

{tensorneko-0.3.4 → tensorneko-0.3.5}/src/tensorneko/evaluation/fid.py

@@ -22,6 +22,34 @@ except ImportError:
 
 
 class FID:
+    """
+    Calculate Fréchet inception distance based on torchmetrics. Require library "torch-fidelity".
+
+    Args:
+        device (``str`` | :class:`~torch.device`, optional): Device to run the metric. Default: ``"cpu"``.
+
+    Example::
+
+        from tensorneko.evaluation import FID
+        fid = FID("cuda")
+        
+        # add predicted and real images
+        fid.add_pred_image("path/to/pred/image1.png")
+        fid.add_pred_image("path/to/pred/image2.png")
+        fid.add_true_image("path/to/true/image1.png")
+        fid.add_true_image("path/to/true/image2.png")
+
+        # add predicted and real videos
+        fid.add_pred_video("path/to/pred/video1.mp4")
+        fid.add_pred_video("path/to/pred/video2.mp4")
+        fid.add_true_video("path/to/true/video1.mp4")
+        fid.add_true_video("path/to/true/video2.mp4")
+
+        # compute FID
+        fid_score = fid.compute(batch_size=128, num_workers=8, progress_bar=True)
+        print(fid_score)
+
+    """
 
     def __init__(self, device: Union[str, Device] = "cpu"):
         self.device = torch.device(device)
@@ -56,14 +84,14 @@ class FID:
     def cuda(self) -> FID:
         return self.to("cuda")
 
-    def compute(self, batch_size=128, num_workers=8, progress_bar: bool = True) -> float:
+    def compute(self, batch_size=128, num_workers=0, progress_bar: bool = False) -> float:
         pred = torch.utils.data.DataLoader(self.pred_data, batch_size=batch_size, num_workers=num_workers)
         true = torch.utils.data.DataLoader(self.true_data, batch_size=batch_size, num_workers=num_workers)
 
         if progress_bar:
             tqdm = import_tqdm_auto().tqdm
-            pred = tqdm(pred, desc="Forward predicted features")
-            true = tqdm(true, desc="Forward ground truth features")
+            pred = tqdm(total=len(pred), desc="Forward predicted features")
+            true = tqdm(total=len(true), desc="Forward ground truth features")
 
         for batch in pred:
             self.fid.update(batch.to(self.device), real=False)
@@ -72,6 +100,11 @@ class FID:
 
         return self.fid.compute().item()
 
+    def reset(self):
+        self.pred_data = _FIDDataset()
+        self.true_data = _FIDDataset()
+        self.fid.reset()
+
 
 @dataclass
 class _FIDEntry:
@@ -104,6 +137,7 @@ class _FIDDataset(IterableDataset):
             raise RuntimeError("Cannot open video file.")
         n_frames = int(cap.get(self.cv2.CAP_PROP_FRAME_COUNT))
         self.length += n_frames
+        cap.release()
 
     @staticmethod
     def _preprocess_image(image: Tensor) -> Tensor:
@@ -130,6 +164,8 @@ class _FIDDataset(IterableDataset):
             frame = self._preprocess_image(frame)
             yield frame
 
+        cap.release()
+
     def __iter__(self):
         for entry in self.content:
             if entry.type == "image":

{tensorneko-0.3.4 → tensorneko-0.3.5}/src/tensorneko/evaluation/iou.py

@@ -70,7 +70,7 @@ def iou_2d(proposal: Union[Tensor, ndarray], target: Union[Tensor, ndarray]) ->
 
     inner_x1 = torch.maximum(proposal_x1, target_x1)
     inner_y1 = torch.maximum(proposal_y1, target_y1)
-    inner_x2 = torch.minimum(proposal_x2, target_y2)
+    inner_x2 = torch.minimum(proposal_x2, target_x2)
     inner_y2 = torch.minimum(proposal_y2, target_y2)
 
     area_proposal = (proposal_x2 - proposal_x1) * (proposal_y2 - proposal_y1)

tensorneko-0.3.5/src/tensorneko/evaluation/secs.py

@@ -0,0 +1,58 @@
+from numpy import ndarray
+from torch import Tensor
+
+from tensorneko_util.util import dispatch, Eval
+
+from tensorneko_util.io import read
+
+
+@Eval.later
+def _secs_encoder():
+    from resemblyzer import VoiceEncoder
+    return VoiceEncoder()
+
+
+@dispatch
+def secs(pred: str, real: str) -> float:
+    from resemblyzer import VoiceEncoder, preprocess_wav
+    pred_audio = preprocess_wav(read.audio(pred).audio[0].numpy())
+    real_audio = preprocess_wav(read.audio(real).audio[0].numpy())
+    return _secs_compute(pred_audio, real_audio)
+
+
+@dispatch
+def secs(pred: Tensor, real: Tensor) -> float:
+    return secs(pred.numpy(), real.numpy())
+
+
+@dispatch
+def secs(pred: ndarray, real: ndarray) -> float:
+    from resemblyzer import VoiceEncoder, preprocess_wav
+    if len(pred.shape) == 2:
+        if pred.shape[0] == 1:
+            pred = pred.squeeze(0)
+        elif pred.shape[1] == 1:
+            pred = pred.squeeze(1)
+        else:
+            raise ValueError("The input audio must be mono.")
+
+    if len(real.shape) == 2:
+        if real.shape[0] == 1:
+            real = real.squeeze(0)
+        elif real.shape[1] == 1:
+            real = real.squeeze(1)
+        else:
+            raise ValueError("The input audio must be mono.")
+
+    pred_audio = preprocess_wav(pred)
+    real_audio = preprocess_wav(real)
+
+    return _secs_compute(pred_audio, real_audio)
+
+
+def _secs_compute(pred_audio: ndarray, real_audio: ndarray) -> float:
+    encoder = _secs_encoder.value
+    real_embed = encoder.embed_utterance(real_audio)
+    pred_embed = encoder.embed_utterance(pred_audio)
+
+    return float((real_embed * pred_embed).sum())

tensorneko-0.3.5/src/tensorneko/version.txt

@@ -0,0 +1 @@
+0.3.5

{tensorneko-0.3.4 → tensorneko-0.3.5/src/tensorneko.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tensorneko
-Version: 0.3.4
+Version: 0.3.5
 Summary: Tensor Neural Engine Kompanion. An util library based on PyTorch and PyTorch Lightning.
 Home-page: https://github.com/ControlNet/tensorneko
 Author: ControlNet
@@ -38,7 +38,7 @@ License-File: LICENSE
 <div align="center">    
     <a href="https://www.python.org/"><img src="https://img.shields.io/pypi/pyversions/tensorneko?style=flat-square"></a>
     <a href="https://pytorch.org/"><img src="https://img.shields.io/badge/PyTorch-%3E%3D1.9.0-EE4C2C?style=flat-square&logo=pytorch"></a>
-    <a href="https://www.pytorchlightning.ai/"><img src="https://img.shields.io/badge/Lightning-2.0.*-792EE5?style=flat-square&logo=lightning"></a>
+    <a href="https://www.pytorchlightning.ai/"><img src="https://img.shields.io/badge/Lightning-2.0.*/2.1.*-792EE5?style=flat-square&logo=lightning"></a>
 </div>
 
 <div align="center">

{tensorneko-0.3.4 → tensorneko-0.3.5}/src/tensorneko.egg-info/SOURCES.txt

@@ -13,6 +13,7 @@ src/tensorneko.egg-info/requires.txt
 src/tensorneko.egg-info/top_level.txt
 src/tensorneko/arch/__init__.py
 src/tensorneko/arch/auto_encoder.py
+src/tensorneko/arch/binary_classifier.py
 src/tensorneko/arch/gan.py
 src/tensorneko/arch/vqvae.py
 src/tensorneko/arch/wgan.py
@@ -27,14 +28,18 @@ src/tensorneko/callback/lr_logger.py
 src/tensorneko/callback/nil_callback.py
 src/tensorneko/callback/system_stats_logger.py
 src/tensorneko/dataset/__init__.py
+src/tensorneko/dataset/list_dataset.py
 src/tensorneko/dataset/nested_dataset.py
 src/tensorneko/dataset/round_robin_dataset.py
+src/tensorneko/dataset/sampler/__init__.py
+src/tensorneko/dataset/sampler/sequential_iter_sampler.py
 src/tensorneko/debug/__init__.py
 src/tensorneko/evaluation/__init__.py
 src/tensorneko/evaluation/enum.py
 src/tensorneko/evaluation/fid.py
 src/tensorneko/evaluation/iou.py
 src/tensorneko/evaluation/psnr.py
+src/tensorneko/evaluation/secs.py
 src/tensorneko/evaluation/ssim.py
 src/tensorneko/io/__init__.py
 src/tensorneko/io/reader.py

{tensorneko-0.3.4 → tensorneko-0.3.5}/src/tensorneko.egg-info/requires.txt

@@ -3,12 +3,12 @@ torchaudio>=0.9.0
 torchvision>=0.10.0
 torchmetrics>=0.7.3
 tensorboard>=2.0.0
-lightning==2.0.*
+lightning<2.2,>=2.0
 pillow>=8.1
 av>=8.0.3
 numpy>=1.20.1
 einops>=0.3.0
-tensorneko_util==0.3.4
+tensorneko_util==0.3.5
 
 [:platform_system == "Windows"]
 pysoundfile>=0.9.0

tensorneko-0.3.4/src/tensorneko/dataset/__init__.py

@@ -1,5 +0,0 @@
-from .round_robin_dataset import RoundRobinDataset
-
-__all__ = [
-    "RoundRobinDataset",
-]

tensorneko-0.3.4/src/tensorneko/version.txt

@@ -1 +0,0 @@
-0.3.4