mirror of
https://github.com/hanzoai/anomalib.git
synced 2026-10-30 18:48:11 +00:00
zeekay
dda5b2a367
The previous commit was meant to remove .github/workflows and removed every tracked path instead — 892 files, the whole repo. This puts back the 874 that were never in scope; the eighteen workflow files stay gone, for the reasons that commit gives. Content is byte-identical to 9d2aec9 outside .github/workflows.
46 lines
1.1 KiB
Python
46 lines
1.1 KiB
Python
"""Tests for the collating DatasetItems into Batches."""
|
|
|
|
# Copyright (C) 2024 Intel Corporation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import torch
|
|
from torchvision.tv_tensors import Image, Mask
|
|
|
|
from anomalib.data.dataclasses.generic import BatchIterateMixin
|
|
|
|
|
|
@dataclass
|
|
class DummyDatasetItem:
|
|
"""Dummy dataset item with image and mask."""
|
|
|
|
image: Image
|
|
mask: Mask
|
|
|
|
|
|
@dataclass
|
|
class DummyBatch(BatchIterateMixin[DummyDatasetItem]):
|
|
"""Dummy batch with image and mask."""
|
|
|
|
item_class = DummyDatasetItem
|
|
image: Image
|
|
mask: Mask
|
|
|
|
|
|
def test_collate_heterogeneous_shapes() -> None:
|
|
"""Test collating items with different shapes."""
|
|
items = [
|
|
DummyDatasetItem(
|
|
image=Image(torch.rand((3, 256, 256))),
|
|
mask=Mask(torch.ones((256, 256))),
|
|
),
|
|
DummyDatasetItem(
|
|
image=Image(torch.rand((3, 224, 224))),
|
|
mask=Mask(torch.ones((224, 224))),
|
|
),
|
|
]
|
|
batch = DummyBatch.collate(items)
|
|
# the collated batch should have the shape of the largest item
|
|
assert batch.image.shape == (2, 3, 256, 256)
|