diff --git a/lib/python3.12/site-packages/torchvision/__init__.py b/lib/python3.12/site-packages/torchvision/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5d06156c25f1dfd34e9f01529e5a6b4bbeda7b42 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/__init__.py @@ -0,0 +1,105 @@ +import os +import warnings +from modulefinder import Module + +import torch + +# Don't re-order these, we need to load the _C extension (done when importing +# .extensions) before entering _meta_registrations. +from .extension import _HAS_OPS # usort:skip +from torchvision import _meta_registrations, datasets, io, models, ops, transforms, utils # usort:skip + +try: + from .version import __version__ # noqa: F401 +except ImportError: + pass + + +# Check if torchvision is being imported within the root folder +if not _HAS_OPS and os.path.dirname(os.path.realpath(__file__)) == os.path.join( + os.path.realpath(os.getcwd()), "torchvision" +): + message = ( + "You are importing torchvision within its own root folder ({}). " + "This is not expected to work and may give errors. Please exit the " + "torchvision project source and relaunch your python interpreter." + ) + warnings.warn(message.format(os.getcwd())) + +_image_backend = "PIL" + +_video_backend = "pyav" + + +def set_image_backend(backend): + """ + Specifies the package used to load images. + + Args: + backend (string): Name of the image backend. one of {'PIL', 'accimage'}. + The :mod:`accimage` package uses the Intel IPP library. It is + generally faster than PIL, but does not support as many operations. + """ + global _image_backend + if backend not in ["PIL", "accimage"]: + raise ValueError(f"Invalid backend '{backend}'. Options are 'PIL' and 'accimage'") + _image_backend = backend + + +def get_image_backend(): + """ + Gets the name of the package used to load images + """ + return _image_backend + + +def set_video_backend(backend): + """ + Specifies the package used to decode videos. + + Args: + backend (string): Name of the video backend. one of {'pyav', 'video_reader'}. + The :mod:`pyav` package uses the 3rd party PyAv library. It is a Pythonic + binding for the FFmpeg libraries. + The :mod:`video_reader` package includes a native C++ implementation on + top of FFMPEG libraries, and a python API of TorchScript custom operator. + It generally decodes faster than :mod:`pyav`, but is perhaps less robust. + + .. note:: + Building with FFMPEG is disabled by default in the latest `main`. If you want to use the 'video_reader' + backend, please compile torchvision from source. + """ + global _video_backend + if backend not in ["pyav", "video_reader", "cuda"]: + raise ValueError("Invalid video backend '%s'. Options are 'pyav', 'video_reader' and 'cuda'" % backend) + if backend == "video_reader" and not io._HAS_CPU_VIDEO_DECODER: + # TODO: better messages + message = "video_reader video backend is not available. Please compile torchvision from source and try again" + raise RuntimeError(message) + elif backend == "cuda" and not io._HAS_GPU_VIDEO_DECODER: + # TODO: better messages + message = "cuda video backend is not available." + raise RuntimeError(message) + else: + _video_backend = backend + + +def get_video_backend(): + """ + Returns the currently active video backend used to decode videos. + + Returns: + str: Name of the video backend. one of {'pyav', 'video_reader'}. + """ + + return _video_backend + + +def _is_tracing(): + return torch._C._get_tracing_state() + + +def disable_beta_transforms_warning(): + # Noop, only exists to avoid breaking existing code. + # See https://github.com/pytorch/vision/issues/7896 + pass diff --git a/lib/python3.12/site-packages/torchvision/_internally_replaced_utils.py b/lib/python3.12/site-packages/torchvision/_internally_replaced_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e0fa72489f1ac8fba771fc7bc20fc80424a71d85 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/_internally_replaced_utils.py @@ -0,0 +1,51 @@ +import importlib.machinery +import os + +from torch.hub import _get_torch_home + + +_HOME = os.path.join(_get_torch_home(), "datasets", "vision") +_USE_SHARDED_DATASETS = False +IN_FBCODE = False + + +def _download_file_from_remote_location(fpath: str, url: str) -> None: + pass + + +def _is_remote_location_available() -> bool: + return False + + +try: + from torch.hub import load_state_dict_from_url # noqa: 401 +except ImportError: + from torch.utils.model_zoo import load_url as load_state_dict_from_url # noqa: 401 + + +def _get_extension_path(lib_name): + + lib_dir = os.path.dirname(__file__) + if os.name == "nt": + # Register the main torchvision library location on the default DLL path + import ctypes + + kernel32 = ctypes.WinDLL("kernel32.dll", use_last_error=True) + with_load_library_flags = hasattr(kernel32, "AddDllDirectory") + prev_error_mode = kernel32.SetErrorMode(0x0001) + + if with_load_library_flags: + kernel32.AddDllDirectory.restype = ctypes.c_void_p + + os.add_dll_directory(lib_dir) + + kernel32.SetErrorMode(prev_error_mode) + + loader_details = (importlib.machinery.ExtensionFileLoader, importlib.machinery.EXTENSION_SUFFIXES) + + extfinder = importlib.machinery.FileFinder(lib_dir, loader_details) + ext_specs = extfinder.find_spec(lib_name) + if ext_specs is None: + raise ImportError + + return ext_specs.origin diff --git a/lib/python3.12/site-packages/torchvision/_meta_registrations.py b/lib/python3.12/site-packages/torchvision/_meta_registrations.py new file mode 100644 index 0000000000000000000000000000000000000000..f75bfb77a7f25a1842509de595f109f232994574 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/_meta_registrations.py @@ -0,0 +1,225 @@ +import functools + +import torch +import torch._custom_ops +import torch.library + +# Ensure that torch.ops.torchvision is visible +import torchvision.extension # noqa: F401 + + +@functools.lru_cache(None) +def get_meta_lib(): + return torch.library.Library("torchvision", "IMPL", "Meta") + + +def register_meta(op_name, overload_name="default"): + def wrapper(fn): + if torchvision.extension._has_ops(): + get_meta_lib().impl(getattr(getattr(torch.ops.torchvision, op_name), overload_name), fn) + return fn + + return wrapper + + +@register_meta("roi_align") +def meta_roi_align(input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio, aligned): + torch._check(rois.size(1) == 5, lambda: "rois must have shape as Tensor[K, 5]") + torch._check( + input.dtype == rois.dtype, + lambda: ( + "Expected tensor for input to have the same type as tensor for rois; " + f"but type {input.dtype} does not equal {rois.dtype}" + ), + ) + num_rois = rois.size(0) + channels = input.size(1) + return input.new_empty((num_rois, channels, pooled_height, pooled_width)) + + +@register_meta("_roi_align_backward") +def meta_roi_align_backward( + grad, rois, spatial_scale, pooled_height, pooled_width, batch_size, channels, height, width, sampling_ratio, aligned +): + torch._check( + grad.dtype == rois.dtype, + lambda: ( + "Expected tensor for grad to have the same type as tensor for rois; " + f"but type {grad.dtype} does not equal {rois.dtype}" + ), + ) + return grad.new_empty((batch_size, channels, height, width)) + + +@register_meta("ps_roi_align") +def meta_ps_roi_align(input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio): + torch._check(rois.size(1) == 5, lambda: "rois must have shape as Tensor[K, 5]") + torch._check( + input.dtype == rois.dtype, + lambda: ( + "Expected tensor for input to have the same type as tensor for rois; " + f"but type {input.dtype} does not equal {rois.dtype}" + ), + ) + channels = input.size(1) + torch._check( + channels % (pooled_height * pooled_width) == 0, + "input channels must be a multiple of pooling height * pooling width", + ) + + num_rois = rois.size(0) + out_size = (num_rois, channels // (pooled_height * pooled_width), pooled_height, pooled_width) + return input.new_empty(out_size), torch.empty(out_size, dtype=torch.int32, device="meta") + + +@register_meta("_ps_roi_align_backward") +def meta_ps_roi_align_backward( + grad, + rois, + channel_mapping, + spatial_scale, + pooled_height, + pooled_width, + sampling_ratio, + batch_size, + channels, + height, + width, +): + torch._check( + grad.dtype == rois.dtype, + lambda: ( + "Expected tensor for grad to have the same type as tensor for rois; " + f"but type {grad.dtype} does not equal {rois.dtype}" + ), + ) + return grad.new_empty((batch_size, channels, height, width)) + + +@register_meta("roi_pool") +def meta_roi_pool(input, rois, spatial_scale, pooled_height, pooled_width): + torch._check(rois.size(1) == 5, lambda: "rois must have shape as Tensor[K, 5]") + torch._check( + input.dtype == rois.dtype, + lambda: ( + "Expected tensor for input to have the same type as tensor for rois; " + f"but type {input.dtype} does not equal {rois.dtype}" + ), + ) + num_rois = rois.size(0) + channels = input.size(1) + out_size = (num_rois, channels, pooled_height, pooled_width) + return input.new_empty(out_size), torch.empty(out_size, device="meta", dtype=torch.int32) + + +@register_meta("_roi_pool_backward") +def meta_roi_pool_backward( + grad, rois, argmax, spatial_scale, pooled_height, pooled_width, batch_size, channels, height, width +): + torch._check( + grad.dtype == rois.dtype, + lambda: ( + "Expected tensor for grad to have the same type as tensor for rois; " + f"but type {grad.dtype} does not equal {rois.dtype}" + ), + ) + return grad.new_empty((batch_size, channels, height, width)) + + +@register_meta("ps_roi_pool") +def meta_ps_roi_pool(input, rois, spatial_scale, pooled_height, pooled_width): + torch._check(rois.size(1) == 5, lambda: "rois must have shape as Tensor[K, 5]") + torch._check( + input.dtype == rois.dtype, + lambda: ( + "Expected tensor for input to have the same type as tensor for rois; " + f"but type {input.dtype} does not equal {rois.dtype}" + ), + ) + channels = input.size(1) + torch._check( + channels % (pooled_height * pooled_width) == 0, + "input channels must be a multiple of pooling height * pooling width", + ) + num_rois = rois.size(0) + out_size = (num_rois, channels // (pooled_height * pooled_width), pooled_height, pooled_width) + return input.new_empty(out_size), torch.empty(out_size, device="meta", dtype=torch.int32) + + +@register_meta("_ps_roi_pool_backward") +def meta_ps_roi_pool_backward( + grad, rois, channel_mapping, spatial_scale, pooled_height, pooled_width, batch_size, channels, height, width +): + torch._check( + grad.dtype == rois.dtype, + lambda: ( + "Expected tensor for grad to have the same type as tensor for rois; " + f"but type {grad.dtype} does not equal {rois.dtype}" + ), + ) + return grad.new_empty((batch_size, channels, height, width)) + + +@torch.library.register_fake("torchvision::nms") +def meta_nms(dets, scores, iou_threshold): + torch._check(dets.dim() == 2, lambda: f"boxes should be a 2d tensor, got {dets.dim()}D") + torch._check(dets.size(1) == 4, lambda: f"boxes should have 4 elements in dimension 1, got {dets.size(1)}") + torch._check(scores.dim() == 1, lambda: f"scores should be a 1d tensor, got {scores.dim()}") + torch._check( + dets.size(0) == scores.size(0), + lambda: f"boxes and scores should have same number of elements in dimension 0, got {dets.size(0)} and {scores.size(0)}", + ) + ctx = torch._custom_ops.get_ctx() + num_to_keep = ctx.create_unbacked_symint() + return dets.new_empty(num_to_keep, dtype=torch.long) + + +@register_meta("deform_conv2d") +def meta_deform_conv2d( + input, + weight, + offset, + mask, + bias, + stride_h, + stride_w, + pad_h, + pad_w, + dil_h, + dil_w, + n_weight_grps, + n_offset_grps, + use_mask, +): + + out_height, out_width = offset.shape[-2:] + out_channels = weight.shape[0] + batch_size = input.shape[0] + return input.new_empty((batch_size, out_channels, out_height, out_width)) + + +@register_meta("_deform_conv2d_backward") +def meta_deform_conv2d_backward( + grad, + input, + weight, + offset, + mask, + bias, + stride_h, + stride_w, + pad_h, + pad_w, + dilation_h, + dilation_w, + groups, + offset_groups, + use_mask, +): + + grad_input = input.new_empty(input.shape) + grad_weight = weight.new_empty(weight.shape) + grad_offset = offset.new_empty(offset.shape) + grad_mask = mask.new_empty(mask.shape) + grad_bias = bias.new_empty(bias.shape) + return grad_input, grad_weight, grad_offset, grad_mask, grad_bias diff --git a/lib/python3.12/site-packages/torchvision/_utils.py b/lib/python3.12/site-packages/torchvision/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b739ef0966e9b6fac4574f3d6f04051799f75a16 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/_utils.py @@ -0,0 +1,32 @@ +import enum +from typing import Sequence, Type, TypeVar + +T = TypeVar("T", bound=enum.Enum) + + +class StrEnumMeta(enum.EnumMeta): + auto = enum.auto + + def from_str(self: Type[T], member: str) -> T: # type: ignore[misc] + try: + return self[member] + except KeyError: + # TODO: use `add_suggestion` from torchvision.prototype.utils._internal to improve the error message as + # soon as it is migrated. + raise ValueError(f"Unknown value '{member}' for {self.__name__}.") from None + + +class StrEnum(enum.Enum, metaclass=StrEnumMeta): + pass + + +def sequence_to_str(seq: Sequence, separate_last: str = "") -> str: + if not seq: + return "" + if len(seq) == 1: + return f"'{seq[0]}'" + + head = "'" + "', '".join([str(item) for item in seq[:-1]]) + "'" + tail = f"{'' if separate_last and len(seq) == 2 else ','} {separate_last}'{seq[-1]}'" + + return head + tail diff --git a/lib/python3.12/site-packages/torchvision/datasets/__init__.py b/lib/python3.12/site-packages/torchvision/datasets/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4005b4a90729c9fe1b811f7388bd8453998d2322 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/__init__.py @@ -0,0 +1,147 @@ +from ._optical_flow import FlyingChairs, FlyingThings3D, HD1K, KittiFlow, Sintel +from ._stereo_matching import ( + CarlaStereo, + CREStereo, + ETH3DStereo, + FallingThingsStereo, + InStereo2k, + Kitti2012Stereo, + Kitti2015Stereo, + Middlebury2014Stereo, + SceneFlowStereo, + SintelStereo, +) +from .caltech import Caltech101, Caltech256 +from .celeba import CelebA +from .cifar import CIFAR10, CIFAR100 +from .cityscapes import Cityscapes +from .clevr import CLEVRClassification +from .coco import CocoCaptions, CocoDetection +from .country211 import Country211 +from .dtd import DTD +from .eurosat import EuroSAT +from .fakedata import FakeData +from .fer2013 import FER2013 +from .fgvc_aircraft import FGVCAircraft +from .flickr import Flickr30k, Flickr8k +from .flowers102 import Flowers102 +from .folder import DatasetFolder, ImageFolder +from .food101 import Food101 +from .gtsrb import GTSRB +from .hmdb51 import HMDB51 +from .imagenet import ImageNet +from .imagenette import Imagenette +from .inaturalist import INaturalist +from .kinetics import Kinetics +from .kitti import Kitti +from .lfw import LFWPairs, LFWPeople +from .lsun import LSUN, LSUNClass +from .mnist import EMNIST, FashionMNIST, KMNIST, MNIST, QMNIST +from .moving_mnist import MovingMNIST +from .omniglot import Omniglot +from .oxford_iiit_pet import OxfordIIITPet +from .pcam import PCAM +from .phototour import PhotoTour +from .places365 import Places365 +from .rendered_sst2 import RenderedSST2 +from .sbd import SBDataset +from .sbu import SBU +from .semeion import SEMEION +from .stanford_cars import StanfordCars +from .stl10 import STL10 +from .sun397 import SUN397 +from .svhn import SVHN +from .ucf101 import UCF101 +from .usps import USPS +from .vision import VisionDataset +from .voc import VOCDetection, VOCSegmentation +from .widerface import WIDERFace + +__all__ = ( + "LSUN", + "LSUNClass", + "ImageFolder", + "DatasetFolder", + "FakeData", + "CocoCaptions", + "CocoDetection", + "CIFAR10", + "CIFAR100", + "EMNIST", + "FashionMNIST", + "QMNIST", + "MNIST", + "KMNIST", + "MovingMNIST", + "StanfordCars", + "STL10", + "SUN397", + "SVHN", + "PhotoTour", + "SEMEION", + "Omniglot", + "SBU", + "Flickr8k", + "Flickr30k", + "Flowers102", + "VOCSegmentation", + "VOCDetection", + "Cityscapes", + "ImageNet", + "Caltech101", + "Caltech256", + "CelebA", + "WIDERFace", + "SBDataset", + "VisionDataset", + "USPS", + "Kinetics", + "HMDB51", + "UCF101", + "Places365", + "Kitti", + "INaturalist", + "LFWPeople", + "LFWPairs", + "KittiFlow", + "Sintel", + "FlyingChairs", + "FlyingThings3D", + "HD1K", + "Food101", + "DTD", + "FER2013", + "GTSRB", + "CLEVRClassification", + "OxfordIIITPet", + "PCAM", + "Country211", + "FGVCAircraft", + "EuroSAT", + "RenderedSST2", + "Kitti2012Stereo", + "Kitti2015Stereo", + "CarlaStereo", + "Middlebury2014Stereo", + "CREStereo", + "FallingThingsStereo", + "SceneFlowStereo", + "SintelStereo", + "InStereo2k", + "ETH3DStereo", + "wrap_dataset_for_transforms_v2", + "Imagenette", +) + + +# We override current module's attributes to handle the import: +# from torchvision.datasets import wrap_dataset_for_transforms_v2 +# without a cyclic error. +# Ref: https://peps.python.org/pep-0562/ +def __getattr__(name): + if name in ("wrap_dataset_for_transforms_v2",): + from torchvision.tv_tensors._dataset_wrapper import wrap_dataset_for_transforms_v2 + + return wrap_dataset_for_transforms_v2 + + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0822fc9d30829c8ee5d8fe890e2f42aa1d334579 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/_optical_flow.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/_optical_flow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dbffa45ad0b4f5174a24f8f395d8860c37059245 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/_optical_flow.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/_stereo_matching.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/_stereo_matching.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..66b749600b9eec9fa02e5de4bf48089474cb87f4 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/_stereo_matching.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/caltech.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/caltech.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9e36b07e192d123d022d86c1666dfef61ae7e2b2 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/caltech.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/celeba.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/celeba.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d32d3d459499cbdf9a2c5cf4f45216393a08be75 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/celeba.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/cifar.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/cifar.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ac84e15064b33871133aa806efb85df986b84053 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/cifar.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/cityscapes.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/cityscapes.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df65c28d01972a534867c27dd325ba747e6d0ef9 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/cityscapes.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/clevr.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/clevr.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b027b837e32f35d2434eaeb46d883390d1d46401 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/clevr.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/coco.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/coco.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b47d6c5fde40f34f3bbec2510dd8ab76577964fc Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/coco.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/country211.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/country211.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..149a2fb0b2dcc23e110985b5b7fb36ba89d1cbfd Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/country211.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/dtd.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/dtd.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d3458b76f5cd3611cc74b15d9e5ed15e59ccd97b Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/dtd.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/eurosat.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/eurosat.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df5690cd148eff3548f92a8ec5ffc93554ef8e6f Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/eurosat.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/fakedata.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/fakedata.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..69ee7652bc655a021b6d5b5fa6ab6078e4f091c5 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/fakedata.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/fer2013.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/fer2013.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..42ec4bc8d4cda6710ec9b8b16a20b9251edd712c Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/fer2013.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/fgvc_aircraft.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/fgvc_aircraft.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a39fabd88efa31868c11e11e86244a33a4314b83 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/fgvc_aircraft.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/flickr.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/flickr.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ed0f3e7158659590566e47d60978083ea8d9b3bb Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/flickr.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/flowers102.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/flowers102.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df57e4ee8ecd9e89b6637e9e9b3c78d97919a643 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/flowers102.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/folder.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/folder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..011c18da646099a631d0a171002b0dfff91c80e3 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/folder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/food101.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/food101.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..09ff6d6b287df8a8e6918f40c9b69028d5e299c1 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/food101.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/gtsrb.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/gtsrb.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2a5f88863d4af40c4e724a0c4dcdabcdcb44f055 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/gtsrb.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/hmdb51.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/hmdb51.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..017fef17708c6ffa6ef4f826e2b1a69c89990240 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/hmdb51.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/imagenet.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/imagenet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6f4e2a92606c144e7db1f818a9450bce7eb503b2 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/imagenet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/imagenette.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/imagenette.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1f63092f2be9c34ded1248906603ec34bf1d4d99 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/imagenette.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/inaturalist.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/inaturalist.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ddc25e0d03dc1ad32c74098582c707d8ea98b2f0 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/inaturalist.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/kinetics.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/kinetics.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1d8f8b7a797bb2d88a4db7f80d8a3ae4a7fa4952 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/kinetics.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/kitti.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/kitti.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..98308502bed3d5783e7b4451c108c18ae874aae2 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/kitti.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/lfw.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/lfw.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..087177aecec5de3639256eaa95fa2f59d93f1fdd Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/lfw.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/lsun.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/lsun.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..25d573cd35fd3de24c0da0836299bb366708e65f Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/lsun.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/mnist.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/mnist.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4d81763a954d424ed971b5d310c45f081525b463 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/mnist.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/moving_mnist.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/moving_mnist.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1f774a61ec98611db4b486eb79e746b85e9d466d Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/moving_mnist.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/omniglot.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/omniglot.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c20fdbada7dcdc181448147fb82d7d2a037b571c Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/omniglot.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/oxford_iiit_pet.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/oxford_iiit_pet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..57c8e16d5615e6953a60988f677e0af938f781f8 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/oxford_iiit_pet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/pcam.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/pcam.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ff67c5e767e4081fab425121abe4f8583003e3c3 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/pcam.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/phototour.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/phototour.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..42a99a3b93dc9d7d11313afd2800c3c5d3ba7576 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/phototour.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/places365.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/places365.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..76978472d0f1850a1027c2b6639f0f13a5ee8e96 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/places365.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/rendered_sst2.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/rendered_sst2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2e754b1f1d49d283d3a5f2a3f00b90e3ffc36be3 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/rendered_sst2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/sbd.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/sbd.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b05b1571702b1f0f3031c9d49d7f0ba170c576bd Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/sbd.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/sbu.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/sbu.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4388b49f207558a14d7cbaf2079e28b541068b4c Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/sbu.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/semeion.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/semeion.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..409cbb9183cc717f0a557f2b034191f49098a948 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/semeion.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/stanford_cars.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/stanford_cars.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7ce26ff7f89ec7885d8e40376d409b2bac718174 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/stanford_cars.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/stl10.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/stl10.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c93de636258246ec302126cd941a1eb9ce38dcf5 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/stl10.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/sun397.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/sun397.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..26768c1af3cc7b32d5e7575aa2019aedcb857595 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/sun397.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/svhn.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/svhn.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2c4eb696e0d1cc7e0f2aeec82f2169d2207b5f3d Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/svhn.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/ucf101.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/ucf101.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f66e5a2b1fe145434424affa7ed1902a9f7d8398 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/ucf101.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/usps.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/usps.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b08815dd1435bf41e26776c5b660fca820492a39 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/usps.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d74240664294d41b7fd406dc35e1cc02341adfb4 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/video_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/video_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..45aba3e86fe4ef0ce56e02ec771fd3a8037bd641 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/video_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/vision.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/vision.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..672f4bf9b2510d069c2f60e99e0549d6b7645d6c Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/vision.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/voc.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/voc.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bbba2960bace577d43789fb40a5ed2c94c6a3590 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/voc.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/__pycache__/widerface.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/widerface.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b5c552bb3f7b21723d8aa901a4497575de68cbb6 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/__pycache__/widerface.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/_optical_flow.py b/lib/python3.12/site-packages/torchvision/datasets/_optical_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..e8d6247f03fadb984686f6ec892f0e9dc2e6a571 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/_optical_flow.py @@ -0,0 +1,490 @@ +import itertools +import os +from abc import ABC, abstractmethod +from glob import glob +from pathlib import Path +from typing import Callable, List, Optional, Tuple, Union + +import numpy as np +import torch +from PIL import Image + +from ..io.image import decode_png, read_file +from .utils import _read_pfm, verify_str_arg +from .vision import VisionDataset + +T1 = Tuple[Image.Image, Image.Image, Optional[np.ndarray], Optional[np.ndarray]] +T2 = Tuple[Image.Image, Image.Image, Optional[np.ndarray]] + + +__all__ = ( + "KittiFlow", + "Sintel", + "FlyingThings3D", + "FlyingChairs", + "HD1K", +) + + +class FlowDataset(ABC, VisionDataset): + # Some datasets like Kitti have a built-in valid_flow_mask, indicating which flow values are valid + # For those we return (img1, img2, flow, valid_flow_mask), and for the rest we return (img1, img2, flow), + # and it's up to whatever consumes the dataset to decide what valid_flow_mask should be. + _has_builtin_flow_mask = False + + def __init__(self, root: Union[str, Path], transforms: Optional[Callable] = None) -> None: + + super().__init__(root=root) + self.transforms = transforms + + self._flow_list: List[str] = [] + self._image_list: List[List[str]] = [] + + def _read_img(self, file_name: str) -> Image.Image: + img = Image.open(file_name) + if img.mode != "RGB": + img = img.convert("RGB") # type: ignore[assignment] + return img + + @abstractmethod + def _read_flow(self, file_name: str): + # Return the flow or a tuple with the flow and the valid_flow_mask if _has_builtin_flow_mask is True + pass + + def __getitem__(self, index: int) -> Union[T1, T2]: + + img1 = self._read_img(self._image_list[index][0]) + img2 = self._read_img(self._image_list[index][1]) + + if self._flow_list: # it will be empty for some dataset when split="test" + flow = self._read_flow(self._flow_list[index]) + if self._has_builtin_flow_mask: + flow, valid_flow_mask = flow + else: + valid_flow_mask = None + else: + flow = valid_flow_mask = None + + if self.transforms is not None: + img1, img2, flow, valid_flow_mask = self.transforms(img1, img2, flow, valid_flow_mask) + + if self._has_builtin_flow_mask or valid_flow_mask is not None: + # The `or valid_flow_mask is not None` part is here because the mask can be generated within a transform + return img1, img2, flow, valid_flow_mask + else: + return img1, img2, flow + + def __len__(self) -> int: + return len(self._image_list) + + def __rmul__(self, v: int) -> torch.utils.data.ConcatDataset: + return torch.utils.data.ConcatDataset([self] * v) + + +class Sintel(FlowDataset): + """`Sintel `_ Dataset for optical flow. + + The dataset is expected to have the following structure: :: + + root + Sintel + testing + clean + scene_1 + scene_2 + ... + final + scene_1 + scene_2 + ... + training + clean + scene_1 + scene_2 + ... + final + scene_1 + scene_2 + ... + flow + scene_1 + scene_2 + ... + + Args: + root (str or ``pathlib.Path``): Root directory of the Sintel Dataset. + split (string, optional): The dataset split, either "train" (default) or "test" + pass_name (string, optional): The pass to use, either "clean" (default), "final", or "both". See link above for + details on the different passes. + transforms (callable, optional): A function/transform that takes in + ``img1, img2, flow, valid_flow_mask`` and returns a transformed version. + ``valid_flow_mask`` is expected for consistency with other datasets which + return a built-in valid mask, such as :class:`~torchvision.datasets.KittiFlow`. + """ + + def __init__( + self, + root: Union[str, Path], + split: str = "train", + pass_name: str = "clean", + transforms: Optional[Callable] = None, + ) -> None: + super().__init__(root=root, transforms=transforms) + + verify_str_arg(split, "split", valid_values=("train", "test")) + verify_str_arg(pass_name, "pass_name", valid_values=("clean", "final", "both")) + passes = ["clean", "final"] if pass_name == "both" else [pass_name] + + root = Path(root) / "Sintel" + flow_root = root / "training" / "flow" + + for pass_name in passes: + split_dir = "training" if split == "train" else split + image_root = root / split_dir / pass_name + for scene in os.listdir(image_root): + image_list = sorted(glob(str(image_root / scene / "*.png"))) + for i in range(len(image_list) - 1): + self._image_list += [[image_list[i], image_list[i + 1]]] + + if split == "train": + self._flow_list += sorted(glob(str(flow_root / scene / "*.flo"))) + + def __getitem__(self, index: int) -> Union[T1, T2]: + """Return example at given index. + + Args: + index(int): The index of the example to retrieve + + Returns: + tuple: A 3-tuple with ``(img1, img2, flow)``. + The flow is a numpy array of shape (2, H, W) and the images are PIL images. + ``flow`` is None if ``split="test"``. + If a valid flow mask is generated within the ``transforms`` parameter, + a 4-tuple with ``(img1, img2, flow, valid_flow_mask)`` is returned. + """ + return super().__getitem__(index) + + def _read_flow(self, file_name: str) -> np.ndarray: + return _read_flo(file_name) + + +class KittiFlow(FlowDataset): + """`KITTI `__ dataset for optical flow (2015). + + The dataset is expected to have the following structure: :: + + root + KittiFlow + testing + image_2 + training + image_2 + flow_occ + + Args: + root (str or ``pathlib.Path``): Root directory of the KittiFlow Dataset. + split (string, optional): The dataset split, either "train" (default) or "test" + transforms (callable, optional): A function/transform that takes in + ``img1, img2, flow, valid_flow_mask`` and returns a transformed version. + """ + + _has_builtin_flow_mask = True + + def __init__(self, root: Union[str, Path], split: str = "train", transforms: Optional[Callable] = None) -> None: + super().__init__(root=root, transforms=transforms) + + verify_str_arg(split, "split", valid_values=("train", "test")) + + root = Path(root) / "KittiFlow" / (split + "ing") + images1 = sorted(glob(str(root / "image_2" / "*_10.png"))) + images2 = sorted(glob(str(root / "image_2" / "*_11.png"))) + + if not images1 or not images2: + raise FileNotFoundError( + "Could not find the Kitti flow images. Please make sure the directory structure is correct." + ) + + for img1, img2 in zip(images1, images2): + self._image_list += [[img1, img2]] + + if split == "train": + self._flow_list = sorted(glob(str(root / "flow_occ" / "*_10.png"))) + + def __getitem__(self, index: int) -> Union[T1, T2]: + """Return example at given index. + + Args: + index(int): The index of the example to retrieve + + Returns: + tuple: A 4-tuple with ``(img1, img2, flow, valid_flow_mask)`` + where ``valid_flow_mask`` is a numpy boolean mask of shape (H, W) + indicating which flow values are valid. The flow is a numpy array of + shape (2, H, W) and the images are PIL images. ``flow`` and ``valid_flow_mask`` are None if + ``split="test"``. + """ + return super().__getitem__(index) + + def _read_flow(self, file_name: str) -> Tuple[np.ndarray, np.ndarray]: + return _read_16bits_png_with_flow_and_valid_mask(file_name) + + +class FlyingChairs(FlowDataset): + """`FlyingChairs `_ Dataset for optical flow. + + You will also need to download the FlyingChairs_train_val.txt file from the dataset page. + + The dataset is expected to have the following structure: :: + + root + FlyingChairs + data + 00001_flow.flo + 00001_img1.ppm + 00001_img2.ppm + ... + FlyingChairs_train_val.txt + + + Args: + root (str or ``pathlib.Path``): Root directory of the FlyingChairs Dataset. + split (string, optional): The dataset split, either "train" (default) or "val" + transforms (callable, optional): A function/transform that takes in + ``img1, img2, flow, valid_flow_mask`` and returns a transformed version. + ``valid_flow_mask`` is expected for consistency with other datasets which + return a built-in valid mask, such as :class:`~torchvision.datasets.KittiFlow`. + """ + + def __init__(self, root: Union[str, Path], split: str = "train", transforms: Optional[Callable] = None) -> None: + super().__init__(root=root, transforms=transforms) + + verify_str_arg(split, "split", valid_values=("train", "val")) + + root = Path(root) / "FlyingChairs" + images = sorted(glob(str(root / "data" / "*.ppm"))) + flows = sorted(glob(str(root / "data" / "*.flo"))) + + split_file_name = "FlyingChairs_train_val.txt" + + if not os.path.exists(root / split_file_name): + raise FileNotFoundError( + "The FlyingChairs_train_val.txt file was not found - please download it from the dataset page (see docstring)." + ) + + split_list = np.loadtxt(str(root / split_file_name), dtype=np.int32) + for i in range(len(flows)): + split_id = split_list[i] + if (split == "train" and split_id == 1) or (split == "val" and split_id == 2): + self._flow_list += [flows[i]] + self._image_list += [[images[2 * i], images[2 * i + 1]]] + + def __getitem__(self, index: int) -> Union[T1, T2]: + """Return example at given index. + + Args: + index(int): The index of the example to retrieve + + Returns: + tuple: A 3-tuple with ``(img1, img2, flow)``. + The flow is a numpy array of shape (2, H, W) and the images are PIL images. + ``flow`` is None if ``split="val"``. + If a valid flow mask is generated within the ``transforms`` parameter, + a 4-tuple with ``(img1, img2, flow, valid_flow_mask)`` is returned. + """ + return super().__getitem__(index) + + def _read_flow(self, file_name: str) -> np.ndarray: + return _read_flo(file_name) + + +class FlyingThings3D(FlowDataset): + """`FlyingThings3D `_ dataset for optical flow. + + The dataset is expected to have the following structure: :: + + root + FlyingThings3D + frames_cleanpass + TEST + TRAIN + frames_finalpass + TEST + TRAIN + optical_flow + TEST + TRAIN + + Args: + root (str or ``pathlib.Path``): Root directory of the intel FlyingThings3D Dataset. + split (string, optional): The dataset split, either "train" (default) or "test" + pass_name (string, optional): The pass to use, either "clean" (default) or "final" or "both". See link above for + details on the different passes. + camera (string, optional): Which camera to return images from. Can be either "left" (default) or "right" or "both". + transforms (callable, optional): A function/transform that takes in + ``img1, img2, flow, valid_flow_mask`` and returns a transformed version. + ``valid_flow_mask`` is expected for consistency with other datasets which + return a built-in valid mask, such as :class:`~torchvision.datasets.KittiFlow`. + """ + + def __init__( + self, + root: Union[str, Path], + split: str = "train", + pass_name: str = "clean", + camera: str = "left", + transforms: Optional[Callable] = None, + ) -> None: + super().__init__(root=root, transforms=transforms) + + verify_str_arg(split, "split", valid_values=("train", "test")) + split = split.upper() + + verify_str_arg(pass_name, "pass_name", valid_values=("clean", "final", "both")) + passes = { + "clean": ["frames_cleanpass"], + "final": ["frames_finalpass"], + "both": ["frames_cleanpass", "frames_finalpass"], + }[pass_name] + + verify_str_arg(camera, "camera", valid_values=("left", "right", "both")) + cameras = ["left", "right"] if camera == "both" else [camera] + + root = Path(root) / "FlyingThings3D" + + directions = ("into_future", "into_past") + for pass_name, camera, direction in itertools.product(passes, cameras, directions): + image_dirs = sorted(glob(str(root / pass_name / split / "*/*"))) + image_dirs = sorted(Path(image_dir) / camera for image_dir in image_dirs) + + flow_dirs = sorted(glob(str(root / "optical_flow" / split / "*/*"))) + flow_dirs = sorted(Path(flow_dir) / direction / camera for flow_dir in flow_dirs) + + if not image_dirs or not flow_dirs: + raise FileNotFoundError( + "Could not find the FlyingThings3D flow images. " + "Please make sure the directory structure is correct." + ) + + for image_dir, flow_dir in zip(image_dirs, flow_dirs): + images = sorted(glob(str(image_dir / "*.png"))) + flows = sorted(glob(str(flow_dir / "*.pfm"))) + for i in range(len(flows) - 1): + if direction == "into_future": + self._image_list += [[images[i], images[i + 1]]] + self._flow_list += [flows[i]] + elif direction == "into_past": + self._image_list += [[images[i + 1], images[i]]] + self._flow_list += [flows[i + 1]] + + def __getitem__(self, index: int) -> Union[T1, T2]: + """Return example at given index. + + Args: + index(int): The index of the example to retrieve + + Returns: + tuple: A 3-tuple with ``(img1, img2, flow)``. + The flow is a numpy array of shape (2, H, W) and the images are PIL images. + ``flow`` is None if ``split="test"``. + If a valid flow mask is generated within the ``transforms`` parameter, + a 4-tuple with ``(img1, img2, flow, valid_flow_mask)`` is returned. + """ + return super().__getitem__(index) + + def _read_flow(self, file_name: str) -> np.ndarray: + return _read_pfm(file_name) + + +class HD1K(FlowDataset): + """`HD1K `__ dataset for optical flow. + + The dataset is expected to have the following structure: :: + + root + hd1k + hd1k_challenge + image_2 + hd1k_flow_gt + flow_occ + hd1k_input + image_2 + + Args: + root (str or ``pathlib.Path``): Root directory of the HD1K Dataset. + split (string, optional): The dataset split, either "train" (default) or "test" + transforms (callable, optional): A function/transform that takes in + ``img1, img2, flow, valid_flow_mask`` and returns a transformed version. + """ + + _has_builtin_flow_mask = True + + def __init__(self, root: Union[str, Path], split: str = "train", transforms: Optional[Callable] = None) -> None: + super().__init__(root=root, transforms=transforms) + + verify_str_arg(split, "split", valid_values=("train", "test")) + + root = Path(root) / "hd1k" + if split == "train": + # There are 36 "sequences" and we don't want seq i to overlap with seq i + 1, so we need this for loop + for seq_idx in range(36): + flows = sorted(glob(str(root / "hd1k_flow_gt" / "flow_occ" / f"{seq_idx:06d}_*.png"))) + images = sorted(glob(str(root / "hd1k_input" / "image_2" / f"{seq_idx:06d}_*.png"))) + for i in range(len(flows) - 1): + self._flow_list += [flows[i]] + self._image_list += [[images[i], images[i + 1]]] + else: + images1 = sorted(glob(str(root / "hd1k_challenge" / "image_2" / "*10.png"))) + images2 = sorted(glob(str(root / "hd1k_challenge" / "image_2" / "*11.png"))) + for image1, image2 in zip(images1, images2): + self._image_list += [[image1, image2]] + + if not self._image_list: + raise FileNotFoundError( + "Could not find the HD1K images. Please make sure the directory structure is correct." + ) + + def _read_flow(self, file_name: str) -> Tuple[np.ndarray, np.ndarray]: + return _read_16bits_png_with_flow_and_valid_mask(file_name) + + def __getitem__(self, index: int) -> Union[T1, T2]: + """Return example at given index. + + Args: + index(int): The index of the example to retrieve + + Returns: + tuple: A 4-tuple with ``(img1, img2, flow, valid_flow_mask)`` where ``valid_flow_mask`` + is a numpy boolean mask of shape (H, W) + indicating which flow values are valid. The flow is a numpy array of + shape (2, H, W) and the images are PIL images. ``flow`` and ``valid_flow_mask`` are None if + ``split="test"``. + """ + return super().__getitem__(index) + + +def _read_flo(file_name: str) -> np.ndarray: + """Read .flo file in Middlebury format""" + # Code adapted from: + # http://stackoverflow.com/questions/28013200/reading-middlebury-flow-files-with-python-bytes-array-numpy + # Everything needs to be in little Endian according to + # https://vision.middlebury.edu/flow/code/flow-code/README.txt + with open(file_name, "rb") as f: + magic = np.fromfile(f, "c", count=4).tobytes() + if magic != b"PIEH": + raise ValueError("Magic number incorrect. Invalid .flo file") + + w = int(np.fromfile(f, " Tuple[np.ndarray, np.ndarray]: + + flow_and_valid = decode_png(read_file(file_name)).to(torch.float32) + flow, valid_flow_mask = flow_and_valid[:2, :, :], flow_and_valid[2, :, :] + flow = (flow - 2**15) / 64 # This conversion is explained somewhere on the kitti archive + valid_flow_mask = valid_flow_mask.bool() + + # For consistency with other datasets, we convert to numpy + return flow.numpy(), valid_flow_mask.numpy() diff --git a/lib/python3.12/site-packages/torchvision/datasets/_stereo_matching.py b/lib/python3.12/site-packages/torchvision/datasets/_stereo_matching.py new file mode 100644 index 0000000000000000000000000000000000000000..09961211cc2b2eecb59ac3878baf60f3c36bd320 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/_stereo_matching.py @@ -0,0 +1,1223 @@ +import functools +import json +import os +import random +import shutil +from abc import ABC, abstractmethod +from glob import glob +from pathlib import Path +from typing import Callable, cast, List, Optional, Tuple, Union + +import numpy as np +from PIL import Image + +from .utils import _read_pfm, download_and_extract_archive, verify_str_arg +from .vision import VisionDataset + +T1 = Tuple[Image.Image, Image.Image, Optional[np.ndarray], np.ndarray] +T2 = Tuple[Image.Image, Image.Image, Optional[np.ndarray]] + +__all__ = () + +_read_pfm_file = functools.partial(_read_pfm, slice_channels=1) + + +class StereoMatchingDataset(ABC, VisionDataset): + """Base interface for Stereo matching datasets""" + + _has_built_in_disparity_mask = False + + def __init__(self, root: Union[str, Path], transforms: Optional[Callable] = None) -> None: + """ + Args: + root(str): Root directory of the dataset. + transforms(callable, optional): A function/transform that takes in Tuples of + (images, disparities, valid_masks) and returns a transformed version of each of them. + images is a Tuple of (``PIL.Image``, ``PIL.Image``) + disparities is a Tuple of (``np.ndarray``, ``np.ndarray``) with shape (1, H, W) + valid_masks is a Tuple of (``np.ndarray``, ``np.ndarray``) with shape (H, W) + In some cases, when a dataset does not provide disparities, the ``disparities`` and + ``valid_masks`` can be Tuples containing None values. + For training splits generally the datasets provide a minimal guarantee of + images: (``PIL.Image``, ``PIL.Image``) + disparities: (``np.ndarray``, ``None``) with shape (1, H, W) + Optionally, based on the dataset, it can return a ``mask`` as well: + valid_masks: (``np.ndarray | None``, ``None``) with shape (H, W) + For some test splits, the datasets provides outputs that look like: + imgaes: (``PIL.Image``, ``PIL.Image``) + disparities: (``None``, ``None``) + Optionally, based on the dataset, it can return a ``mask`` as well: + valid_masks: (``None``, ``None``) + """ + super().__init__(root=root) + self.transforms = transforms + + self._images = [] # type: ignore + self._disparities = [] # type: ignore + + def _read_img(self, file_path: Union[str, Path]) -> Image.Image: + img = Image.open(file_path) + if img.mode != "RGB": + img = img.convert("RGB") # type: ignore [assignment] + return img + + def _scan_pairs( + self, + paths_left_pattern: str, + paths_right_pattern: Optional[str] = None, + ) -> List[Tuple[str, Optional[str]]]: + + left_paths = list(sorted(glob(paths_left_pattern))) + + right_paths: List[Union[None, str]] + if paths_right_pattern: + right_paths = list(sorted(glob(paths_right_pattern))) + else: + right_paths = list(None for _ in left_paths) + + if not left_paths: + raise FileNotFoundError(f"Could not find any files matching the patterns: {paths_left_pattern}") + + if not right_paths: + raise FileNotFoundError(f"Could not find any files matching the patterns: {paths_right_pattern}") + + if len(left_paths) != len(right_paths): + raise ValueError( + f"Found {len(left_paths)} left files but {len(right_paths)} right files using:\n " + f"left pattern: {paths_left_pattern}\n" + f"right pattern: {paths_right_pattern}\n" + ) + + paths = list((left, right) for left, right in zip(left_paths, right_paths)) + return paths + + @abstractmethod + def _read_disparity(self, file_path: str) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]: + # function that returns a disparity map and an occlusion map + pass + + def __getitem__(self, index: int) -> Union[T1, T2]: + """Return example at given index. + + Args: + index(int): The index of the example to retrieve + + Returns: + tuple: A 3 or 4-tuple with ``(img_left, img_right, disparity, Optional[valid_mask])`` where ``valid_mask`` + can be a numpy boolean mask of shape (H, W) if the dataset provides a file + indicating which disparity pixels are valid. The disparity is a numpy array of + shape (1, H, W) and the images are PIL images. ``disparity`` is None for + datasets on which for ``split="test"`` the authors did not provide annotations. + """ + img_left = self._read_img(self._images[index][0]) + img_right = self._read_img(self._images[index][1]) + + dsp_map_left, valid_mask_left = self._read_disparity(self._disparities[index][0]) + dsp_map_right, valid_mask_right = self._read_disparity(self._disparities[index][1]) + + imgs = (img_left, img_right) + dsp_maps = (dsp_map_left, dsp_map_right) + valid_masks = (valid_mask_left, valid_mask_right) + + if self.transforms is not None: + ( + imgs, + dsp_maps, + valid_masks, + ) = self.transforms(imgs, dsp_maps, valid_masks) + + if self._has_built_in_disparity_mask or valid_masks[0] is not None: + return imgs[0], imgs[1], dsp_maps[0], cast(np.ndarray, valid_masks[0]) + else: + return imgs[0], imgs[1], dsp_maps[0] + + def __len__(self) -> int: + return len(self._images) + + +class CarlaStereo(StereoMatchingDataset): + """ + Carla simulator data linked in the `CREStereo github repo `_. + + The dataset is expected to have the following structure: :: + + root + carla-highres + trainingF + scene1 + img0.png + img1.png + disp0GT.pfm + disp1GT.pfm + calib.txt + scene2 + img0.png + img1.png + disp0GT.pfm + disp1GT.pfm + calib.txt + ... + + Args: + root (str or ``pathlib.Path``): Root directory where `carla-highres` is located. + transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version. + """ + + def __init__(self, root: Union[str, Path], transforms: Optional[Callable] = None) -> None: + super().__init__(root, transforms) + + root = Path(root) / "carla-highres" + + left_image_pattern = str(root / "trainingF" / "*" / "im0.png") + right_image_pattern = str(root / "trainingF" / "*" / "im1.png") + imgs = self._scan_pairs(left_image_pattern, right_image_pattern) + self._images = imgs + + left_disparity_pattern = str(root / "trainingF" / "*" / "disp0GT.pfm") + right_disparity_pattern = str(root / "trainingF" / "*" / "disp1GT.pfm") + disparities = self._scan_pairs(left_disparity_pattern, right_disparity_pattern) + self._disparities = disparities + + def _read_disparity(self, file_path: str) -> Tuple[np.ndarray, None]: + disparity_map = _read_pfm_file(file_path) + disparity_map = np.abs(disparity_map) # ensure that the disparity is positive + valid_mask = None + return disparity_map, valid_mask + + def __getitem__(self, index: int) -> T1: + """Return example at given index. + + Args: + index(int): The index of the example to retrieve + + Returns: + tuple: A 3-tuple with ``(img_left, img_right, disparity)``. + The disparity is a numpy array of shape (1, H, W) and the images are PIL images. + If a ``valid_mask`` is generated within the ``transforms`` parameter, + a 4-tuple with ``(img_left, img_right, disparity, valid_mask)`` is returned. + """ + return cast(T1, super().__getitem__(index)) + + +class Kitti2012Stereo(StereoMatchingDataset): + """ + KITTI dataset from the `2012 stereo evaluation benchmark `_. + Uses the RGB images for consistency with KITTI 2015. + + The dataset is expected to have the following structure: :: + + root + Kitti2012 + testing + colored_0 + 1_10.png + 2_10.png + ... + colored_1 + 1_10.png + 2_10.png + ... + training + colored_0 + 1_10.png + 2_10.png + ... + colored_1 + 1_10.png + 2_10.png + ... + disp_noc + 1.png + 2.png + ... + calib + + Args: + root (str or ``pathlib.Path``): Root directory where `Kitti2012` is located. + split (string, optional): The dataset split of scenes, either "train" (default) or "test". + transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version. + """ + + _has_built_in_disparity_mask = True + + def __init__(self, root: Union[str, Path], split: str = "train", transforms: Optional[Callable] = None) -> None: + super().__init__(root, transforms) + + verify_str_arg(split, "split", valid_values=("train", "test")) + + root = Path(root) / "Kitti2012" / (split + "ing") + + left_img_pattern = str(root / "colored_0" / "*_10.png") + right_img_pattern = str(root / "colored_1" / "*_10.png") + self._images = self._scan_pairs(left_img_pattern, right_img_pattern) + + if split == "train": + disparity_pattern = str(root / "disp_noc" / "*.png") + self._disparities = self._scan_pairs(disparity_pattern, None) + else: + self._disparities = list((None, None) for _ in self._images) + + def _read_disparity(self, file_path: str) -> Tuple[Optional[np.ndarray], None]: + # test split has no disparity maps + if file_path is None: + return None, None + + disparity_map = np.asarray(Image.open(file_path)) / 256.0 + # unsqueeze the disparity map into (C, H, W) format + disparity_map = disparity_map[None, :, :] + valid_mask = None + return disparity_map, valid_mask + + def __getitem__(self, index: int) -> T1: + """Return example at given index. + + Args: + index(int): The index of the example to retrieve + + Returns: + tuple: A 4-tuple with ``(img_left, img_right, disparity, valid_mask)``. + The disparity is a numpy array of shape (1, H, W) and the images are PIL images. + ``valid_mask`` is implicitly ``None`` if the ``transforms`` parameter does not + generate a valid mask. + Both ``disparity`` and ``valid_mask`` are ``None`` if the dataset split is test. + """ + return cast(T1, super().__getitem__(index)) + + +class Kitti2015Stereo(StereoMatchingDataset): + """ + KITTI dataset from the `2015 stereo evaluation benchmark `_. + + The dataset is expected to have the following structure: :: + + root + Kitti2015 + testing + image_2 + img1.png + img2.png + ... + image_3 + img1.png + img2.png + ... + training + image_2 + img1.png + img2.png + ... + image_3 + img1.png + img2.png + ... + disp_occ_0 + img1.png + img2.png + ... + disp_occ_1 + img1.png + img2.png + ... + calib + + Args: + root (str or ``pathlib.Path``): Root directory where `Kitti2015` is located. + split (string, optional): The dataset split of scenes, either "train" (default) or "test". + transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version. + """ + + _has_built_in_disparity_mask = True + + def __init__(self, root: Union[str, Path], split: str = "train", transforms: Optional[Callable] = None) -> None: + super().__init__(root, transforms) + + verify_str_arg(split, "split", valid_values=("train", "test")) + + root = Path(root) / "Kitti2015" / (split + "ing") + left_img_pattern = str(root / "image_2" / "*.png") + right_img_pattern = str(root / "image_3" / "*.png") + self._images = self._scan_pairs(left_img_pattern, right_img_pattern) + + if split == "train": + left_disparity_pattern = str(root / "disp_occ_0" / "*.png") + right_disparity_pattern = str(root / "disp_occ_1" / "*.png") + self._disparities = self._scan_pairs(left_disparity_pattern, right_disparity_pattern) + else: + self._disparities = list((None, None) for _ in self._images) + + def _read_disparity(self, file_path: str) -> Tuple[Optional[np.ndarray], None]: + # test split has no disparity maps + if file_path is None: + return None, None + + disparity_map = np.asarray(Image.open(file_path)) / 256.0 + # unsqueeze the disparity map into (C, H, W) format + disparity_map = disparity_map[None, :, :] + valid_mask = None + return disparity_map, valid_mask + + def __getitem__(self, index: int) -> T1: + """Return example at given index. + + Args: + index(int): The index of the example to retrieve + + Returns: + tuple: A 4-tuple with ``(img_left, img_right, disparity, valid_mask)``. + The disparity is a numpy array of shape (1, H, W) and the images are PIL images. + ``valid_mask`` is implicitly ``None`` if the ``transforms`` parameter does not + generate a valid mask. + Both ``disparity`` and ``valid_mask`` are ``None`` if the dataset split is test. + """ + return cast(T1, super().__getitem__(index)) + + +class Middlebury2014Stereo(StereoMatchingDataset): + """Publicly available scenes from the Middlebury dataset `2014 version `. + + The dataset mostly follows the original format, without containing the ambient subdirectories. : :: + + root + Middlebury2014 + train + scene1-{perfect,imperfect} + calib.txt + im{0,1}.png + im1E.png + im1L.png + disp{0,1}.pfm + disp{0,1}-n.png + disp{0,1}-sd.pfm + disp{0,1}y.pfm + scene2-{perfect,imperfect} + calib.txt + im{0,1}.png + im1E.png + im1L.png + disp{0,1}.pfm + disp{0,1}-n.png + disp{0,1}-sd.pfm + disp{0,1}y.pfm + ... + additional + scene1-{perfect,imperfect} + calib.txt + im{0,1}.png + im1E.png + im1L.png + disp{0,1}.pfm + disp{0,1}-n.png + disp{0,1}-sd.pfm + disp{0,1}y.pfm + ... + test + scene1 + calib.txt + im{0,1}.png + scene2 + calib.txt + im{0,1}.png + ... + + Args: + root (str or ``pathlib.Path``): Root directory of the Middleburry 2014 Dataset. + split (string, optional): The dataset split of scenes, either "train" (default), "test", or "additional" + use_ambient_views (boolean, optional): Whether to use different expose or lightning views when possible. + The dataset samples with equal probability between ``[im1.png, im1E.png, im1L.png]``. + calibration (string, optional): Whether or not to use the calibrated (default) or uncalibrated scenes. + transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version. + download (boolean, optional): Whether or not to download the dataset in the ``root`` directory. + """ + + splits = { + "train": [ + "Adirondack", + "Jadeplant", + "Motorcycle", + "Piano", + "Pipes", + "Playroom", + "Playtable", + "Recycle", + "Shelves", + "Vintage", + ], + "additional": [ + "Backpack", + "Bicycle1", + "Cable", + "Classroom1", + "Couch", + "Flowers", + "Mask", + "Shopvac", + "Sticks", + "Storage", + "Sword1", + "Sword2", + "Umbrella", + ], + "test": [ + "Plants", + "Classroom2E", + "Classroom2", + "Australia", + "DjembeL", + "CrusadeP", + "Crusade", + "Hoops", + "Bicycle2", + "Staircase", + "Newkuba", + "AustraliaP", + "Djembe", + "Livingroom", + "Computer", + ], + } + + _has_built_in_disparity_mask = True + + def __init__( + self, + root: Union[str, Path], + split: str = "train", + calibration: Optional[str] = "perfect", + use_ambient_views: bool = False, + transforms: Optional[Callable] = None, + download: bool = False, + ) -> None: + super().__init__(root, transforms) + + verify_str_arg(split, "split", valid_values=("train", "test", "additional")) + self.split = split + + if calibration: + verify_str_arg(calibration, "calibration", valid_values=("perfect", "imperfect", "both", None)) # type: ignore + if split == "test": + raise ValueError("Split 'test' has only no calibration settings, please set `calibration=None`.") + else: + if split != "test": + raise ValueError( + f"Split '{split}' has calibration settings, however None was provided as an argument." + f"\nSetting calibration to 'perfect' for split '{split}'. Available calibration settings are: 'perfect', 'imperfect', 'both'.", + ) + + if download: + self._download_dataset(root) + + root = Path(root) / "Middlebury2014" + + if not os.path.exists(root / split): + raise FileNotFoundError(f"The {split} directory was not found in the provided root directory") + + split_scenes = self.splits[split] + # check that the provided root folder contains the scene splits + if not any( + # using startswith to account for perfect / imperfect calibrartion + scene.startswith(s) + for scene in os.listdir(root / split) + for s in split_scenes + ): + raise FileNotFoundError(f"Provided root folder does not contain any scenes from the {split} split.") + + calibrartion_suffixes = { + None: [""], + "perfect": ["-perfect"], + "imperfect": ["-imperfect"], + "both": ["-perfect", "-imperfect"], + }[calibration] + + for calibration_suffix in calibrartion_suffixes: + scene_pattern = "*" + calibration_suffix + left_img_pattern = str(root / split / scene_pattern / "im0.png") + right_img_pattern = str(root / split / scene_pattern / "im1.png") + self._images += self._scan_pairs(left_img_pattern, right_img_pattern) + + if split == "test": + self._disparities = list((None, None) for _ in self._images) + else: + left_dispartity_pattern = str(root / split / scene_pattern / "disp0.pfm") + right_dispartity_pattern = str(root / split / scene_pattern / "disp1.pfm") + self._disparities += self._scan_pairs(left_dispartity_pattern, right_dispartity_pattern) + + self.use_ambient_views = use_ambient_views + + def _read_img(self, file_path: Union[str, Path]) -> Image.Image: + """ + Function that reads either the original right image or an augmented view when ``use_ambient_views`` is True. + When ``use_ambient_views`` is True, the dataset will return at random one of ``[im1.png, im1E.png, im1L.png]`` + as the right image. + """ + ambient_file_paths: List[Union[str, Path]] # make mypy happy + + if not isinstance(file_path, Path): + file_path = Path(file_path) + + if file_path.name == "im1.png" and self.use_ambient_views: + base_path = file_path.parent + # initialize sampleable container + ambient_file_paths = list(base_path / view_name for view_name in ["im1E.png", "im1L.png"]) + # double check that we're not going to try to read from an invalid file path + ambient_file_paths = list(filter(lambda p: os.path.exists(p), ambient_file_paths)) + # keep the original image as an option as well for uniform sampling between base views + ambient_file_paths.append(file_path) + file_path = random.choice(ambient_file_paths) # type: ignore + return super()._read_img(file_path) + + def _read_disparity(self, file_path: str) -> Union[Tuple[None, None], Tuple[np.ndarray, np.ndarray]]: + # test split has not disparity maps + if file_path is None: + return None, None + + disparity_map = _read_pfm_file(file_path) + disparity_map = np.abs(disparity_map) # ensure that the disparity is positive + disparity_map[disparity_map == np.inf] = 0 # remove infinite disparities + valid_mask = (disparity_map > 0).squeeze(0) # mask out invalid disparities + return disparity_map, valid_mask + + def _download_dataset(self, root: Union[str, Path]) -> None: + base_url = "https://vision.middlebury.edu/stereo/data/scenes2014/zip" + # train and additional splits have 2 different calibration settings + root = Path(root) / "Middlebury2014" + split_name = self.split + + if split_name != "test": + for split_scene in self.splits[split_name]: + split_root = root / split_name + for calibration in ["perfect", "imperfect"]: + scene_name = f"{split_scene}-{calibration}" + scene_url = f"{base_url}/{scene_name}.zip" + # download the scene only if it doesn't exist + if not (split_root / scene_name).exists(): + download_and_extract_archive( + url=scene_url, + filename=f"{scene_name}.zip", + download_root=str(split_root), + remove_finished=True, + ) + else: + os.makedirs(root / "test") + if any(s not in os.listdir(root / "test") for s in self.splits["test"]): + # test split is downloaded from a different location + test_set_url = "https://vision.middlebury.edu/stereo/submit3/zip/MiddEval3-data-F.zip" + # the unzip is going to produce a directory MiddEval3 with two subdirectories trainingF and testF + # we want to move the contents from testF into the directory + download_and_extract_archive(url=test_set_url, download_root=str(root), remove_finished=True) + for scene_dir, scene_names, _ in os.walk(str(root / "MiddEval3/testF")): + for scene in scene_names: + scene_dst_dir = root / "test" + scene_src_dir = Path(scene_dir) / scene + os.makedirs(scene_dst_dir, exist_ok=True) + shutil.move(str(scene_src_dir), str(scene_dst_dir)) + + # cleanup MiddEval3 directory + shutil.rmtree(str(root / "MiddEval3")) + + def __getitem__(self, index: int) -> T2: + """Return example at given index. + + Args: + index(int): The index of the example to retrieve + + Returns: + tuple: A 4-tuple with ``(img_left, img_right, disparity, valid_mask)``. + The disparity is a numpy array of shape (1, H, W) and the images are PIL images. + ``valid_mask`` is implicitly ``None`` for `split=test`. + """ + return cast(T2, super().__getitem__(index)) + + +class CREStereo(StereoMatchingDataset): + """Synthetic dataset used in training the `CREStereo `_ architecture. + Dataset details on the official paper `repo `_. + + The dataset is expected to have the following structure: :: + + root + CREStereo + tree + img1_left.jpg + img1_right.jpg + img1_left.disp.jpg + img1_right.disp.jpg + img2_left.jpg + img2_right.jpg + img2_left.disp.jpg + img2_right.disp.jpg + ... + shapenet + img1_left.jpg + img1_right.jpg + img1_left.disp.jpg + img1_right.disp.jpg + ... + reflective + img1_left.jpg + img1_right.jpg + img1_left.disp.jpg + img1_right.disp.jpg + ... + hole + img1_left.jpg + img1_right.jpg + img1_left.disp.jpg + img1_right.disp.jpg + ... + + Args: + root (str): Root directory of the dataset. + transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version. + """ + + _has_built_in_disparity_mask = True + + def __init__( + self, + root: Union[str, Path], + transforms: Optional[Callable] = None, + ) -> None: + super().__init__(root, transforms) + + root = Path(root) / "CREStereo" + + dirs = ["shapenet", "reflective", "tree", "hole"] + + for s in dirs: + left_image_pattern = str(root / s / "*_left.jpg") + right_image_pattern = str(root / s / "*_right.jpg") + imgs = self._scan_pairs(left_image_pattern, right_image_pattern) + self._images += imgs + + left_disparity_pattern = str(root / s / "*_left.disp.png") + right_disparity_pattern = str(root / s / "*_right.disp.png") + disparities = self._scan_pairs(left_disparity_pattern, right_disparity_pattern) + self._disparities += disparities + + def _read_disparity(self, file_path: str) -> Tuple[np.ndarray, None]: + disparity_map = np.asarray(Image.open(file_path), dtype=np.float32) + # unsqueeze the disparity map into (C, H, W) format + disparity_map = disparity_map[None, :, :] / 32.0 + valid_mask = None + return disparity_map, valid_mask + + def __getitem__(self, index: int) -> T1: + """Return example at given index. + + Args: + index(int): The index of the example to retrieve + + Returns: + tuple: A 4-tuple with ``(img_left, img_right, disparity, valid_mask)``. + The disparity is a numpy array of shape (1, H, W) and the images are PIL images. + ``valid_mask`` is implicitly ``None`` if the ``transforms`` parameter does not + generate a valid mask. + """ + return cast(T1, super().__getitem__(index)) + + +class FallingThingsStereo(StereoMatchingDataset): + """`FallingThings `_ dataset. + + The dataset is expected to have the following structure: :: + + root + FallingThings + single + dir1 + scene1 + _object_settings.json + _camera_settings.json + image1.left.depth.png + image1.right.depth.png + image1.left.jpg + image1.right.jpg + image2.left.depth.png + image2.right.depth.png + image2.left.jpg + image2.right + ... + scene2 + ... + mixed + scene1 + _object_settings.json + _camera_settings.json + image1.left.depth.png + image1.right.depth.png + image1.left.jpg + image1.right.jpg + image2.left.depth.png + image2.right.depth.png + image2.left.jpg + image2.right + ... + scene2 + ... + + Args: + root (str or ``pathlib.Path``): Root directory where FallingThings is located. + variant (string): Which variant to use. Either "single", "mixed", or "both". + transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version. + """ + + def __init__(self, root: Union[str, Path], variant: str = "single", transforms: Optional[Callable] = None) -> None: + super().__init__(root, transforms) + + root = Path(root) / "FallingThings" + + verify_str_arg(variant, "variant", valid_values=("single", "mixed", "both")) + + variants = { + "single": ["single"], + "mixed": ["mixed"], + "both": ["single", "mixed"], + }[variant] + + split_prefix = { + "single": Path("*") / "*", + "mixed": Path("*"), + } + + for s in variants: + left_img_pattern = str(root / s / split_prefix[s] / "*.left.jpg") + right_img_pattern = str(root / s / split_prefix[s] / "*.right.jpg") + self._images += self._scan_pairs(left_img_pattern, right_img_pattern) + + left_disparity_pattern = str(root / s / split_prefix[s] / "*.left.depth.png") + right_disparity_pattern = str(root / s / split_prefix[s] / "*.right.depth.png") + self._disparities += self._scan_pairs(left_disparity_pattern, right_disparity_pattern) + + def _read_disparity(self, file_path: str) -> Tuple[np.ndarray, None]: + # (H, W) image + depth = np.asarray(Image.open(file_path)) + # as per https://research.nvidia.com/sites/default/files/pubs/2018-06_Falling-Things/readme_0.txt + # in order to extract disparity from depth maps + camera_settings_path = Path(file_path).parent / "_camera_settings.json" + with open(camera_settings_path, "r") as f: + # inverse of depth-from-disparity equation: depth = (baseline * focal) / (disparity * pixel_constant) + intrinsics = json.load(f) + focal = intrinsics["camera_settings"][0]["intrinsic_settings"]["fx"] + baseline, pixel_constant = 6, 100 # pixel constant is inverted + disparity_map = (baseline * focal * pixel_constant) / depth.astype(np.float32) + # unsqueeze disparity to (C, H, W) + disparity_map = disparity_map[None, :, :] + valid_mask = None + return disparity_map, valid_mask + + def __getitem__(self, index: int) -> T1: + """Return example at given index. + + Args: + index(int): The index of the example to retrieve + + Returns: + tuple: A 3-tuple with ``(img_left, img_right, disparity)``. + The disparity is a numpy array of shape (1, H, W) and the images are PIL images. + If a ``valid_mask`` is generated within the ``transforms`` parameter, + a 4-tuple with ``(img_left, img_right, disparity, valid_mask)`` is returned. + """ + return cast(T1, super().__getitem__(index)) + + +class SceneFlowStereo(StereoMatchingDataset): + """Dataset interface for `Scene Flow `_ datasets. + This interface provides access to the `FlyingThings3D, `Monkaa` and `Driving` datasets. + + The dataset is expected to have the following structure: :: + + root + SceneFlow + Monkaa + frames_cleanpass + scene1 + left + img1.png + img2.png + right + img1.png + img2.png + scene2 + left + img1.png + img2.png + right + img1.png + img2.png + frames_finalpass + scene1 + left + img1.png + img2.png + right + img1.png + img2.png + ... + ... + disparity + scene1 + left + img1.pfm + img2.pfm + right + img1.pfm + img2.pfm + FlyingThings3D + ... + ... + + Args: + root (str or ``pathlib.Path``): Root directory where SceneFlow is located. + variant (string): Which dataset variant to user, "FlyingThings3D" (default), "Monkaa" or "Driving". + pass_name (string): Which pass to use, "clean" (default), "final" or "both". + transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version. + + """ + + def __init__( + self, + root: Union[str, Path], + variant: str = "FlyingThings3D", + pass_name: str = "clean", + transforms: Optional[Callable] = None, + ) -> None: + super().__init__(root, transforms) + + root = Path(root) / "SceneFlow" + + verify_str_arg(variant, "variant", valid_values=("FlyingThings3D", "Driving", "Monkaa")) + verify_str_arg(pass_name, "pass_name", valid_values=("clean", "final", "both")) + + passes = { + "clean": ["frames_cleanpass"], + "final": ["frames_finalpass"], + "both": ["frames_cleanpass", "frames_finalpass"], + }[pass_name] + + root = root / variant + + prefix_directories = { + "Monkaa": Path("*"), + "FlyingThings3D": Path("*") / "*" / "*", + "Driving": Path("*") / "*" / "*", + } + + for p in passes: + left_image_pattern = str(root / p / prefix_directories[variant] / "left" / "*.png") + right_image_pattern = str(root / p / prefix_directories[variant] / "right" / "*.png") + self._images += self._scan_pairs(left_image_pattern, right_image_pattern) + + left_disparity_pattern = str(root / "disparity" / prefix_directories[variant] / "left" / "*.pfm") + right_disparity_pattern = str(root / "disparity" / prefix_directories[variant] / "right" / "*.pfm") + self._disparities += self._scan_pairs(left_disparity_pattern, right_disparity_pattern) + + def _read_disparity(self, file_path: str) -> Tuple[np.ndarray, None]: + disparity_map = _read_pfm_file(file_path) + disparity_map = np.abs(disparity_map) # ensure that the disparity is positive + valid_mask = None + return disparity_map, valid_mask + + def __getitem__(self, index: int) -> T1: + """Return example at given index. + + Args: + index(int): The index of the example to retrieve + + Returns: + tuple: A 3-tuple with ``(img_left, img_right, disparity)``. + The disparity is a numpy array of shape (1, H, W) and the images are PIL images. + If a ``valid_mask`` is generated within the ``transforms`` parameter, + a 4-tuple with ``(img_left, img_right, disparity, valid_mask)`` is returned. + """ + return cast(T1, super().__getitem__(index)) + + +class SintelStereo(StereoMatchingDataset): + """Sintel `Stereo Dataset `_. + + The dataset is expected to have the following structure: :: + + root + Sintel + training + final_left + scene1 + img1.png + img2.png + ... + ... + final_right + scene2 + img1.png + img2.png + ... + ... + disparities + scene1 + img1.png + img2.png + ... + ... + occlusions + scene1 + img1.png + img2.png + ... + ... + outofframe + scene1 + img1.png + img2.png + ... + ... + + Args: + root (str or ``pathlib.Path``): Root directory where Sintel Stereo is located. + pass_name (string): The name of the pass to use, either "final", "clean" or "both". + transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version. + """ + + _has_built_in_disparity_mask = True + + def __init__(self, root: Union[str, Path], pass_name: str = "final", transforms: Optional[Callable] = None) -> None: + super().__init__(root, transforms) + + verify_str_arg(pass_name, "pass_name", valid_values=("final", "clean", "both")) + + root = Path(root) / "Sintel" + pass_names = { + "final": ["final"], + "clean": ["clean"], + "both": ["final", "clean"], + }[pass_name] + + for p in pass_names: + left_img_pattern = str(root / "training" / f"{p}_left" / "*" / "*.png") + right_img_pattern = str(root / "training" / f"{p}_right" / "*" / "*.png") + self._images += self._scan_pairs(left_img_pattern, right_img_pattern) + + disparity_pattern = str(root / "training" / "disparities" / "*" / "*.png") + self._disparities += self._scan_pairs(disparity_pattern, None) + + def _get_occlussion_mask_paths(self, file_path: str) -> Tuple[str, str]: + # helper function to get the occlusion mask paths + # a path will look like .../.../.../training/disparities/scene1/img1.png + # we want to get something like .../.../.../training/occlusions/scene1/img1.png + fpath = Path(file_path) + basename = fpath.name + scenedir = fpath.parent + # the parent of the scenedir is actually the disparity dir + sampledir = scenedir.parent.parent + + occlusion_path = str(sampledir / "occlusions" / scenedir.name / basename) + outofframe_path = str(sampledir / "outofframe" / scenedir.name / basename) + + if not os.path.exists(occlusion_path): + raise FileNotFoundError(f"Occlusion mask {occlusion_path} does not exist") + + if not os.path.exists(outofframe_path): + raise FileNotFoundError(f"Out of frame mask {outofframe_path} does not exist") + + return occlusion_path, outofframe_path + + def _read_disparity(self, file_path: str) -> Union[Tuple[None, None], Tuple[np.ndarray, np.ndarray]]: + if file_path is None: + return None, None + + # disparity decoding as per Sintel instructions in the README provided with the dataset + disparity_map = np.asarray(Image.open(file_path), dtype=np.float32) + r, g, b = np.split(disparity_map, 3, axis=-1) + disparity_map = r * 4 + g / (2**6) + b / (2**14) + # reshape into (C, H, W) format + disparity_map = np.transpose(disparity_map, (2, 0, 1)) + # find the appropriate file paths + occlued_mask_path, out_of_frame_mask_path = self._get_occlussion_mask_paths(file_path) + # occlusion masks + valid_mask = np.asarray(Image.open(occlued_mask_path)) == 0 + # out of frame masks + off_mask = np.asarray(Image.open(out_of_frame_mask_path)) == 0 + # combine the masks together + valid_mask = np.logical_and(off_mask, valid_mask) + return disparity_map, valid_mask + + def __getitem__(self, index: int) -> T2: + """Return example at given index. + + Args: + index(int): The index of the example to retrieve + + Returns: + tuple: A 4-tuple with ``(img_left, img_right, disparity, valid_mask)`` is returned. + The disparity is a numpy array of shape (1, H, W) and the images are PIL images whilst + the valid_mask is a numpy array of shape (H, W). + """ + return cast(T2, super().__getitem__(index)) + + +class InStereo2k(StereoMatchingDataset): + """`InStereo2k `_ dataset. + + The dataset is expected to have the following structure: :: + + root + InStereo2k + train + scene1 + left.png + right.png + left_disp.png + right_disp.png + ... + scene2 + ... + test + scene1 + left.png + right.png + left_disp.png + right_disp.png + ... + scene2 + ... + + Args: + root (str or ``pathlib.Path``): Root directory where InStereo2k is located. + split (string): Either "train" or "test". + transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version. + """ + + def __init__(self, root: Union[str, Path], split: str = "train", transforms: Optional[Callable] = None) -> None: + super().__init__(root, transforms) + + root = Path(root) / "InStereo2k" / split + + verify_str_arg(split, "split", valid_values=("train", "test")) + + left_img_pattern = str(root / "*" / "left.png") + right_img_pattern = str(root / "*" / "right.png") + self._images = self._scan_pairs(left_img_pattern, right_img_pattern) + + left_disparity_pattern = str(root / "*" / "left_disp.png") + right_disparity_pattern = str(root / "*" / "right_disp.png") + self._disparities = self._scan_pairs(left_disparity_pattern, right_disparity_pattern) + + def _read_disparity(self, file_path: str) -> Tuple[np.ndarray, None]: + disparity_map = np.asarray(Image.open(file_path), dtype=np.float32) + # unsqueeze disparity to (C, H, W) + disparity_map = disparity_map[None, :, :] / 1024.0 + valid_mask = None + return disparity_map, valid_mask + + def __getitem__(self, index: int) -> T1: + """Return example at given index. + + Args: + index(int): The index of the example to retrieve + + Returns: + tuple: A 3-tuple with ``(img_left, img_right, disparity)``. + The disparity is a numpy array of shape (1, H, W) and the images are PIL images. + If a ``valid_mask`` is generated within the ``transforms`` parameter, + a 4-tuple with ``(img_left, img_right, disparity, valid_mask)`` is returned. + """ + return cast(T1, super().__getitem__(index)) + + +class ETH3DStereo(StereoMatchingDataset): + """ETH3D `Low-Res Two-View `_ dataset. + + The dataset is expected to have the following structure: :: + + root + ETH3D + two_view_training + scene1 + im1.png + im0.png + images.txt + cameras.txt + calib.txt + scene2 + im1.png + im0.png + images.txt + cameras.txt + calib.txt + ... + two_view_training_gt + scene1 + disp0GT.pfm + mask0nocc.png + scene2 + disp0GT.pfm + mask0nocc.png + ... + two_view_testing + scene1 + im1.png + im0.png + images.txt + cameras.txt + calib.txt + scene2 + im1.png + im0.png + images.txt + cameras.txt + calib.txt + ... + + Args: + root (str or ``pathlib.Path``): Root directory of the ETH3D Dataset. + split (string, optional): The dataset split of scenes, either "train" (default) or "test". + transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version. + """ + + _has_built_in_disparity_mask = True + + def __init__(self, root: Union[str, Path], split: str = "train", transforms: Optional[Callable] = None) -> None: + super().__init__(root, transforms) + + verify_str_arg(split, "split", valid_values=("train", "test")) + + root = Path(root) / "ETH3D" + + img_dir = "two_view_training" if split == "train" else "two_view_test" + anot_dir = "two_view_training_gt" + + left_img_pattern = str(root / img_dir / "*" / "im0.png") + right_img_pattern = str(root / img_dir / "*" / "im1.png") + self._images = self._scan_pairs(left_img_pattern, right_img_pattern) + + if split == "test": + self._disparities = list((None, None) for _ in self._images) + else: + disparity_pattern = str(root / anot_dir / "*" / "disp0GT.pfm") + self._disparities = self._scan_pairs(disparity_pattern, None) + + def _read_disparity(self, file_path: str) -> Union[Tuple[None, None], Tuple[np.ndarray, np.ndarray]]: + # test split has no disparity maps + if file_path is None: + return None, None + + disparity_map = _read_pfm_file(file_path) + disparity_map = np.abs(disparity_map) # ensure that the disparity is positive + mask_path = Path(file_path).parent / "mask0nocc.png" + valid_mask = Image.open(mask_path) + valid_mask = np.asarray(valid_mask).astype(bool) + return disparity_map, valid_mask + + def __getitem__(self, index: int) -> T2: + """Return example at given index. + + Args: + index(int): The index of the example to retrieve + + Returns: + tuple: A 4-tuple with ``(img_left, img_right, disparity, valid_mask)``. + The disparity is a numpy array of shape (1, H, W) and the images are PIL images. + ``valid_mask`` is implicitly ``None`` if the ``transforms`` parameter does not + generate a valid mask. + Both ``disparity`` and ``valid_mask`` are ``None`` if the dataset split is test. + """ + return cast(T2, super().__getitem__(index)) diff --git a/lib/python3.12/site-packages/torchvision/datasets/caltech.py b/lib/python3.12/site-packages/torchvision/datasets/caltech.py new file mode 100644 index 0000000000000000000000000000000000000000..b152c425c77c78e24eb6f17772fc54ed31329e6e --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/caltech.py @@ -0,0 +1,240 @@ +import os +import os.path +from pathlib import Path +from typing import Any, Callable, List, Optional, Tuple, Union + +from PIL import Image + +from .utils import download_and_extract_archive, verify_str_arg +from .vision import VisionDataset + + +class Caltech101(VisionDataset): + """`Caltech 101 `_ Dataset. + + .. warning:: + + This class needs `scipy `_ to load target files from `.mat` format. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset where directory + ``caltech101`` exists or will be saved to if download is set to True. + target_type (string or list, optional): Type of target to use, ``category`` or + ``annotation``. Can also be a list to output a tuple with all specified + target types. ``category`` represents the target class, and + ``annotation`` is a list of points from a hand-generated outline. + Defaults to ``category``. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + + .. warning:: + + To download the dataset `gdown `_ is required. + """ + + def __init__( + self, + root: Union[str, Path], + target_type: Union[List[str], str] = "category", + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + super().__init__(os.path.join(root, "caltech101"), transform=transform, target_transform=target_transform) + os.makedirs(self.root, exist_ok=True) + if isinstance(target_type, str): + target_type = [target_type] + self.target_type = [verify_str_arg(t, "target_type", ("category", "annotation")) for t in target_type] + + if download: + self.download() + + if not self._check_integrity(): + raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") + + self.categories = sorted(os.listdir(os.path.join(self.root, "101_ObjectCategories"))) + self.categories.remove("BACKGROUND_Google") # this is not a real class + + # For some reason, the category names in "101_ObjectCategories" and + # "Annotations" do not always match. This is a manual map between the + # two. Defaults to using same name, since most names are fine. + name_map = { + "Faces": "Faces_2", + "Faces_easy": "Faces_3", + "Motorbikes": "Motorbikes_16", + "airplanes": "Airplanes_Side_2", + } + self.annotation_categories = list(map(lambda x: name_map[x] if x in name_map else x, self.categories)) + + self.index: List[int] = [] + self.y = [] + for (i, c) in enumerate(self.categories): + n = len(os.listdir(os.path.join(self.root, "101_ObjectCategories", c))) + self.index.extend(range(1, n + 1)) + self.y.extend(n * [i]) + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: (image, target) where the type of target specified by target_type. + """ + import scipy.io + + img = Image.open( + os.path.join( + self.root, + "101_ObjectCategories", + self.categories[self.y[index]], + f"image_{self.index[index]:04d}.jpg", + ) + ) + + target: Any = [] + for t in self.target_type: + if t == "category": + target.append(self.y[index]) + elif t == "annotation": + data = scipy.io.loadmat( + os.path.join( + self.root, + "Annotations", + self.annotation_categories[self.y[index]], + f"annotation_{self.index[index]:04d}.mat", + ) + ) + target.append(data["obj_contour"]) + target = tuple(target) if len(target) > 1 else target[0] + + if self.transform is not None: + img = self.transform(img) + + if self.target_transform is not None: + target = self.target_transform(target) + + return img, target + + def _check_integrity(self) -> bool: + # can be more robust and check hash of files + return os.path.exists(os.path.join(self.root, "101_ObjectCategories")) + + def __len__(self) -> int: + return len(self.index) + + def download(self) -> None: + if self._check_integrity(): + return + + download_and_extract_archive( + "https://drive.google.com/file/d/137RyRjvTBkBiIfeYBNZBtViDHQ6_Ewsp", + self.root, + filename="101_ObjectCategories.tar.gz", + md5="b224c7392d521a49829488ab0f1120d9", + ) + download_and_extract_archive( + "https://drive.google.com/file/d/175kQy3UsZ0wUEHZjqkUDdNVssr7bgh_m", + self.root, + filename="Annotations.tar", + md5="6f83eeb1f24d99cab4eb377263132c91", + ) + + def extra_repr(self) -> str: + return "Target type: {target_type}".format(**self.__dict__) + + +class Caltech256(VisionDataset): + """`Caltech 256 `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset where directory + ``caltech256`` exists or will be saved to if download is set to True. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + """ + + def __init__( + self, + root: str, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + super().__init__(os.path.join(root, "caltech256"), transform=transform, target_transform=target_transform) + os.makedirs(self.root, exist_ok=True) + + if download: + self.download() + + if not self._check_integrity(): + raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") + + self.categories = sorted(os.listdir(os.path.join(self.root, "256_ObjectCategories"))) + self.index: List[int] = [] + self.y = [] + for (i, c) in enumerate(self.categories): + n = len( + [ + item + for item in os.listdir(os.path.join(self.root, "256_ObjectCategories", c)) + if item.endswith(".jpg") + ] + ) + self.index.extend(range(1, n + 1)) + self.y.extend(n * [i]) + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: (image, target) where target is index of the target class. + """ + img = Image.open( + os.path.join( + self.root, + "256_ObjectCategories", + self.categories[self.y[index]], + f"{self.y[index] + 1:03d}_{self.index[index]:04d}.jpg", + ) + ) + + target = self.y[index] + + if self.transform is not None: + img = self.transform(img) + + if self.target_transform is not None: + target = self.target_transform(target) + + return img, target + + def _check_integrity(self) -> bool: + # can be more robust and check hash of files + return os.path.exists(os.path.join(self.root, "256_ObjectCategories")) + + def __len__(self) -> int: + return len(self.index) + + def download(self) -> None: + if self._check_integrity(): + return + + download_and_extract_archive( + "https://drive.google.com/file/d/1r6o0pSROcV1_VwT4oSjA2FBUSCWGuxLK", + self.root, + filename="256_ObjectCategories.tar", + md5="67b4f42ca05d46448c6bb8ecd2220f6d", + ) diff --git a/lib/python3.12/site-packages/torchvision/datasets/celeba.py b/lib/python3.12/site-packages/torchvision/datasets/celeba.py new file mode 100644 index 0000000000000000000000000000000000000000..c15120af5a55a88ba215ae255b643e42577b51a4 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/celeba.py @@ -0,0 +1,193 @@ +import csv +import os +from collections import namedtuple +from pathlib import Path +from typing import Any, Callable, List, Optional, Tuple, Union + +import PIL +import torch + +from .utils import check_integrity, download_file_from_google_drive, extract_archive, verify_str_arg +from .vision import VisionDataset + +CSV = namedtuple("CSV", ["header", "index", "data"]) + + +class CelebA(VisionDataset): + """`Large-scale CelebFaces Attributes (CelebA) Dataset `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory where images are downloaded to. + split (string): One of {'train', 'valid', 'test', 'all'}. + Accordingly dataset is selected. + target_type (string or list, optional): Type of target to use, ``attr``, ``identity``, ``bbox``, + or ``landmarks``. Can also be a list to output a tuple with all specified target types. + The targets represent: + + - ``attr`` (Tensor shape=(40,) dtype=int): binary (0, 1) labels for attributes + - ``identity`` (int): label for each person (data points with the same identity are the same person) + - ``bbox`` (Tensor shape=(4,) dtype=int): bounding box (x, y, width, height) + - ``landmarks`` (Tensor shape=(10,) dtype=int): landmark points (lefteye_x, lefteye_y, righteye_x, + righteye_y, nose_x, nose_y, leftmouth_x, leftmouth_y, rightmouth_x, rightmouth_y) + + Defaults to ``attr``. If empty, ``None`` will be returned as target. + + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.PILToTensor`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + + .. warning:: + + To download the dataset `gdown `_ is required. + """ + + base_folder = "celeba" + # There currently does not appear to be an easy way to extract 7z in python (without introducing additional + # dependencies). The "in-the-wild" (not aligned+cropped) images are only in 7z, so they are not available + # right now. + file_list = [ + # File ID MD5 Hash Filename + ("0B7EVK8r0v71pZjFTYXZWM3FlRnM", "00d2c5bc6d35e252742224ab0c1e8fcb", "img_align_celeba.zip"), + # ("0B7EVK8r0v71pbWNEUjJKdDQ3dGc","b6cd7e93bc7a96c2dc33f819aa3ac651", "img_align_celeba_png.7z"), + # ("0B7EVK8r0v71peklHb0pGdDl6R28", "b6cd7e93bc7a96c2dc33f819aa3ac651", "img_celeba.7z"), + ("0B7EVK8r0v71pblRyaVFSWGxPY0U", "75e246fa4810816ffd6ee81facbd244c", "list_attr_celeba.txt"), + ("1_ee_0u7vcNLOfNLegJRHmolfH5ICW-XS", "32bd1bd63d3c78cd57e08160ec5ed1e2", "identity_CelebA.txt"), + ("0B7EVK8r0v71pbThiMVRxWXZ4dU0", "00566efa6fedff7a56946cd1c10f1c16", "list_bbox_celeba.txt"), + ("0B7EVK8r0v71pd0FJY3Blby1HUTQ", "cc24ecafdb5b50baae59b03474781f8c", "list_landmarks_align_celeba.txt"), + # ("0B7EVK8r0v71pTzJIdlJWdHczRlU", "063ee6ddb681f96bc9ca28c6febb9d1a", "list_landmarks_celeba.txt"), + ("0B7EVK8r0v71pY0NSMzRuSXJEVkk", "d32c9cbf5e040fd4025c592c306e6668", "list_eval_partition.txt"), + ] + + def __init__( + self, + root: Union[str, Path], + split: str = "train", + target_type: Union[List[str], str] = "attr", + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + super().__init__(root, transform=transform, target_transform=target_transform) + self.split = split + if isinstance(target_type, list): + self.target_type = target_type + else: + self.target_type = [target_type] + + if not self.target_type and self.target_transform is not None: + raise RuntimeError("target_transform is specified but target_type is empty") + + if download: + self.download() + + if not self._check_integrity(): + raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") + + split_map = { + "train": 0, + "valid": 1, + "test": 2, + "all": None, + } + split_ = split_map[verify_str_arg(split.lower(), "split", ("train", "valid", "test", "all"))] + splits = self._load_csv("list_eval_partition.txt") + identity = self._load_csv("identity_CelebA.txt") + bbox = self._load_csv("list_bbox_celeba.txt", header=1) + landmarks_align = self._load_csv("list_landmarks_align_celeba.txt", header=1) + attr = self._load_csv("list_attr_celeba.txt", header=1) + + mask = slice(None) if split_ is None else (splits.data == split_).squeeze() + + if mask == slice(None): # if split == "all" + self.filename = splits.index + else: + self.filename = [splits.index[i] for i in torch.squeeze(torch.nonzero(mask))] # type: ignore[arg-type] + self.identity = identity.data[mask] + self.bbox = bbox.data[mask] + self.landmarks_align = landmarks_align.data[mask] + self.attr = attr.data[mask] + # map from {-1, 1} to {0, 1} + self.attr = torch.div(self.attr + 1, 2, rounding_mode="floor") + self.attr_names = attr.header + + def _load_csv( + self, + filename: str, + header: Optional[int] = None, + ) -> CSV: + with open(os.path.join(self.root, self.base_folder, filename)) as csv_file: + data = list(csv.reader(csv_file, delimiter=" ", skipinitialspace=True)) + + if header is not None: + headers = data[header] + data = data[header + 1 :] + else: + headers = [] + + indices = [row[0] for row in data] + data = [row[1:] for row in data] + data_int = [list(map(int, i)) for i in data] + + return CSV(headers, indices, torch.tensor(data_int)) + + def _check_integrity(self) -> bool: + for (_, md5, filename) in self.file_list: + fpath = os.path.join(self.root, self.base_folder, filename) + _, ext = os.path.splitext(filename) + # Allow original archive to be deleted (zip and 7z) + # Only need the extracted images + if ext not in [".zip", ".7z"] and not check_integrity(fpath, md5): + return False + + # Should check a hash of the images + return os.path.isdir(os.path.join(self.root, self.base_folder, "img_align_celeba")) + + def download(self) -> None: + if self._check_integrity(): + return + + for (file_id, md5, filename) in self.file_list: + download_file_from_google_drive(file_id, os.path.join(self.root, self.base_folder), filename, md5) + + extract_archive(os.path.join(self.root, self.base_folder, "img_align_celeba.zip")) + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + X = PIL.Image.open(os.path.join(self.root, self.base_folder, "img_align_celeba", self.filename[index])) + + target: Any = [] + for t in self.target_type: + if t == "attr": + target.append(self.attr[index, :]) + elif t == "identity": + target.append(self.identity[index, 0]) + elif t == "bbox": + target.append(self.bbox[index, :]) + elif t == "landmarks": + target.append(self.landmarks_align[index, :]) + else: + # TODO: refactor with utils.verify_str_arg + raise ValueError(f'Target type "{t}" is not recognized.') + + if self.transform is not None: + X = self.transform(X) + + if target: + target = tuple(target) if len(target) > 1 else target[0] + + if self.target_transform is not None: + target = self.target_transform(target) + else: + target = None + + return X, target + + def __len__(self) -> int: + return len(self.attr) + + def extra_repr(self) -> str: + lines = ["Target type: {target_type}", "Split: {split}"] + return "\n".join(lines).format(**self.__dict__) diff --git a/lib/python3.12/site-packages/torchvision/datasets/cifar.py b/lib/python3.12/site-packages/torchvision/datasets/cifar.py new file mode 100644 index 0000000000000000000000000000000000000000..0f425d76c5763d3bc98c7c3cbb6b9241b6385622 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/cifar.py @@ -0,0 +1,167 @@ +import os.path +import pickle +from pathlib import Path +from typing import Any, Callable, Optional, Tuple, Union + +import numpy as np +from PIL import Image + +from .utils import check_integrity, download_and_extract_archive +from .vision import VisionDataset + + +class CIFAR10(VisionDataset): + """`CIFAR10 `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset where directory + ``cifar-10-batches-py`` exists or will be saved to if download is set to True. + train (bool, optional): If True, creates dataset from training set, otherwise + creates from test set. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + + """ + + base_folder = "cifar-10-batches-py" + url = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz" + filename = "cifar-10-python.tar.gz" + tgz_md5 = "c58f30108f718f92721af3b95e74349a" + train_list = [ + ["data_batch_1", "c99cafc152244af753f735de768cd75f"], + ["data_batch_2", "d4bba439e000b95fd0a9bffe97cbabec"], + ["data_batch_3", "54ebc095f3ab1f0389bbae665268c751"], + ["data_batch_4", "634d18415352ddfa80567beed471001a"], + ["data_batch_5", "482c414d41f54cd18b22e5b47cb7c3cb"], + ] + + test_list = [ + ["test_batch", "40351d587109b95175f43aff81a1287e"], + ] + meta = { + "filename": "batches.meta", + "key": "label_names", + "md5": "5ff9c542aee3614f3951f8cda6e48888", + } + + def __init__( + self, + root: Union[str, Path], + train: bool = True, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + + super().__init__(root, transform=transform, target_transform=target_transform) + + self.train = train # training set or test set + + if download: + self.download() + + if not self._check_integrity(): + raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") + + if self.train: + downloaded_list = self.train_list + else: + downloaded_list = self.test_list + + self.data: Any = [] + self.targets = [] + + # now load the picked numpy arrays + for file_name, checksum in downloaded_list: + file_path = os.path.join(self.root, self.base_folder, file_name) + with open(file_path, "rb") as f: + entry = pickle.load(f, encoding="latin1") + self.data.append(entry["data"]) + if "labels" in entry: + self.targets.extend(entry["labels"]) + else: + self.targets.extend(entry["fine_labels"]) + + self.data = np.vstack(self.data).reshape(-1, 3, 32, 32) + self.data = self.data.transpose((0, 2, 3, 1)) # convert to HWC + + self._load_meta() + + def _load_meta(self) -> None: + path = os.path.join(self.root, self.base_folder, self.meta["filename"]) + if not check_integrity(path, self.meta["md5"]): + raise RuntimeError("Dataset metadata file not found or corrupted. You can use download=True to download it") + with open(path, "rb") as infile: + data = pickle.load(infile, encoding="latin1") + self.classes = data[self.meta["key"]] + self.class_to_idx = {_class: i for i, _class in enumerate(self.classes)} + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: (image, target) where target is index of the target class. + """ + img, target = self.data[index], self.targets[index] + + # doing this so that it is consistent with all other datasets + # to return a PIL Image + img = Image.fromarray(img) + + if self.transform is not None: + img = self.transform(img) + + if self.target_transform is not None: + target = self.target_transform(target) + + return img, target + + def __len__(self) -> int: + return len(self.data) + + def _check_integrity(self) -> bool: + for filename, md5 in self.train_list + self.test_list: + fpath = os.path.join(self.root, self.base_folder, filename) + if not check_integrity(fpath, md5): + return False + return True + + def download(self) -> None: + if self._check_integrity(): + return + download_and_extract_archive(self.url, self.root, filename=self.filename, md5=self.tgz_md5) + + def extra_repr(self) -> str: + split = "Train" if self.train is True else "Test" + return f"Split: {split}" + + +class CIFAR100(CIFAR10): + """`CIFAR100 `_ Dataset. + + This is a subclass of the `CIFAR10` Dataset. + """ + + base_folder = "cifar-100-python" + url = "https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz" + filename = "cifar-100-python.tar.gz" + tgz_md5 = "eb9058c3a382ffc7106e4002c42a8d85" + train_list = [ + ["train", "16019d7e3df5f24257cddd939b257f8d"], + ] + + test_list = [ + ["test", "f0ef6b0ae62326f3e7ffdfab6717acfc"], + ] + meta = { + "filename": "meta", + "key": "fine_label_names", + "md5": "7973b15100ade9c7d40fb424638fde48", + } diff --git a/lib/python3.12/site-packages/torchvision/datasets/cityscapes.py b/lib/python3.12/site-packages/torchvision/datasets/cityscapes.py new file mode 100644 index 0000000000000000000000000000000000000000..97a47c07bebba250fd598c8b48a95a8e1af290fd --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/cityscapes.py @@ -0,0 +1,222 @@ +import json +import os +from collections import namedtuple +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +from PIL import Image + +from .utils import extract_archive, iterable_to_str, verify_str_arg +from .vision import VisionDataset + + +class Cityscapes(VisionDataset): + """`Cityscapes `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset where directory ``leftImg8bit`` + and ``gtFine`` or ``gtCoarse`` are located. + split (string, optional): The image split to use, ``train``, ``test`` or ``val`` if mode="fine" + otherwise ``train``, ``train_extra`` or ``val`` + mode (string, optional): The quality mode to use, ``fine`` or ``coarse`` + target_type (string or list, optional): Type of target to use, ``instance``, ``semantic``, ``polygon`` + or ``color``. Can also be a list to output a tuple with all specified target types. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + transforms (callable, optional): A function/transform that takes input sample and its target as entry + and returns a transformed version. + + Examples: + + Get semantic segmentation target + + .. code-block:: python + + dataset = Cityscapes('./data/cityscapes', split='train', mode='fine', + target_type='semantic') + + img, smnt = dataset[0] + + Get multiple targets + + .. code-block:: python + + dataset = Cityscapes('./data/cityscapes', split='train', mode='fine', + target_type=['instance', 'color', 'polygon']) + + img, (inst, col, poly) = dataset[0] + + Validate on the "coarse" set + + .. code-block:: python + + dataset = Cityscapes('./data/cityscapes', split='val', mode='coarse', + target_type='semantic') + + img, smnt = dataset[0] + """ + + # Based on https://github.com/mcordts/cityscapesScripts + CityscapesClass = namedtuple( + "CityscapesClass", + ["name", "id", "train_id", "category", "category_id", "has_instances", "ignore_in_eval", "color"], + ) + + classes = [ + CityscapesClass("unlabeled", 0, 255, "void", 0, False, True, (0, 0, 0)), + CityscapesClass("ego vehicle", 1, 255, "void", 0, False, True, (0, 0, 0)), + CityscapesClass("rectification border", 2, 255, "void", 0, False, True, (0, 0, 0)), + CityscapesClass("out of roi", 3, 255, "void", 0, False, True, (0, 0, 0)), + CityscapesClass("static", 4, 255, "void", 0, False, True, (0, 0, 0)), + CityscapesClass("dynamic", 5, 255, "void", 0, False, True, (111, 74, 0)), + CityscapesClass("ground", 6, 255, "void", 0, False, True, (81, 0, 81)), + CityscapesClass("road", 7, 0, "flat", 1, False, False, (128, 64, 128)), + CityscapesClass("sidewalk", 8, 1, "flat", 1, False, False, (244, 35, 232)), + CityscapesClass("parking", 9, 255, "flat", 1, False, True, (250, 170, 160)), + CityscapesClass("rail track", 10, 255, "flat", 1, False, True, (230, 150, 140)), + CityscapesClass("building", 11, 2, "construction", 2, False, False, (70, 70, 70)), + CityscapesClass("wall", 12, 3, "construction", 2, False, False, (102, 102, 156)), + CityscapesClass("fence", 13, 4, "construction", 2, False, False, (190, 153, 153)), + CityscapesClass("guard rail", 14, 255, "construction", 2, False, True, (180, 165, 180)), + CityscapesClass("bridge", 15, 255, "construction", 2, False, True, (150, 100, 100)), + CityscapesClass("tunnel", 16, 255, "construction", 2, False, True, (150, 120, 90)), + CityscapesClass("pole", 17, 5, "object", 3, False, False, (153, 153, 153)), + CityscapesClass("polegroup", 18, 255, "object", 3, False, True, (153, 153, 153)), + CityscapesClass("traffic light", 19, 6, "object", 3, False, False, (250, 170, 30)), + CityscapesClass("traffic sign", 20, 7, "object", 3, False, False, (220, 220, 0)), + CityscapesClass("vegetation", 21, 8, "nature", 4, False, False, (107, 142, 35)), + CityscapesClass("terrain", 22, 9, "nature", 4, False, False, (152, 251, 152)), + CityscapesClass("sky", 23, 10, "sky", 5, False, False, (70, 130, 180)), + CityscapesClass("person", 24, 11, "human", 6, True, False, (220, 20, 60)), + CityscapesClass("rider", 25, 12, "human", 6, True, False, (255, 0, 0)), + CityscapesClass("car", 26, 13, "vehicle", 7, True, False, (0, 0, 142)), + CityscapesClass("truck", 27, 14, "vehicle", 7, True, False, (0, 0, 70)), + CityscapesClass("bus", 28, 15, "vehicle", 7, True, False, (0, 60, 100)), + CityscapesClass("caravan", 29, 255, "vehicle", 7, True, True, (0, 0, 90)), + CityscapesClass("trailer", 30, 255, "vehicle", 7, True, True, (0, 0, 110)), + CityscapesClass("train", 31, 16, "vehicle", 7, True, False, (0, 80, 100)), + CityscapesClass("motorcycle", 32, 17, "vehicle", 7, True, False, (0, 0, 230)), + CityscapesClass("bicycle", 33, 18, "vehicle", 7, True, False, (119, 11, 32)), + CityscapesClass("license plate", -1, -1, "vehicle", 7, False, True, (0, 0, 142)), + ] + + def __init__( + self, + root: Union[str, Path], + split: str = "train", + mode: str = "fine", + target_type: Union[List[str], str] = "instance", + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + transforms: Optional[Callable] = None, + ) -> None: + super().__init__(root, transforms, transform, target_transform) + self.mode = "gtFine" if mode == "fine" else "gtCoarse" + self.images_dir = os.path.join(self.root, "leftImg8bit", split) + self.targets_dir = os.path.join(self.root, self.mode, split) + self.target_type = target_type + self.split = split + self.images = [] + self.targets = [] + + verify_str_arg(mode, "mode", ("fine", "coarse")) + if mode == "fine": + valid_modes = ("train", "test", "val") + else: + valid_modes = ("train", "train_extra", "val") + msg = "Unknown value '{}' for argument split if mode is '{}'. Valid values are {{{}}}." + msg = msg.format(split, mode, iterable_to_str(valid_modes)) + verify_str_arg(split, "split", valid_modes, msg) + + if not isinstance(target_type, list): + self.target_type = [target_type] + [ + verify_str_arg(value, "target_type", ("instance", "semantic", "polygon", "color")) + for value in self.target_type + ] + + if not os.path.isdir(self.images_dir) or not os.path.isdir(self.targets_dir): + + if split == "train_extra": + image_dir_zip = os.path.join(self.root, "leftImg8bit_trainextra.zip") + else: + image_dir_zip = os.path.join(self.root, "leftImg8bit_trainvaltest.zip") + + if self.mode == "gtFine": + target_dir_zip = os.path.join(self.root, f"{self.mode}_trainvaltest.zip") + elif self.mode == "gtCoarse": + target_dir_zip = os.path.join(self.root, f"{self.mode}.zip") + + if os.path.isfile(image_dir_zip) and os.path.isfile(target_dir_zip): + extract_archive(from_path=image_dir_zip, to_path=self.root) + extract_archive(from_path=target_dir_zip, to_path=self.root) + else: + raise RuntimeError( + "Dataset not found or incomplete. Please make sure all required folders for the" + ' specified "split" and "mode" are inside the "root" directory' + ) + + for city in os.listdir(self.images_dir): + img_dir = os.path.join(self.images_dir, city) + target_dir = os.path.join(self.targets_dir, city) + for file_name in os.listdir(img_dir): + target_types = [] + for t in self.target_type: + target_name = "{}_{}".format( + file_name.split("_leftImg8bit")[0], self._get_target_suffix(self.mode, t) + ) + target_types.append(os.path.join(target_dir, target_name)) + + self.images.append(os.path.join(img_dir, file_name)) + self.targets.append(target_types) + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + Returns: + tuple: (image, target) where target is a tuple of all target types if target_type is a list with more + than one item. Otherwise, target is a json object if target_type="polygon", else the image segmentation. + """ + + image = Image.open(self.images[index]).convert("RGB") + + targets: Any = [] + for i, t in enumerate(self.target_type): + if t == "polygon": + target = self._load_json(self.targets[index][i]) + else: + target = Image.open(self.targets[index][i]) # type: ignore[assignment] + + targets.append(target) + + target = tuple(targets) if len(targets) > 1 else targets[0] # type: ignore[assignment] + + if self.transforms is not None: + image, target = self.transforms(image, target) + + return image, target + + def __len__(self) -> int: + return len(self.images) + + def extra_repr(self) -> str: + lines = ["Split: {split}", "Mode: {mode}", "Type: {target_type}"] + return "\n".join(lines).format(**self.__dict__) + + def _load_json(self, path: str) -> Dict[str, Any]: + with open(path) as file: + data = json.load(file) + return data + + def _get_target_suffix(self, mode: str, target_type: str) -> str: + if target_type == "instance": + return f"{mode}_instanceIds.png" + elif target_type == "semantic": + return f"{mode}_labelIds.png" + elif target_type == "color": + return f"{mode}_color.png" + else: + return f"{mode}_polygons.json" diff --git a/lib/python3.12/site-packages/torchvision/datasets/clevr.py b/lib/python3.12/site-packages/torchvision/datasets/clevr.py new file mode 100644 index 0000000000000000000000000000000000000000..328eb7d79da70c3607b86ded512021f901119d1b --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/clevr.py @@ -0,0 +1,88 @@ +import json +import pathlib +from typing import Any, Callable, List, Optional, Tuple, Union +from urllib.parse import urlparse + +from PIL import Image + +from .utils import download_and_extract_archive, verify_str_arg +from .vision import VisionDataset + + +class CLEVRClassification(VisionDataset): + """`CLEVR `_ classification dataset. + + The number of objects in a scene are used as label. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset where directory ``root/clevr`` exists or will be saved to if download is + set to True. + split (string, optional): The dataset split, supports ``"train"`` (default), ``"val"``, or ``"test"``. + transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed + version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in them target and transforms it. + download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If + dataset is already downloaded, it is not downloaded again. + """ + + _URL = "https://dl.fbaipublicfiles.com/clevr/CLEVR_v1.0.zip" + _MD5 = "b11922020e72d0cd9154779b2d3d07d2" + + def __init__( + self, + root: Union[str, pathlib.Path], + split: str = "train", + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + self._split = verify_str_arg(split, "split", ("train", "val", "test")) + super().__init__(root, transform=transform, target_transform=target_transform) + self._base_folder = pathlib.Path(self.root) / "clevr" + self._data_folder = self._base_folder / pathlib.Path(urlparse(self._URL).path).stem + + if download: + self._download() + + if not self._check_exists(): + raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") + + self._image_files = sorted(self._data_folder.joinpath("images", self._split).glob("*")) + + self._labels: List[Optional[int]] + if self._split != "test": + with open(self._data_folder / "scenes" / f"CLEVR_{self._split}_scenes.json") as file: + content = json.load(file) + num_objects = {scene["image_filename"]: len(scene["objects"]) for scene in content["scenes"]} + self._labels = [num_objects[image_file.name] for image_file in self._image_files] + else: + self._labels = [None] * len(self._image_files) + + def __len__(self) -> int: + return len(self._image_files) + + def __getitem__(self, idx: int) -> Tuple[Any, Any]: + image_file = self._image_files[idx] + label = self._labels[idx] + + image = Image.open(image_file).convert("RGB") + + if self.transform: + image = self.transform(image) + + if self.target_transform: + label = self.target_transform(label) + + return image, label + + def _check_exists(self) -> bool: + return self._data_folder.exists() and self._data_folder.is_dir() + + def _download(self) -> None: + if self._check_exists(): + return + + download_and_extract_archive(self._URL, str(self._base_folder), md5=self._MD5) + + def extra_repr(self) -> str: + return f"split={self._split}" diff --git a/lib/python3.12/site-packages/torchvision/datasets/coco.py b/lib/python3.12/site-packages/torchvision/datasets/coco.py new file mode 100644 index 0000000000000000000000000000000000000000..f3b7be798b2457de56416f230f82a49373fbe941 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/coco.py @@ -0,0 +1,109 @@ +import os.path +from pathlib import Path +from typing import Any, Callable, List, Optional, Tuple, Union + +from PIL import Image + +from .vision import VisionDataset + + +class CocoDetection(VisionDataset): + """`MS Coco Detection `_ Dataset. + + It requires the `COCO API to be installed `_. + + Args: + root (str or ``pathlib.Path``): Root directory where images are downloaded to. + annFile (string): Path to json annotation file. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.PILToTensor`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + transforms (callable, optional): A function/transform that takes input sample and its target as entry + and returns a transformed version. + """ + + def __init__( + self, + root: Union[str, Path], + annFile: str, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + transforms: Optional[Callable] = None, + ) -> None: + super().__init__(root, transforms, transform, target_transform) + from pycocotools.coco import COCO + + self.coco = COCO(annFile) + self.ids = list(sorted(self.coco.imgs.keys())) + + def _load_image(self, id: int) -> Image.Image: + path = self.coco.loadImgs(id)[0]["file_name"] + return Image.open(os.path.join(self.root, path)).convert("RGB") + + def _load_target(self, id: int) -> List[Any]: + return self.coco.loadAnns(self.coco.getAnnIds(id)) + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + + if not isinstance(index, int): + raise ValueError(f"Index must be of type integer, got {type(index)} instead.") + + id = self.ids[index] + image = self._load_image(id) + target = self._load_target(id) + + if self.transforms is not None: + image, target = self.transforms(image, target) + + return image, target + + def __len__(self) -> int: + return len(self.ids) + + +class CocoCaptions(CocoDetection): + """`MS Coco Captions `_ Dataset. + + It requires the `COCO API to be installed `_. + + Args: + root (str or ``pathlib.Path``): Root directory where images are downloaded to. + annFile (string): Path to json annotation file. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.PILToTensor`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + transforms (callable, optional): A function/transform that takes input sample and its target as entry + and returns a transformed version. + + Example: + + .. code:: python + + import torchvision.datasets as dset + import torchvision.transforms as transforms + cap = dset.CocoCaptions(root = 'dir where images are', + annFile = 'json annotation file', + transform=transforms.PILToTensor()) + + print('Number of samples: ', len(cap)) + img, target = cap[3] # load 4th sample + + print("Image Size: ", img.size()) + print(target) + + Output: :: + + Number of samples: 82783 + Image Size: (3L, 427L, 640L) + [u'A plane emitting smoke stream flying over a mountain.', + u'A plane darts across a bright blue sky behind a mountain covered in snow', + u'A plane leaves a contrail above the snowy mountain top.', + u'A mountain that has a plane flying overheard in the distance.', + u'A mountain view with a plume of smoke in the background'] + + """ + + def _load_target(self, id: int) -> List[str]: + return [ann["caption"] for ann in super()._load_target(id)] diff --git a/lib/python3.12/site-packages/torchvision/datasets/country211.py b/lib/python3.12/site-packages/torchvision/datasets/country211.py new file mode 100644 index 0000000000000000000000000000000000000000..a0f82ee1226670ae274cc09dc0697dd4c51fb1c2 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/country211.py @@ -0,0 +1,58 @@ +from pathlib import Path +from typing import Callable, Optional, Union + +from .folder import ImageFolder +from .utils import download_and_extract_archive, verify_str_arg + + +class Country211(ImageFolder): + """`The Country211 Data Set `_ from OpenAI. + + This dataset was built by filtering the images from the YFCC100m dataset + that have GPS coordinate corresponding to a ISO-3166 country code. The + dataset is balanced by sampling 150 train images, 50 validation images, and + 100 test images for each country. + + Args: + root (str or ``pathlib.Path``): Root directory of the dataset. + split (string, optional): The dataset split, supports ``"train"`` (default), ``"valid"`` and ``"test"``. + transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed + version. E.g, ``transforms.RandomCrop``. + target_transform (callable, optional): A function/transform that takes in the target and transforms it. + download (bool, optional): If True, downloads the dataset from the internet and puts it into + ``root/country211/``. If dataset is already downloaded, it is not downloaded again. + """ + + _URL = "https://openaipublic.azureedge.net/clip/data/country211.tgz" + _MD5 = "84988d7644798601126c29e9877aab6a" + + def __init__( + self, + root: Union[str, Path], + split: str = "train", + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + self._split = verify_str_arg(split, "split", ("train", "valid", "test")) + + root = Path(root).expanduser() + self.root = str(root) + self._base_folder = root / "country211" + + if download: + self._download() + + if not self._check_exists(): + raise RuntimeError("Dataset not found. You can use download=True to download it") + + super().__init__(str(self._base_folder / self._split), transform=transform, target_transform=target_transform) + self.root = str(root) + + def _check_exists(self) -> bool: + return self._base_folder.exists() and self._base_folder.is_dir() + + def _download(self) -> None: + if self._check_exists(): + return + download_and_extract_archive(self._URL, download_root=self.root, md5=self._MD5) diff --git a/lib/python3.12/site-packages/torchvision/datasets/dtd.py b/lib/python3.12/site-packages/torchvision/datasets/dtd.py new file mode 100644 index 0000000000000000000000000000000000000000..71c556bd201b37b0622df050e8c9dadd5f32f4e0 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/dtd.py @@ -0,0 +1,100 @@ +import os +import pathlib +from typing import Any, Callable, Optional, Tuple, Union + +import PIL.Image + +from .utils import download_and_extract_archive, verify_str_arg +from .vision import VisionDataset + + +class DTD(VisionDataset): + """`Describable Textures Dataset (DTD) `_. + + Args: + root (str or ``pathlib.Path``): Root directory of the dataset. + split (string, optional): The dataset split, supports ``"train"`` (default), ``"val"``, or ``"test"``. + partition (int, optional): The dataset partition. Should be ``1 <= partition <= 10``. Defaults to ``1``. + + .. note:: + + The partition only changes which split each image belongs to. Thus, regardless of the selected + partition, combining all splits will result in all images. + + transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed + version. E.g, ``transforms.RandomCrop``. + target_transform (callable, optional): A function/transform that takes in the target and transforms it. + download (bool, optional): If True, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. Default is False. + """ + + _URL = "https://www.robots.ox.ac.uk/~vgg/data/dtd/download/dtd-r1.0.1.tar.gz" + _MD5 = "fff73e5086ae6bdbea199a49dfb8a4c1" + + def __init__( + self, + root: Union[str, pathlib.Path], + split: str = "train", + partition: int = 1, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + self._split = verify_str_arg(split, "split", ("train", "val", "test")) + if not isinstance(partition, int) and not (1 <= partition <= 10): + raise ValueError( + f"Parameter 'partition' should be an integer with `1 <= partition <= 10`, " + f"but got {partition} instead" + ) + self._partition = partition + + super().__init__(root, transform=transform, target_transform=target_transform) + self._base_folder = pathlib.Path(self.root) / type(self).__name__.lower() + self._data_folder = self._base_folder / "dtd" + self._meta_folder = self._data_folder / "labels" + self._images_folder = self._data_folder / "images" + + if download: + self._download() + + if not self._check_exists(): + raise RuntimeError("Dataset not found. You can use download=True to download it") + + self._image_files = [] + classes = [] + with open(self._meta_folder / f"{self._split}{self._partition}.txt") as file: + for line in file: + cls, name = line.strip().split("/") + self._image_files.append(self._images_folder.joinpath(cls, name)) + classes.append(cls) + + self.classes = sorted(set(classes)) + self.class_to_idx = dict(zip(self.classes, range(len(self.classes)))) + self._labels = [self.class_to_idx[cls] for cls in classes] + + def __len__(self) -> int: + return len(self._image_files) + + def __getitem__(self, idx: int) -> Tuple[Any, Any]: + image_file, label = self._image_files[idx], self._labels[idx] + image = PIL.Image.open(image_file).convert("RGB") + + if self.transform: + image = self.transform(image) + + if self.target_transform: + label = self.target_transform(label) + + return image, label + + def extra_repr(self) -> str: + return f"split={self._split}, partition={self._partition}" + + def _check_exists(self) -> bool: + return os.path.exists(self._data_folder) and os.path.isdir(self._data_folder) + + def _download(self) -> None: + if self._check_exists(): + return + download_and_extract_archive(self._URL, download_root=str(self._base_folder), md5=self._MD5) diff --git a/lib/python3.12/site-packages/torchvision/datasets/eurosat.py b/lib/python3.12/site-packages/torchvision/datasets/eurosat.py new file mode 100644 index 0000000000000000000000000000000000000000..c6571d2abab11d457f0c0b9ff910c7a58efbcf87 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/eurosat.py @@ -0,0 +1,62 @@ +import os +from pathlib import Path +from typing import Callable, Optional, Union + +from .folder import ImageFolder +from .utils import download_and_extract_archive + + +class EuroSAT(ImageFolder): + """RGB version of the `EuroSAT `_ Dataset. + + For the MS version of the dataset, see + `TorchGeo `__. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset where ``root/eurosat`` exists. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + download (bool, optional): If True, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. Default is False. + """ + + def __init__( + self, + root: Union[str, Path], + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + self.root = os.path.expanduser(root) + self._base_folder = os.path.join(self.root, "eurosat") + self._data_folder = os.path.join(self._base_folder, "2750") + + if download: + self.download() + + if not self._check_exists(): + raise RuntimeError("Dataset not found. You can use download=True to download it") + + super().__init__(self._data_folder, transform=transform, target_transform=target_transform) + self.root = os.path.expanduser(root) + + def __len__(self) -> int: + return len(self.samples) + + def _check_exists(self) -> bool: + return os.path.exists(self._data_folder) + + def download(self) -> None: + + if self._check_exists(): + return + + os.makedirs(self._base_folder, exist_ok=True) + download_and_extract_archive( + "https://huggingface.co/datasets/torchgeo/eurosat/resolve/c877bcd43f099cd0196738f714544e355477f3fd/EuroSAT.zip", + download_root=self._base_folder, + md5="c8fa014336c82ac7804f0398fcb19387", + ) diff --git a/lib/python3.12/site-packages/torchvision/datasets/fakedata.py b/lib/python3.12/site-packages/torchvision/datasets/fakedata.py new file mode 100644 index 0000000000000000000000000000000000000000..af26a8579e5d954bde9fc06966d3518b7ca76c1c --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/fakedata.py @@ -0,0 +1,67 @@ +from typing import Any, Callable, Optional, Tuple + +import torch + +from .. import transforms +from .vision import VisionDataset + + +class FakeData(VisionDataset): + """A fake dataset that returns randomly generated images and returns them as PIL images + + Args: + size (int, optional): Size of the dataset. Default: 1000 images + image_size(tuple, optional): Size if the returned images. Default: (3, 224, 224) + num_classes(int, optional): Number of classes in the dataset. Default: 10 + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + random_offset (int): Offsets the index-based random seed used to + generate each image. Default: 0 + + """ + + def __init__( + self, + size: int = 1000, + image_size: Tuple[int, int, int] = (3, 224, 224), + num_classes: int = 10, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + random_offset: int = 0, + ) -> None: + super().__init__(transform=transform, target_transform=target_transform) + self.size = size + self.num_classes = num_classes + self.image_size = image_size + self.random_offset = random_offset + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: (image, target) where target is class_index of the target class. + """ + # create random image that is consistent with the index id + if index >= len(self): + raise IndexError(f"{self.__class__.__name__} index out of range") + rng_state = torch.get_rng_state() + torch.manual_seed(index + self.random_offset) + img = torch.randn(*self.image_size) + target = torch.randint(0, self.num_classes, size=(1,), dtype=torch.long)[0] + torch.set_rng_state(rng_state) + + # convert to PIL Image + img = transforms.ToPILImage()(img) + if self.transform is not None: + img = self.transform(img) + if self.target_transform is not None: + target = self.target_transform(target) + + return img, target.item() + + def __len__(self) -> int: + return self.size diff --git a/lib/python3.12/site-packages/torchvision/datasets/fer2013.py b/lib/python3.12/site-packages/torchvision/datasets/fer2013.py new file mode 100644 index 0000000000000000000000000000000000000000..3afda07846b6d18b1cd43c37f7bb8747b8b172b5 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/fer2013.py @@ -0,0 +1,120 @@ +import csv +import pathlib +from typing import Any, Callable, Optional, Tuple, Union + +import torch +from PIL import Image + +from .utils import check_integrity, verify_str_arg +from .vision import VisionDataset + + +class FER2013(VisionDataset): + """`FER2013 + `_ Dataset. + + .. note:: + This dataset can return test labels only if ``fer2013.csv`` OR + ``icml_face_data.csv`` are present in ``root/fer2013/``. If only + ``train.csv`` and ``test.csv`` are present, the test labels are set to + ``None``. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset where directory + ``root/fer2013`` exists. This directory may contain either + ``fer2013.csv``, ``icml_face_data.csv``, or both ``train.csv`` and + ``test.csv``. Precendence is given in that order, i.e. if + ``fer2013.csv`` is present then the rest of the files will be + ignored. All these (combinations of) files contain the same data and + are supported for convenience, but only ``fer2013.csv`` and + ``icml_face_data.csv`` are able to return non-None test labels. + split (string, optional): The dataset split, supports ``"train"`` (default), or ``"test"``. + transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed + version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the target and transforms it. + """ + + _RESOURCES = { + "train": ("train.csv", "3f0dfb3d3fd99c811a1299cb947e3131"), + "test": ("test.csv", "b02c2298636a634e8c2faabbf3ea9a23"), + # The fer2013.csv and icml_face_data.csv files contain both train and + # tests instances, and unlike test.csv they contain the labels for the + # test instances. We give these 2 files precedence over train.csv and + # test.csv. And yes, they both contain the same data, but with different + # column names (note the spaces) and ordering: + # $ head -n 1 fer2013.csv icml_face_data.csv train.csv test.csv + # ==> fer2013.csv <== + # emotion,pixels,Usage + # + # ==> icml_face_data.csv <== + # emotion, Usage, pixels + # + # ==> train.csv <== + # emotion,pixels + # + # ==> test.csv <== + # pixels + "fer": ("fer2013.csv", "f8428a1edbd21e88f42c73edd2a14f95"), + "icml": ("icml_face_data.csv", "b114b9e04e6949e5fe8b6a98b3892b1d"), + } + + def __init__( + self, + root: Union[str, pathlib.Path], + split: str = "train", + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + ) -> None: + self._split = verify_str_arg(split, "split", ("train", "test")) + super().__init__(root, transform=transform, target_transform=target_transform) + + base_folder = pathlib.Path(self.root) / "fer2013" + use_fer_file = (base_folder / self._RESOURCES["fer"][0]).exists() + use_icml_file = not use_fer_file and (base_folder / self._RESOURCES["icml"][0]).exists() + file_name, md5 = self._RESOURCES["fer" if use_fer_file else "icml" if use_icml_file else self._split] + data_file = base_folder / file_name + if not check_integrity(str(data_file), md5=md5): + raise RuntimeError( + f"{file_name} not found in {base_folder} or corrupted. " + f"You can download it from " + f"https://www.kaggle.com/c/challenges-in-representation-learning-facial-expression-recognition-challenge" + ) + + pixels_key = " pixels" if use_icml_file else "pixels" + usage_key = " Usage" if use_icml_file else "Usage" + + def get_img(row): + return torch.tensor([int(idx) for idx in row[pixels_key].split()], dtype=torch.uint8).reshape(48, 48) + + def get_label(row): + if use_fer_file or use_icml_file or self._split == "train": + return int(row["emotion"]) + else: + return None + + with open(data_file, "r", newline="") as file: + rows = (row for row in csv.DictReader(file)) + + if use_fer_file or use_icml_file: + valid_keys = ("Training",) if self._split == "train" else ("PublicTest", "PrivateTest") + rows = (row for row in rows if row[usage_key] in valid_keys) + + self._samples = [(get_img(row), get_label(row)) for row in rows] + + def __len__(self) -> int: + return len(self._samples) + + def __getitem__(self, idx: int) -> Tuple[Any, Any]: + image_tensor, target = self._samples[idx] + image = Image.fromarray(image_tensor.numpy()) + + if self.transform is not None: + image = self.transform(image) + + if self.target_transform is not None: + target = self.target_transform(target) + + return image, target + + def extra_repr(self) -> str: + return f"split={self._split}" diff --git a/lib/python3.12/site-packages/torchvision/datasets/fgvc_aircraft.py b/lib/python3.12/site-packages/torchvision/datasets/fgvc_aircraft.py new file mode 100644 index 0000000000000000000000000000000000000000..bbf4e970a787556e634f6c2eeb64ed4cd706fa2b --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/fgvc_aircraft.py @@ -0,0 +1,115 @@ +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any, Callable, Optional, Tuple, Union + +import PIL.Image + +from .utils import download_and_extract_archive, verify_str_arg +from .vision import VisionDataset + + +class FGVCAircraft(VisionDataset): + """`FGVC Aircraft `_ Dataset. + + The dataset contains 10,000 images of aircraft, with 100 images for each of 100 + different aircraft model variants, most of which are airplanes. + Aircraft models are organized in a three-levels hierarchy. The three levels, from + finer to coarser, are: + + - ``variant``, e.g. Boeing 737-700. A variant collapses all the models that are visually + indistinguishable into one class. The dataset comprises 100 different variants. + - ``family``, e.g. Boeing 737. The dataset comprises 70 different families. + - ``manufacturer``, e.g. Boeing. The dataset comprises 30 different manufacturers. + + Args: + root (str or ``pathlib.Path``): Root directory of the FGVC Aircraft dataset. + split (string, optional): The dataset split, supports ``train``, ``val``, + ``trainval`` and ``test``. + annotation_level (str, optional): The annotation level, supports ``variant``, + ``family`` and ``manufacturer``. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + download (bool, optional): If True, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + """ + + _URL = "https://www.robots.ox.ac.uk/~vgg/data/fgvc-aircraft/archives/fgvc-aircraft-2013b.tar.gz" + + def __init__( + self, + root: Union[str, Path], + split: str = "trainval", + annotation_level: str = "variant", + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + super().__init__(root, transform=transform, target_transform=target_transform) + self._split = verify_str_arg(split, "split", ("train", "val", "trainval", "test")) + self._annotation_level = verify_str_arg( + annotation_level, "annotation_level", ("variant", "family", "manufacturer") + ) + + self._data_path = os.path.join(self.root, "fgvc-aircraft-2013b") + if download: + self._download() + + if not self._check_exists(): + raise RuntimeError("Dataset not found. You can use download=True to download it") + + annotation_file = os.path.join( + self._data_path, + "data", + { + "variant": "variants.txt", + "family": "families.txt", + "manufacturer": "manufacturers.txt", + }[self._annotation_level], + ) + with open(annotation_file, "r") as f: + self.classes = [line.strip() for line in f] + + self.class_to_idx = dict(zip(self.classes, range(len(self.classes)))) + + image_data_folder = os.path.join(self._data_path, "data", "images") + labels_file = os.path.join(self._data_path, "data", f"images_{self._annotation_level}_{self._split}.txt") + + self._image_files = [] + self._labels = [] + + with open(labels_file, "r") as f: + for line in f: + image_name, label_name = line.strip().split(" ", 1) + self._image_files.append(os.path.join(image_data_folder, f"{image_name}.jpg")) + self._labels.append(self.class_to_idx[label_name]) + + def __len__(self) -> int: + return len(self._image_files) + + def __getitem__(self, idx: int) -> Tuple[Any, Any]: + image_file, label = self._image_files[idx], self._labels[idx] + image = PIL.Image.open(image_file).convert("RGB") + + if self.transform: + image = self.transform(image) + + if self.target_transform: + label = self.target_transform(label) + + return image, label + + def _download(self) -> None: + """ + Download the FGVC Aircraft dataset archive and extract it under root. + """ + if self._check_exists(): + return + download_and_extract_archive(self._URL, self.root) + + def _check_exists(self) -> bool: + return os.path.exists(self._data_path) and os.path.isdir(self._data_path) diff --git a/lib/python3.12/site-packages/torchvision/datasets/flickr.py b/lib/python3.12/site-packages/torchvision/datasets/flickr.py new file mode 100644 index 0000000000000000000000000000000000000000..1021309db0540bcd32c7bf35668b4c43ce586a4d --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/flickr.py @@ -0,0 +1,167 @@ +import glob +import os +from collections import defaultdict +from html.parser import HTMLParser +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +from PIL import Image + +from .vision import VisionDataset + + +class Flickr8kParser(HTMLParser): + """Parser for extracting captions from the Flickr8k dataset web page.""" + + def __init__(self, root: Union[str, Path]) -> None: + super().__init__() + + self.root = root + + # Data structure to store captions + self.annotations: Dict[str, List[str]] = {} + + # State variables + self.in_table = False + self.current_tag: Optional[str] = None + self.current_img: Optional[str] = None + + def handle_starttag(self, tag: str, attrs: List[Tuple[str, Optional[str]]]) -> None: + self.current_tag = tag + + if tag == "table": + self.in_table = True + + def handle_endtag(self, tag: str) -> None: + self.current_tag = None + + if tag == "table": + self.in_table = False + + def handle_data(self, data: str) -> None: + if self.in_table: + if data == "Image Not Found": + self.current_img = None + elif self.current_tag == "a": + img_id = data.split("/")[-2] + img_id = os.path.join(self.root, img_id + "_*.jpg") + img_id = glob.glob(img_id)[0] + self.current_img = img_id + self.annotations[img_id] = [] + elif self.current_tag == "li" and self.current_img: + img_id = self.current_img + self.annotations[img_id].append(data.strip()) + + +class Flickr8k(VisionDataset): + """`Flickr8k Entities `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory where images are downloaded to. + ann_file (string): Path to annotation file. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.PILToTensor`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + """ + + def __init__( + self, + root: Union[str, Path], + ann_file: str, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + ) -> None: + super().__init__(root, transform=transform, target_transform=target_transform) + self.ann_file = os.path.expanduser(ann_file) + + # Read annotations and store in a dict + parser = Flickr8kParser(self.root) + with open(self.ann_file) as fh: + parser.feed(fh.read()) + self.annotations = parser.annotations + + self.ids = list(sorted(self.annotations.keys())) + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: Tuple (image, target). target is a list of captions for the image. + """ + img_id = self.ids[index] + + # Image + img = Image.open(img_id).convert("RGB") + if self.transform is not None: + img = self.transform(img) + + # Captions + target = self.annotations[img_id] + if self.target_transform is not None: + target = self.target_transform(target) + + return img, target + + def __len__(self) -> int: + return len(self.ids) + + +class Flickr30k(VisionDataset): + """`Flickr30k Entities `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory where images are downloaded to. + ann_file (string): Path to annotation file. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.PILToTensor`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + """ + + def __init__( + self, + root: str, + ann_file: str, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + ) -> None: + super().__init__(root, transform=transform, target_transform=target_transform) + self.ann_file = os.path.expanduser(ann_file) + + # Read annotations and store in a dict + self.annotations = defaultdict(list) + with open(self.ann_file) as fh: + for line in fh: + img_id, caption = line.strip().split("\t") + self.annotations[img_id[:-2]].append(caption) + + self.ids = list(sorted(self.annotations.keys())) + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: Tuple (image, target). target is a list of captions for the image. + """ + img_id = self.ids[index] + + # Image + filename = os.path.join(self.root, img_id) + img = Image.open(filename).convert("RGB") + if self.transform is not None: + img = self.transform(img) + + # Captions + target = self.annotations[img_id] + if self.target_transform is not None: + target = self.target_transform(target) + + return img, target + + def __len__(self) -> int: + return len(self.ids) diff --git a/lib/python3.12/site-packages/torchvision/datasets/flowers102.py b/lib/python3.12/site-packages/torchvision/datasets/flowers102.py new file mode 100644 index 0000000000000000000000000000000000000000..07f403702f5f0e06f3d890713f693571190546c6 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/flowers102.py @@ -0,0 +1,114 @@ +from pathlib import Path +from typing import Any, Callable, Optional, Tuple, Union + +import PIL.Image + +from .utils import check_integrity, download_and_extract_archive, download_url, verify_str_arg +from .vision import VisionDataset + + +class Flowers102(VisionDataset): + """`Oxford 102 Flower `_ Dataset. + + .. warning:: + + This class needs `scipy `_ to load target files from `.mat` format. + + Oxford 102 Flower is an image classification dataset consisting of 102 flower categories. The + flowers were chosen to be flowers commonly occurring in the United Kingdom. Each class consists of + between 40 and 258 images. + + The images have large scale, pose and light variations. In addition, there are categories that + have large variations within the category, and several very similar categories. + + Args: + root (str or ``pathlib.Path``): Root directory of the dataset. + split (string, optional): The dataset split, supports ``"train"`` (default), ``"val"``, or ``"test"``. + transform (callable, optional): A function/transform that takes in a PIL image and returns a + transformed version. E.g, ``transforms.RandomCrop``. + target_transform (callable, optional): A function/transform that takes in the target and transforms it. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + """ + + _download_url_prefix = "https://www.robots.ox.ac.uk/~vgg/data/flowers/102/" + _file_dict = { # filename, md5 + "image": ("102flowers.tgz", "52808999861908f626f3c1f4e79d11fa"), + "label": ("imagelabels.mat", "e0620be6f572b9609742df49c70aed4d"), + "setid": ("setid.mat", "a5357ecc9cb78c4bef273ce3793fc85c"), + } + _splits_map = {"train": "trnid", "val": "valid", "test": "tstid"} + + def __init__( + self, + root: Union[str, Path], + split: str = "train", + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + super().__init__(root, transform=transform, target_transform=target_transform) + self._split = verify_str_arg(split, "split", ("train", "val", "test")) + self._base_folder = Path(self.root) / "flowers-102" + self._images_folder = self._base_folder / "jpg" + + if download: + self.download() + + if not self._check_integrity(): + raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") + + from scipy.io import loadmat + + set_ids = loadmat(self._base_folder / self._file_dict["setid"][0], squeeze_me=True) + image_ids = set_ids[self._splits_map[self._split]].tolist() + + labels = loadmat(self._base_folder / self._file_dict["label"][0], squeeze_me=True) + image_id_to_label = dict(enumerate((labels["labels"] - 1).tolist(), 1)) + + self._labels = [] + self._image_files = [] + for image_id in image_ids: + self._labels.append(image_id_to_label[image_id]) + self._image_files.append(self._images_folder / f"image_{image_id:05d}.jpg") + + def __len__(self) -> int: + return len(self._image_files) + + def __getitem__(self, idx: int) -> Tuple[Any, Any]: + image_file, label = self._image_files[idx], self._labels[idx] + image = PIL.Image.open(image_file).convert("RGB") + + if self.transform: + image = self.transform(image) + + if self.target_transform: + label = self.target_transform(label) + + return image, label + + def extra_repr(self) -> str: + return f"split={self._split}" + + def _check_integrity(self): + if not (self._images_folder.exists() and self._images_folder.is_dir()): + return False + + for id in ["label", "setid"]: + filename, md5 = self._file_dict[id] + if not check_integrity(str(self._base_folder / filename), md5): + return False + return True + + def download(self): + if self._check_integrity(): + return + download_and_extract_archive( + f"{self._download_url_prefix}{self._file_dict['image'][0]}", + str(self._base_folder), + md5=self._file_dict["image"][1], + ) + for id in ["label", "setid"]: + filename, md5 = self._file_dict[id] + download_url(self._download_url_prefix + filename, str(self._base_folder), md5=md5) diff --git a/lib/python3.12/site-packages/torchvision/datasets/folder.py b/lib/python3.12/site-packages/torchvision/datasets/folder.py new file mode 100644 index 0000000000000000000000000000000000000000..8f2f65c7b615aa12ec4cdb2bd472d521a6880ad7 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/folder.py @@ -0,0 +1,337 @@ +import os +import os.path +from pathlib import Path +from typing import Any, Callable, cast, Dict, List, Optional, Tuple, Union + +from PIL import Image + +from .vision import VisionDataset + + +def has_file_allowed_extension(filename: str, extensions: Union[str, Tuple[str, ...]]) -> bool: + """Checks if a file is an allowed extension. + + Args: + filename (string): path to a file + extensions (tuple of strings): extensions to consider (lowercase) + + Returns: + bool: True if the filename ends with one of given extensions + """ + return filename.lower().endswith(extensions if isinstance(extensions, str) else tuple(extensions)) + + +def is_image_file(filename: str) -> bool: + """Checks if a file is an allowed image extension. + + Args: + filename (string): path to a file + + Returns: + bool: True if the filename ends with a known image extension + """ + return has_file_allowed_extension(filename, IMG_EXTENSIONS) + + +def find_classes(directory: Union[str, Path]) -> Tuple[List[str], Dict[str, int]]: + """Finds the class folders in a dataset. + + See :class:`DatasetFolder` for details. + """ + classes = sorted(entry.name for entry in os.scandir(directory) if entry.is_dir()) + if not classes: + raise FileNotFoundError(f"Couldn't find any class folder in {directory}.") + + class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)} + return classes, class_to_idx + + +def make_dataset( + directory: Union[str, Path], + class_to_idx: Optional[Dict[str, int]] = None, + extensions: Optional[Union[str, Tuple[str, ...]]] = None, + is_valid_file: Optional[Callable[[str], bool]] = None, + allow_empty: bool = False, +) -> List[Tuple[str, int]]: + """Generates a list of samples of a form (path_to_sample, class). + + See :class:`DatasetFolder` for details. + + Note: The class_to_idx parameter is here optional and will use the logic of the ``find_classes`` function + by default. + """ + directory = os.path.expanduser(directory) + + if class_to_idx is None: + _, class_to_idx = find_classes(directory) + elif not class_to_idx: + raise ValueError("'class_to_index' must have at least one entry to collect any samples.") + + both_none = extensions is None and is_valid_file is None + both_something = extensions is not None and is_valid_file is not None + if both_none or both_something: + raise ValueError("Both extensions and is_valid_file cannot be None or not None at the same time") + + if extensions is not None: + + def is_valid_file(x: str) -> bool: + return has_file_allowed_extension(x, extensions) # type: ignore[arg-type] + + is_valid_file = cast(Callable[[str], bool], is_valid_file) + + instances = [] + available_classes = set() + for target_class in sorted(class_to_idx.keys()): + class_index = class_to_idx[target_class] + target_dir = os.path.join(directory, target_class) + if not os.path.isdir(target_dir): + continue + for root, _, fnames in sorted(os.walk(target_dir, followlinks=True)): + for fname in sorted(fnames): + path = os.path.join(root, fname) + if is_valid_file(path): + item = path, class_index + instances.append(item) + + if target_class not in available_classes: + available_classes.add(target_class) + + empty_classes = set(class_to_idx.keys()) - available_classes + if empty_classes and not allow_empty: + msg = f"Found no valid file for the classes {', '.join(sorted(empty_classes))}. " + if extensions is not None: + msg += f"Supported extensions are: {extensions if isinstance(extensions, str) else ', '.join(extensions)}" + raise FileNotFoundError(msg) + + return instances + + +class DatasetFolder(VisionDataset): + """A generic data loader. + + This default directory structure can be customized by overriding the + :meth:`find_classes` method. + + Args: + root (str or ``pathlib.Path``): Root directory path. + loader (callable): A function to load a sample given its path. + extensions (tuple[string]): A list of allowed extensions. + both extensions and is_valid_file should not be passed. + transform (callable, optional): A function/transform that takes in + a sample and returns a transformed version. + E.g, ``transforms.RandomCrop`` for images. + target_transform (callable, optional): A function/transform that takes + in the target and transforms it. + is_valid_file (callable, optional): A function that takes path of a file + and check if the file is a valid file (used to check of corrupt files) + both extensions and is_valid_file should not be passed. + allow_empty(bool, optional): If True, empty folders are considered to be valid classes. + An error is raised on empty folders if False (default). + + Attributes: + classes (list): List of the class names sorted alphabetically. + class_to_idx (dict): Dict with items (class_name, class_index). + samples (list): List of (sample path, class_index) tuples + targets (list): The class_index value for each image in the dataset + """ + + def __init__( + self, + root: Union[str, Path], + loader: Callable[[str], Any], + extensions: Optional[Tuple[str, ...]] = None, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + is_valid_file: Optional[Callable[[str], bool]] = None, + allow_empty: bool = False, + ) -> None: + super().__init__(root, transform=transform, target_transform=target_transform) + classes, class_to_idx = self.find_classes(self.root) + samples = self.make_dataset( + self.root, + class_to_idx=class_to_idx, + extensions=extensions, + is_valid_file=is_valid_file, + allow_empty=allow_empty, + ) + + self.loader = loader + self.extensions = extensions + + self.classes = classes + self.class_to_idx = class_to_idx + self.samples = samples + self.targets = [s[1] for s in samples] + + @staticmethod + def make_dataset( + directory: Union[str, Path], + class_to_idx: Dict[str, int], + extensions: Optional[Tuple[str, ...]] = None, + is_valid_file: Optional[Callable[[str], bool]] = None, + allow_empty: bool = False, + ) -> List[Tuple[str, int]]: + """Generates a list of samples of a form (path_to_sample, class). + + This can be overridden to e.g. read files from a compressed zip file instead of from the disk. + + Args: + directory (str): root dataset directory, corresponding to ``self.root``. + class_to_idx (Dict[str, int]): Dictionary mapping class name to class index. + extensions (optional): A list of allowed extensions. + Either extensions or is_valid_file should be passed. Defaults to None. + is_valid_file (optional): A function that takes path of a file + and checks if the file is a valid file + (used to check of corrupt files) both extensions and + is_valid_file should not be passed. Defaults to None. + allow_empty(bool, optional): If True, empty folders are considered to be valid classes. + An error is raised on empty folders if False (default). + + Raises: + ValueError: In case ``class_to_idx`` is empty. + ValueError: In case ``extensions`` and ``is_valid_file`` are None or both are not None. + FileNotFoundError: In case no valid file was found for any class. + + Returns: + List[Tuple[str, int]]: samples of a form (path_to_sample, class) + """ + if class_to_idx is None: + # prevent potential bug since make_dataset() would use the class_to_idx logic of the + # find_classes() function, instead of using that of the find_classes() method, which + # is potentially overridden and thus could have a different logic. + raise ValueError("The class_to_idx parameter cannot be None.") + return make_dataset( + directory, class_to_idx, extensions=extensions, is_valid_file=is_valid_file, allow_empty=allow_empty + ) + + def find_classes(self, directory: Union[str, Path]) -> Tuple[List[str], Dict[str, int]]: + """Find the class folders in a dataset structured as follows:: + + directory/ + ├── class_x + │ ├── xxx.ext + │ ├── xxy.ext + │ └── ... + │ └── xxz.ext + └── class_y + ├── 123.ext + ├── nsdf3.ext + └── ... + └── asd932_.ext + + This method can be overridden to only consider + a subset of classes, or to adapt to a different dataset directory structure. + + Args: + directory(str): Root directory path, corresponding to ``self.root`` + + Raises: + FileNotFoundError: If ``dir`` has no class folders. + + Returns: + (Tuple[List[str], Dict[str, int]]): List of all classes and dictionary mapping each class to an index. + """ + return find_classes(directory) + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: (sample, target) where target is class_index of the target class. + """ + path, target = self.samples[index] + sample = self.loader(path) + if self.transform is not None: + sample = self.transform(sample) + if self.target_transform is not None: + target = self.target_transform(target) + + return sample, target + + def __len__(self) -> int: + return len(self.samples) + + +IMG_EXTENSIONS = (".jpg", ".jpeg", ".png", ".ppm", ".bmp", ".pgm", ".tif", ".tiff", ".webp") + + +def pil_loader(path: str) -> Image.Image: + # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835) + with open(path, "rb") as f: + img = Image.open(f) + return img.convert("RGB") + + +# TODO: specify the return type +def accimage_loader(path: str) -> Any: + import accimage + + try: + return accimage.Image(path) + except OSError: + # Potentially a decoding problem, fall back to PIL.Image + return pil_loader(path) + + +def default_loader(path: str) -> Any: + from torchvision import get_image_backend + + if get_image_backend() == "accimage": + return accimage_loader(path) + else: + return pil_loader(path) + + +class ImageFolder(DatasetFolder): + """A generic data loader where the images are arranged in this way by default: :: + + root/dog/xxx.png + root/dog/xxy.png + root/dog/[...]/xxz.png + + root/cat/123.png + root/cat/nsdf3.png + root/cat/[...]/asd932_.png + + This class inherits from :class:`~torchvision.datasets.DatasetFolder` so + the same methods can be overridden to customize the dataset. + + Args: + root (str or ``pathlib.Path``): Root directory path. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + loader (callable, optional): A function to load an image given its path. + is_valid_file (callable, optional): A function that takes path of an Image file + and check if the file is a valid file (used to check of corrupt files) + allow_empty(bool, optional): If True, empty folders are considered to be valid classes. + An error is raised on empty folders if False (default). + + Attributes: + classes (list): List of the class names sorted alphabetically. + class_to_idx (dict): Dict with items (class_name, class_index). + imgs (list): List of (image path, class_index) tuples + """ + + def __init__( + self, + root: Union[str, Path], + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + loader: Callable[[str], Any] = default_loader, + is_valid_file: Optional[Callable[[str], bool]] = None, + allow_empty: bool = False, + ): + super().__init__( + root, + loader, + IMG_EXTENSIONS if is_valid_file is None else None, + transform=transform, + target_transform=target_transform, + is_valid_file=is_valid_file, + allow_empty=allow_empty, + ) + self.imgs = self.samples diff --git a/lib/python3.12/site-packages/torchvision/datasets/food101.py b/lib/python3.12/site-packages/torchvision/datasets/food101.py new file mode 100644 index 0000000000000000000000000000000000000000..f734787c1bf638867ad8f05bc52f953e00954060 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/food101.py @@ -0,0 +1,93 @@ +import json +from pathlib import Path +from typing import Any, Callable, Optional, Tuple, Union + +import PIL.Image + +from .utils import download_and_extract_archive, verify_str_arg +from .vision import VisionDataset + + +class Food101(VisionDataset): + """`The Food-101 Data Set `_. + + The Food-101 is a challenging data set of 101 food categories with 101,000 images. + For each class, 250 manually reviewed test images are provided as well as 750 training images. + On purpose, the training images were not cleaned, and thus still contain some amount of noise. + This comes mostly in the form of intense colors and sometimes wrong labels. All images were + rescaled to have a maximum side length of 512 pixels. + + + Args: + root (str or ``pathlib.Path``): Root directory of the dataset. + split (string, optional): The dataset split, supports ``"train"`` (default) and ``"test"``. + transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed + version. E.g, ``transforms.RandomCrop``. + target_transform (callable, optional): A function/transform that takes in the target and transforms it. + download (bool, optional): If True, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. Default is False. + """ + + _URL = "http://data.vision.ee.ethz.ch/cvl/food-101.tar.gz" + _MD5 = "85eeb15f3717b99a5da872d97d918f87" + + def __init__( + self, + root: Union[str, Path], + split: str = "train", + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + super().__init__(root, transform=transform, target_transform=target_transform) + self._split = verify_str_arg(split, "split", ("train", "test")) + self._base_folder = Path(self.root) / "food-101" + self._meta_folder = self._base_folder / "meta" + self._images_folder = self._base_folder / "images" + + if download: + self._download() + + if not self._check_exists(): + raise RuntimeError("Dataset not found. You can use download=True to download it") + + self._labels = [] + self._image_files = [] + with open(self._meta_folder / f"{split}.json") as f: + metadata = json.loads(f.read()) + + self.classes = sorted(metadata.keys()) + self.class_to_idx = dict(zip(self.classes, range(len(self.classes)))) + + for class_label, im_rel_paths in metadata.items(): + self._labels += [self.class_to_idx[class_label]] * len(im_rel_paths) + self._image_files += [ + self._images_folder.joinpath(*f"{im_rel_path}.jpg".split("/")) for im_rel_path in im_rel_paths + ] + + def __len__(self) -> int: + return len(self._image_files) + + def __getitem__(self, idx: int) -> Tuple[Any, Any]: + image_file, label = self._image_files[idx], self._labels[idx] + image = PIL.Image.open(image_file).convert("RGB") + + if self.transform: + image = self.transform(image) + + if self.target_transform: + label = self.target_transform(label) + + return image, label + + def extra_repr(self) -> str: + return f"split={self._split}" + + def _check_exists(self) -> bool: + return all(folder.exists() and folder.is_dir() for folder in (self._meta_folder, self._images_folder)) + + def _download(self) -> None: + if self._check_exists(): + return + download_and_extract_archive(self._URL, download_root=self.root, md5=self._MD5) diff --git a/lib/python3.12/site-packages/torchvision/datasets/gtsrb.py b/lib/python3.12/site-packages/torchvision/datasets/gtsrb.py new file mode 100644 index 0000000000000000000000000000000000000000..a3d012c70b22fd8209534a01a51fa9978c705d00 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/gtsrb.py @@ -0,0 +1,103 @@ +import csv +import pathlib +from typing import Any, Callable, Optional, Tuple, Union + +import PIL + +from .folder import make_dataset +from .utils import download_and_extract_archive, verify_str_arg +from .vision import VisionDataset + + +class GTSRB(VisionDataset): + """`German Traffic Sign Recognition Benchmark (GTSRB) `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of the dataset. + split (string, optional): The dataset split, supports ``"train"`` (default), or ``"test"``. + transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed + version. E.g, ``transforms.RandomCrop``. + target_transform (callable, optional): A function/transform that takes in the target and transforms it. + download (bool, optional): If True, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + """ + + def __init__( + self, + root: Union[str, pathlib.Path], + split: str = "train", + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + + super().__init__(root, transform=transform, target_transform=target_transform) + + self._split = verify_str_arg(split, "split", ("train", "test")) + self._base_folder = pathlib.Path(root) / "gtsrb" + self._target_folder = ( + self._base_folder / "GTSRB" / ("Training" if self._split == "train" else "Final_Test/Images") + ) + + if download: + self.download() + + if not self._check_exists(): + raise RuntimeError("Dataset not found. You can use download=True to download it") + + if self._split == "train": + samples = make_dataset(str(self._target_folder), extensions=(".ppm",)) + else: + with open(self._base_folder / "GT-final_test.csv") as csv_file: + samples = [ + (str(self._target_folder / row["Filename"]), int(row["ClassId"])) + for row in csv.DictReader(csv_file, delimiter=";", skipinitialspace=True) + ] + + self._samples = samples + self.transform = transform + self.target_transform = target_transform + + def __len__(self) -> int: + return len(self._samples) + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + + path, target = self._samples[index] + sample = PIL.Image.open(path).convert("RGB") + + if self.transform is not None: + sample = self.transform(sample) + + if self.target_transform is not None: + target = self.target_transform(target) + + return sample, target + + def _check_exists(self) -> bool: + return self._target_folder.is_dir() + + def download(self) -> None: + if self._check_exists(): + return + + base_url = "https://sid.erda.dk/public/archives/daaeac0d7ce1152aea9b61d9f1e19370/" + + if self._split == "train": + download_and_extract_archive( + f"{base_url}GTSRB-Training_fixed.zip", + download_root=str(self._base_folder), + md5="513f3c79a4c5141765e10e952eaa2478", + ) + else: + download_and_extract_archive( + f"{base_url}GTSRB_Final_Test_Images.zip", + download_root=str(self._base_folder), + md5="c7e4e6327067d32654124b0fe9e82185", + ) + download_and_extract_archive( + f"{base_url}GTSRB_Final_Test_GT.zip", + download_root=str(self._base_folder), + md5="fe31e9c9270bbcd7b84b7f21a9d9d9e5", + ) diff --git a/lib/python3.12/site-packages/torchvision/datasets/hmdb51.py b/lib/python3.12/site-packages/torchvision/datasets/hmdb51.py new file mode 100644 index 0000000000000000000000000000000000000000..8377e40d57ca7c2b29f98704b67bd03793a40836 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/hmdb51.py @@ -0,0 +1,152 @@ +import glob +import os +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +from torch import Tensor + +from .folder import find_classes, make_dataset +from .video_utils import VideoClips +from .vision import VisionDataset + + +class HMDB51(VisionDataset): + """ + `HMDB51 `_ + dataset. + + HMDB51 is an action recognition video dataset. + This dataset consider every video as a collection of video clips of fixed size, specified + by ``frames_per_clip``, where the step in frames between each clip is given by + ``step_between_clips``. + + To give an example, for 2 videos with 10 and 15 frames respectively, if ``frames_per_clip=5`` + and ``step_between_clips=5``, the dataset size will be (2 + 3) = 5, where the first two + elements will come from video 1, and the next three elements from video 2. + Note that we drop clips which do not have exactly ``frames_per_clip`` elements, so not all + frames in a video might be present. + + Internally, it uses a VideoClips object to handle clip creation. + + Args: + root (str or ``pathlib.Path``): Root directory of the HMDB51 Dataset. + annotation_path (str): Path to the folder containing the split files. + frames_per_clip (int): Number of frames in a clip. + step_between_clips (int): Number of frames between each clip. + fold (int, optional): Which fold to use. Should be between 1 and 3. + train (bool, optional): If ``True``, creates a dataset from the train split, + otherwise from the ``test`` split. + transform (callable, optional): A function/transform that takes in a TxHxWxC video + and returns a transformed version. + output_format (str, optional): The format of the output video tensors (before transforms). + Can be either "THWC" (default) or "TCHW". + + Returns: + tuple: A 3-tuple with the following entries: + + - video (Tensor[T, H, W, C] or Tensor[T, C, H, W]): The `T` video frames + - audio(Tensor[K, L]): the audio frames, where `K` is the number of channels + and `L` is the number of points + - label (int): class of the video clip + """ + + data_url = "https://serre-lab.clps.brown.edu/wp-content/uploads/2013/10/hmdb51_org.rar" + splits = { + "url": "https://serre-lab.clps.brown.edu/wp-content/uploads/2013/10/test_train_splits.rar", + "md5": "15e67781e70dcfbdce2d7dbb9b3344b5", + } + TRAIN_TAG = 1 + TEST_TAG = 2 + + def __init__( + self, + root: Union[str, Path], + annotation_path: str, + frames_per_clip: int, + step_between_clips: int = 1, + frame_rate: Optional[int] = None, + fold: int = 1, + train: bool = True, + transform: Optional[Callable] = None, + _precomputed_metadata: Optional[Dict[str, Any]] = None, + num_workers: int = 1, + _video_width: int = 0, + _video_height: int = 0, + _video_min_dimension: int = 0, + _audio_samples: int = 0, + output_format: str = "THWC", + ) -> None: + super().__init__(root) + if fold not in (1, 2, 3): + raise ValueError(f"fold should be between 1 and 3, got {fold}") + + extensions = ("avi",) + self.classes, class_to_idx = find_classes(self.root) + self.samples = make_dataset( + self.root, + class_to_idx, + extensions, + ) + + video_paths = [path for (path, _) in self.samples] + video_clips = VideoClips( + video_paths, + frames_per_clip, + step_between_clips, + frame_rate, + _precomputed_metadata, + num_workers=num_workers, + _video_width=_video_width, + _video_height=_video_height, + _video_min_dimension=_video_min_dimension, + _audio_samples=_audio_samples, + output_format=output_format, + ) + # we bookkeep the full version of video clips because we want to be able + # to return the metadata of full version rather than the subset version of + # video clips + self.full_video_clips = video_clips + self.fold = fold + self.train = train + self.indices = self._select_fold(video_paths, annotation_path, fold, train) + self.video_clips = video_clips.subset(self.indices) + self.transform = transform + + @property + def metadata(self) -> Dict[str, Any]: + return self.full_video_clips.metadata + + def _select_fold(self, video_list: List[str], annotations_dir: str, fold: int, train: bool) -> List[int]: + target_tag = self.TRAIN_TAG if train else self.TEST_TAG + split_pattern_name = f"*test_split{fold}.txt" + split_pattern_path = os.path.join(annotations_dir, split_pattern_name) + annotation_paths = glob.glob(split_pattern_path) + selected_files = set() + for filepath in annotation_paths: + with open(filepath) as fid: + lines = fid.readlines() + for line in lines: + video_filename, tag_string = line.split() + tag = int(tag_string) + if tag == target_tag: + selected_files.add(video_filename) + + indices = [] + for video_index, video_path in enumerate(video_list): + if os.path.basename(video_path) in selected_files: + indices.append(video_index) + + return indices + + def __len__(self) -> int: + return self.video_clips.num_clips() + + def __getitem__(self, idx: int) -> Tuple[Tensor, Tensor, int]: + video, audio, _, video_idx = self.video_clips.get_clip(idx) + sample_index = self.indices[video_idx] + _, class_index = self.samples[sample_index] + + if self.transform is not None: + video = self.transform(video) + + return video, audio, class_index diff --git a/lib/python3.12/site-packages/torchvision/datasets/imagenet.py b/lib/python3.12/site-packages/torchvision/datasets/imagenet.py new file mode 100644 index 0000000000000000000000000000000000000000..d7caf328d2bb1adc149548daf83ee49c2a8459c3 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/imagenet.py @@ -0,0 +1,219 @@ +import os +import shutil +import tempfile +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Dict, Iterator, List, Optional, Tuple, Union + +import torch + +from .folder import ImageFolder +from .utils import check_integrity, extract_archive, verify_str_arg + +ARCHIVE_META = { + "train": ("ILSVRC2012_img_train.tar", "1d675b47d978889d74fa0da5fadfb00e"), + "val": ("ILSVRC2012_img_val.tar", "29b22e2961454d5413ddabcf34fc5622"), + "devkit": ("ILSVRC2012_devkit_t12.tar.gz", "fa75699e90414af021442c21a62c3abf"), +} + +META_FILE = "meta.bin" + + +class ImageNet(ImageFolder): + """`ImageNet `_ 2012 Classification Dataset. + + .. note:: + Before using this class, it is required to download ImageNet 2012 dataset from + `here `_ and + place the files ``ILSVRC2012_devkit_t12.tar.gz`` and ``ILSVRC2012_img_train.tar`` + or ``ILSVRC2012_img_val.tar`` based on ``split`` in the root directory. + + Args: + root (str or ``pathlib.Path``): Root directory of the ImageNet Dataset. + split (string, optional): The dataset split, supports ``train``, or ``val``. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + loader (callable, optional): A function to load an image given its path. + + Attributes: + classes (list): List of the class name tuples. + class_to_idx (dict): Dict with items (class_name, class_index). + wnids (list): List of the WordNet IDs. + wnid_to_idx (dict): Dict with items (wordnet_id, class_index). + imgs (list): List of (image path, class_index) tuples + targets (list): The class_index value for each image in the dataset + """ + + def __init__(self, root: Union[str, Path], split: str = "train", **kwargs: Any) -> None: + root = self.root = os.path.expanduser(root) + self.split = verify_str_arg(split, "split", ("train", "val")) + + self.parse_archives() + wnid_to_classes = load_meta_file(self.root)[0] + + super().__init__(self.split_folder, **kwargs) + self.root = root + + self.wnids = self.classes + self.wnid_to_idx = self.class_to_idx + self.classes = [wnid_to_classes[wnid] for wnid in self.wnids] + self.class_to_idx = {cls: idx for idx, clss in enumerate(self.classes) for cls in clss} + + def parse_archives(self) -> None: + if not check_integrity(os.path.join(self.root, META_FILE)): + parse_devkit_archive(self.root) + + if not os.path.isdir(self.split_folder): + if self.split == "train": + parse_train_archive(self.root) + elif self.split == "val": + parse_val_archive(self.root) + + @property + def split_folder(self) -> str: + return os.path.join(self.root, self.split) + + def extra_repr(self) -> str: + return "Split: {split}".format(**self.__dict__) + + +def load_meta_file(root: Union[str, Path], file: Optional[str] = None) -> Tuple[Dict[str, str], List[str]]: + if file is None: + file = META_FILE + file = os.path.join(root, file) + + if check_integrity(file): + return torch.load(file, weights_only=True) + else: + msg = ( + "The meta file {} is not present in the root directory or is corrupted. " + "This file is automatically created by the ImageNet dataset." + ) + raise RuntimeError(msg.format(file, root)) + + +def _verify_archive(root: Union[str, Path], file: str, md5: str) -> None: + if not check_integrity(os.path.join(root, file), md5): + msg = ( + "The archive {} is not present in the root directory or is corrupted. " + "You need to download it externally and place it in {}." + ) + raise RuntimeError(msg.format(file, root)) + + +def parse_devkit_archive(root: Union[str, Path], file: Optional[str] = None) -> None: + """Parse the devkit archive of the ImageNet2012 classification dataset and save + the meta information in a binary file. + + Args: + root (str or ``pathlib.Path``): Root directory containing the devkit archive + file (str, optional): Name of devkit archive. Defaults to + 'ILSVRC2012_devkit_t12.tar.gz' + """ + import scipy.io as sio + + def parse_meta_mat(devkit_root: str) -> Tuple[Dict[int, str], Dict[str, Tuple[str, ...]]]: + metafile = os.path.join(devkit_root, "data", "meta.mat") + meta = sio.loadmat(metafile, squeeze_me=True)["synsets"] + nums_children = list(zip(*meta))[4] + meta = [meta[idx] for idx, num_children in enumerate(nums_children) if num_children == 0] + idcs, wnids, classes = list(zip(*meta))[:3] + classes = [tuple(clss.split(", ")) for clss in classes] + idx_to_wnid = {idx: wnid for idx, wnid in zip(idcs, wnids)} + wnid_to_classes = {wnid: clss for wnid, clss in zip(wnids, classes)} + return idx_to_wnid, wnid_to_classes + + def parse_val_groundtruth_txt(devkit_root: str) -> List[int]: + file = os.path.join(devkit_root, "data", "ILSVRC2012_validation_ground_truth.txt") + with open(file) as txtfh: + val_idcs = txtfh.readlines() + return [int(val_idx) for val_idx in val_idcs] + + @contextmanager + def get_tmp_dir() -> Iterator[str]: + tmp_dir = tempfile.mkdtemp() + try: + yield tmp_dir + finally: + shutil.rmtree(tmp_dir) + + archive_meta = ARCHIVE_META["devkit"] + if file is None: + file = archive_meta[0] + md5 = archive_meta[1] + + _verify_archive(root, file, md5) + + with get_tmp_dir() as tmp_dir: + extract_archive(os.path.join(root, file), tmp_dir) + + devkit_root = os.path.join(tmp_dir, "ILSVRC2012_devkit_t12") + idx_to_wnid, wnid_to_classes = parse_meta_mat(devkit_root) + val_idcs = parse_val_groundtruth_txt(devkit_root) + val_wnids = [idx_to_wnid[idx] for idx in val_idcs] + + torch.save((wnid_to_classes, val_wnids), os.path.join(root, META_FILE)) + + +def parse_train_archive(root: Union[str, Path], file: Optional[str] = None, folder: str = "train") -> None: + """Parse the train images archive of the ImageNet2012 classification dataset and + prepare it for usage with the ImageNet dataset. + + Args: + root (str or ``pathlib.Path``): Root directory containing the train images archive + file (str, optional): Name of train images archive. Defaults to + 'ILSVRC2012_img_train.tar' + folder (str, optional): Optional name for train images folder. Defaults to + 'train' + """ + archive_meta = ARCHIVE_META["train"] + if file is None: + file = archive_meta[0] + md5 = archive_meta[1] + + _verify_archive(root, file, md5) + + train_root = os.path.join(root, folder) + extract_archive(os.path.join(root, file), train_root) + + archives = [os.path.join(train_root, archive) for archive in os.listdir(train_root)] + for archive in archives: + extract_archive(archive, os.path.splitext(archive)[0], remove_finished=True) + + +def parse_val_archive( + root: Union[str, Path], file: Optional[str] = None, wnids: Optional[List[str]] = None, folder: str = "val" +) -> None: + """Parse the validation images archive of the ImageNet2012 classification dataset + and prepare it for usage with the ImageNet dataset. + + Args: + root (str or ``pathlib.Path``): Root directory containing the validation images archive + file (str, optional): Name of validation images archive. Defaults to + 'ILSVRC2012_img_val.tar' + wnids (list, optional): List of WordNet IDs of the validation images. If None + is given, the IDs are loaded from the meta file in the root directory + folder (str, optional): Optional name for validation images folder. Defaults to + 'val' + """ + archive_meta = ARCHIVE_META["val"] + if file is None: + file = archive_meta[0] + md5 = archive_meta[1] + if wnids is None: + wnids = load_meta_file(root)[1] + + _verify_archive(root, file, md5) + + val_root = os.path.join(root, folder) + extract_archive(os.path.join(root, file), val_root) + + images = sorted(os.path.join(val_root, image) for image in os.listdir(val_root)) + + for wnid in set(wnids): + os.mkdir(os.path.join(val_root, wnid)) + + for wnid, img_file in zip(wnids, images): + shutil.move(img_file, os.path.join(val_root, wnid, os.path.basename(img_file))) diff --git a/lib/python3.12/site-packages/torchvision/datasets/imagenette.py b/lib/python3.12/site-packages/torchvision/datasets/imagenette.py new file mode 100644 index 0000000000000000000000000000000000000000..0b27f3b25e51b1f87ccb17cebf7f20a539c0a5bf --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/imagenette.py @@ -0,0 +1,101 @@ +from pathlib import Path +from typing import Any, Callable, Optional, Tuple, Union + +from PIL import Image + +from .folder import find_classes, make_dataset +from .utils import download_and_extract_archive, verify_str_arg +from .vision import VisionDataset + + +class Imagenette(VisionDataset): + """`Imagenette `_ image classification dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of the Imagenette dataset. + split (string, optional): The dataset split. Supports ``"train"`` (default), and ``"val"``. + size (string, optional): The image size. Supports ``"full"`` (default), ``"320px"``, and ``"160px"``. + download (bool, optional): If ``True``, downloads the dataset components and places them in ``root``. Already + downloaded archives are not downloaded again. + transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed + version, e.g. ``transforms.RandomCrop``. + target_transform (callable, optional): A function/transform that takes in the target and transforms it. + + Attributes: + classes (list): List of the class name tuples. + class_to_idx (dict): Dict with items (class name, class index). + wnids (list): List of the WordNet IDs. + wnid_to_idx (dict): Dict with items (WordNet ID, class index). + """ + + _ARCHIVES = { + "full": ("https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz", "fe2fc210e6bb7c5664d602c3cd71e612"), + "320px": ("https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-320.tgz", "3df6f0d01a2c9592104656642f5e78a3"), + "160px": ("https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-160.tgz", "e793b78cc4c9e9a4ccc0c1155377a412"), + } + _WNID_TO_CLASS = { + "n01440764": ("tench", "Tinca tinca"), + "n02102040": ("English springer", "English springer spaniel"), + "n02979186": ("cassette player",), + "n03000684": ("chain saw", "chainsaw"), + "n03028079": ("church", "church building"), + "n03394916": ("French horn", "horn"), + "n03417042": ("garbage truck", "dustcart"), + "n03425413": ("gas pump", "gasoline pump", "petrol pump", "island dispenser"), + "n03445777": ("golf ball",), + "n03888257": ("parachute", "chute"), + } + + def __init__( + self, + root: Union[str, Path], + split: str = "train", + size: str = "full", + download=False, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + ) -> None: + super().__init__(root, transform=transform, target_transform=target_transform) + + self._split = verify_str_arg(split, "split", ["train", "val"]) + self._size = verify_str_arg(size, "size", ["full", "320px", "160px"]) + + self._url, self._md5 = self._ARCHIVES[self._size] + self._size_root = Path(self.root) / Path(self._url).stem + self._image_root = str(self._size_root / self._split) + + if download: + self._download() + elif not self._check_exists(): + raise RuntimeError("Dataset not found. You can use download=True to download it.") + + self.wnids, self.wnid_to_idx = find_classes(self._image_root) + self.classes = [self._WNID_TO_CLASS[wnid] for wnid in self.wnids] + self.class_to_idx = { + class_name: idx for wnid, idx in self.wnid_to_idx.items() for class_name in self._WNID_TO_CLASS[wnid] + } + self._samples = make_dataset(self._image_root, self.wnid_to_idx, extensions=".jpeg") + + def _check_exists(self) -> bool: + return self._size_root.exists() + + def _download(self): + if self._check_exists(): + return + + download_and_extract_archive(self._url, self.root, md5=self._md5) + + def __getitem__(self, idx: int) -> Tuple[Any, Any]: + path, label = self._samples[idx] + image = Image.open(path).convert("RGB") + + if self.transform is not None: + image = self.transform(image) + + if self.target_transform is not None: + label = self.target_transform(label) + + return image, label + + def __len__(self) -> int: + return len(self._samples) diff --git a/lib/python3.12/site-packages/torchvision/datasets/inaturalist.py b/lib/python3.12/site-packages/torchvision/datasets/inaturalist.py new file mode 100644 index 0000000000000000000000000000000000000000..e041d41f4a2322580607123cd33e5c3e23ac46df --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/inaturalist.py @@ -0,0 +1,238 @@ +import os +import os.path +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +from PIL import Image + +from .utils import download_and_extract_archive, verify_str_arg +from .vision import VisionDataset + +CATEGORIES_2021 = ["kingdom", "phylum", "class", "order", "family", "genus"] + +DATASET_URLS = { + "2017": "https://ml-inat-competition-datasets.s3.amazonaws.com/2017/train_val_images.tar.gz", + "2018": "https://ml-inat-competition-datasets.s3.amazonaws.com/2018/train_val2018.tar.gz", + "2019": "https://ml-inat-competition-datasets.s3.amazonaws.com/2019/train_val2019.tar.gz", + "2021_train": "https://ml-inat-competition-datasets.s3.amazonaws.com/2021/train.tar.gz", + "2021_train_mini": "https://ml-inat-competition-datasets.s3.amazonaws.com/2021/train_mini.tar.gz", + "2021_valid": "https://ml-inat-competition-datasets.s3.amazonaws.com/2021/val.tar.gz", +} + +DATASET_MD5 = { + "2017": "7c784ea5e424efaec655bd392f87301f", + "2018": "b1c6952ce38f31868cc50ea72d066cc3", + "2019": "c60a6e2962c9b8ccbd458d12c8582644", + "2021_train": "e0526d53c7f7b2e3167b2b43bb2690ed", + "2021_train_mini": "db6ed8330e634445efc8fec83ae81442", + "2021_valid": "f6f6e0e242e3d4c9569ba56400938afc", +} + + +class INaturalist(VisionDataset): + """`iNaturalist `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset where the image files are stored. + This class does not require/use annotation files. + version (string, optional): Which version of the dataset to download/use. One of + '2017', '2018', '2019', '2021_train', '2021_train_mini', '2021_valid'. + Default: `2021_train`. + target_type (string or list, optional): Type of target to use, for 2021 versions, one of: + + - ``full``: the full category (species) + - ``kingdom``: e.g. "Animalia" + - ``phylum``: e.g. "Arthropoda" + - ``class``: e.g. "Insecta" + - ``order``: e.g. "Coleoptera" + - ``family``: e.g. "Cleridae" + - ``genus``: e.g. "Trichodes" + + for 2017-2019 versions, one of: + + - ``full``: the full (numeric) category + - ``super``: the super category, e.g. "Amphibians" + + Can also be a list to output a tuple with all specified target types. + Defaults to ``full``. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + """ + + def __init__( + self, + root: Union[str, Path], + version: str = "2021_train", + target_type: Union[List[str], str] = "full", + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + self.version = verify_str_arg(version, "version", DATASET_URLS.keys()) + + super().__init__(os.path.join(root, version), transform=transform, target_transform=target_transform) + + os.makedirs(root, exist_ok=True) + if download: + self.download() + + if not self._check_exists(): + raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") + + self.all_categories: List[str] = [] + + # map: category type -> name of category -> index + self.categories_index: Dict[str, Dict[str, int]] = {} + + # list indexed by category id, containing mapping from category type -> index + self.categories_map: List[Dict[str, int]] = [] + + if not isinstance(target_type, list): + target_type = [target_type] + if self.version[:4] == "2021": + self.target_type = [verify_str_arg(t, "target_type", ("full", *CATEGORIES_2021)) for t in target_type] + self._init_2021() + else: + self.target_type = [verify_str_arg(t, "target_type", ("full", "super")) for t in target_type] + self._init_pre2021() + + # index of all files: (full category id, filename) + self.index: List[Tuple[int, str]] = [] + + for dir_index, dir_name in enumerate(self.all_categories): + files = os.listdir(os.path.join(self.root, dir_name)) + for fname in files: + self.index.append((dir_index, fname)) + + def _init_2021(self) -> None: + """Initialize based on 2021 layout""" + + self.all_categories = sorted(os.listdir(self.root)) + + # map: category type -> name of category -> index + self.categories_index = {k: {} for k in CATEGORIES_2021} + + for dir_index, dir_name in enumerate(self.all_categories): + pieces = dir_name.split("_") + if len(pieces) != 8: + raise RuntimeError(f"Unexpected category name {dir_name}, wrong number of pieces") + if pieces[0] != f"{dir_index:05d}": + raise RuntimeError(f"Unexpected category id {pieces[0]}, expecting {dir_index:05d}") + cat_map = {} + for cat, name in zip(CATEGORIES_2021, pieces[1:7]): + if name in self.categories_index[cat]: + cat_id = self.categories_index[cat][name] + else: + cat_id = len(self.categories_index[cat]) + self.categories_index[cat][name] = cat_id + cat_map[cat] = cat_id + self.categories_map.append(cat_map) + + def _init_pre2021(self) -> None: + """Initialize based on 2017-2019 layout""" + + # map: category type -> name of category -> index + self.categories_index = {"super": {}} + + cat_index = 0 + super_categories = sorted(os.listdir(self.root)) + for sindex, scat in enumerate(super_categories): + self.categories_index["super"][scat] = sindex + subcategories = sorted(os.listdir(os.path.join(self.root, scat))) + for subcat in subcategories: + if self.version == "2017": + # this version does not use ids as directory names + subcat_i = cat_index + cat_index += 1 + else: + try: + subcat_i = int(subcat) + except ValueError: + raise RuntimeError(f"Unexpected non-numeric dir name: {subcat}") + if subcat_i >= len(self.categories_map): + old_len = len(self.categories_map) + self.categories_map.extend([{}] * (subcat_i - old_len + 1)) + self.all_categories.extend([""] * (subcat_i - old_len + 1)) + if self.categories_map[subcat_i]: + raise RuntimeError(f"Duplicate category {subcat}") + self.categories_map[subcat_i] = {"super": sindex} + self.all_categories[subcat_i] = os.path.join(scat, subcat) + + # validate the dictionary + for cindex, c in enumerate(self.categories_map): + if not c: + raise RuntimeError(f"Missing category {cindex}") + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: (image, target) where the type of target specified by target_type. + """ + + cat_id, fname = self.index[index] + img = Image.open(os.path.join(self.root, self.all_categories[cat_id], fname)) + + target: Any = [] + for t in self.target_type: + if t == "full": + target.append(cat_id) + else: + target.append(self.categories_map[cat_id][t]) + target = tuple(target) if len(target) > 1 else target[0] + + if self.transform is not None: + img = self.transform(img) + + if self.target_transform is not None: + target = self.target_transform(target) + + return img, target + + def __len__(self) -> int: + return len(self.index) + + def category_name(self, category_type: str, category_id: int) -> str: + """ + Args: + category_type(str): one of "full", "kingdom", "phylum", "class", "order", "family", "genus" or "super" + category_id(int): an index (class id) from this category + + Returns: + the name of the category + """ + if category_type == "full": + return self.all_categories[category_id] + else: + if category_type not in self.categories_index: + raise ValueError(f"Invalid category type '{category_type}'") + else: + for name, id in self.categories_index[category_type].items(): + if id == category_id: + return name + raise ValueError(f"Invalid category id {category_id} for {category_type}") + + def _check_exists(self) -> bool: + return os.path.exists(self.root) and len(os.listdir(self.root)) > 0 + + def download(self) -> None: + if self._check_exists(): + return + + base_root = os.path.dirname(self.root) + + download_and_extract_archive( + DATASET_URLS[self.version], base_root, filename=f"{self.version}.tgz", md5=DATASET_MD5[self.version] + ) + + orig_dir_name = os.path.join(base_root, os.path.basename(DATASET_URLS[self.version]).rstrip(".tar.gz")) + if not os.path.exists(orig_dir_name): + raise RuntimeError(f"Unable to find downloaded files at {orig_dir_name}") + os.rename(orig_dir_name, self.root) diff --git a/lib/python3.12/site-packages/torchvision/datasets/kinetics.py b/lib/python3.12/site-packages/torchvision/datasets/kinetics.py new file mode 100644 index 0000000000000000000000000000000000000000..773d9f68ca9f849dcf30ac8f368212107f6a3cae --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/kinetics.py @@ -0,0 +1,237 @@ +import csv +import os +import urllib +from functools import partial +from multiprocessing import Pool +from os import path +from pathlib import Path +from typing import Any, Callable, Dict, Optional, Tuple, Union + +from torch import Tensor + +from .folder import find_classes, make_dataset +from .utils import check_integrity, download_and_extract_archive, download_url, verify_str_arg +from .video_utils import VideoClips +from .vision import VisionDataset + + +def _dl_wrap(tarpath: Union[str, Path], videopath: Union[str, Path], line: str) -> None: + download_and_extract_archive(line, tarpath, videopath) + + +class Kinetics(VisionDataset): + """`Generic Kinetics `_ + dataset. + + Kinetics-400/600/700 are action recognition video datasets. + This dataset consider every video as a collection of video clips of fixed size, specified + by ``frames_per_clip``, where the step in frames between each clip is given by + ``step_between_clips``. + + To give an example, for 2 videos with 10 and 15 frames respectively, if ``frames_per_clip=5`` + and ``step_between_clips=5``, the dataset size will be (2 + 3) = 5, where the first two + elements will come from video 1, and the next three elements from video 2. + Note that we drop clips which do not have exactly ``frames_per_clip`` elements, so not all + frames in a video might be present. + + Args: + root (str or ``pathlib.Path``): Root directory of the Kinetics Dataset. + Directory should be structured as follows: + .. code:: + + root/ + ├── split + │ ├── class1 + │ │ ├── vid1.mp4 + │ │ ├── vid2.mp4 + │ │ ├── vid3.mp4 + │ │ ├── ... + │ ├── class2 + │ │ ├── vidx.mp4 + │ │ └── ... + + Note: split is appended automatically using the split argument. + frames_per_clip (int): number of frames in a clip + num_classes (int): select between Kinetics-400 (default), Kinetics-600, and Kinetics-700 + split (str): split of the dataset to consider; supports ``"train"`` (default) ``"val"`` ``"test"`` + frame_rate (float): If omitted, interpolate different frame rate for each clip. + step_between_clips (int): number of frames between each clip + transform (callable, optional): A function/transform that takes in a TxHxWxC video + and returns a transformed version. + download (bool): Download the official version of the dataset to root folder. + num_workers (int): Use multiple workers for VideoClips creation + num_download_workers (int): Use multiprocessing in order to speed up download. + output_format (str, optional): The format of the output video tensors (before transforms). + Can be either "THWC" or "TCHW" (default). + Note that in most other utils and datasets, the default is actually "THWC". + + Returns: + tuple: A 3-tuple with the following entries: + + - video (Tensor[T, C, H, W] or Tensor[T, H, W, C]): the `T` video frames in torch.uint8 tensor + - audio(Tensor[K, L]): the audio frames, where `K` is the number of channels + and `L` is the number of points in torch.float tensor + - label (int): class of the video clip + + Raises: + RuntimeError: If ``download is True`` and the video archives are already extracted. + """ + + _TAR_URLS = { + "400": "https://s3.amazonaws.com/kinetics/400/{split}/k400_{split}_path.txt", + "600": "https://s3.amazonaws.com/kinetics/600/{split}/k600_{split}_path.txt", + "700": "https://s3.amazonaws.com/kinetics/700_2020/{split}/k700_2020_{split}_path.txt", + } + _ANNOTATION_URLS = { + "400": "https://s3.amazonaws.com/kinetics/400/annotations/{split}.csv", + "600": "https://s3.amazonaws.com/kinetics/600/annotations/{split}.csv", + "700": "https://s3.amazonaws.com/kinetics/700_2020/annotations/{split}.csv", + } + + def __init__( + self, + root: Union[str, Path], + frames_per_clip: int, + num_classes: str = "400", + split: str = "train", + frame_rate: Optional[int] = None, + step_between_clips: int = 1, + transform: Optional[Callable] = None, + extensions: Tuple[str, ...] = ("avi", "mp4"), + download: bool = False, + num_download_workers: int = 1, + num_workers: int = 1, + _precomputed_metadata: Optional[Dict[str, Any]] = None, + _video_width: int = 0, + _video_height: int = 0, + _video_min_dimension: int = 0, + _audio_samples: int = 0, + _audio_channels: int = 0, + _legacy: bool = False, + output_format: str = "TCHW", + ) -> None: + + # TODO: support test + self.num_classes = verify_str_arg(num_classes, arg="num_classes", valid_values=["400", "600", "700"]) + self.extensions = extensions + self.num_download_workers = num_download_workers + + self.root = root + self._legacy = _legacy + + if _legacy: + self.split_folder = root + self.split = "unknown" + output_format = "THWC" + if download: + raise ValueError("Cannot download the videos using legacy_structure.") + else: + self.split_folder = path.join(root, split) + self.split = verify_str_arg(split, arg="split", valid_values=["train", "val", "test"]) + + if download: + self.download_and_process_videos() + + super().__init__(self.root) + + self.classes, class_to_idx = find_classes(self.split_folder) + self.samples = make_dataset(self.split_folder, class_to_idx, extensions, is_valid_file=None) + video_list = [x[0] for x in self.samples] + self.video_clips = VideoClips( + video_list, + frames_per_clip, + step_between_clips, + frame_rate, + _precomputed_metadata, + num_workers=num_workers, + _video_width=_video_width, + _video_height=_video_height, + _video_min_dimension=_video_min_dimension, + _audio_samples=_audio_samples, + _audio_channels=_audio_channels, + output_format=output_format, + ) + self.transform = transform + + def download_and_process_videos(self) -> None: + """Downloads all the videos to the _root_ folder in the expected format.""" + self._download_videos() + self._make_ds_structure() + + def _download_videos(self) -> None: + """download tarballs containing the video to "tars" folder and extract them into the _split_ folder where + split is one of the official dataset splits. + + Raises: + RuntimeError: if download folder exists, break to prevent downloading entire dataset again. + """ + if path.exists(self.split_folder): + return + tar_path = path.join(self.root, "tars") + file_list_path = path.join(self.root, "files") + + split_url = self._TAR_URLS[self.num_classes].format(split=self.split) + split_url_filepath = path.join(file_list_path, path.basename(split_url)) + if not check_integrity(split_url_filepath): + download_url(split_url, file_list_path) + with open(split_url_filepath) as file: + list_video_urls = [urllib.parse.quote(line, safe="/,:") for line in file.read().splitlines()] + + if self.num_download_workers == 1: + for line in list_video_urls: + download_and_extract_archive(line, tar_path, self.split_folder) + else: + part = partial(_dl_wrap, tar_path, self.split_folder) + poolproc = Pool(self.num_download_workers) + poolproc.map(part, list_video_urls) + + def _make_ds_structure(self) -> None: + """move videos from + split_folder/ + ├── clip1.avi + ├── clip2.avi + + to the correct format as described below: + split_folder/ + ├── class1 + │ ├── clip1.avi + + """ + annotation_path = path.join(self.root, "annotations") + if not check_integrity(path.join(annotation_path, f"{self.split}.csv")): + download_url(self._ANNOTATION_URLS[self.num_classes].format(split=self.split), annotation_path) + annotations = path.join(annotation_path, f"{self.split}.csv") + + file_fmtstr = "{ytid}_{start:06}_{end:06}.mp4" + with open(annotations) as csvfile: + reader = csv.DictReader(csvfile) + for row in reader: + f = file_fmtstr.format( + ytid=row["youtube_id"], + start=int(row["time_start"]), + end=int(row["time_end"]), + ) + label = row["label"].replace(" ", "_").replace("'", "").replace("(", "").replace(")", "") + os.makedirs(path.join(self.split_folder, label), exist_ok=True) + downloaded_file = path.join(self.split_folder, f) + if path.isfile(downloaded_file): + os.replace( + downloaded_file, + path.join(self.split_folder, label, f), + ) + + @property + def metadata(self) -> Dict[str, Any]: + return self.video_clips.metadata + + def __len__(self) -> int: + return self.video_clips.num_clips() + + def __getitem__(self, idx: int) -> Tuple[Tensor, Tensor, int]: + video, audio, info, video_idx = self.video_clips.get_clip(idx) + label = self.samples[video_idx][1] + + if self.transform is not None: + video = self.transform(video) + + return video, audio, label diff --git a/lib/python3.12/site-packages/torchvision/datasets/kitti.py b/lib/python3.12/site-packages/torchvision/datasets/kitti.py new file mode 100644 index 0000000000000000000000000000000000000000..69e603c76f22bb6387a4ab6e42bb4d92f49a0bee --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/kitti.py @@ -0,0 +1,158 @@ +import csv +import os +from pathlib import Path +from typing import Any, Callable, List, Optional, Tuple, Union + +from PIL import Image + +from .utils import download_and_extract_archive +from .vision import VisionDataset + + +class Kitti(VisionDataset): + """`KITTI `_ Dataset. + + It corresponds to the "left color images of object" dataset, for object detection. + + Args: + root (str or ``pathlib.Path``): Root directory where images are downloaded to. + Expects the following folder structure if download=False: + + .. code:: + + + └── Kitti + └─ raw + ├── training + | ├── image_2 + | └── label_2 + └── testing + └── image_2 + train (bool, optional): Use ``train`` split if true, else ``test`` split. + Defaults to ``train``. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.PILToTensor`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + transforms (callable, optional): A function/transform that takes input sample + and its target as entry and returns a transformed version. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + + """ + + data_url = "https://s3.eu-central-1.amazonaws.com/avg-kitti/" + resources = [ + "data_object_image_2.zip", + "data_object_label_2.zip", + ] + image_dir_name = "image_2" + labels_dir_name = "label_2" + + def __init__( + self, + root: Union[str, Path], + train: bool = True, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + transforms: Optional[Callable] = None, + download: bool = False, + ): + super().__init__( + root, + transform=transform, + target_transform=target_transform, + transforms=transforms, + ) + self.images = [] + self.targets = [] + self.train = train + self._location = "training" if self.train else "testing" + + if download: + self.download() + if not self._check_exists(): + raise RuntimeError("Dataset not found. You may use download=True to download it.") + + image_dir = os.path.join(self._raw_folder, self._location, self.image_dir_name) + if self.train: + labels_dir = os.path.join(self._raw_folder, self._location, self.labels_dir_name) + for img_file in os.listdir(image_dir): + self.images.append(os.path.join(image_dir, img_file)) + if self.train: + self.targets.append(os.path.join(labels_dir, f"{img_file.split('.')[0]}.txt")) + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """Get item at a given index. + + Args: + index (int): Index + Returns: + tuple: (image, target), where + target is a list of dictionaries with the following keys: + + - type: str + - truncated: float + - occluded: int + - alpha: float + - bbox: float[4] + - dimensions: float[3] + - locations: float[3] + - rotation_y: float + + """ + image = Image.open(self.images[index]) + target = self._parse_target(index) if self.train else None + if self.transforms: + image, target = self.transforms(image, target) + return image, target + + def _parse_target(self, index: int) -> List: + target = [] + with open(self.targets[index]) as inp: + content = csv.reader(inp, delimiter=" ") + for line in content: + target.append( + { + "type": line[0], + "truncated": float(line[1]), + "occluded": int(line[2]), + "alpha": float(line[3]), + "bbox": [float(x) for x in line[4:8]], + "dimensions": [float(x) for x in line[8:11]], + "location": [float(x) for x in line[11:14]], + "rotation_y": float(line[14]), + } + ) + return target + + def __len__(self) -> int: + return len(self.images) + + @property + def _raw_folder(self) -> str: + return os.path.join(self.root, self.__class__.__name__, "raw") + + def _check_exists(self) -> bool: + """Check if the data directory exists.""" + folders = [self.image_dir_name] + if self.train: + folders.append(self.labels_dir_name) + return all(os.path.isdir(os.path.join(self._raw_folder, self._location, fname)) for fname in folders) + + def download(self) -> None: + """Download the KITTI data if it doesn't exist already.""" + + if self._check_exists(): + return + + os.makedirs(self._raw_folder, exist_ok=True) + + # download files + for fname in self.resources: + download_and_extract_archive( + url=f"{self.data_url}{fname}", + download_root=self._raw_folder, + filename=fname, + ) diff --git a/lib/python3.12/site-packages/torchvision/datasets/lfw.py b/lib/python3.12/site-packages/torchvision/datasets/lfw.py new file mode 100644 index 0000000000000000000000000000000000000000..18374fc3c9bea77e6631ce23bd67d20a5642082c --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/lfw.py @@ -0,0 +1,255 @@ +import os +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +from PIL import Image + +from .utils import check_integrity, download_and_extract_archive, download_url, verify_str_arg +from .vision import VisionDataset + + +class _LFW(VisionDataset): + + base_folder = "lfw-py" + download_url_prefix = "http://vis-www.cs.umass.edu/lfw/" + + file_dict = { + "original": ("lfw", "lfw.tgz", "a17d05bd522c52d84eca14327a23d494"), + "funneled": ("lfw_funneled", "lfw-funneled.tgz", "1b42dfed7d15c9b2dd63d5e5840c86ad"), + "deepfunneled": ("lfw-deepfunneled", "lfw-deepfunneled.tgz", "68331da3eb755a505a502b5aacb3c201"), + } + checksums = { + "pairs.txt": "9f1ba174e4e1c508ff7cdf10ac338a7d", + "pairsDevTest.txt": "5132f7440eb68cf58910c8a45a2ac10b", + "pairsDevTrain.txt": "4f27cbf15b2da4a85c1907eb4181ad21", + "people.txt": "450f0863dd89e85e73936a6d71a3474b", + "peopleDevTest.txt": "e4bf5be0a43b5dcd9dc5ccfcb8fb19c5", + "peopleDevTrain.txt": "54eaac34beb6d042ed3a7d883e247a21", + "lfw-names.txt": "a6d0a479bd074669f656265a6e693f6d", + } + annot_file = {"10fold": "", "train": "DevTrain", "test": "DevTest"} + names = "lfw-names.txt" + + def __init__( + self, + root: Union[str, Path], + split: str, + image_set: str, + view: str, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + super().__init__(os.path.join(root, self.base_folder), transform=transform, target_transform=target_transform) + + self.image_set = verify_str_arg(image_set.lower(), "image_set", self.file_dict.keys()) + images_dir, self.filename, self.md5 = self.file_dict[self.image_set] + + self.view = verify_str_arg(view.lower(), "view", ["people", "pairs"]) + self.split = verify_str_arg(split.lower(), "split", ["10fold", "train", "test"]) + self.labels_file = f"{self.view}{self.annot_file[self.split]}.txt" + self.data: List[Any] = [] + + if download: + self.download() + + if not self._check_integrity(): + raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") + + self.images_dir = os.path.join(self.root, images_dir) + + def _loader(self, path: str) -> Image.Image: + with open(path, "rb") as f: + img = Image.open(f) + return img.convert("RGB") + + def _check_integrity(self) -> bool: + st1 = check_integrity(os.path.join(self.root, self.filename), self.md5) + st2 = check_integrity(os.path.join(self.root, self.labels_file), self.checksums[self.labels_file]) + if not st1 or not st2: + return False + if self.view == "people": + return check_integrity(os.path.join(self.root, self.names), self.checksums[self.names]) + return True + + def download(self) -> None: + if self._check_integrity(): + return + url = f"{self.download_url_prefix}{self.filename}" + download_and_extract_archive(url, self.root, filename=self.filename, md5=self.md5) + download_url(f"{self.download_url_prefix}{self.labels_file}", self.root) + if self.view == "people": + download_url(f"{self.download_url_prefix}{self.names}", self.root) + + def _get_path(self, identity: str, no: Union[int, str]) -> str: + return os.path.join(self.images_dir, identity, f"{identity}_{int(no):04d}.jpg") + + def extra_repr(self) -> str: + return f"Alignment: {self.image_set}\nSplit: {self.split}" + + def __len__(self) -> int: + return len(self.data) + + +class LFWPeople(_LFW): + """`LFW `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset where directory + ``lfw-py`` exists or will be saved to if download is set to True. + split (string, optional): The image split to use. Can be one of ``train``, ``test``, + ``10fold`` (default). + image_set (str, optional): Type of image funneling to use, ``original``, ``funneled`` or + ``deepfunneled``. Defaults to ``funneled``. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomRotation`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + + """ + + def __init__( + self, + root: str, + split: str = "10fold", + image_set: str = "funneled", + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + super().__init__(root, split, image_set, "people", transform, target_transform, download) + + self.class_to_idx = self._get_classes() + self.data, self.targets = self._get_people() + + def _get_people(self) -> Tuple[List[str], List[int]]: + data, targets = [], [] + with open(os.path.join(self.root, self.labels_file)) as f: + lines = f.readlines() + n_folds, s = (int(lines[0]), 1) if self.split == "10fold" else (1, 0) + + for fold in range(n_folds): + n_lines = int(lines[s]) + people = [line.strip().split("\t") for line in lines[s + 1 : s + n_lines + 1]] + s += n_lines + 1 + for i, (identity, num_imgs) in enumerate(people): + for num in range(1, int(num_imgs) + 1): + img = self._get_path(identity, num) + data.append(img) + targets.append(self.class_to_idx[identity]) + + return data, targets + + def _get_classes(self) -> Dict[str, int]: + with open(os.path.join(self.root, self.names)) as f: + lines = f.readlines() + names = [line.strip().split()[0] for line in lines] + class_to_idx = {name: i for i, name in enumerate(names)} + return class_to_idx + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: Tuple (image, target) where target is the identity of the person. + """ + img = self._loader(self.data[index]) + target = self.targets[index] + + if self.transform is not None: + img = self.transform(img) + + if self.target_transform is not None: + target = self.target_transform(target) + + return img, target + + def extra_repr(self) -> str: + return super().extra_repr() + f"\nClasses (identities): {len(self.class_to_idx)}" + + +class LFWPairs(_LFW): + """`LFW `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset where directory + ``lfw-py`` exists or will be saved to if download is set to True. + split (string, optional): The image split to use. Can be one of ``train``, ``test``, + ``10fold``. Defaults to ``10fold``. + image_set (str, optional): Type of image funneling to use, ``original``, ``funneled`` or + ``deepfunneled``. Defaults to ``funneled``. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomRotation`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + + """ + + def __init__( + self, + root: str, + split: str = "10fold", + image_set: str = "funneled", + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + super().__init__(root, split, image_set, "pairs", transform, target_transform, download) + + self.pair_names, self.data, self.targets = self._get_pairs(self.images_dir) + + def _get_pairs(self, images_dir: str) -> Tuple[List[Tuple[str, str]], List[Tuple[str, str]], List[int]]: + pair_names, data, targets = [], [], [] + with open(os.path.join(self.root, self.labels_file)) as f: + lines = f.readlines() + if self.split == "10fold": + n_folds, n_pairs = lines[0].split("\t") + n_folds, n_pairs = int(n_folds), int(n_pairs) + else: + n_folds, n_pairs = 1, int(lines[0]) + s = 1 + + for fold in range(n_folds): + matched_pairs = [line.strip().split("\t") for line in lines[s : s + n_pairs]] + unmatched_pairs = [line.strip().split("\t") for line in lines[s + n_pairs : s + (2 * n_pairs)]] + s += 2 * n_pairs + for pair in matched_pairs: + img1, img2, same = self._get_path(pair[0], pair[1]), self._get_path(pair[0], pair[2]), 1 + pair_names.append((pair[0], pair[0])) + data.append((img1, img2)) + targets.append(same) + for pair in unmatched_pairs: + img1, img2, same = self._get_path(pair[0], pair[1]), self._get_path(pair[2], pair[3]), 0 + pair_names.append((pair[0], pair[2])) + data.append((img1, img2)) + targets.append(same) + + return pair_names, data, targets + + def __getitem__(self, index: int) -> Tuple[Any, Any, int]: + """ + Args: + index (int): Index + + Returns: + tuple: (image1, image2, target) where target is `0` for different indentities and `1` for same identities. + """ + img1, img2 = self.data[index] + img1, img2 = self._loader(img1), self._loader(img2) + target = self.targets[index] + + if self.transform is not None: + img1, img2 = self.transform(img1), self.transform(img2) + + if self.target_transform is not None: + target = self.target_transform(target) + + return img1, img2, target diff --git a/lib/python3.12/site-packages/torchvision/datasets/lsun.py b/lib/python3.12/site-packages/torchvision/datasets/lsun.py new file mode 100644 index 0000000000000000000000000000000000000000..61d40eee221dc1656e00b4927c64de76129156f6 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/lsun.py @@ -0,0 +1,168 @@ +import io +import os.path +import pickle +import string +from collections.abc import Iterable +from pathlib import Path +from typing import Any, Callable, cast, List, Optional, Tuple, Union + +from PIL import Image + +from .utils import iterable_to_str, verify_str_arg +from .vision import VisionDataset + + +class LSUNClass(VisionDataset): + def __init__( + self, root: str, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None + ) -> None: + import lmdb + + super().__init__(root, transform=transform, target_transform=target_transform) + + self.env = lmdb.open(root, max_readers=1, readonly=True, lock=False, readahead=False, meminit=False) + with self.env.begin(write=False) as txn: + self.length = txn.stat()["entries"] + cache_file = "_cache_" + "".join(c for c in root if c in string.ascii_letters) + if os.path.isfile(cache_file): + self.keys = pickle.load(open(cache_file, "rb")) + else: + with self.env.begin(write=False) as txn: + self.keys = [key for key in txn.cursor().iternext(keys=True, values=False)] + pickle.dump(self.keys, open(cache_file, "wb")) + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + img, target = None, None + env = self.env + with env.begin(write=False) as txn: + imgbuf = txn.get(self.keys[index]) + + buf = io.BytesIO() + buf.write(imgbuf) + buf.seek(0) + img = Image.open(buf).convert("RGB") + + if self.transform is not None: + img = self.transform(img) + + if self.target_transform is not None: + target = self.target_transform(target) + + return img, target + + def __len__(self) -> int: + return self.length + + +class LSUN(VisionDataset): + """`LSUN `_ dataset. + + You will need to install the ``lmdb`` package to use this dataset: run + ``pip install lmdb`` + + Args: + root (str or ``pathlib.Path``): Root directory for the database files. + classes (string or list): One of {'train', 'val', 'test'} or a list of + categories to load. e,g. ['bedroom_train', 'church_outdoor_train']. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + """ + + def __init__( + self, + root: Union[str, Path], + classes: Union[str, List[str]] = "train", + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + ) -> None: + super().__init__(root, transform=transform, target_transform=target_transform) + self.classes = self._verify_classes(classes) + + # for each class, create an LSUNClassDataset + self.dbs = [] + for c in self.classes: + self.dbs.append(LSUNClass(root=os.path.join(root, f"{c}_lmdb"), transform=transform)) + + self.indices = [] + count = 0 + for db in self.dbs: + count += len(db) + self.indices.append(count) + + self.length = count + + def _verify_classes(self, classes: Union[str, List[str]]) -> List[str]: + categories = [ + "bedroom", + "bridge", + "church_outdoor", + "classroom", + "conference_room", + "dining_room", + "kitchen", + "living_room", + "restaurant", + "tower", + ] + dset_opts = ["train", "val", "test"] + + try: + classes = cast(str, classes) + verify_str_arg(classes, "classes", dset_opts) + if classes == "test": + classes = [classes] + else: + classes = [c + "_" + classes for c in categories] + except ValueError: + if not isinstance(classes, Iterable): + msg = "Expected type str or Iterable for argument classes, but got type {}." + raise ValueError(msg.format(type(classes))) + + classes = list(classes) + msg_fmtstr_type = "Expected type str for elements in argument classes, but got type {}." + for c in classes: + verify_str_arg(c, custom_msg=msg_fmtstr_type.format(type(c))) + c_short = c.split("_") + category, dset_opt = "_".join(c_short[:-1]), c_short[-1] + + msg_fmtstr = "Unknown value '{}' for {}. Valid values are {{{}}}." + msg = msg_fmtstr.format(category, "LSUN class", iterable_to_str(categories)) + verify_str_arg(category, valid_values=categories, custom_msg=msg) + + msg = msg_fmtstr.format(dset_opt, "postfix", iterable_to_str(dset_opts)) + verify_str_arg(dset_opt, valid_values=dset_opts, custom_msg=msg) + + return classes + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: Tuple (image, target) where target is the index of the target category. + """ + target = 0 + sub = 0 + for ind in self.indices: + if index < ind: + break + target += 1 + sub = ind + + db = self.dbs[target] + index = index - sub + + if self.target_transform is not None: + target = self.target_transform(target) + + img, _ = db[index] + return img, target + + def __len__(self) -> int: + return self.length + + def extra_repr(self) -> str: + return "Classes: {classes}".format(**self.__dict__) diff --git a/lib/python3.12/site-packages/torchvision/datasets/mnist.py b/lib/python3.12/site-packages/torchvision/datasets/mnist.py new file mode 100644 index 0000000000000000000000000000000000000000..fd1455535298440467e5750c981529c7de50e493 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/mnist.py @@ -0,0 +1,560 @@ +import codecs +import os +import os.path +import shutil +import string +import sys +import warnings +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from urllib.error import URLError + +import numpy as np +import torch +from PIL import Image + +from .utils import _flip_byte_order, check_integrity, download_and_extract_archive, extract_archive, verify_str_arg +from .vision import VisionDataset + + +class MNIST(VisionDataset): + """`MNIST `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset where ``MNIST/raw/train-images-idx3-ubyte`` + and ``MNIST/raw/t10k-images-idx3-ubyte`` exist. + train (bool, optional): If True, creates dataset from ``train-images-idx3-ubyte``, + otherwise from ``t10k-images-idx3-ubyte``. + download (bool, optional): If True, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + """ + + mirrors = [ + "http://yann.lecun.com/exdb/mnist/", + "https://ossci-datasets.s3.amazonaws.com/mnist/", + ] + + resources = [ + ("train-images-idx3-ubyte.gz", "f68b3c2dcbeaaa9fbdd348bbdeb94873"), + ("train-labels-idx1-ubyte.gz", "d53e105ee54ea40749a09fcbcd1e9432"), + ("t10k-images-idx3-ubyte.gz", "9fb629c4189551a2d022fa330f9573f3"), + ("t10k-labels-idx1-ubyte.gz", "ec29112dd5afa0611ce80d1b7f02629c"), + ] + + training_file = "training.pt" + test_file = "test.pt" + classes = [ + "0 - zero", + "1 - one", + "2 - two", + "3 - three", + "4 - four", + "5 - five", + "6 - six", + "7 - seven", + "8 - eight", + "9 - nine", + ] + + @property + def train_labels(self): + warnings.warn("train_labels has been renamed targets") + return self.targets + + @property + def test_labels(self): + warnings.warn("test_labels has been renamed targets") + return self.targets + + @property + def train_data(self): + warnings.warn("train_data has been renamed data") + return self.data + + @property + def test_data(self): + warnings.warn("test_data has been renamed data") + return self.data + + def __init__( + self, + root: Union[str, Path], + train: bool = True, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + super().__init__(root, transform=transform, target_transform=target_transform) + self.train = train # training set or test set + + if self._check_legacy_exist(): + self.data, self.targets = self._load_legacy_data() + return + + if download: + self.download() + + if not self._check_exists(): + raise RuntimeError("Dataset not found. You can use download=True to download it") + + self.data, self.targets = self._load_data() + + def _check_legacy_exist(self): + processed_folder_exists = os.path.exists(self.processed_folder) + if not processed_folder_exists: + return False + + return all( + check_integrity(os.path.join(self.processed_folder, file)) for file in (self.training_file, self.test_file) + ) + + def _load_legacy_data(self): + # This is for BC only. We no longer cache the data in a custom binary, but simply read from the raw data + # directly. + data_file = self.training_file if self.train else self.test_file + return torch.load(os.path.join(self.processed_folder, data_file), weights_only=True) + + def _load_data(self): + image_file = f"{'train' if self.train else 't10k'}-images-idx3-ubyte" + data = read_image_file(os.path.join(self.raw_folder, image_file)) + + label_file = f"{'train' if self.train else 't10k'}-labels-idx1-ubyte" + targets = read_label_file(os.path.join(self.raw_folder, label_file)) + + return data, targets + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: (image, target) where target is index of the target class. + """ + img, target = self.data[index], int(self.targets[index]) + + # doing this so that it is consistent with all other datasets + # to return a PIL Image + img = Image.fromarray(img.numpy(), mode="L") + + if self.transform is not None: + img = self.transform(img) + + if self.target_transform is not None: + target = self.target_transform(target) + + return img, target + + def __len__(self) -> int: + return len(self.data) + + @property + def raw_folder(self) -> str: + return os.path.join(self.root, self.__class__.__name__, "raw") + + @property + def processed_folder(self) -> str: + return os.path.join(self.root, self.__class__.__name__, "processed") + + @property + def class_to_idx(self) -> Dict[str, int]: + return {_class: i for i, _class in enumerate(self.classes)} + + def _check_exists(self) -> bool: + return all( + check_integrity(os.path.join(self.raw_folder, os.path.splitext(os.path.basename(url))[0])) + for url, _ in self.resources + ) + + def download(self) -> None: + """Download the MNIST data if it doesn't exist already.""" + + if self._check_exists(): + return + + os.makedirs(self.raw_folder, exist_ok=True) + + # download files + for filename, md5 in self.resources: + errors = [] + for mirror in self.mirrors: + url = f"{mirror}{filename}" + try: + download_and_extract_archive(url, download_root=self.raw_folder, filename=filename, md5=md5) + except URLError as e: + errors.append(e) + continue + break + else: + s = f"Error downloading {filename}:\n" + for mirror, err in zip(self.mirrors, errors): + s += f"Tried {mirror}, got:\n{str(err)}\n" + raise RuntimeError(s) + + def extra_repr(self) -> str: + split = "Train" if self.train is True else "Test" + return f"Split: {split}" + + +class FashionMNIST(MNIST): + """`Fashion-MNIST `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset where ``FashionMNIST/raw/train-images-idx3-ubyte`` + and ``FashionMNIST/raw/t10k-images-idx3-ubyte`` exist. + train (bool, optional): If True, creates dataset from ``train-images-idx3-ubyte``, + otherwise from ``t10k-images-idx3-ubyte``. + download (bool, optional): If True, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + """ + + mirrors = ["http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/"] + + resources = [ + ("train-images-idx3-ubyte.gz", "8d4fb7e6c68d591d4c3dfef9ec88bf0d"), + ("train-labels-idx1-ubyte.gz", "25c81989df183df01b3e8a0aad5dffbe"), + ("t10k-images-idx3-ubyte.gz", "bef4ecab320f06d8554ea6380940ec79"), + ("t10k-labels-idx1-ubyte.gz", "bb300cfdad3c16e7a12a480ee83cd310"), + ] + classes = ["T-shirt/top", "Trouser", "Pullover", "Dress", "Coat", "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"] + + +class KMNIST(MNIST): + """`Kuzushiji-MNIST `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset where ``KMNIST/raw/train-images-idx3-ubyte`` + and ``KMNIST/raw/t10k-images-idx3-ubyte`` exist. + train (bool, optional): If True, creates dataset from ``train-images-idx3-ubyte``, + otherwise from ``t10k-images-idx3-ubyte``. + download (bool, optional): If True, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + """ + + mirrors = ["http://codh.rois.ac.jp/kmnist/dataset/kmnist/"] + + resources = [ + ("train-images-idx3-ubyte.gz", "bdb82020997e1d708af4cf47b453dcf7"), + ("train-labels-idx1-ubyte.gz", "e144d726b3acfaa3e44228e80efcd344"), + ("t10k-images-idx3-ubyte.gz", "5c965bf0a639b31b8f53240b1b52f4d7"), + ("t10k-labels-idx1-ubyte.gz", "7320c461ea6c1c855c0b718fb2a4b134"), + ] + classes = ["o", "ki", "su", "tsu", "na", "ha", "ma", "ya", "re", "wo"] + + +class EMNIST(MNIST): + """`EMNIST `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset where ``EMNIST/raw/train-images-idx3-ubyte`` + and ``EMNIST/raw/t10k-images-idx3-ubyte`` exist. + split (string): The dataset has 6 different splits: ``byclass``, ``bymerge``, + ``balanced``, ``letters``, ``digits`` and ``mnist``. This argument specifies + which one to use. + train (bool, optional): If True, creates dataset from ``training.pt``, + otherwise from ``test.pt``. + download (bool, optional): If True, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + """ + + url = "https://biometrics.nist.gov/cs_links/EMNIST/gzip.zip" + md5 = "58c8d27c78d21e728a6bc7b3cc06412e" + splits = ("byclass", "bymerge", "balanced", "letters", "digits", "mnist") + # Merged Classes assumes Same structure for both uppercase and lowercase version + _merged_classes = {"c", "i", "j", "k", "l", "m", "o", "p", "s", "u", "v", "w", "x", "y", "z"} + _all_classes = set(string.digits + string.ascii_letters) + classes_split_dict = { + "byclass": sorted(list(_all_classes)), + "bymerge": sorted(list(_all_classes - _merged_classes)), + "balanced": sorted(list(_all_classes - _merged_classes)), + "letters": ["N/A"] + list(string.ascii_lowercase), + "digits": list(string.digits), + "mnist": list(string.digits), + } + + def __init__(self, root: Union[str, Path], split: str, **kwargs: Any) -> None: + self.split = verify_str_arg(split, "split", self.splits) + self.training_file = self._training_file(split) + self.test_file = self._test_file(split) + super().__init__(root, **kwargs) + self.classes = self.classes_split_dict[self.split] + + @staticmethod + def _training_file(split) -> str: + return f"training_{split}.pt" + + @staticmethod + def _test_file(split) -> str: + return f"test_{split}.pt" + + @property + def _file_prefix(self) -> str: + return f"emnist-{self.split}-{'train' if self.train else 'test'}" + + @property + def images_file(self) -> str: + return os.path.join(self.raw_folder, f"{self._file_prefix}-images-idx3-ubyte") + + @property + def labels_file(self) -> str: + return os.path.join(self.raw_folder, f"{self._file_prefix}-labels-idx1-ubyte") + + def _load_data(self): + return read_image_file(self.images_file), read_label_file(self.labels_file) + + def _check_exists(self) -> bool: + return all(check_integrity(file) for file in (self.images_file, self.labels_file)) + + def download(self) -> None: + """Download the EMNIST data if it doesn't exist already.""" + + if self._check_exists(): + return + + os.makedirs(self.raw_folder, exist_ok=True) + + download_and_extract_archive(self.url, download_root=self.raw_folder, md5=self.md5) + gzip_folder = os.path.join(self.raw_folder, "gzip") + for gzip_file in os.listdir(gzip_folder): + if gzip_file.endswith(".gz"): + extract_archive(os.path.join(gzip_folder, gzip_file), self.raw_folder) + shutil.rmtree(gzip_folder) + + +class QMNIST(MNIST): + """`QMNIST `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset whose ``raw`` + subdir contains binary files of the datasets. + what (string,optional): Can be 'train', 'test', 'test10k', + 'test50k', or 'nist' for respectively the mnist compatible + training set, the 60k qmnist testing set, the 10k qmnist + examples that match the mnist testing set, the 50k + remaining qmnist testing examples, or all the nist + digits. The default is to select 'train' or 'test' + according to the compatibility argument 'train'. + compat (bool,optional): A boolean that says whether the target + for each example is class number (for compatibility with + the MNIST dataloader) or a torch vector containing the + full qmnist information. Default=True. + download (bool, optional): If True, downloads the dataset from + the internet and puts it in root directory. If dataset is + already downloaded, it is not downloaded again. + transform (callable, optional): A function/transform that + takes in a PIL image and returns a transformed + version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform + that takes in the target and transforms it. + train (bool,optional,compatibility): When argument 'what' is + not specified, this boolean decides whether to load the + training set or the testing set. Default: True. + """ + + subsets = {"train": "train", "test": "test", "test10k": "test", "test50k": "test", "nist": "nist"} + resources: Dict[str, List[Tuple[str, str]]] = { # type: ignore[assignment] + "train": [ + ( + "https://raw.githubusercontent.com/facebookresearch/qmnist/master/qmnist-train-images-idx3-ubyte.gz", + "ed72d4157d28c017586c42bc6afe6370", + ), + ( + "https://raw.githubusercontent.com/facebookresearch/qmnist/master/qmnist-train-labels-idx2-int.gz", + "0058f8dd561b90ffdd0f734c6a30e5e4", + ), + ], + "test": [ + ( + "https://raw.githubusercontent.com/facebookresearch/qmnist/master/qmnist-test-images-idx3-ubyte.gz", + "1394631089c404de565df7b7aeaf9412", + ), + ( + "https://raw.githubusercontent.com/facebookresearch/qmnist/master/qmnist-test-labels-idx2-int.gz", + "5b5b05890a5e13444e108efe57b788aa", + ), + ], + "nist": [ + ( + "https://raw.githubusercontent.com/facebookresearch/qmnist/master/xnist-images-idx3-ubyte.xz", + "7f124b3b8ab81486c9d8c2749c17f834", + ), + ( + "https://raw.githubusercontent.com/facebookresearch/qmnist/master/xnist-labels-idx2-int.xz", + "5ed0e788978e45d4a8bd4b7caec3d79d", + ), + ], + } + classes = [ + "0 - zero", + "1 - one", + "2 - two", + "3 - three", + "4 - four", + "5 - five", + "6 - six", + "7 - seven", + "8 - eight", + "9 - nine", + ] + + def __init__( + self, root: Union[str, Path], what: Optional[str] = None, compat: bool = True, train: bool = True, **kwargs: Any + ) -> None: + if what is None: + what = "train" if train else "test" + self.what = verify_str_arg(what, "what", tuple(self.subsets.keys())) + self.compat = compat + self.data_file = what + ".pt" + self.training_file = self.data_file + self.test_file = self.data_file + super().__init__(root, train, **kwargs) + + @property + def images_file(self) -> str: + (url, _), _ = self.resources[self.subsets[self.what]] + return os.path.join(self.raw_folder, os.path.splitext(os.path.basename(url))[0]) + + @property + def labels_file(self) -> str: + _, (url, _) = self.resources[self.subsets[self.what]] + return os.path.join(self.raw_folder, os.path.splitext(os.path.basename(url))[0]) + + def _check_exists(self) -> bool: + return all(check_integrity(file) for file in (self.images_file, self.labels_file)) + + def _load_data(self): + data = read_sn3_pascalvincent_tensor(self.images_file) + if data.dtype != torch.uint8: + raise TypeError(f"data should be of dtype torch.uint8 instead of {data.dtype}") + if data.ndimension() != 3: + raise ValueError("data should have 3 dimensions instead of {data.ndimension()}") + + targets = read_sn3_pascalvincent_tensor(self.labels_file).long() + if targets.ndimension() != 2: + raise ValueError(f"targets should have 2 dimensions instead of {targets.ndimension()}") + + if self.what == "test10k": + data = data[0:10000, :, :].clone() + targets = targets[0:10000, :].clone() + elif self.what == "test50k": + data = data[10000:, :, :].clone() + targets = targets[10000:, :].clone() + + return data, targets + + def download(self) -> None: + """Download the QMNIST data if it doesn't exist already. + Note that we only download what has been asked for (argument 'what'). + """ + if self._check_exists(): + return + + os.makedirs(self.raw_folder, exist_ok=True) + split = self.resources[self.subsets[self.what]] + + for url, md5 in split: + download_and_extract_archive(url, self.raw_folder, md5=md5) + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + # redefined to handle the compat flag + img, target = self.data[index], self.targets[index] + img = Image.fromarray(img.numpy(), mode="L") + if self.transform is not None: + img = self.transform(img) + if self.compat: + target = int(target[0]) + if self.target_transform is not None: + target = self.target_transform(target) + return img, target + + def extra_repr(self) -> str: + return f"Split: {self.what}" + + +def get_int(b: bytes) -> int: + return int(codecs.encode(b, "hex"), 16) + + +SN3_PASCALVINCENT_TYPEMAP = { + 8: torch.uint8, + 9: torch.int8, + 11: torch.int16, + 12: torch.int32, + 13: torch.float32, + 14: torch.float64, +} + + +def read_sn3_pascalvincent_tensor(path: str, strict: bool = True) -> torch.Tensor: + """Read a SN3 file in "Pascal Vincent" format (Lush file 'libidx/idx-io.lsh'). + Argument may be a filename, compressed filename, or file object. + """ + # read + with open(path, "rb") as f: + data = f.read() + + # parse + if sys.byteorder == "little": + magic = get_int(data[0:4]) + nd = magic % 256 + ty = magic // 256 + else: + nd = get_int(data[0:1]) + ty = get_int(data[1:2]) + get_int(data[2:3]) * 256 + get_int(data[3:4]) * 256 * 256 + + assert 1 <= nd <= 3 + assert 8 <= ty <= 14 + torch_type = SN3_PASCALVINCENT_TYPEMAP[ty] + s = [get_int(data[4 * (i + 1) : 4 * (i + 2)]) for i in range(nd)] + + if sys.byteorder == "big": + for i in range(len(s)): + s[i] = int.from_bytes(s[i].to_bytes(4, byteorder="little"), byteorder="big", signed=False) + + parsed = torch.frombuffer(bytearray(data), dtype=torch_type, offset=(4 * (nd + 1))) + + # The MNIST format uses the big endian byte order, while `torch.frombuffer` uses whatever the system uses. In case + # that is little endian and the dtype has more than one byte, we need to flip them. + if sys.byteorder == "little" and parsed.element_size() > 1: + parsed = _flip_byte_order(parsed) + + assert parsed.shape[0] == np.prod(s) or not strict + return parsed.view(*s) + + +def read_label_file(path: str) -> torch.Tensor: + x = read_sn3_pascalvincent_tensor(path, strict=False) + if x.dtype != torch.uint8: + raise TypeError(f"x should be of dtype torch.uint8 instead of {x.dtype}") + if x.ndimension() != 1: + raise ValueError(f"x should have 1 dimension instead of {x.ndimension()}") + return x.long() + + +def read_image_file(path: str) -> torch.Tensor: + x = read_sn3_pascalvincent_tensor(path, strict=False) + if x.dtype != torch.uint8: + raise TypeError(f"x should be of dtype torch.uint8 instead of {x.dtype}") + if x.ndimension() != 3: + raise ValueError(f"x should have 3 dimension instead of {x.ndimension()}") + return x diff --git a/lib/python3.12/site-packages/torchvision/datasets/moving_mnist.py b/lib/python3.12/site-packages/torchvision/datasets/moving_mnist.py new file mode 100644 index 0000000000000000000000000000000000000000..48715de4e8dea86fec71eb09ea98f3ff1b9d6fda --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/moving_mnist.py @@ -0,0 +1,94 @@ +import os.path +from pathlib import Path +from typing import Callable, Optional, Union + +import numpy as np +import torch +from torchvision.datasets.utils import download_url, verify_str_arg +from torchvision.datasets.vision import VisionDataset + + +class MovingMNIST(VisionDataset): + """`MovingMNIST `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset where ``MovingMNIST/mnist_test_seq.npy`` exists. + split (string, optional): The dataset split, supports ``None`` (default), ``"train"`` and ``"test"``. + If ``split=None``, the full data is returned. + split_ratio (int, optional): The split ratio of number of frames. If ``split="train"``, the first split + frames ``data[:, :split_ratio]`` is returned. If ``split="test"``, the last split frames ``data[:, split_ratio:]`` + is returned. If ``split=None``, this parameter is ignored and the all frames data is returned. + transform (callable, optional): A function/transform that takes in a torch Tensor + and returns a transformed version. E.g, ``transforms.RandomCrop`` + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + """ + + _URL = "http://www.cs.toronto.edu/~nitish/unsupervised_video/mnist_test_seq.npy" + + def __init__( + self, + root: Union[str, Path], + split: Optional[str] = None, + split_ratio: int = 10, + download: bool = False, + transform: Optional[Callable] = None, + ) -> None: + super().__init__(root, transform=transform) + + self._base_folder = os.path.join(self.root, self.__class__.__name__) + self._filename = self._URL.split("/")[-1] + + if split is not None: + verify_str_arg(split, "split", ("train", "test")) + self.split = split + + if not isinstance(split_ratio, int): + raise TypeError(f"`split_ratio` should be an integer, but got {type(split_ratio)}") + elif not (1 <= split_ratio <= 19): + raise ValueError(f"`split_ratio` should be `1 <= split_ratio <= 19`, but got {split_ratio} instead.") + self.split_ratio = split_ratio + + if download: + self.download() + + if not self._check_exists(): + raise RuntimeError("Dataset not found. You can use download=True to download it.") + + data = torch.from_numpy(np.load(os.path.join(self._base_folder, self._filename))) + if self.split == "train": + data = data[: self.split_ratio] + elif self.split == "test": + data = data[self.split_ratio :] + self.data = data.transpose(0, 1).unsqueeze(2).contiguous() + + def __getitem__(self, idx: int) -> torch.Tensor: + """ + Args: + idx (int): Index + Returns: + torch.Tensor: Video frames (torch Tensor[T, C, H, W]). The `T` is the number of frames. + """ + data = self.data[idx] + if self.transform is not None: + data = self.transform(data) + + return data + + def __len__(self) -> int: + return len(self.data) + + def _check_exists(self) -> bool: + return os.path.exists(os.path.join(self._base_folder, self._filename)) + + def download(self) -> None: + if self._check_exists(): + return + + download_url( + url=self._URL, + root=self._base_folder, + filename=self._filename, + md5="be083ec986bfe91a449d63653c411eb2", + ) diff --git a/lib/python3.12/site-packages/torchvision/datasets/omniglot.py b/lib/python3.12/site-packages/torchvision/datasets/omniglot.py new file mode 100644 index 0000000000000000000000000000000000000000..c3434a7245616178b5a0aab9efe6643cc33bff14 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/omniglot.py @@ -0,0 +1,102 @@ +from os.path import join +from pathlib import Path +from typing import Any, Callable, List, Optional, Tuple, Union + +from PIL import Image + +from .utils import check_integrity, download_and_extract_archive, list_dir, list_files +from .vision import VisionDataset + + +class Omniglot(VisionDataset): + """`Omniglot `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset where directory + ``omniglot-py`` exists. + background (bool, optional): If True, creates dataset from the "background" set, otherwise + creates from the "evaluation" set. This terminology is defined by the authors. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + download (bool, optional): If true, downloads the dataset zip files from the internet and + puts it in root directory. If the zip files are already downloaded, they are not + downloaded again. + """ + + folder = "omniglot-py" + download_url_prefix = "https://raw.githubusercontent.com/brendenlake/omniglot/master/python" + zips_md5 = { + "images_background": "68d2efa1b9178cc56df9314c21c6e718", + "images_evaluation": "6b91aef0f799c5bb55b94e3f2daec811", + } + + def __init__( + self, + root: Union[str, Path], + background: bool = True, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + super().__init__(join(root, self.folder), transform=transform, target_transform=target_transform) + self.background = background + + if download: + self.download() + + if not self._check_integrity(): + raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") + + self.target_folder = join(self.root, self._get_target_folder()) + self._alphabets = list_dir(self.target_folder) + self._characters: List[str] = sum( + ([join(a, c) for c in list_dir(join(self.target_folder, a))] for a in self._alphabets), [] + ) + self._character_images = [ + [(image, idx) for image in list_files(join(self.target_folder, character), ".png")] + for idx, character in enumerate(self._characters) + ] + self._flat_character_images: List[Tuple[str, int]] = sum(self._character_images, []) + + def __len__(self) -> int: + return len(self._flat_character_images) + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: (image, target) where target is index of the target character class. + """ + image_name, character_class = self._flat_character_images[index] + image_path = join(self.target_folder, self._characters[character_class], image_name) + image = Image.open(image_path, mode="r").convert("L") + + if self.transform: + image = self.transform(image) + + if self.target_transform: + character_class = self.target_transform(character_class) + + return image, character_class + + def _check_integrity(self) -> bool: + zip_filename = self._get_target_folder() + if not check_integrity(join(self.root, zip_filename + ".zip"), self.zips_md5[zip_filename]): + return False + return True + + def download(self) -> None: + if self._check_integrity(): + return + + filename = self._get_target_folder() + zip_filename = filename + ".zip" + url = self.download_url_prefix + "/" + zip_filename + download_and_extract_archive(url, self.root, filename=zip_filename, md5=self.zips_md5[filename]) + + def _get_target_folder(self) -> str: + return "images_background" if self.background else "images_evaluation" diff --git a/lib/python3.12/site-packages/torchvision/datasets/oxford_iiit_pet.py b/lib/python3.12/site-packages/torchvision/datasets/oxford_iiit_pet.py new file mode 100644 index 0000000000000000000000000000000000000000..1d6d990fdf90106ef765d9ff47d1b760ce26d5d0 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/oxford_iiit_pet.py @@ -0,0 +1,132 @@ +import os +import os.path +import pathlib +from typing import Any, Callable, Optional, Sequence, Tuple, Union + +from PIL import Image + +from .utils import download_and_extract_archive, verify_str_arg +from .vision import VisionDataset + + +class OxfordIIITPet(VisionDataset): + """`Oxford-IIIT Pet Dataset `_. + + Args: + root (str or ``pathlib.Path``): Root directory of the dataset. + split (string, optional): The dataset split, supports ``"trainval"`` (default) or ``"test"``. + target_types (string, sequence of strings, optional): Types of target to use. Can be ``category`` (default) or + ``segmentation``. Can also be a list to output a tuple with all specified target types. The types represent: + + - ``category`` (int): Label for one of the 37 pet categories. + - ``binary-category`` (int): Binary label for cat or dog. + - ``segmentation`` (PIL image): Segmentation trimap of the image. + + If empty, ``None`` will be returned as target. + + transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed + version. E.g, ``transforms.RandomCrop``. + target_transform (callable, optional): A function/transform that takes in the target and transforms it. + download (bool, optional): If True, downloads the dataset from the internet and puts it into + ``root/oxford-iiit-pet``. If dataset is already downloaded, it is not downloaded again. + """ + + _RESOURCES = ( + ("https://www.robots.ox.ac.uk/~vgg/data/pets/data/images.tar.gz", "5c4f3ee8e5d25df40f4fd59a7f44e54c"), + ("https://www.robots.ox.ac.uk/~vgg/data/pets/data/annotations.tar.gz", "95a8c909bbe2e81eed6a22bccdf3f68f"), + ) + _VALID_TARGET_TYPES = ("category", "binary-category", "segmentation") + + def __init__( + self, + root: Union[str, pathlib.Path], + split: str = "trainval", + target_types: Union[Sequence[str], str] = "category", + transforms: Optional[Callable] = None, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ): + self._split = verify_str_arg(split, "split", ("trainval", "test")) + if isinstance(target_types, str): + target_types = [target_types] + self._target_types = [ + verify_str_arg(target_type, "target_types", self._VALID_TARGET_TYPES) for target_type in target_types + ] + + super().__init__(root, transforms=transforms, transform=transform, target_transform=target_transform) + self._base_folder = pathlib.Path(self.root) / "oxford-iiit-pet" + self._images_folder = self._base_folder / "images" + self._anns_folder = self._base_folder / "annotations" + self._segs_folder = self._anns_folder / "trimaps" + + if download: + self._download() + + if not self._check_exists(): + raise RuntimeError("Dataset not found. You can use download=True to download it") + + image_ids = [] + self._labels = [] + self._bin_labels = [] + with open(self._anns_folder / f"{self._split}.txt") as file: + for line in file: + image_id, label, bin_label, _ = line.strip().split() + image_ids.append(image_id) + self._labels.append(int(label) - 1) + self._bin_labels.append(int(bin_label) - 1) + + self.bin_classes = ["Cat", "Dog"] + self.classes = [ + " ".join(part.title() for part in raw_cls.split("_")) + for raw_cls, _ in sorted( + {(image_id.rsplit("_", 1)[0], label) for image_id, label in zip(image_ids, self._labels)}, + key=lambda image_id_and_label: image_id_and_label[1], + ) + ] + self.bin_class_to_idx = dict(zip(self.bin_classes, range(len(self.bin_classes)))) + self.class_to_idx = dict(zip(self.classes, range(len(self.classes)))) + + self._images = [self._images_folder / f"{image_id}.jpg" for image_id in image_ids] + self._segs = [self._segs_folder / f"{image_id}.png" for image_id in image_ids] + + def __len__(self) -> int: + return len(self._images) + + def __getitem__(self, idx: int) -> Tuple[Any, Any]: + image = Image.open(self._images[idx]).convert("RGB") + + target: Any = [] + for target_type in self._target_types: + if target_type == "category": + target.append(self._labels[idx]) + elif target_type == "binary-category": + target.append(self._bin_labels[idx]) + else: # target_type == "segmentation" + target.append(Image.open(self._segs[idx])) + + if not target: + target = None + elif len(target) == 1: + target = target[0] + else: + target = tuple(target) + + if self.transforms: + image, target = self.transforms(image, target) + + return image, target + + def _check_exists(self) -> bool: + for folder in (self._images_folder, self._anns_folder): + if not (os.path.exists(folder) and os.path.isdir(folder)): + return False + else: + return True + + def _download(self) -> None: + if self._check_exists(): + return + + for url, md5 in self._RESOURCES: + download_and_extract_archive(url, download_root=str(self._base_folder), md5=md5) diff --git a/lib/python3.12/site-packages/torchvision/datasets/pcam.py b/lib/python3.12/site-packages/torchvision/datasets/pcam.py new file mode 100644 index 0000000000000000000000000000000000000000..8849e0ea39dd1bfb76920bc8afc7f6ced597cda4 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/pcam.py @@ -0,0 +1,134 @@ +import pathlib +from typing import Any, Callable, Optional, Tuple, Union + +from PIL import Image + +from .utils import _decompress, download_file_from_google_drive, verify_str_arg +from .vision import VisionDataset + + +class PCAM(VisionDataset): + """`PCAM Dataset `_. + + The PatchCamelyon dataset is a binary classification dataset with 327,680 + color images (96px x 96px), extracted from histopathologic scans of lymph node + sections. Each image is annotated with a binary label indicating presence of + metastatic tissue. + + This dataset requires the ``h5py`` package which you can install with ``pip install h5py``. + + Args: + root (str or ``pathlib.Path``): Root directory of the dataset. + split (string, optional): The dataset split, supports ``"train"`` (default), ``"test"`` or ``"val"``. + transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed + version. E.g, ``transforms.RandomCrop``. + target_transform (callable, optional): A function/transform that takes in the target and transforms it. + download (bool, optional): If True, downloads the dataset from the internet and puts it into ``root/pcam``. If + dataset is already downloaded, it is not downloaded again. + + .. warning:: + + To download the dataset `gdown `_ is required. + """ + + _FILES = { + "train": { + "images": ( + "camelyonpatch_level_2_split_train_x.h5", # Data file name + "1Ka0XfEMiwgCYPdTI-vv6eUElOBnKFKQ2", # Google Drive ID + "1571f514728f59376b705fc836ff4b63", # md5 hash + ), + "targets": ( + "camelyonpatch_level_2_split_train_y.h5", + "1269yhu3pZDP8UYFQs-NYs3FPwuK-nGSG", + "35c2d7259d906cfc8143347bb8e05be7", + ), + }, + "test": { + "images": ( + "camelyonpatch_level_2_split_test_x.h5", + "1qV65ZqZvWzuIVthK8eVDhIwrbnsJdbg_", + "d8c2d60d490dbd479f8199bdfa0cf6ec", + ), + "targets": ( + "camelyonpatch_level_2_split_test_y.h5", + "17BHrSrwWKjYsOgTMmoqrIjDy6Fa2o_gP", + "60a7035772fbdb7f34eb86d4420cf66a", + ), + }, + "val": { + "images": ( + "camelyonpatch_level_2_split_valid_x.h5", + "1hgshYGWK8V-eGRy8LToWJJgDU_rXWVJ3", + "d5b63470df7cfa627aeec8b9dc0c066e", + ), + "targets": ( + "camelyonpatch_level_2_split_valid_y.h5", + "1bH8ZRbhSVAhScTS0p9-ZzGnX91cHT3uO", + "2b85f58b927af9964a4c15b8f7e8f179", + ), + }, + } + + def __init__( + self, + root: Union[str, pathlib.Path], + split: str = "train", + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ): + try: + import h5py + + self.h5py = h5py + except ImportError: + raise RuntimeError( + "h5py is not found. This dataset needs to have h5py installed: please run pip install h5py" + ) + + self._split = verify_str_arg(split, "split", ("train", "test", "val")) + + super().__init__(root, transform=transform, target_transform=target_transform) + self._base_folder = pathlib.Path(self.root) / "pcam" + + if download: + self._download() + + if not self._check_exists(): + raise RuntimeError("Dataset not found. You can use download=True to download it") + + def __len__(self) -> int: + images_file = self._FILES[self._split]["images"][0] + with self.h5py.File(self._base_folder / images_file) as images_data: + return images_data["x"].shape[0] + + def __getitem__(self, idx: int) -> Tuple[Any, Any]: + images_file = self._FILES[self._split]["images"][0] + with self.h5py.File(self._base_folder / images_file) as images_data: + image = Image.fromarray(images_data["x"][idx]).convert("RGB") + + targets_file = self._FILES[self._split]["targets"][0] + with self.h5py.File(self._base_folder / targets_file) as targets_data: + target = int(targets_data["y"][idx, 0, 0, 0]) # shape is [num_images, 1, 1, 1] + + if self.transform: + image = self.transform(image) + if self.target_transform: + target = self.target_transform(target) + + return image, target + + def _check_exists(self) -> bool: + images_file = self._FILES[self._split]["images"][0] + targets_file = self._FILES[self._split]["targets"][0] + return all(self._base_folder.joinpath(h5_file).exists() for h5_file in (images_file, targets_file)) + + def _download(self) -> None: + if self._check_exists(): + return + + for file_name, file_id, md5 in self._FILES[self._split].values(): + archive_name = file_name + ".gz" + download_file_from_google_drive(file_id, str(self._base_folder), filename=archive_name, md5=md5) + _decompress(str(self._base_folder / archive_name)) diff --git a/lib/python3.12/site-packages/torchvision/datasets/phototour.py b/lib/python3.12/site-packages/torchvision/datasets/phototour.py new file mode 100644 index 0000000000000000000000000000000000000000..9511f0626b40da10bb4d60158ea9091676fd150d --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/phototour.py @@ -0,0 +1,230 @@ +import os +from pathlib import Path +from typing import Any, Callable, List, Optional, Tuple, Union + +import numpy as np +import torch +from PIL import Image + +from .utils import download_url +from .vision import VisionDataset + + +class PhotoTour(VisionDataset): + """`Multi-view Stereo Correspondence `_ Dataset. + + .. note:: + + We only provide the newer version of the dataset, since the authors state that it + + is more suitable for training descriptors based on difference of Gaussian, or Harris corners, as the + patches are centred on real interest point detections, rather than being projections of 3D points as is the + case in the old dataset. + + The original dataset is available under http://phototour.cs.washington.edu/patches/default.htm. + + + Args: + root (str or ``pathlib.Path``): Root directory where images are. + name (string): Name of the dataset to load. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + + """ + + urls = { + "notredame_harris": [ + "http://matthewalunbrown.com/patchdata/notredame_harris.zip", + "notredame_harris.zip", + "69f8c90f78e171349abdf0307afefe4d", + ], + "yosemite_harris": [ + "http://matthewalunbrown.com/patchdata/yosemite_harris.zip", + "yosemite_harris.zip", + "a73253d1c6fbd3ba2613c45065c00d46", + ], + "liberty_harris": [ + "http://matthewalunbrown.com/patchdata/liberty_harris.zip", + "liberty_harris.zip", + "c731fcfb3abb4091110d0ae8c7ba182c", + ], + "notredame": [ + "http://icvl.ee.ic.ac.uk/vbalnt/notredame.zip", + "notredame.zip", + "509eda8535847b8c0a90bbb210c83484", + ], + "yosemite": ["http://icvl.ee.ic.ac.uk/vbalnt/yosemite.zip", "yosemite.zip", "533b2e8eb7ede31be40abc317b2fd4f0"], + "liberty": ["http://icvl.ee.ic.ac.uk/vbalnt/liberty.zip", "liberty.zip", "fdd9152f138ea5ef2091746689176414"], + } + means = { + "notredame": 0.4854, + "yosemite": 0.4844, + "liberty": 0.4437, + "notredame_harris": 0.4854, + "yosemite_harris": 0.4844, + "liberty_harris": 0.4437, + } + stds = { + "notredame": 0.1864, + "yosemite": 0.1818, + "liberty": 0.2019, + "notredame_harris": 0.1864, + "yosemite_harris": 0.1818, + "liberty_harris": 0.2019, + } + lens = { + "notredame": 468159, + "yosemite": 633587, + "liberty": 450092, + "liberty_harris": 379587, + "yosemite_harris": 450912, + "notredame_harris": 325295, + } + image_ext = "bmp" + info_file = "info.txt" + matches_files = "m50_100000_100000_0.txt" + + def __init__( + self, + root: Union[str, Path], + name: str, + train: bool = True, + transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + super().__init__(root, transform=transform) + self.name = name + self.data_dir = os.path.join(self.root, name) + self.data_down = os.path.join(self.root, f"{name}.zip") + self.data_file = os.path.join(self.root, f"{name}.pt") + + self.train = train + self.mean = self.means[name] + self.std = self.stds[name] + + if download: + self.download() + + if not self._check_datafile_exists(): + self.cache() + + # load the serialized data + self.data, self.labels, self.matches = torch.load(self.data_file, weights_only=True) + + def __getitem__(self, index: int) -> Union[torch.Tensor, Tuple[Any, Any, torch.Tensor]]: + """ + Args: + index (int): Index + + Returns: + tuple: (data1, data2, matches) + """ + if self.train: + data = self.data[index] + if self.transform is not None: + data = self.transform(data) + return data + m = self.matches[index] + data1, data2 = self.data[m[0]], self.data[m[1]] + if self.transform is not None: + data1 = self.transform(data1) + data2 = self.transform(data2) + return data1, data2, m[2] + + def __len__(self) -> int: + return len(self.data if self.train else self.matches) + + def _check_datafile_exists(self) -> bool: + return os.path.exists(self.data_file) + + def _check_downloaded(self) -> bool: + return os.path.exists(self.data_dir) + + def download(self) -> None: + if self._check_datafile_exists(): + return + + if not self._check_downloaded(): + # download files + url = self.urls[self.name][0] + filename = self.urls[self.name][1] + md5 = self.urls[self.name][2] + fpath = os.path.join(self.root, filename) + + download_url(url, self.root, filename, md5) + + import zipfile + + with zipfile.ZipFile(fpath, "r") as z: + z.extractall(self.data_dir) + + os.unlink(fpath) + + def cache(self) -> None: + # process and save as torch files + + dataset = ( + read_image_file(self.data_dir, self.image_ext, self.lens[self.name]), + read_info_file(self.data_dir, self.info_file), + read_matches_files(self.data_dir, self.matches_files), + ) + + with open(self.data_file, "wb") as f: + torch.save(dataset, f) + + def extra_repr(self) -> str: + split = "Train" if self.train is True else "Test" + return f"Split: {split}" + + +def read_image_file(data_dir: str, image_ext: str, n: int) -> torch.Tensor: + """Return a Tensor containing the patches""" + + def PIL2array(_img: Image.Image) -> np.ndarray: + """Convert PIL image type to numpy 2D array""" + return np.array(_img.getdata(), dtype=np.uint8).reshape(64, 64) + + def find_files(_data_dir: str, _image_ext: str) -> List[str]: + """Return a list with the file names of the images containing the patches""" + files = [] + # find those files with the specified extension + for file_dir in os.listdir(_data_dir): + if file_dir.endswith(_image_ext): + files.append(os.path.join(_data_dir, file_dir)) + return sorted(files) # sort files in ascend order to keep relations + + patches = [] + list_files = find_files(data_dir, image_ext) + + for fpath in list_files: + img = Image.open(fpath) + for y in range(0, img.height, 64): + for x in range(0, img.width, 64): + patch = img.crop((x, y, x + 64, y + 64)) + patches.append(PIL2array(patch)) + return torch.ByteTensor(np.array(patches[:n])) + + +def read_info_file(data_dir: str, info_file: str) -> torch.Tensor: + """Return a Tensor containing the list of labels + Read the file and keep only the ID of the 3D point. + """ + with open(os.path.join(data_dir, info_file)) as f: + labels = [int(line.split()[0]) for line in f] + return torch.LongTensor(labels) + + +def read_matches_files(data_dir: str, matches_file: str) -> torch.Tensor: + """Return a Tensor containing the ground truth matches + Read the file and keep only 3D point ID. + Matches are represented with a 1, non matches with a 0. + """ + matches = [] + with open(os.path.join(data_dir, matches_file)) as f: + for line in f: + line_split = line.split() + matches.append([int(line_split[0]), int(line_split[3]), int(line_split[1] == line_split[4])]) + return torch.LongTensor(matches) diff --git a/lib/python3.12/site-packages/torchvision/datasets/places365.py b/lib/python3.12/site-packages/torchvision/datasets/places365.py new file mode 100644 index 0000000000000000000000000000000000000000..a120e0e217a246f660bda0d05aa97e34fce8ea74 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/places365.py @@ -0,0 +1,168 @@ +import os +from os import path +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from urllib.parse import urljoin + +from .folder import default_loader +from .utils import check_integrity, download_and_extract_archive, verify_str_arg +from .vision import VisionDataset + + +class Places365(VisionDataset): + r"""`Places365 `_ classification dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of the Places365 dataset. + split (string, optional): The dataset split. Can be one of ``train-standard`` (default), ``train-challenge``, + ``val``. + small (bool, optional): If ``True``, uses the small images, i.e. resized to 256 x 256 pixels, instead of the + high resolution ones. + download (bool, optional): If ``True``, downloads the dataset components and places them in ``root``. Already + downloaded archives are not downloaded again. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + loader (callable, optional): A function to load an image given its path. + + Attributes: + classes (list): List of the class names. + class_to_idx (dict): Dict with items (class_name, class_index). + imgs (list): List of (image path, class_index) tuples + targets (list): The class_index value for each image in the dataset + + Raises: + RuntimeError: If ``download is False`` and the meta files, i.e. the devkit, are not present or corrupted. + RuntimeError: If ``download is True`` and the image archive is already extracted. + """ + _SPLITS = ("train-standard", "train-challenge", "val") + _BASE_URL = "http://data.csail.mit.edu/places/places365/" + # {variant: (archive, md5)} + _DEVKIT_META = { + "standard": ("filelist_places365-standard.tar", "35a0585fee1fa656440f3ab298f8479c"), + "challenge": ("filelist_places365-challenge.tar", "70a8307e459c3de41690a7c76c931734"), + } + # (file, md5) + _CATEGORIES_META = ("categories_places365.txt", "06c963b85866bd0649f97cb43dd16673") + # {split: (file, md5)} + _FILE_LIST_META = { + "train-standard": ("places365_train_standard.txt", "30f37515461640559006b8329efbed1a"), + "train-challenge": ("places365_train_challenge.txt", "b2931dc997b8c33c27e7329c073a6b57"), + "val": ("places365_val.txt", "e9f2fd57bfd9d07630173f4e8708e4b1"), + } + # {(split, small): (file, md5)} + _IMAGES_META = { + ("train-standard", False): ("train_large_places365standard.tar", "67e186b496a84c929568076ed01a8aa1"), + ("train-challenge", False): ("train_large_places365challenge.tar", "605f18e68e510c82b958664ea134545f"), + ("val", False): ("val_large.tar", "9b71c4993ad89d2d8bcbdc4aef38042f"), + ("train-standard", True): ("train_256_places365standard.tar", "53ca1c756c3d1e7809517cc47c5561c5"), + ("train-challenge", True): ("train_256_places365challenge.tar", "741915038a5e3471ec7332404dfb64ef"), + ("val", True): ("val_256.tar", "e27b17d8d44f4af9a78502beb927f808"), + } + + def __init__( + self, + root: Union[str, Path], + split: str = "train-standard", + small: bool = False, + download: bool = False, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + loader: Callable[[str], Any] = default_loader, + ) -> None: + super().__init__(root, transform=transform, target_transform=target_transform) + + self.split = self._verify_split(split) + self.small = small + self.loader = loader + + self.classes, self.class_to_idx = self.load_categories(download) + self.imgs, self.targets = self.load_file_list(download) + + if download: + self.download_images() + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + file, target = self.imgs[index] + image = self.loader(file) + + if self.transforms is not None: + image, target = self.transforms(image, target) + + return image, target + + def __len__(self) -> int: + return len(self.imgs) + + @property + def variant(self) -> str: + return "challenge" if "challenge" in self.split else "standard" + + @property + def images_dir(self) -> str: + size = "256" if self.small else "large" + if self.split.startswith("train"): + dir = f"data_{size}_{self.variant}" + else: + dir = f"{self.split}_{size}" + return path.join(self.root, dir) + + def load_categories(self, download: bool = True) -> Tuple[List[str], Dict[str, int]]: + def process(line: str) -> Tuple[str, int]: + cls, idx = line.split() + return cls, int(idx) + + file, md5 = self._CATEGORIES_META + file = path.join(self.root, file) + if not self._check_integrity(file, md5, download): + self.download_devkit() + + with open(file) as fh: + class_to_idx = dict(process(line) for line in fh) + + return sorted(class_to_idx.keys()), class_to_idx + + def load_file_list(self, download: bool = True) -> Tuple[List[Tuple[str, int]], List[int]]: + def process(line: str, sep="/") -> Tuple[str, int]: + image, idx = line.split() + return path.join(self.images_dir, image.lstrip(sep).replace(sep, os.sep)), int(idx) + + file, md5 = self._FILE_LIST_META[self.split] + file = path.join(self.root, file) + if not self._check_integrity(file, md5, download): + self.download_devkit() + + with open(file) as fh: + images = [process(line) for line in fh] + + _, targets = zip(*images) + return images, list(targets) + + def download_devkit(self) -> None: + file, md5 = self._DEVKIT_META[self.variant] + download_and_extract_archive(urljoin(self._BASE_URL, file), self.root, md5=md5) + + def download_images(self) -> None: + if path.exists(self.images_dir): + return + + file, md5 = self._IMAGES_META[(self.split, self.small)] + download_and_extract_archive(urljoin(self._BASE_URL, file), self.root, md5=md5) + + if self.split.startswith("train"): + os.rename(self.images_dir.rsplit("_", 1)[0], self.images_dir) + + def extra_repr(self) -> str: + return "\n".join(("Split: {split}", "Small: {small}")).format(**self.__dict__) + + def _verify_split(self, split: str) -> str: + return verify_str_arg(split, "split", self._SPLITS) + + def _check_integrity(self, file: str, md5: str, download: bool) -> bool: + integrity = check_integrity(file, md5=md5) + if not integrity and not download: + raise RuntimeError( + f"The file {file} does not exist or is corrupted. You can set download=True to download it." + ) + return integrity diff --git a/lib/python3.12/site-packages/torchvision/datasets/rendered_sst2.py b/lib/python3.12/site-packages/torchvision/datasets/rendered_sst2.py new file mode 100644 index 0000000000000000000000000000000000000000..48b0ddfc4fb3394505c1125e67cc026b3fd14fc8 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/rendered_sst2.py @@ -0,0 +1,86 @@ +from pathlib import Path +from typing import Any, Callable, Optional, Tuple, Union + +import PIL.Image + +from .folder import make_dataset +from .utils import download_and_extract_archive, verify_str_arg +from .vision import VisionDataset + + +class RenderedSST2(VisionDataset): + """`The Rendered SST2 Dataset `_. + + Rendered SST2 is an image classification dataset used to evaluate the models capability on optical + character recognition. This dataset was generated by rendering sentences in the Standford Sentiment + Treebank v2 dataset. + + This dataset contains two classes (positive and negative) and is divided in three splits: a train + split containing 6920 images (3610 positive and 3310 negative), a validation split containing 872 images + (444 positive and 428 negative), and a test split containing 1821 images (909 positive and 912 negative). + + Args: + root (str or ``pathlib.Path``): Root directory of the dataset. + split (string, optional): The dataset split, supports ``"train"`` (default), `"val"` and ``"test"``. + transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed + version. E.g, ``transforms.RandomCrop``. + target_transform (callable, optional): A function/transform that takes in the target and transforms it. + download (bool, optional): If True, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. Default is False. + """ + + _URL = "https://openaipublic.azureedge.net/clip/data/rendered-sst2.tgz" + _MD5 = "2384d08e9dcfa4bd55b324e610496ee5" + + def __init__( + self, + root: Union[str, Path], + split: str = "train", + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + super().__init__(root, transform=transform, target_transform=target_transform) + self._split = verify_str_arg(split, "split", ("train", "val", "test")) + self._split_to_folder = {"train": "train", "val": "valid", "test": "test"} + self._base_folder = Path(self.root) / "rendered-sst2" + self.classes = ["negative", "positive"] + self.class_to_idx = {"negative": 0, "positive": 1} + + if download: + self._download() + + if not self._check_exists(): + raise RuntimeError("Dataset not found. You can use download=True to download it") + + self._samples = make_dataset(str(self._base_folder / self._split_to_folder[self._split]), extensions=("png",)) + + def __len__(self) -> int: + return len(self._samples) + + def __getitem__(self, idx: int) -> Tuple[Any, Any]: + image_file, label = self._samples[idx] + image = PIL.Image.open(image_file).convert("RGB") + + if self.transform: + image = self.transform(image) + + if self.target_transform: + label = self.target_transform(label) + + return image, label + + def extra_repr(self) -> str: + return f"split={self._split}" + + def _check_exists(self) -> bool: + for class_label in set(self.classes): + if not (self._base_folder / self._split_to_folder[self._split] / class_label).is_dir(): + return False + return True + + def _download(self) -> None: + if self._check_exists(): + return + download_and_extract_archive(self._URL, download_root=self.root, md5=self._MD5) diff --git a/lib/python3.12/site-packages/torchvision/datasets/samplers/__init__.py b/lib/python3.12/site-packages/torchvision/datasets/samplers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..58b2d2abd936d885221174d194a633a8e413935f --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/samplers/__init__.py @@ -0,0 +1,3 @@ +from .clip_sampler import DistributedSampler, RandomClipSampler, UniformClipSampler + +__all__ = ("DistributedSampler", "UniformClipSampler", "RandomClipSampler") diff --git a/lib/python3.12/site-packages/torchvision/datasets/samplers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/samplers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1c898c46caf0ea47040c50f993836bf8dd0fb926 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/samplers/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/samplers/__pycache__/clip_sampler.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/datasets/samplers/__pycache__/clip_sampler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..38f4785c6d672eb275273b5c77ccacce72041b51 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/datasets/samplers/__pycache__/clip_sampler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/datasets/samplers/clip_sampler.py b/lib/python3.12/site-packages/torchvision/datasets/samplers/clip_sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..026c3d75d3b8acd5d0240d8e537e608a92cddfdb --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/samplers/clip_sampler.py @@ -0,0 +1,172 @@ +import math +from typing import cast, Iterator, List, Optional, Sized, Union + +import torch +import torch.distributed as dist +from torch.utils.data import Sampler +from torchvision.datasets.video_utils import VideoClips + + +class DistributedSampler(Sampler): + """ + Extension of DistributedSampler, as discussed in + https://github.com/pytorch/pytorch/issues/23430 + + Example: + dataset: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13] + num_replicas: 4 + shuffle: False + + when group_size = 1 + RANK | shard_dataset + ========================= + rank_0 | [0, 4, 8, 12] + rank_1 | [1, 5, 9, 13] + rank_2 | [2, 6, 10, 0] + rank_3 | [3, 7, 11, 1] + + when group_size = 2 + + RANK | shard_dataset + ========================= + rank_0 | [0, 1, 8, 9] + rank_1 | [2, 3, 10, 11] + rank_2 | [4, 5, 12, 13] + rank_3 | [6, 7, 0, 1] + + """ + + def __init__( + self, + dataset: Sized, + num_replicas: Optional[int] = None, + rank: Optional[int] = None, + shuffle: bool = False, + group_size: int = 1, + ) -> None: + if num_replicas is None: + if not dist.is_available(): + raise RuntimeError("Requires distributed package to be available") + num_replicas = dist.get_world_size() + if rank is None: + if not dist.is_available(): + raise RuntimeError("Requires distributed package to be available") + rank = dist.get_rank() + if len(dataset) % group_size != 0: + raise ValueError( + f"dataset length must be a multiplier of group size dataset length: {len(dataset)}, group size: {group_size}" + ) + self.dataset = dataset + self.group_size = group_size + self.num_replicas = num_replicas + self.rank = rank + self.epoch = 0 + dataset_group_length = len(dataset) // group_size + self.num_group_samples = int(math.ceil(dataset_group_length * 1.0 / self.num_replicas)) + self.num_samples = self.num_group_samples * group_size + self.total_size = self.num_samples * self.num_replicas + self.shuffle = shuffle + + def __iter__(self) -> Iterator[int]: + # deterministically shuffle based on epoch + g = torch.Generator() + g.manual_seed(self.epoch) + indices: Union[torch.Tensor, List[int]] + if self.shuffle: + indices = torch.randperm(len(self.dataset), generator=g).tolist() + else: + indices = list(range(len(self.dataset))) + + # add extra samples to make it evenly divisible + indices += indices[: (self.total_size - len(indices))] + assert len(indices) == self.total_size + + total_group_size = self.total_size // self.group_size + indices = torch.reshape(torch.LongTensor(indices), (total_group_size, self.group_size)) + + # subsample + indices = indices[self.rank : total_group_size : self.num_replicas, :] + indices = torch.reshape(indices, (-1,)).tolist() + assert len(indices) == self.num_samples + + if isinstance(self.dataset, Sampler): + orig_indices = list(iter(self.dataset)) + indices = [orig_indices[i] for i in indices] + + return iter(indices) + + def __len__(self) -> int: + return self.num_samples + + def set_epoch(self, epoch: int) -> None: + self.epoch = epoch + + +class UniformClipSampler(Sampler): + """ + Sample `num_video_clips_per_video` clips for each video, equally spaced. + When number of unique clips in the video is fewer than num_video_clips_per_video, + repeat the clips until `num_video_clips_per_video` clips are collected + + Args: + video_clips (VideoClips): video clips to sample from + num_clips_per_video (int): number of clips to be sampled per video + """ + + def __init__(self, video_clips: VideoClips, num_clips_per_video: int) -> None: + if not isinstance(video_clips, VideoClips): + raise TypeError(f"Expected video_clips to be an instance of VideoClips, got {type(video_clips)}") + self.video_clips = video_clips + self.num_clips_per_video = num_clips_per_video + + def __iter__(self) -> Iterator[int]: + idxs = [] + s = 0 + # select num_clips_per_video for each video, uniformly spaced + for c in self.video_clips.clips: + length = len(c) + if length == 0: + # corner case where video decoding fails + continue + + sampled = torch.linspace(s, s + length - 1, steps=self.num_clips_per_video).floor().to(torch.int64) + s += length + idxs.append(sampled) + return iter(cast(List[int], torch.cat(idxs).tolist())) + + def __len__(self) -> int: + return sum(self.num_clips_per_video for c in self.video_clips.clips if len(c) > 0) + + +class RandomClipSampler(Sampler): + """ + Samples at most `max_video_clips_per_video` clips for each video randomly + + Args: + video_clips (VideoClips): video clips to sample from + max_clips_per_video (int): maximum number of clips to be sampled per video + """ + + def __init__(self, video_clips: VideoClips, max_clips_per_video: int) -> None: + if not isinstance(video_clips, VideoClips): + raise TypeError(f"Expected video_clips to be an instance of VideoClips, got {type(video_clips)}") + self.video_clips = video_clips + self.max_clips_per_video = max_clips_per_video + + def __iter__(self) -> Iterator[int]: + idxs = [] + s = 0 + # select at most max_clips_per_video for each video, randomly + for c in self.video_clips.clips: + length = len(c) + size = min(length, self.max_clips_per_video) + sampled = torch.randperm(length)[:size] + s + s += length + idxs.append(sampled) + idxs_ = torch.cat(idxs) + # shuffle all clips randomly + perm = torch.randperm(len(idxs_)) + return iter(idxs_[perm].tolist()) + + def __len__(self) -> int: + return sum(min(len(c), self.max_clips_per_video) for c in self.video_clips.clips) diff --git a/lib/python3.12/site-packages/torchvision/datasets/sbd.py b/lib/python3.12/site-packages/torchvision/datasets/sbd.py new file mode 100644 index 0000000000000000000000000000000000000000..4b9ccb75eb9bc6c5c2e9af6fbb4b5868e8488923 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/sbd.py @@ -0,0 +1,126 @@ +import os +import shutil +from pathlib import Path +from typing import Any, Callable, Optional, Tuple, Union + +import numpy as np +from PIL import Image + +from .utils import download_and_extract_archive, download_url, verify_str_arg +from .vision import VisionDataset + + +class SBDataset(VisionDataset): + """`Semantic Boundaries Dataset `_ + + The SBD currently contains annotations from 11355 images taken from the PASCAL VOC 2011 dataset. + + .. note :: + + Please note that the train and val splits included with this dataset are different from + the splits in the PASCAL VOC dataset. In particular some "train" images might be part of + VOC2012 val. + If you are interested in testing on VOC 2012 val, then use `image_set='train_noval'`, + which excludes all val images. + + .. warning:: + + This class needs `scipy `_ to load target files from `.mat` format. + + Args: + root (str or ``pathlib.Path``): Root directory of the Semantic Boundaries Dataset + image_set (string, optional): Select the image_set to use, ``train``, ``val`` or ``train_noval``. + Image set ``train_noval`` excludes VOC 2012 val images. + mode (string, optional): Select target type. Possible values 'boundaries' or 'segmentation'. + In case of 'boundaries', the target is an array of shape `[num_classes, H, W]`, + where `num_classes=20`. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + transforms (callable, optional): A function/transform that takes input sample and its target as entry + and returns a transformed version. Input sample is PIL image and target is a numpy array + if `mode='boundaries'` or PIL image if `mode='segmentation'`. + """ + + url = "https://www2.eecs.berkeley.edu/Research/Projects/CS/vision/grouping/semantic_contours/benchmark.tgz" + md5 = "82b4d87ceb2ed10f6038a1cba92111cb" + filename = "benchmark.tgz" + + voc_train_url = "https://www.cs.cornell.edu/~bharathh/train_noval.txt" + voc_split_filename = "train_noval.txt" + voc_split_md5 = "79bff800c5f0b1ec6b21080a3c066722" + + def __init__( + self, + root: Union[str, Path], + image_set: str = "train", + mode: str = "boundaries", + download: bool = False, + transforms: Optional[Callable] = None, + ) -> None: + + try: + from scipy.io import loadmat + + self._loadmat = loadmat + except ImportError: + raise RuntimeError("Scipy is not found. This dataset needs to have scipy installed: pip install scipy") + + super().__init__(root, transforms) + self.image_set = verify_str_arg(image_set, "image_set", ("train", "val", "train_noval")) + self.mode = verify_str_arg(mode, "mode", ("segmentation", "boundaries")) + self.num_classes = 20 + + sbd_root = self.root + image_dir = os.path.join(sbd_root, "img") + mask_dir = os.path.join(sbd_root, "cls") + + if download: + download_and_extract_archive(self.url, self.root, filename=self.filename, md5=self.md5) + extracted_ds_root = os.path.join(self.root, "benchmark_RELEASE", "dataset") + for f in ["cls", "img", "inst", "train.txt", "val.txt"]: + old_path = os.path.join(extracted_ds_root, f) + shutil.move(old_path, sbd_root) + if self.image_set == "train_noval": + # Note: this is failing as of June 2024 https://github.com/pytorch/vision/issues/8471 + download_url(self.voc_train_url, sbd_root, self.voc_split_filename, self.voc_split_md5) + + if not os.path.isdir(sbd_root): + raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") + + split_f = os.path.join(sbd_root, image_set.rstrip("\n") + ".txt") + + with open(os.path.join(split_f)) as fh: + file_names = [x.strip() for x in fh.readlines()] + + self.images = [os.path.join(image_dir, x + ".jpg") for x in file_names] + self.masks = [os.path.join(mask_dir, x + ".mat") for x in file_names] + + self._get_target = self._get_segmentation_target if self.mode == "segmentation" else self._get_boundaries_target + + def _get_segmentation_target(self, filepath: str) -> Image.Image: + mat = self._loadmat(filepath) + return Image.fromarray(mat["GTcls"][0]["Segmentation"][0]) + + def _get_boundaries_target(self, filepath: str) -> np.ndarray: + mat = self._loadmat(filepath) + return np.concatenate( + [np.expand_dims(mat["GTcls"][0]["Boundaries"][0][i][0].toarray(), axis=0) for i in range(self.num_classes)], + axis=0, + ) + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + img = Image.open(self.images[index]).convert("RGB") + target = self._get_target(self.masks[index]) + + if self.transforms is not None: + img, target = self.transforms(img, target) + + return img, target + + def __len__(self) -> int: + return len(self.images) + + def extra_repr(self) -> str: + lines = ["Image set: {image_set}", "Mode: {mode}"] + return "\n".join(lines).format(**self.__dict__) diff --git a/lib/python3.12/site-packages/torchvision/datasets/sbu.py b/lib/python3.12/site-packages/torchvision/datasets/sbu.py new file mode 100644 index 0000000000000000000000000000000000000000..b5f46101e07ed146f4a2c6392cf2df119413d50f --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/sbu.py @@ -0,0 +1,109 @@ +import os +from pathlib import Path +from typing import Any, Callable, Optional, Tuple, Union + +from PIL import Image + +from .utils import check_integrity, download_and_extract_archive, download_url +from .vision import VisionDataset + + +class SBU(VisionDataset): + """`SBU Captioned Photo `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset where tarball + ``SBUCaptionedPhotoDataset.tar.gz`` exists. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + download (bool, optional): If True, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + """ + + url = "https://www.cs.rice.edu/~vo9/sbucaptions/SBUCaptionedPhotoDataset.tar.gz" + filename = "SBUCaptionedPhotoDataset.tar.gz" + md5_checksum = "9aec147b3488753cf758b4d493422285" + + def __init__( + self, + root: Union[str, Path], + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = True, + ) -> None: + super().__init__(root, transform=transform, target_transform=target_transform) + + if download: + self.download() + + if not self._check_integrity(): + raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") + + # Read the caption for each photo + self.photos = [] + self.captions = [] + + file1 = os.path.join(self.root, "dataset", "SBU_captioned_photo_dataset_urls.txt") + file2 = os.path.join(self.root, "dataset", "SBU_captioned_photo_dataset_captions.txt") + + for line1, line2 in zip(open(file1), open(file2)): + url = line1.rstrip() + photo = os.path.basename(url) + filename = os.path.join(self.root, "dataset", photo) + if os.path.exists(filename): + caption = line2.rstrip() + self.photos.append(photo) + self.captions.append(caption) + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: (image, target) where target is a caption for the photo. + """ + filename = os.path.join(self.root, "dataset", self.photos[index]) + img = Image.open(filename).convert("RGB") + if self.transform is not None: + img = self.transform(img) + + target = self.captions[index] + if self.target_transform is not None: + target = self.target_transform(target) + + return img, target + + def __len__(self) -> int: + """The number of photos in the dataset.""" + return len(self.photos) + + def _check_integrity(self) -> bool: + """Check the md5 checksum of the downloaded tarball.""" + root = self.root + fpath = os.path.join(root, self.filename) + if not check_integrity(fpath, self.md5_checksum): + return False + return True + + def download(self) -> None: + """Download and extract the tarball, and download each individual photo.""" + + if self._check_integrity(): + return + + download_and_extract_archive(self.url, self.root, self.root, self.filename, self.md5_checksum) + + # Download individual photos + with open(os.path.join(self.root, "dataset", "SBU_captioned_photo_dataset_urls.txt")) as fh: + for line in fh: + url = line.rstrip() + try: + download_url(url, os.path.join(self.root, "dataset")) + except OSError: + # The images point to public images on Flickr. + # Note: Images might be removed by users at anytime. + pass diff --git a/lib/python3.12/site-packages/torchvision/datasets/semeion.py b/lib/python3.12/site-packages/torchvision/datasets/semeion.py new file mode 100644 index 0000000000000000000000000000000000000000..71485b14dcf99dc7b5064c77582c5657a7d1571d --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/semeion.py @@ -0,0 +1,91 @@ +import os.path +from pathlib import Path +from typing import Any, Callable, Optional, Tuple, Union + +import numpy as np +from PIL import Image + +from .utils import check_integrity, download_url +from .vision import VisionDataset + + +class SEMEION(VisionDataset): + r"""`SEMEION `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset where directory + ``semeion.py`` exists. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + + """ + url = "http://archive.ics.uci.edu/ml/machine-learning-databases/semeion/semeion.data" + filename = "semeion.data" + md5_checksum = "cb545d371d2ce14ec121470795a77432" + + def __init__( + self, + root: Union[str, Path], + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = True, + ) -> None: + super().__init__(root, transform=transform, target_transform=target_transform) + + if download: + self.download() + + if not self._check_integrity(): + raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") + + fp = os.path.join(self.root, self.filename) + data = np.loadtxt(fp) + # convert value to 8 bit unsigned integer + # color (white #255) the pixels + self.data = (data[:, :256] * 255).astype("uint8") + self.data = np.reshape(self.data, (-1, 16, 16)) + self.labels = np.nonzero(data[:, 256:])[1] + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: (image, target) where target is index of the target class. + """ + img, target = self.data[index], int(self.labels[index]) + + # doing this so that it is consistent with all other datasets + # to return a PIL Image + img = Image.fromarray(img, mode="L") + + if self.transform is not None: + img = self.transform(img) + + if self.target_transform is not None: + target = self.target_transform(target) + + return img, target + + def __len__(self) -> int: + return len(self.data) + + def _check_integrity(self) -> bool: + root = self.root + fpath = os.path.join(root, self.filename) + if not check_integrity(fpath, self.md5_checksum): + return False + return True + + def download(self) -> None: + if self._check_integrity(): + return + + root = self.root + download_url(self.url, root, self.filename, self.md5_checksum) diff --git a/lib/python3.12/site-packages/torchvision/datasets/stanford_cars.py b/lib/python3.12/site-packages/torchvision/datasets/stanford_cars.py new file mode 100644 index 0000000000000000000000000000000000000000..6264de82eb7a567d10572d5d733011dd84e08d5b --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/stanford_cars.py @@ -0,0 +1,114 @@ +import pathlib +from typing import Any, Callable, Optional, Tuple, Union + +from PIL import Image + +from .utils import verify_str_arg +from .vision import VisionDataset + + +class StanfordCars(VisionDataset): + """Stanford Cars Dataset + + The Cars dataset contains 16,185 images of 196 classes of cars. The data is + split into 8,144 training images and 8,041 testing images, where each class + has been split roughly in a 50-50 split + + The original URL is https://ai.stanford.edu/~jkrause/cars/car_dataset.html, but it is broken. + Follow the instructions in ``download`` argument to obtain and use the dataset offline. + + .. note:: + + This class needs `scipy `_ to load target files from `.mat` format. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset + split (string, optional): The dataset split, supports ``"train"`` (default) or ``"test"``. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + download (bool, optional): This parameter exists for backward compatibility but it does not + download the dataset, since the original URL is not available anymore. The dataset + seems to be available on Kaggle so you can try to manually download and configure it using + `these instructions `_, + or use an integrated + `dataset on Kaggle `_. + In both cases, first download and configure the dataset locally, and use the dataset with + ``"download=False"``. + """ + + def __init__( + self, + root: Union[str, pathlib.Path], + split: str = "train", + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + + try: + import scipy.io as sio + except ImportError: + raise RuntimeError("Scipy is not found. This dataset needs to have scipy installed: pip install scipy") + + super().__init__(root, transform=transform, target_transform=target_transform) + + self._split = verify_str_arg(split, "split", ("train", "test")) + self._base_folder = pathlib.Path(root) / "stanford_cars" + devkit = self._base_folder / "devkit" + + if self._split == "train": + self._annotations_mat_path = devkit / "cars_train_annos.mat" + self._images_base_path = self._base_folder / "cars_train" + else: + self._annotations_mat_path = self._base_folder / "cars_test_annos_withlabels.mat" + self._images_base_path = self._base_folder / "cars_test" + + if download: + self.download() + + if not self._check_exists(): + raise RuntimeError( + "Dataset not found. Try to manually download following the instructions in " + "https://github.com/pytorch/vision/issues/7545#issuecomment-1631441616." + ) + + self._samples = [ + ( + str(self._images_base_path / annotation["fname"]), + annotation["class"] - 1, # Original target mapping starts from 1, hence -1 + ) + for annotation in sio.loadmat(self._annotations_mat_path, squeeze_me=True)["annotations"] + ] + + self.classes = sio.loadmat(str(devkit / "cars_meta.mat"), squeeze_me=True)["class_names"].tolist() + self.class_to_idx = {cls: i for i, cls in enumerate(self.classes)} + + def __len__(self) -> int: + return len(self._samples) + + def __getitem__(self, idx: int) -> Tuple[Any, Any]: + """Returns pil_image and class_id for given index""" + image_path, target = self._samples[idx] + pil_image = Image.open(image_path).convert("RGB") + + if self.transform is not None: + pil_image = self.transform(pil_image) + if self.target_transform is not None: + target = self.target_transform(target) + return pil_image, target + + def _check_exists(self) -> bool: + if not (self._base_folder / "devkit").is_dir(): + return False + + return self._annotations_mat_path.exists() and self._images_base_path.is_dir() + + def download(self): + raise ValueError( + "The original URL is broken so the StanfordCars dataset is not available for automatic " + "download anymore. You can try to download it manually following " + "https://github.com/pytorch/vision/issues/7545#issuecomment-1631441616, " + "and set download=False to avoid this error." + ) diff --git a/lib/python3.12/site-packages/torchvision/datasets/stl10.py b/lib/python3.12/site-packages/torchvision/datasets/stl10.py new file mode 100644 index 0000000000000000000000000000000000000000..39c058e497d4d950d57b9b027907b45f0bc46a58 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/stl10.py @@ -0,0 +1,174 @@ +import os.path +from pathlib import Path +from typing import Any, Callable, cast, Optional, Tuple, Union + +import numpy as np +from PIL import Image + +from .utils import check_integrity, download_and_extract_archive, verify_str_arg +from .vision import VisionDataset + + +class STL10(VisionDataset): + """`STL10 `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset where directory + ``stl10_binary`` exists. + split (string): One of {'train', 'test', 'unlabeled', 'train+unlabeled'}. + Accordingly, dataset is selected. + folds (int, optional): One of {0-9} or None. + For training, loads one of the 10 pre-defined folds of 1k samples for the + standard evaluation procedure. If no value is passed, loads the 5k samples. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + """ + + base_folder = "stl10_binary" + url = "http://ai.stanford.edu/~acoates/stl10/stl10_binary.tar.gz" + filename = "stl10_binary.tar.gz" + tgz_md5 = "91f7769df0f17e558f3565bffb0c7dfb" + class_names_file = "class_names.txt" + folds_list_file = "fold_indices.txt" + train_list = [ + ["train_X.bin", "918c2871b30a85fa023e0c44e0bee87f"], + ["train_y.bin", "5a34089d4802c674881badbb80307741"], + ["unlabeled_X.bin", "5242ba1fed5e4be9e1e742405eb56ca4"], + ] + + test_list = [["test_X.bin", "7f263ba9f9e0b06b93213547f721ac82"], ["test_y.bin", "36f9794fa4beb8a2c72628de14fa638e"]] + splits = ("train", "train+unlabeled", "unlabeled", "test") + + def __init__( + self, + root: Union[str, Path], + split: str = "train", + folds: Optional[int] = None, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + super().__init__(root, transform=transform, target_transform=target_transform) + self.split = verify_str_arg(split, "split", self.splits) + self.folds = self._verify_folds(folds) + + if download: + self.download() + elif not self._check_integrity(): + raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") + + # now load the picked numpy arrays + self.labels: Optional[np.ndarray] + if self.split == "train": + self.data, self.labels = self.__loadfile(self.train_list[0][0], self.train_list[1][0]) + self.labels = cast(np.ndarray, self.labels) + self.__load_folds(folds) + + elif self.split == "train+unlabeled": + self.data, self.labels = self.__loadfile(self.train_list[0][0], self.train_list[1][0]) + self.labels = cast(np.ndarray, self.labels) + self.__load_folds(folds) + unlabeled_data, _ = self.__loadfile(self.train_list[2][0]) + self.data = np.concatenate((self.data, unlabeled_data)) + self.labels = np.concatenate((self.labels, np.asarray([-1] * unlabeled_data.shape[0]))) + + elif self.split == "unlabeled": + self.data, _ = self.__loadfile(self.train_list[2][0]) + self.labels = np.asarray([-1] * self.data.shape[0]) + else: # self.split == 'test': + self.data, self.labels = self.__loadfile(self.test_list[0][0], self.test_list[1][0]) + + class_file = os.path.join(self.root, self.base_folder, self.class_names_file) + if os.path.isfile(class_file): + with open(class_file) as f: + self.classes = f.read().splitlines() + + def _verify_folds(self, folds: Optional[int]) -> Optional[int]: + if folds is None: + return folds + elif isinstance(folds, int): + if folds in range(10): + return folds + msg = "Value for argument folds should be in the range [0, 10), but got {}." + raise ValueError(msg.format(folds)) + else: + msg = "Expected type None or int for argument folds, but got type {}." + raise ValueError(msg.format(type(folds))) + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: (image, target) where target is index of the target class. + """ + target: Optional[int] + if self.labels is not None: + img, target = self.data[index], int(self.labels[index]) + else: + img, target = self.data[index], None + + # doing this so that it is consistent with all other datasets + # to return a PIL Image + img = Image.fromarray(np.transpose(img, (1, 2, 0))) + + if self.transform is not None: + img = self.transform(img) + + if self.target_transform is not None: + target = self.target_transform(target) + + return img, target + + def __len__(self) -> int: + return self.data.shape[0] + + def __loadfile(self, data_file: str, labels_file: Optional[str] = None) -> Tuple[np.ndarray, Optional[np.ndarray]]: + labels = None + if labels_file: + path_to_labels = os.path.join(self.root, self.base_folder, labels_file) + with open(path_to_labels, "rb") as f: + labels = np.fromfile(f, dtype=np.uint8) - 1 # 0-based + + path_to_data = os.path.join(self.root, self.base_folder, data_file) + with open(path_to_data, "rb") as f: + # read whole file in uint8 chunks + everything = np.fromfile(f, dtype=np.uint8) + images = np.reshape(everything, (-1, 3, 96, 96)) + images = np.transpose(images, (0, 1, 3, 2)) + + return images, labels + + def _check_integrity(self) -> bool: + for filename, md5 in self.train_list + self.test_list: + fpath = os.path.join(self.root, self.base_folder, filename) + if not check_integrity(fpath, md5): + return False + return True + + def download(self) -> None: + if self._check_integrity(): + return + download_and_extract_archive(self.url, self.root, filename=self.filename, md5=self.tgz_md5) + self._check_integrity() + + def extra_repr(self) -> str: + return "Split: {split}".format(**self.__dict__) + + def __load_folds(self, folds: Optional[int]) -> None: + # loads one of the folds if specified + if folds is None: + return + path_to_folds = os.path.join(self.root, self.base_folder, self.folds_list_file) + with open(path_to_folds) as f: + str_idx = f.read().splitlines()[folds] + list_idx = np.fromstring(str_idx, dtype=np.int64, sep=" ") + self.data = self.data[list_idx, :, :, :] + if self.labels is not None: + self.labels = self.labels[list_idx] diff --git a/lib/python3.12/site-packages/torchvision/datasets/sun397.py b/lib/python3.12/site-packages/torchvision/datasets/sun397.py new file mode 100644 index 0000000000000000000000000000000000000000..4db0a3cf237376f9a59e06f770f90554abfd87e2 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/sun397.py @@ -0,0 +1,76 @@ +from pathlib import Path +from typing import Any, Callable, Optional, Tuple, Union + +import PIL.Image + +from .utils import download_and_extract_archive +from .vision import VisionDataset + + +class SUN397(VisionDataset): + """`The SUN397 Data Set `_. + + The SUN397 or Scene UNderstanding (SUN) is a dataset for scene recognition consisting of + 397 categories with 108'754 images. + + Args: + root (str or ``pathlib.Path``): Root directory of the dataset. + transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed + version. E.g, ``transforms.RandomCrop``. + target_transform (callable, optional): A function/transform that takes in the target and transforms it. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + """ + + _DATASET_URL = "http://vision.princeton.edu/projects/2010/SUN/SUN397.tar.gz" + _DATASET_MD5 = "8ca2778205c41d23104230ba66911c7a" + + def __init__( + self, + root: Union[str, Path], + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + super().__init__(root, transform=transform, target_transform=target_transform) + self._data_dir = Path(self.root) / "SUN397" + + if download: + self._download() + + if not self._check_exists(): + raise RuntimeError("Dataset not found. You can use download=True to download it") + + with open(self._data_dir / "ClassName.txt") as f: + self.classes = [c[3:].strip() for c in f] + + self.class_to_idx = dict(zip(self.classes, range(len(self.classes)))) + self._image_files = list(self._data_dir.rglob("sun_*.jpg")) + + self._labels = [ + self.class_to_idx["/".join(path.relative_to(self._data_dir).parts[1:-1])] for path in self._image_files + ] + + def __len__(self) -> int: + return len(self._image_files) + + def __getitem__(self, idx: int) -> Tuple[Any, Any]: + image_file, label = self._image_files[idx], self._labels[idx] + image = PIL.Image.open(image_file).convert("RGB") + + if self.transform: + image = self.transform(image) + + if self.target_transform: + label = self.target_transform(label) + + return image, label + + def _check_exists(self) -> bool: + return self._data_dir.is_dir() + + def _download(self) -> None: + if self._check_exists(): + return + download_and_extract_archive(self._DATASET_URL, download_root=self.root, md5=self._DATASET_MD5) diff --git a/lib/python3.12/site-packages/torchvision/datasets/svhn.py b/lib/python3.12/site-packages/torchvision/datasets/svhn.py new file mode 100644 index 0000000000000000000000000000000000000000..5d20d7db7e3ccc1d77c5949830cf7b610d88355f --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/svhn.py @@ -0,0 +1,130 @@ +import os.path +from pathlib import Path +from typing import Any, Callable, Optional, Tuple, Union + +import numpy as np +from PIL import Image + +from .utils import check_integrity, download_url, verify_str_arg +from .vision import VisionDataset + + +class SVHN(VisionDataset): + """`SVHN `_ Dataset. + Note: The SVHN dataset assigns the label `10` to the digit `0`. However, in this Dataset, + we assign the label `0` to the digit `0` to be compatible with PyTorch loss functions which + expect the class labels to be in the range `[0, C-1]` + + .. warning:: + + This class needs `scipy `_ to load data from `.mat` format. + + Args: + root (str or ``pathlib.Path``): Root directory of the dataset where the data is stored. + split (string): One of {'train', 'test', 'extra'}. + Accordingly dataset is selected. 'extra' is Extra training set. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + + """ + + split_list = { + "train": [ + "http://ufldl.stanford.edu/housenumbers/train_32x32.mat", + "train_32x32.mat", + "e26dedcc434d2e4c54c9b2d4a06d8373", + ], + "test": [ + "http://ufldl.stanford.edu/housenumbers/test_32x32.mat", + "test_32x32.mat", + "eb5a983be6a315427106f1b164d9cef3", + ], + "extra": [ + "http://ufldl.stanford.edu/housenumbers/extra_32x32.mat", + "extra_32x32.mat", + "a93ce644f1a588dc4d68dda5feec44a7", + ], + } + + def __init__( + self, + root: Union[str, Path], + split: str = "train", + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + super().__init__(root, transform=transform, target_transform=target_transform) + self.split = verify_str_arg(split, "split", tuple(self.split_list.keys())) + self.url = self.split_list[split][0] + self.filename = self.split_list[split][1] + self.file_md5 = self.split_list[split][2] + + if download: + self.download() + + if not self._check_integrity(): + raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") + + # import here rather than at top of file because this is + # an optional dependency for torchvision + import scipy.io as sio + + # reading(loading) mat file as array + loaded_mat = sio.loadmat(os.path.join(self.root, self.filename)) + + self.data = loaded_mat["X"] + # loading from the .mat file gives an np.ndarray of type np.uint8 + # converting to np.int64, so that we have a LongTensor after + # the conversion from the numpy array + # the squeeze is needed to obtain a 1D tensor + self.labels = loaded_mat["y"].astype(np.int64).squeeze() + + # the svhn dataset assigns the class label "10" to the digit 0 + # this makes it inconsistent with several loss functions + # which expect the class labels to be in the range [0, C-1] + np.place(self.labels, self.labels == 10, 0) + self.data = np.transpose(self.data, (3, 2, 0, 1)) + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: (image, target) where target is index of the target class. + """ + img, target = self.data[index], int(self.labels[index]) + + # doing this so that it is consistent with all other datasets + # to return a PIL Image + img = Image.fromarray(np.transpose(img, (1, 2, 0))) + + if self.transform is not None: + img = self.transform(img) + + if self.target_transform is not None: + target = self.target_transform(target) + + return img, target + + def __len__(self) -> int: + return len(self.data) + + def _check_integrity(self) -> bool: + root = self.root + md5 = self.split_list[self.split][2] + fpath = os.path.join(root, self.filename) + return check_integrity(fpath, md5) + + def download(self) -> None: + md5 = self.split_list[self.split][2] + download_url(self.url, self.root, self.filename, md5) + + def extra_repr(self) -> str: + return "Split: {split}".format(**self.__dict__) diff --git a/lib/python3.12/site-packages/torchvision/datasets/ucf101.py b/lib/python3.12/site-packages/torchvision/datasets/ucf101.py new file mode 100644 index 0000000000000000000000000000000000000000..935f8ad41c78e6086444b72ed1f5ae1f6de34bad --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/ucf101.py @@ -0,0 +1,131 @@ +import os +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +from torch import Tensor + +from .folder import find_classes, make_dataset +from .video_utils import VideoClips +from .vision import VisionDataset + + +class UCF101(VisionDataset): + """ + `UCF101 `_ dataset. + + UCF101 is an action recognition video dataset. + This dataset consider every video as a collection of video clips of fixed size, specified + by ``frames_per_clip``, where the step in frames between each clip is given by + ``step_between_clips``. The dataset itself can be downloaded from the dataset website; + annotations that ``annotation_path`` should be pointing to can be downloaded from `here + `_. + + To give an example, for 2 videos with 10 and 15 frames respectively, if ``frames_per_clip=5`` + and ``step_between_clips=5``, the dataset size will be (2 + 3) = 5, where the first two + elements will come from video 1, and the next three elements from video 2. + Note that we drop clips which do not have exactly ``frames_per_clip`` elements, so not all + frames in a video might be present. + + Internally, it uses a VideoClips object to handle clip creation. + + Args: + root (str or ``pathlib.Path``): Root directory of the UCF101 Dataset. + annotation_path (str): path to the folder containing the split files; + see docstring above for download instructions of these files + frames_per_clip (int): number of frames in a clip. + step_between_clips (int, optional): number of frames between each clip. + fold (int, optional): which fold to use. Should be between 1 and 3. + train (bool, optional): if ``True``, creates a dataset from the train split, + otherwise from the ``test`` split. + transform (callable, optional): A function/transform that takes in a TxHxWxC video + and returns a transformed version. + output_format (str, optional): The format of the output video tensors (before transforms). + Can be either "THWC" (default) or "TCHW". + + Returns: + tuple: A 3-tuple with the following entries: + + - video (Tensor[T, H, W, C] or Tensor[T, C, H, W]): The `T` video frames + - audio(Tensor[K, L]): the audio frames, where `K` is the number of channels + and `L` is the number of points + - label (int): class of the video clip + """ + + def __init__( + self, + root: Union[str, Path], + annotation_path: str, + frames_per_clip: int, + step_between_clips: int = 1, + frame_rate: Optional[int] = None, + fold: int = 1, + train: bool = True, + transform: Optional[Callable] = None, + _precomputed_metadata: Optional[Dict[str, Any]] = None, + num_workers: int = 1, + _video_width: int = 0, + _video_height: int = 0, + _video_min_dimension: int = 0, + _audio_samples: int = 0, + output_format: str = "THWC", + ) -> None: + super().__init__(root) + if not 1 <= fold <= 3: + raise ValueError(f"fold should be between 1 and 3, got {fold}") + + extensions = ("avi",) + self.fold = fold + self.train = train + + self.classes, class_to_idx = find_classes(self.root) + self.samples = make_dataset(self.root, class_to_idx, extensions, is_valid_file=None) + video_list = [x[0] for x in self.samples] + video_clips = VideoClips( + video_list, + frames_per_clip, + step_between_clips, + frame_rate, + _precomputed_metadata, + num_workers=num_workers, + _video_width=_video_width, + _video_height=_video_height, + _video_min_dimension=_video_min_dimension, + _audio_samples=_audio_samples, + output_format=output_format, + ) + # we bookkeep the full version of video clips because we want to be able + # to return the metadata of full version rather than the subset version of + # video clips + self.full_video_clips = video_clips + self.indices = self._select_fold(video_list, annotation_path, fold, train) + self.video_clips = video_clips.subset(self.indices) + self.transform = transform + + @property + def metadata(self) -> Dict[str, Any]: + return self.full_video_clips.metadata + + def _select_fold(self, video_list: List[str], annotation_path: str, fold: int, train: bool) -> List[int]: + name = "train" if train else "test" + name = f"{name}list{fold:02d}.txt" + f = os.path.join(annotation_path, name) + selected_files = set() + with open(f) as fid: + data = fid.readlines() + data = [x.strip().split(" ")[0] for x in data] + data = [os.path.join(self.root, *x.split("/")) for x in data] + selected_files.update(data) + indices = [i for i in range(len(video_list)) if video_list[i] in selected_files] + return indices + + def __len__(self) -> int: + return self.video_clips.num_clips() + + def __getitem__(self, idx: int) -> Tuple[Tensor, Tensor, int]: + video, audio, info, video_idx = self.video_clips.get_clip(idx) + label = self.samples[self.indices[video_idx]][1] + + if self.transform is not None: + video = self.transform(video) + + return video, audio, label diff --git a/lib/python3.12/site-packages/torchvision/datasets/usps.py b/lib/python3.12/site-packages/torchvision/datasets/usps.py new file mode 100644 index 0000000000000000000000000000000000000000..9c681e79f6c3dc8fb6567f90b741ac7839298931 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/usps.py @@ -0,0 +1,96 @@ +import os +from pathlib import Path +from typing import Any, Callable, Optional, Tuple, Union + +import numpy as np +from PIL import Image + +from .utils import download_url +from .vision import VisionDataset + + +class USPS(VisionDataset): + """`USPS `_ Dataset. + The data-format is : [label [index:value ]*256 \\n] * num_lines, where ``label`` lies in ``[1, 10]``. + The value for each pixel lies in ``[-1, 1]``. Here we transform the ``label`` into ``[0, 9]`` + and make pixel values in ``[0, 255]``. + + Args: + root (str or ``pathlib.Path``): Root directory of dataset to store``USPS`` data files. + train (bool, optional): If True, creates dataset from ``usps.bz2``, + otherwise from ``usps.t.bz2``. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + + """ + + split_list = { + "train": [ + "https://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/multiclass/usps.bz2", + "usps.bz2", + "ec16c51db3855ca6c91edd34d0e9b197", + ], + "test": [ + "https://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/multiclass/usps.t.bz2", + "usps.t.bz2", + "8ea070ee2aca1ac39742fdd1ef5ed118", + ], + } + + def __init__( + self, + root: Union[str, Path], + train: bool = True, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + super().__init__(root, transform=transform, target_transform=target_transform) + split = "train" if train else "test" + url, filename, checksum = self.split_list[split] + full_path = os.path.join(self.root, filename) + + if download and not os.path.exists(full_path): + download_url(url, self.root, filename, md5=checksum) + + import bz2 + + with bz2.open(full_path) as fp: + raw_data = [line.decode().split() for line in fp.readlines()] + tmp_list = [[x.split(":")[-1] for x in data[1:]] for data in raw_data] + imgs = np.asarray(tmp_list, dtype=np.float32).reshape((-1, 16, 16)) + imgs = ((imgs + 1) / 2 * 255).astype(dtype=np.uint8) + targets = [int(d[0]) - 1 for d in raw_data] + + self.data = imgs + self.targets = targets + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: (image, target) where target is index of the target class. + """ + img, target = self.data[index], int(self.targets[index]) + + # doing this so that it is consistent with all other datasets + # to return a PIL Image + img = Image.fromarray(img, mode="L") + + if self.transform is not None: + img = self.transform(img) + + if self.target_transform is not None: + target = self.target_transform(target) + + return img, target + + def __len__(self) -> int: + return len(self.data) diff --git a/lib/python3.12/site-packages/torchvision/datasets/utils.py b/lib/python3.12/site-packages/torchvision/datasets/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8bf310896c7efbdb5693156860afbf9977ebf4bb --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/utils.py @@ -0,0 +1,471 @@ +import bz2 +import gzip +import hashlib +import lzma +import os +import os.path +import pathlib +import re +import sys +import tarfile +import urllib +import urllib.error +import urllib.request +import zipfile +from typing import Any, Callable, Dict, IO, Iterable, List, Optional, Tuple, TypeVar, Union +from urllib.parse import urlparse + +import numpy as np +import torch +from torch.utils.model_zoo import tqdm + +from .._internally_replaced_utils import _download_file_from_remote_location, _is_remote_location_available + +USER_AGENT = "pytorch/vision" + + +def _urlretrieve(url: str, filename: Union[str, pathlib.Path], chunk_size: int = 1024 * 32) -> None: + with urllib.request.urlopen(urllib.request.Request(url, headers={"User-Agent": USER_AGENT})) as response: + with open(filename, "wb") as fh, tqdm(total=response.length, unit="B", unit_scale=True) as pbar: + while chunk := response.read(chunk_size): + fh.write(chunk) + pbar.update(len(chunk)) + + +def calculate_md5(fpath: Union[str, pathlib.Path], chunk_size: int = 1024 * 1024) -> str: + # Setting the `usedforsecurity` flag does not change anything about the functionality, but indicates that we are + # not using the MD5 checksum for cryptography. This enables its usage in restricted environments like FIPS. Without + # it torchvision.datasets is unusable in these environments since we perform a MD5 check everywhere. + if sys.version_info >= (3, 9): + md5 = hashlib.md5(usedforsecurity=False) + else: + md5 = hashlib.md5() + with open(fpath, "rb") as f: + while chunk := f.read(chunk_size): + md5.update(chunk) + return md5.hexdigest() + + +def check_md5(fpath: Union[str, pathlib.Path], md5: str, **kwargs: Any) -> bool: + return md5 == calculate_md5(fpath, **kwargs) + + +def check_integrity(fpath: Union[str, pathlib.Path], md5: Optional[str] = None) -> bool: + if not os.path.isfile(fpath): + return False + if md5 is None: + return True + return check_md5(fpath, md5) + + +def _get_redirect_url(url: str, max_hops: int = 3) -> str: + initial_url = url + headers = {"Method": "HEAD", "User-Agent": USER_AGENT} + + for _ in range(max_hops + 1): + with urllib.request.urlopen(urllib.request.Request(url, headers=headers)) as response: + if response.url == url or response.url is None: + return url + + url = response.url + else: + raise RecursionError( + f"Request to {initial_url} exceeded {max_hops} redirects. The last redirect points to {url}." + ) + + +def _get_google_drive_file_id(url: str) -> Optional[str]: + parts = urlparse(url) + + if re.match(r"(drive|docs)[.]google[.]com", parts.netloc) is None: + return None + + match = re.match(r"/file/d/(?P[^/]*)", parts.path) + if match is None: + return None + + return match.group("id") + + +def download_url( + url: str, + root: Union[str, pathlib.Path], + filename: Optional[Union[str, pathlib.Path]] = None, + md5: Optional[str] = None, + max_redirect_hops: int = 3, +) -> None: + """Download a file from a url and place it in root. + + Args: + url (str): URL to download file from + root (str): Directory to place downloaded file in + filename (str, optional): Name to save the file under. If None, use the basename of the URL + md5 (str, optional): MD5 checksum of the download. If None, do not check + max_redirect_hops (int, optional): Maximum number of redirect hops allowed + """ + root = os.path.expanduser(root) + if not filename: + filename = os.path.basename(url) + fpath = os.fspath(os.path.join(root, filename)) + + os.makedirs(root, exist_ok=True) + + # check if file is already present locally + if check_integrity(fpath, md5): + return + + if _is_remote_location_available(): + _download_file_from_remote_location(fpath, url) + else: + # expand redirect chain if needed + url = _get_redirect_url(url, max_hops=max_redirect_hops) + + # check if file is located on Google Drive + file_id = _get_google_drive_file_id(url) + if file_id is not None: + return download_file_from_google_drive(file_id, root, filename, md5) + + # download the file + try: + _urlretrieve(url, fpath) + except (urllib.error.URLError, OSError) as e: # type: ignore[attr-defined] + if url[:5] == "https": + url = url.replace("https:", "http:") + _urlretrieve(url, fpath) + else: + raise e + + # check integrity of downloaded file + if not check_integrity(fpath, md5): + raise RuntimeError("File not found or corrupted.") + + +def list_dir(root: Union[str, pathlib.Path], prefix: bool = False) -> List[str]: + """List all directories at a given root + + Args: + root (str): Path to directory whose folders need to be listed + prefix (bool, optional): If true, prepends the path to each result, otherwise + only returns the name of the directories found + """ + root = os.path.expanduser(root) + directories = [p for p in os.listdir(root) if os.path.isdir(os.path.join(root, p))] + if prefix is True: + directories = [os.path.join(root, d) for d in directories] + return directories + + +def list_files(root: Union[str, pathlib.Path], suffix: str, prefix: bool = False) -> List[str]: + """List all files ending with a suffix at a given root + + Args: + root (str): Path to directory whose folders need to be listed + suffix (str or tuple): Suffix of the files to match, e.g. '.png' or ('.jpg', '.png'). + It uses the Python "str.endswith" method and is passed directly + prefix (bool, optional): If true, prepends the path to each result, otherwise + only returns the name of the files found + """ + root = os.path.expanduser(root) + files = [p for p in os.listdir(root) if os.path.isfile(os.path.join(root, p)) and p.endswith(suffix)] + if prefix is True: + files = [os.path.join(root, d) for d in files] + return files + + +def download_file_from_google_drive( + file_id: str, + root: Union[str, pathlib.Path], + filename: Optional[Union[str, pathlib.Path]] = None, + md5: Optional[str] = None, +): + """Download a Google Drive file from and place it in root. + + Args: + file_id (str): id of file to be downloaded + root (str): Directory to place downloaded file in + filename (str, optional): Name to save the file under. If None, use the id of the file. + md5 (str, optional): MD5 checksum of the download. If None, do not check + """ + try: + import gdown + except ModuleNotFoundError: + raise RuntimeError( + "To download files from GDrive, 'gdown' is required. You can install it with 'pip install gdown'." + ) + + root = os.path.expanduser(root) + if not filename: + filename = file_id + fpath = os.fspath(os.path.join(root, filename)) + + os.makedirs(root, exist_ok=True) + + if check_integrity(fpath, md5): + return + + gdown.download(id=file_id, output=fpath, quiet=False, user_agent=USER_AGENT) + + if not check_integrity(fpath, md5): + raise RuntimeError("File not found or corrupted.") + + +def _extract_tar( + from_path: Union[str, pathlib.Path], to_path: Union[str, pathlib.Path], compression: Optional[str] +) -> None: + with tarfile.open(from_path, f"r:{compression[1:]}" if compression else "r") as tar: + tar.extractall(to_path) + + +_ZIP_COMPRESSION_MAP: Dict[str, int] = { + ".bz2": zipfile.ZIP_BZIP2, + ".xz": zipfile.ZIP_LZMA, +} + + +def _extract_zip( + from_path: Union[str, pathlib.Path], to_path: Union[str, pathlib.Path], compression: Optional[str] +) -> None: + with zipfile.ZipFile( + from_path, "r", compression=_ZIP_COMPRESSION_MAP[compression] if compression else zipfile.ZIP_STORED + ) as zip: + zip.extractall(to_path) + + +_ARCHIVE_EXTRACTORS: Dict[str, Callable[[Union[str, pathlib.Path], Union[str, pathlib.Path], Optional[str]], None]] = { + ".tar": _extract_tar, + ".zip": _extract_zip, +} +_COMPRESSED_FILE_OPENERS: Dict[str, Callable[..., IO]] = { + ".bz2": bz2.open, + ".gz": gzip.open, + ".xz": lzma.open, +} +_FILE_TYPE_ALIASES: Dict[str, Tuple[Optional[str], Optional[str]]] = { + ".tbz": (".tar", ".bz2"), + ".tbz2": (".tar", ".bz2"), + ".tgz": (".tar", ".gz"), +} + + +def _detect_file_type(file: Union[str, pathlib.Path]) -> Tuple[str, Optional[str], Optional[str]]: + """Detect the archive type and/or compression of a file. + + Args: + file (str): the filename + + Returns: + (tuple): tuple of suffix, archive type, and compression + + Raises: + RuntimeError: if file has no suffix or suffix is not supported + """ + suffixes = pathlib.Path(file).suffixes + if not suffixes: + raise RuntimeError( + f"File '{file}' has no suffixes that could be used to detect the archive type and compression." + ) + suffix = suffixes[-1] + + # check if the suffix is a known alias + if suffix in _FILE_TYPE_ALIASES: + return (suffix, *_FILE_TYPE_ALIASES[suffix]) + + # check if the suffix is an archive type + if suffix in _ARCHIVE_EXTRACTORS: + return suffix, suffix, None + + # check if the suffix is a compression + if suffix in _COMPRESSED_FILE_OPENERS: + # check for suffix hierarchy + if len(suffixes) > 1: + suffix2 = suffixes[-2] + + # check if the suffix2 is an archive type + if suffix2 in _ARCHIVE_EXTRACTORS: + return suffix2 + suffix, suffix2, suffix + + return suffix, None, suffix + + valid_suffixes = sorted(set(_FILE_TYPE_ALIASES) | set(_ARCHIVE_EXTRACTORS) | set(_COMPRESSED_FILE_OPENERS)) + raise RuntimeError(f"Unknown compression or archive type: '{suffix}'.\nKnown suffixes are: '{valid_suffixes}'.") + + +def _decompress( + from_path: Union[str, pathlib.Path], + to_path: Optional[Union[str, pathlib.Path]] = None, + remove_finished: bool = False, +) -> pathlib.Path: + r"""Decompress a file. + + The compression is automatically detected from the file name. + + Args: + from_path (str): Path to the file to be decompressed. + to_path (str): Path to the decompressed file. If omitted, ``from_path`` without compression extension is used. + remove_finished (bool): If ``True``, remove the file after the extraction. + + Returns: + (str): Path to the decompressed file. + """ + suffix, archive_type, compression = _detect_file_type(from_path) + if not compression: + raise RuntimeError(f"Couldn't detect a compression from suffix {suffix}.") + + if to_path is None: + to_path = pathlib.Path(os.fspath(from_path).replace(suffix, archive_type if archive_type is not None else "")) + + # We don't need to check for a missing key here, since this was already done in _detect_file_type() + compressed_file_opener = _COMPRESSED_FILE_OPENERS[compression] + + with compressed_file_opener(from_path, "rb") as rfh, open(to_path, "wb") as wfh: + wfh.write(rfh.read()) + + if remove_finished: + os.remove(from_path) + + return pathlib.Path(to_path) + + +def extract_archive( + from_path: Union[str, pathlib.Path], + to_path: Optional[Union[str, pathlib.Path]] = None, + remove_finished: bool = False, +) -> Union[str, pathlib.Path]: + """Extract an archive. + + The archive type and a possible compression is automatically detected from the file name. If the file is compressed + but not an archive the call is dispatched to :func:`decompress`. + + Args: + from_path (str): Path to the file to be extracted. + to_path (str): Path to the directory the file will be extracted to. If omitted, the directory of the file is + used. + remove_finished (bool): If ``True``, remove the file after the extraction. + + Returns: + (str): Path to the directory the file was extracted to. + """ + + def path_or_str(ret_path: pathlib.Path) -> Union[str, pathlib.Path]: + if isinstance(from_path, str): + return os.fspath(ret_path) + else: + return ret_path + + if to_path is None: + to_path = os.path.dirname(from_path) + + suffix, archive_type, compression = _detect_file_type(from_path) + if not archive_type: + ret_path = _decompress( + from_path, + os.path.join(to_path, os.path.basename(from_path).replace(suffix, "")), + remove_finished=remove_finished, + ) + return path_or_str(ret_path) + + # We don't need to check for a missing key here, since this was already done in _detect_file_type() + extractor = _ARCHIVE_EXTRACTORS[archive_type] + + extractor(from_path, to_path, compression) + if remove_finished: + os.remove(from_path) + + return path_or_str(pathlib.Path(to_path)) + + +def download_and_extract_archive( + url: str, + download_root: Union[str, pathlib.Path], + extract_root: Optional[Union[str, pathlib.Path]] = None, + filename: Optional[Union[str, pathlib.Path]] = None, + md5: Optional[str] = None, + remove_finished: bool = False, +) -> None: + download_root = os.path.expanduser(download_root) + if extract_root is None: + extract_root = download_root + if not filename: + filename = os.path.basename(url) + + download_url(url, download_root, filename, md5) + + archive = os.path.join(download_root, filename) + extract_archive(archive, extract_root, remove_finished) + + +def iterable_to_str(iterable: Iterable) -> str: + return "'" + "', '".join([str(item) for item in iterable]) + "'" + + +T = TypeVar("T", str, bytes) + + +def verify_str_arg( + value: T, + arg: Optional[str] = None, + valid_values: Optional[Iterable[T]] = None, + custom_msg: Optional[str] = None, +) -> T: + if not isinstance(value, str): + if arg is None: + msg = "Expected type str, but got type {type}." + else: + msg = "Expected type str for argument {arg}, but got type {type}." + msg = msg.format(type=type(value), arg=arg) + raise ValueError(msg) + + if valid_values is None: + return value + + if value not in valid_values: + if custom_msg is not None: + msg = custom_msg + else: + msg = "Unknown value '{value}' for argument {arg}. Valid values are {{{valid_values}}}." + msg = msg.format(value=value, arg=arg, valid_values=iterable_to_str(valid_values)) + raise ValueError(msg) + + return value + + +def _read_pfm(file_name: Union[str, pathlib.Path], slice_channels: int = 2) -> np.ndarray: + """Read file in .pfm format. Might contain either 1 or 3 channels of data. + + Args: + file_name (str): Path to the file. + slice_channels (int): Number of channels to slice out of the file. + Useful for reading different data formats stored in .pfm files: Optical Flows, Stereo Disparity Maps, etc. + """ + + with open(file_name, "rb") as f: + header = f.readline().rstrip() + if header not in [b"PF", b"Pf"]: + raise ValueError("Invalid PFM file") + + dim_match = re.match(rb"^(\d+)\s(\d+)\s$", f.readline()) + if not dim_match: + raise Exception("Malformed PFM header.") + w, h = (int(dim) for dim in dim_match.groups()) + + scale = float(f.readline().rstrip()) + if scale < 0: # little-endian + endian = "<" + scale = -scale + else: + endian = ">" # big-endian + + data = np.fromfile(f, dtype=endian + "f") + + pfm_channels = 3 if header == b"PF" else 1 + + data = data.reshape(h, w, pfm_channels).transpose(2, 0, 1) + data = np.flip(data, axis=1) # flip on h dimension + data = data[:slice_channels, :, :] + return data.astype(np.float32) + + +def _flip_byte_order(t: torch.Tensor) -> torch.Tensor: + return ( + t.contiguous().view(torch.uint8).view(*t.shape, t.element_size()).flip(-1).view(*t.shape[:-1], -1).view(t.dtype) + ) diff --git a/lib/python3.12/site-packages/torchvision/datasets/video_utils.py b/lib/python3.12/site-packages/torchvision/datasets/video_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a412bc5841cca891e6808215a6bd7d658c69e0f3 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/video_utils.py @@ -0,0 +1,419 @@ +import bisect +import math +import warnings +from fractions import Fraction +from typing import Any, Callable, cast, Dict, List, Optional, Tuple, TypeVar, Union + +import torch +from torchvision.io import _probe_video_from_file, _read_video_from_file, read_video, read_video_timestamps + +from .utils import tqdm + +T = TypeVar("T") + + +def pts_convert(pts: int, timebase_from: Fraction, timebase_to: Fraction, round_func: Callable = math.floor) -> int: + """convert pts between different time bases + Args: + pts: presentation timestamp, float + timebase_from: original timebase. Fraction + timebase_to: new timebase. Fraction + round_func: rounding function. + """ + new_pts = Fraction(pts, 1) * timebase_from / timebase_to + return round_func(new_pts) + + +def unfold(tensor: torch.Tensor, size: int, step: int, dilation: int = 1) -> torch.Tensor: + """ + similar to tensor.unfold, but with the dilation + and specialized for 1d tensors + + Returns all consecutive windows of `size` elements, with + `step` between windows. The distance between each element + in a window is given by `dilation`. + """ + if tensor.dim() != 1: + raise ValueError(f"tensor should have 1 dimension instead of {tensor.dim()}") + o_stride = tensor.stride(0) + numel = tensor.numel() + new_stride = (step * o_stride, dilation * o_stride) + new_size = ((numel - (dilation * (size - 1) + 1)) // step + 1, size) + if new_size[0] < 1: + new_size = (0, size) + return torch.as_strided(tensor, new_size, new_stride) + + +class _VideoTimestampsDataset: + """ + Dataset used to parallelize the reading of the timestamps + of a list of videos, given their paths in the filesystem. + + Used in VideoClips and defined at top level, so it can be + pickled when forking. + """ + + def __init__(self, video_paths: List[str]) -> None: + self.video_paths = video_paths + + def __len__(self) -> int: + return len(self.video_paths) + + def __getitem__(self, idx: int) -> Tuple[List[int], Optional[float]]: + return read_video_timestamps(self.video_paths[idx]) + + +def _collate_fn(x: T) -> T: + """ + Dummy collate function to be used with _VideoTimestampsDataset + """ + return x + + +class VideoClips: + """ + Given a list of video files, computes all consecutive subvideos of size + `clip_length_in_frames`, where the distance between each subvideo in the + same video is defined by `frames_between_clips`. + If `frame_rate` is specified, it will also resample all the videos to have + the same frame rate, and the clips will refer to this frame rate. + + Creating this instance the first time is time-consuming, as it needs to + decode all the videos in `video_paths`. It is recommended that you + cache the results after instantiation of the class. + + Recreating the clips for different clip lengths is fast, and can be done + with the `compute_clips` method. + + Args: + video_paths (List[str]): paths to the video files + clip_length_in_frames (int): size of a clip in number of frames + frames_between_clips (int): step (in frames) between each clip + frame_rate (float, optional): if specified, it will resample the video + so that it has `frame_rate`, and then the clips will be defined + on the resampled video + num_workers (int): how many subprocesses to use for data loading. + 0 means that the data will be loaded in the main process. (default: 0) + output_format (str): The format of the output video tensors. Can be either "THWC" (default) or "TCHW". + """ + + def __init__( + self, + video_paths: List[str], + clip_length_in_frames: int = 16, + frames_between_clips: int = 1, + frame_rate: Optional[float] = None, + _precomputed_metadata: Optional[Dict[str, Any]] = None, + num_workers: int = 0, + _video_width: int = 0, + _video_height: int = 0, + _video_min_dimension: int = 0, + _video_max_dimension: int = 0, + _audio_samples: int = 0, + _audio_channels: int = 0, + output_format: str = "THWC", + ) -> None: + + self.video_paths = video_paths + self.num_workers = num_workers + + # these options are not valid for pyav backend + self._video_width = _video_width + self._video_height = _video_height + self._video_min_dimension = _video_min_dimension + self._video_max_dimension = _video_max_dimension + self._audio_samples = _audio_samples + self._audio_channels = _audio_channels + self.output_format = output_format.upper() + if self.output_format not in ("THWC", "TCHW"): + raise ValueError(f"output_format should be either 'THWC' or 'TCHW', got {output_format}.") + + if _precomputed_metadata is None: + self._compute_frame_pts() + else: + self._init_from_metadata(_precomputed_metadata) + self.compute_clips(clip_length_in_frames, frames_between_clips, frame_rate) + + def _compute_frame_pts(self) -> None: + self.video_pts = [] # len = num_videos. Each entry is a tensor of shape (num_frames_in_video,) + self.video_fps: List[float] = [] # len = num_videos + + # strategy: use a DataLoader to parallelize read_video_timestamps + # so need to create a dummy dataset first + import torch.utils.data + + dl: torch.utils.data.DataLoader = torch.utils.data.DataLoader( + _VideoTimestampsDataset(self.video_paths), # type: ignore[arg-type] + batch_size=16, + num_workers=self.num_workers, + collate_fn=_collate_fn, + ) + + with tqdm(total=len(dl)) as pbar: + for batch in dl: + pbar.update(1) + batch_pts, batch_fps = list(zip(*batch)) + # we need to specify dtype=torch.long because for empty list, + # torch.as_tensor will use torch.float as default dtype. This + # happens when decoding fails and no pts is returned in the list. + batch_pts = [torch.as_tensor(pts, dtype=torch.long) for pts in batch_pts] + self.video_pts.extend(batch_pts) + self.video_fps.extend(batch_fps) + + def _init_from_metadata(self, metadata: Dict[str, Any]) -> None: + self.video_paths = metadata["video_paths"] + assert len(self.video_paths) == len(metadata["video_pts"]) + self.video_pts = metadata["video_pts"] + assert len(self.video_paths) == len(metadata["video_fps"]) + self.video_fps = metadata["video_fps"] + + @property + def metadata(self) -> Dict[str, Any]: + _metadata = { + "video_paths": self.video_paths, + "video_pts": self.video_pts, + "video_fps": self.video_fps, + } + return _metadata + + def subset(self, indices: List[int]) -> "VideoClips": + video_paths = [self.video_paths[i] for i in indices] + video_pts = [self.video_pts[i] for i in indices] + video_fps = [self.video_fps[i] for i in indices] + metadata = { + "video_paths": video_paths, + "video_pts": video_pts, + "video_fps": video_fps, + } + return type(self)( + video_paths, + clip_length_in_frames=self.num_frames, + frames_between_clips=self.step, + frame_rate=self.frame_rate, + _precomputed_metadata=metadata, + num_workers=self.num_workers, + _video_width=self._video_width, + _video_height=self._video_height, + _video_min_dimension=self._video_min_dimension, + _video_max_dimension=self._video_max_dimension, + _audio_samples=self._audio_samples, + _audio_channels=self._audio_channels, + output_format=self.output_format, + ) + + @staticmethod + def compute_clips_for_video( + video_pts: torch.Tensor, num_frames: int, step: int, fps: Optional[float], frame_rate: Optional[float] = None + ) -> Tuple[torch.Tensor, Union[List[slice], torch.Tensor]]: + if fps is None: + # if for some reason the video doesn't have fps (because doesn't have a video stream) + # set the fps to 1. The value doesn't matter, because video_pts is empty anyway + fps = 1 + if frame_rate is None: + frame_rate = fps + total_frames = len(video_pts) * frame_rate / fps + _idxs = VideoClips._resample_video_idx(int(math.floor(total_frames)), fps, frame_rate) + video_pts = video_pts[_idxs] + clips = unfold(video_pts, num_frames, step) + if not clips.numel(): + warnings.warn( + "There aren't enough frames in the current video to get a clip for the given clip length and " + "frames between clips. The video (and potentially others) will be skipped." + ) + idxs: Union[List[slice], torch.Tensor] + if isinstance(_idxs, slice): + idxs = [_idxs] * len(clips) + else: + idxs = unfold(_idxs, num_frames, step) + return clips, idxs + + def compute_clips(self, num_frames: int, step: int, frame_rate: Optional[float] = None) -> None: + """ + Compute all consecutive sequences of clips from video_pts. + Always returns clips of size `num_frames`, meaning that the + last few frames in a video can potentially be dropped. + + Args: + num_frames (int): number of frames for the clip + step (int): distance between two clips + frame_rate (int, optional): The frame rate + """ + self.num_frames = num_frames + self.step = step + self.frame_rate = frame_rate + self.clips = [] + self.resampling_idxs = [] + for video_pts, fps in zip(self.video_pts, self.video_fps): + clips, idxs = self.compute_clips_for_video(video_pts, num_frames, step, fps, frame_rate) + self.clips.append(clips) + self.resampling_idxs.append(idxs) + clip_lengths = torch.as_tensor([len(v) for v in self.clips]) + self.cumulative_sizes = clip_lengths.cumsum(0).tolist() + + def __len__(self) -> int: + return self.num_clips() + + def num_videos(self) -> int: + return len(self.video_paths) + + def num_clips(self) -> int: + """ + Number of subclips that are available in the video list. + """ + return self.cumulative_sizes[-1] + + def get_clip_location(self, idx: int) -> Tuple[int, int]: + """ + Converts a flattened representation of the indices into a video_idx, clip_idx + representation. + """ + video_idx = bisect.bisect_right(self.cumulative_sizes, idx) + if video_idx == 0: + clip_idx = idx + else: + clip_idx = idx - self.cumulative_sizes[video_idx - 1] + return video_idx, clip_idx + + @staticmethod + def _resample_video_idx(num_frames: int, original_fps: float, new_fps: float) -> Union[slice, torch.Tensor]: + step = original_fps / new_fps + if step.is_integer(): + # optimization: if step is integer, don't need to perform + # advanced indexing + step = int(step) + return slice(None, None, step) + idxs = torch.arange(num_frames, dtype=torch.float32) * step + idxs = idxs.floor().to(torch.int64) + return idxs + + def get_clip(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, Any], int]: + """ + Gets a subclip from a list of videos. + + Args: + idx (int): index of the subclip. Must be between 0 and num_clips(). + + Returns: + video (Tensor) + audio (Tensor) + info (Dict) + video_idx (int): index of the video in `video_paths` + """ + if idx >= self.num_clips(): + raise IndexError(f"Index {idx} out of range ({self.num_clips()} number of clips)") + video_idx, clip_idx = self.get_clip_location(idx) + video_path = self.video_paths[video_idx] + clip_pts = self.clips[video_idx][clip_idx] + + from torchvision import get_video_backend + + backend = get_video_backend() + + if backend == "pyav": + # check for invalid options + if self._video_width != 0: + raise ValueError("pyav backend doesn't support _video_width != 0") + if self._video_height != 0: + raise ValueError("pyav backend doesn't support _video_height != 0") + if self._video_min_dimension != 0: + raise ValueError("pyav backend doesn't support _video_min_dimension != 0") + if self._video_max_dimension != 0: + raise ValueError("pyav backend doesn't support _video_max_dimension != 0") + if self._audio_samples != 0: + raise ValueError("pyav backend doesn't support _audio_samples != 0") + + if backend == "pyav": + start_pts = clip_pts[0].item() + end_pts = clip_pts[-1].item() + video, audio, info = read_video(video_path, start_pts, end_pts) + else: + _info = _probe_video_from_file(video_path) + video_fps = _info.video_fps + audio_fps = None + + video_start_pts = cast(int, clip_pts[0].item()) + video_end_pts = cast(int, clip_pts[-1].item()) + + audio_start_pts, audio_end_pts = 0, -1 + audio_timebase = Fraction(0, 1) + video_timebase = Fraction(_info.video_timebase.numerator, _info.video_timebase.denominator) + if _info.has_audio: + audio_timebase = Fraction(_info.audio_timebase.numerator, _info.audio_timebase.denominator) + audio_start_pts = pts_convert(video_start_pts, video_timebase, audio_timebase, math.floor) + audio_end_pts = pts_convert(video_end_pts, video_timebase, audio_timebase, math.ceil) + audio_fps = _info.audio_sample_rate + video, audio, _ = _read_video_from_file( + video_path, + video_width=self._video_width, + video_height=self._video_height, + video_min_dimension=self._video_min_dimension, + video_max_dimension=self._video_max_dimension, + video_pts_range=(video_start_pts, video_end_pts), + video_timebase=video_timebase, + audio_samples=self._audio_samples, + audio_channels=self._audio_channels, + audio_pts_range=(audio_start_pts, audio_end_pts), + audio_timebase=audio_timebase, + ) + + info = {"video_fps": video_fps} + if audio_fps is not None: + info["audio_fps"] = audio_fps + + if self.frame_rate is not None: + resampling_idx = self.resampling_idxs[video_idx][clip_idx] + if isinstance(resampling_idx, torch.Tensor): + resampling_idx = resampling_idx - resampling_idx[0] + video = video[resampling_idx] + info["video_fps"] = self.frame_rate + assert len(video) == self.num_frames, f"{video.shape} x {self.num_frames}" + + if self.output_format == "TCHW": + # [T,H,W,C] --> [T,C,H,W] + video = video.permute(0, 3, 1, 2) + + return video, audio, info, video_idx + + def __getstate__(self) -> Dict[str, Any]: + video_pts_sizes = [len(v) for v in self.video_pts] + # To be back-compatible, we convert data to dtype torch.long as needed + # because for empty list, in legacy implementation, torch.as_tensor will + # use torch.float as default dtype. This happens when decoding fails and + # no pts is returned in the list. + video_pts = [x.to(torch.int64) for x in self.video_pts] + # video_pts can be an empty list if no frames have been decoded + if video_pts: + video_pts = torch.cat(video_pts) # type: ignore[assignment] + # avoid bug in https://github.com/pytorch/pytorch/issues/32351 + # TODO: Revert it once the bug is fixed. + video_pts = video_pts.numpy() # type: ignore[attr-defined] + + # make a copy of the fields of self + d = self.__dict__.copy() + d["video_pts_sizes"] = video_pts_sizes + d["video_pts"] = video_pts + # delete the following attributes to reduce the size of dictionary. They + # will be re-computed in "__setstate__()" + del d["clips"] + del d["resampling_idxs"] + del d["cumulative_sizes"] + + # for backwards-compatibility + d["_version"] = 2 + return d + + def __setstate__(self, d: Dict[str, Any]) -> None: + # for backwards-compatibility + if "_version" not in d: + self.__dict__ = d + return + + video_pts = torch.as_tensor(d["video_pts"], dtype=torch.int64) + video_pts = torch.split(video_pts, d["video_pts_sizes"], dim=0) + # don't need this info anymore + del d["video_pts_sizes"] + + d["video_pts"] = video_pts + self.__dict__ = d + # recompute attributes "clips", "resampling_idxs" and other derivative ones + self.compute_clips(self.num_frames, self.step, self.frame_rate) diff --git a/lib/python3.12/site-packages/torchvision/datasets/vision.py b/lib/python3.12/site-packages/torchvision/datasets/vision.py new file mode 100644 index 0000000000000000000000000000000000000000..e524c67e263abd9ec9ed5b86aa55203e1f3ba144 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/vision.py @@ -0,0 +1,111 @@ +import os +from pathlib import Path +from typing import Any, Callable, List, Optional, Tuple, Union + +import torch.utils.data as data + +from ..utils import _log_api_usage_once + + +class VisionDataset(data.Dataset): + """ + Base Class For making datasets which are compatible with torchvision. + It is necessary to override the ``__getitem__`` and ``__len__`` method. + + Args: + root (string, optional): Root directory of dataset. Only used for `__repr__`. + transforms (callable, optional): A function/transforms that takes in + an image and a label and returns the transformed versions of both. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + + .. note:: + + :attr:`transforms` and the combination of :attr:`transform` and :attr:`target_transform` are mutually exclusive. + """ + + _repr_indent = 4 + + def __init__( + self, + root: Union[str, Path] = None, # type: ignore[assignment] + transforms: Optional[Callable] = None, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + ) -> None: + _log_api_usage_once(self) + if isinstance(root, str): + root = os.path.expanduser(root) + self.root = root + + has_transforms = transforms is not None + has_separate_transform = transform is not None or target_transform is not None + if has_transforms and has_separate_transform: + raise ValueError("Only transforms or transform/target_transform can be passed as argument") + + # for backwards-compatibility + self.transform = transform + self.target_transform = target_transform + + if has_separate_transform: + transforms = StandardTransform(transform, target_transform) + self.transforms = transforms + + def __getitem__(self, index: int) -> Any: + """ + Args: + index (int): Index + + Returns: + (Any): Sample and meta data, optionally transformed by the respective transforms. + """ + raise NotImplementedError + + def __len__(self) -> int: + raise NotImplementedError + + def __repr__(self) -> str: + head = "Dataset " + self.__class__.__name__ + body = [f"Number of datapoints: {self.__len__()}"] + if self.root is not None: + body.append(f"Root location: {self.root}") + body += self.extra_repr().splitlines() + if hasattr(self, "transforms") and self.transforms is not None: + body += [repr(self.transforms)] + lines = [head] + [" " * self._repr_indent + line for line in body] + return "\n".join(lines) + + def _format_transform_repr(self, transform: Callable, head: str) -> List[str]: + lines = transform.__repr__().splitlines() + return [f"{head}{lines[0]}"] + ["{}{}".format(" " * len(head), line) for line in lines[1:]] + + def extra_repr(self) -> str: + return "" + + +class StandardTransform: + def __init__(self, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None) -> None: + self.transform = transform + self.target_transform = target_transform + + def __call__(self, input: Any, target: Any) -> Tuple[Any, Any]: + if self.transform is not None: + input = self.transform(input) + if self.target_transform is not None: + target = self.target_transform(target) + return input, target + + def _format_transform_repr(self, transform: Callable, head: str) -> List[str]: + lines = transform.__repr__().splitlines() + return [f"{head}{lines[0]}"] + ["{}{}".format(" " * len(head), line) for line in lines[1:]] + + def __repr__(self) -> str: + body = [self.__class__.__name__] + if self.transform is not None: + body += self._format_transform_repr(self.transform, "Transform: ") + if self.target_transform is not None: + body += self._format_transform_repr(self.target_transform, "Target transform: ") + + return "\n".join(body) diff --git a/lib/python3.12/site-packages/torchvision/datasets/voc.py b/lib/python3.12/site-packages/torchvision/datasets/voc.py new file mode 100644 index 0000000000000000000000000000000000000000..0f0e84c84fa159e942b6da20fb9b2f651816ac06 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/voc.py @@ -0,0 +1,224 @@ +import collections +import os +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from xml.etree.ElementTree import Element as ET_Element + +try: + from defusedxml.ElementTree import parse as ET_parse +except ImportError: + from xml.etree.ElementTree import parse as ET_parse + +from PIL import Image + +from .utils import download_and_extract_archive, verify_str_arg +from .vision import VisionDataset + +DATASET_YEAR_DICT = { + "2012": { + "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar", + "filename": "VOCtrainval_11-May-2012.tar", + "md5": "6cd6e144f989b92b3379bac3b3de84fd", + "base_dir": os.path.join("VOCdevkit", "VOC2012"), + }, + "2011": { + "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2011/VOCtrainval_25-May-2011.tar", + "filename": "VOCtrainval_25-May-2011.tar", + "md5": "6c3384ef61512963050cb5d687e5bf1e", + "base_dir": os.path.join("TrainVal", "VOCdevkit", "VOC2011"), + }, + "2010": { + "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2010/VOCtrainval_03-May-2010.tar", + "filename": "VOCtrainval_03-May-2010.tar", + "md5": "da459979d0c395079b5c75ee67908abb", + "base_dir": os.path.join("VOCdevkit", "VOC2010"), + }, + "2009": { + "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2009/VOCtrainval_11-May-2009.tar", + "filename": "VOCtrainval_11-May-2009.tar", + "md5": "a3e00b113cfcfebf17e343f59da3caa1", + "base_dir": os.path.join("VOCdevkit", "VOC2009"), + }, + "2008": { + "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2008/VOCtrainval_14-Jul-2008.tar", + "filename": "VOCtrainval_11-May-2012.tar", + "md5": "2629fa636546599198acfcfbfcf1904a", + "base_dir": os.path.join("VOCdevkit", "VOC2008"), + }, + "2007": { + "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar", + "filename": "VOCtrainval_06-Nov-2007.tar", + "md5": "c52e279531787c972589f7e41ab4ae64", + "base_dir": os.path.join("VOCdevkit", "VOC2007"), + }, + "2007-test": { + "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar", + "filename": "VOCtest_06-Nov-2007.tar", + "md5": "b6e924de25625d8de591ea690078ad9f", + "base_dir": os.path.join("VOCdevkit", "VOC2007"), + }, +} + + +class _VOCBase(VisionDataset): + _SPLITS_DIR: str + _TARGET_DIR: str + _TARGET_FILE_EXT: str + + def __init__( + self, + root: Union[str, Path], + year: str = "2012", + image_set: str = "train", + download: bool = False, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + transforms: Optional[Callable] = None, + ): + super().__init__(root, transforms, transform, target_transform) + + self.year = verify_str_arg(year, "year", valid_values=[str(yr) for yr in range(2007, 2013)]) + + valid_image_sets = ["train", "trainval", "val"] + if year == "2007": + valid_image_sets.append("test") + self.image_set = verify_str_arg(image_set, "image_set", valid_image_sets) + + key = "2007-test" if year == "2007" and image_set == "test" else year + dataset_year_dict = DATASET_YEAR_DICT[key] + + self.url = dataset_year_dict["url"] + self.filename = dataset_year_dict["filename"] + self.md5 = dataset_year_dict["md5"] + + base_dir = dataset_year_dict["base_dir"] + voc_root = os.path.join(self.root, base_dir) + + if download: + download_and_extract_archive(self.url, self.root, filename=self.filename, md5=self.md5) + + if not os.path.isdir(voc_root): + raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") + + splits_dir = os.path.join(voc_root, "ImageSets", self._SPLITS_DIR) + split_f = os.path.join(splits_dir, image_set.rstrip("\n") + ".txt") + with open(os.path.join(split_f)) as f: + file_names = [x.strip() for x in f.readlines()] + + image_dir = os.path.join(voc_root, "JPEGImages") + self.images = [os.path.join(image_dir, x + ".jpg") for x in file_names] + + target_dir = os.path.join(voc_root, self._TARGET_DIR) + self.targets = [os.path.join(target_dir, x + self._TARGET_FILE_EXT) for x in file_names] + + assert len(self.images) == len(self.targets) + + def __len__(self) -> int: + return len(self.images) + + +class VOCSegmentation(_VOCBase): + """`Pascal VOC `_ Segmentation Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of the VOC Dataset. + year (string, optional): The dataset year, supports years ``"2007"`` to ``"2012"``. + image_set (string, optional): Select the image_set to use, ``"train"``, ``"trainval"`` or ``"val"``. If + ``year=="2007"``, can also be ``"test"``. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + transforms (callable, optional): A function/transform that takes input sample and its target as entry + and returns a transformed version. + """ + + _SPLITS_DIR = "Segmentation" + _TARGET_DIR = "SegmentationClass" + _TARGET_FILE_EXT = ".png" + + @property + def masks(self) -> List[str]: + return self.targets + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: (image, target) where target is the image segmentation. + """ + img = Image.open(self.images[index]).convert("RGB") + target = Image.open(self.masks[index]) + + if self.transforms is not None: + img, target = self.transforms(img, target) + + return img, target + + +class VOCDetection(_VOCBase): + """`Pascal VOC `_ Detection Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory of the VOC Dataset. + year (string, optional): The dataset year, supports years ``"2007"`` to ``"2012"``. + image_set (string, optional): Select the image_set to use, ``"train"``, ``"trainval"`` or ``"val"``. If + ``year=="2007"``, can also be ``"test"``. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + (default: alphabetic indexing of VOC's 20 classes). + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, required): A function/transform that takes in the + target and transforms it. + transforms (callable, optional): A function/transform that takes input sample and its target as entry + and returns a transformed version. + """ + + _SPLITS_DIR = "Main" + _TARGET_DIR = "Annotations" + _TARGET_FILE_EXT = ".xml" + + @property + def annotations(self) -> List[str]: + return self.targets + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: (image, target) where target is a dictionary of the XML tree. + """ + img = Image.open(self.images[index]).convert("RGB") + target = self.parse_voc_xml(ET_parse(self.annotations[index]).getroot()) + + if self.transforms is not None: + img, target = self.transforms(img, target) + + return img, target + + @staticmethod + def parse_voc_xml(node: ET_Element) -> Dict[str, Any]: + voc_dict: Dict[str, Any] = {} + children = list(node) + if children: + def_dic: Dict[str, Any] = collections.defaultdict(list) + for dc in map(VOCDetection.parse_voc_xml, children): + for ind, v in dc.items(): + def_dic[ind].append(v) + if node.tag == "annotation": + def_dic["object"] = [def_dic["object"]] + voc_dict = {node.tag: {ind: v[0] if len(v) == 1 else v for ind, v in def_dic.items()}} + if node.text: + text = node.text.strip() + if not children: + voc_dict[node.tag] = text + return voc_dict diff --git a/lib/python3.12/site-packages/torchvision/datasets/widerface.py b/lib/python3.12/site-packages/torchvision/datasets/widerface.py new file mode 100644 index 0000000000000000000000000000000000000000..b451ebe25b9c1999ae53e3df40e55e009d64942c --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/datasets/widerface.py @@ -0,0 +1,196 @@ +import os +from os.path import abspath, expanduser +from pathlib import Path + +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import torch +from PIL import Image + +from .utils import download_and_extract_archive, download_file_from_google_drive, extract_archive, verify_str_arg +from .vision import VisionDataset + + +class WIDERFace(VisionDataset): + """`WIDERFace `_ Dataset. + + Args: + root (str or ``pathlib.Path``): Root directory where images and annotations are downloaded to. + Expects the following folder structure if download=False: + + .. code:: + + + └── widerface + ├── wider_face_split ('wider_face_split.zip' if compressed) + ├── WIDER_train ('WIDER_train.zip' if compressed) + ├── WIDER_val ('WIDER_val.zip' if compressed) + └── WIDER_test ('WIDER_test.zip' if compressed) + split (string): The dataset split to use. One of {``train``, ``val``, ``test``}. + Defaults to ``train``. + transform (callable, optional): A function/transform that takes in a PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + + .. warning:: + + To download the dataset `gdown `_ is required. + + """ + + BASE_FOLDER = "widerface" + FILE_LIST = [ + # File ID MD5 Hash Filename + ("15hGDLhsx8bLgLcIRD5DhYt5iBxnjNF1M", "3fedf70df600953d25982bcd13d91ba2", "WIDER_train.zip"), + ("1GUCogbp16PMGa39thoMMeWxp7Rp5oM8Q", "dfa7d7e790efa35df3788964cf0bbaea", "WIDER_val.zip"), + ("1HIfDbVEWKmsYKJZm4lchTBDLW5N7dY5T", "e5d8f4248ed24c334bbd12f49c29dd40", "WIDER_test.zip"), + ] + ANNOTATIONS_FILE = ( + "http://shuoyang1213.me/WIDERFACE/support/bbx_annotation/wider_face_split.zip", + "0e3767bcf0e326556d407bf5bff5d27c", + "wider_face_split.zip", + ) + + def __init__( + self, + root: Union[str, Path], + split: str = "train", + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + super().__init__( + root=os.path.join(root, self.BASE_FOLDER), transform=transform, target_transform=target_transform + ) + # check arguments + self.split = verify_str_arg(split, "split", ("train", "val", "test")) + + if download: + self.download() + + if not self._check_integrity(): + raise RuntimeError("Dataset not found or corrupted. You can use download=True to download and prepare it") + + self.img_info: List[Dict[str, Union[str, Dict[str, torch.Tensor]]]] = [] + if self.split in ("train", "val"): + self.parse_train_val_annotations_file() + else: + self.parse_test_annotations_file() + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: (image, target) where target is a dict of annotations for all faces in the image. + target=None for the test split. + """ + + # stay consistent with other datasets and return a PIL Image + img = Image.open(self.img_info[index]["img_path"]) # type: ignore[arg-type] + + if self.transform is not None: + img = self.transform(img) + + target = None if self.split == "test" else self.img_info[index]["annotations"] + if self.target_transform is not None: + target = self.target_transform(target) + + return img, target + + def __len__(self) -> int: + return len(self.img_info) + + def extra_repr(self) -> str: + lines = ["Split: {split}"] + return "\n".join(lines).format(**self.__dict__) + + def parse_train_val_annotations_file(self) -> None: + filename = "wider_face_train_bbx_gt.txt" if self.split == "train" else "wider_face_val_bbx_gt.txt" + filepath = os.path.join(self.root, "wider_face_split", filename) + + with open(filepath) as f: + lines = f.readlines() + file_name_line, num_boxes_line, box_annotation_line = True, False, False + num_boxes, box_counter = 0, 0 + labels = [] + for line in lines: + line = line.rstrip() + if file_name_line: + img_path = os.path.join(self.root, "WIDER_" + self.split, "images", line) + img_path = abspath(expanduser(img_path)) + file_name_line = False + num_boxes_line = True + elif num_boxes_line: + num_boxes = int(line) + num_boxes_line = False + box_annotation_line = True + elif box_annotation_line: + box_counter += 1 + line_split = line.split(" ") + line_values = [int(x) for x in line_split] + labels.append(line_values) + if box_counter >= num_boxes: + box_annotation_line = False + file_name_line = True + labels_tensor = torch.tensor(labels) + self.img_info.append( + { + "img_path": img_path, + "annotations": { + "bbox": labels_tensor[:, 0:4].clone(), # x, y, width, height + "blur": labels_tensor[:, 4].clone(), + "expression": labels_tensor[:, 5].clone(), + "illumination": labels_tensor[:, 6].clone(), + "occlusion": labels_tensor[:, 7].clone(), + "pose": labels_tensor[:, 8].clone(), + "invalid": labels_tensor[:, 9].clone(), + }, + } + ) + box_counter = 0 + labels.clear() + else: + raise RuntimeError(f"Error parsing annotation file {filepath}") + + def parse_test_annotations_file(self) -> None: + filepath = os.path.join(self.root, "wider_face_split", "wider_face_test_filelist.txt") + filepath = abspath(expanduser(filepath)) + with open(filepath) as f: + lines = f.readlines() + for line in lines: + line = line.rstrip() + img_path = os.path.join(self.root, "WIDER_test", "images", line) + img_path = abspath(expanduser(img_path)) + self.img_info.append({"img_path": img_path}) + + def _check_integrity(self) -> bool: + # Allow original archive to be deleted (zip). Only need the extracted images + all_files = self.FILE_LIST.copy() + all_files.append(self.ANNOTATIONS_FILE) + for (_, md5, filename) in all_files: + file, ext = os.path.splitext(filename) + extracted_dir = os.path.join(self.root, file) + if not os.path.exists(extracted_dir): + return False + return True + + def download(self) -> None: + if self._check_integrity(): + return + + # download and extract image data + for (file_id, md5, filename) in self.FILE_LIST: + download_file_from_google_drive(file_id, self.root, filename, md5) + filepath = os.path.join(self.root, filename) + extract_archive(filepath) + + # download and extract annotation files + download_and_extract_archive( + url=self.ANNOTATIONS_FILE[0], download_root=self.root, md5=self.ANNOTATIONS_FILE[1] + ) diff --git a/lib/python3.12/site-packages/torchvision/extension.py b/lib/python3.12/site-packages/torchvision/extension.py new file mode 100644 index 0000000000000000000000000000000000000000..67801056e88b44d40bc2d382d62c389bf4ef039e --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/extension.py @@ -0,0 +1,92 @@ +import os +import sys + +import torch + +from ._internally_replaced_utils import _get_extension_path + + +_HAS_OPS = False + + +def _has_ops(): + return False + + +try: + # On Windows Python-3.8.x has `os.add_dll_directory` call, + # which is called to configure dll search path. + # To find cuda related dlls we need to make sure the + # conda environment/bin path is configured Please take a look: + # https://stackoverflow.com/questions/59330863/cant-import-dll-module-in-python + # Please note: if some path can't be added using add_dll_directory we simply ignore this path + if os.name == "nt" and sys.version_info < (3, 9): + env_path = os.environ["PATH"] + path_arr = env_path.split(";") + for path in path_arr: + if os.path.exists(path): + try: + os.add_dll_directory(path) # type: ignore[attr-defined] + except Exception: + pass + + lib_path = _get_extension_path("_C") + torch.ops.load_library(lib_path) + _HAS_OPS = True + + def _has_ops(): # noqa: F811 + return True + +except (ImportError, OSError): + pass + + +def _assert_has_ops(): + if not _has_ops(): + raise RuntimeError( + "Couldn't load custom C++ ops. This can happen if your PyTorch and " + "torchvision versions are incompatible, or if you had errors while compiling " + "torchvision from source. For further information on the compatible versions, check " + "https://github.com/pytorch/vision#installation for the compatibility matrix. " + "Please check your PyTorch version with torch.__version__ and your torchvision " + "version with torchvision.__version__ and verify if they are compatible, and if not " + "please reinstall torchvision so that it matches your PyTorch install." + ) + + +def _check_cuda_version(): + """ + Make sure that CUDA versions match between the pytorch install and torchvision install + """ + if not _HAS_OPS: + return -1 + from torch.version import cuda as torch_version_cuda + + _version = torch.ops.torchvision._cuda_version() + if _version != -1 and torch_version_cuda is not None: + tv_version = str(_version) + if int(tv_version) < 10000: + tv_major = int(tv_version[0]) + tv_minor = int(tv_version[2]) + else: + tv_major = int(tv_version[0:2]) + tv_minor = int(tv_version[3]) + t_version = torch_version_cuda.split(".") + t_major = int(t_version[0]) + t_minor = int(t_version[1]) + if t_major != tv_major: + raise RuntimeError( + "Detected that PyTorch and torchvision were compiled with different CUDA major versions. " + f"PyTorch has CUDA Version={t_major}.{t_minor} and torchvision has " + f"CUDA Version={tv_major}.{tv_minor}. " + "Please reinstall the torchvision that matches your PyTorch install." + ) + return _version + + +def _load_library(lib_name): + lib_path = _get_extension_path(lib_name) + torch.ops.load_library(lib_path) + + +_check_cuda_version() diff --git a/lib/python3.12/site-packages/torchvision/models/__init__.py b/lib/python3.12/site-packages/torchvision/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6ea0a1f7178b6ca03776d58c17411a8ff483f8b2 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/__init__.py @@ -0,0 +1,23 @@ +from .alexnet import * +from .convnext import * +from .densenet import * +from .efficientnet import * +from .googlenet import * +from .inception import * +from .mnasnet import * +from .mobilenet import * +from .regnet import * +from .resnet import * +from .shufflenetv2 import * +from .squeezenet import * +from .vgg import * +from .vision_transformer import * +from .swin_transformer import * +from .maxvit import * +from . import detection, optical_flow, quantization, segmentation, video + +# The Weights and WeightsEnum are developer-facing utils that we make public for +# downstream libs like torchgeo https://github.com/pytorch/vision/issues/7094 +# TODO: we could / should document them publicly, but it's not clear where, as +# they're not intended for end users. +from ._api import get_model, get_model_builder, get_model_weights, get_weight, list_models, Weights, WeightsEnum diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4b719ca1896223ff2140181a9f2b006788591c4e Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/_api.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/_api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..87bf0368eadb08a4b926f43eb71b13bffd87356d Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/_api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/_meta.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/_meta.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c53049f803fbb7e5b603f0cec8b85c7d98f13cc5 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/_meta.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b5d29a8cf03906c9435e5c1b530154c2e0109c7d Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/alexnet.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/alexnet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2b4424b54837a39a14d484f5a074e1b7a97e276e Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/alexnet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/convnext.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/convnext.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5e18a087b9851168756192bb3a5b0362524e6a3b Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/convnext.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/densenet.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/densenet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0d7a00cad4f4c7f49615785ffa58493f9dd56532 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/densenet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/efficientnet.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/efficientnet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..56810a8e7c2402039af55efa241fd13b72795d48 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/efficientnet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/feature_extraction.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/feature_extraction.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e71bf6f21364ac495d1582c17cdc39c0d153a5d5 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/feature_extraction.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/googlenet.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/googlenet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9e601caf496921f04d93d36b4684d641aa15a7c1 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/googlenet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/inception.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/inception.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bcb332cb0e62ea601639f7406b7a212c75a57430 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/inception.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/maxvit.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/maxvit.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..088bdc8d5b095a753807fa4ec19b5744da7a0287 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/maxvit.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/mnasnet.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/mnasnet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4f37675ec237f4c448432da63780887e5cddb0ae Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/mnasnet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/mobilenet.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/mobilenet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..259e47bafca1a8bf829244d2a92a3af5061f2550 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/mobilenet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/mobilenetv2.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/mobilenetv2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6536610219a32a560a3f01ba1f6cf770ee86f747 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/mobilenetv2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/mobilenetv3.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/mobilenetv3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cef22ccc1392ded32508839d7ffde3aa60fedbe3 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/mobilenetv3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/regnet.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/regnet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5920160c452b0c5bfcaad74ad88e9c08148b4686 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/regnet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/resnet.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/resnet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..32af30407510a5f28c03e0da959a16f968584ea0 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/resnet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/shufflenetv2.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/shufflenetv2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e871dba222faf37ebe9a2aebacfad75535f11147 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/shufflenetv2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/squeezenet.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/squeezenet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..80bda17bf890e7402f12bb52762c0eb1ca68e266 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/squeezenet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/swin_transformer.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/swin_transformer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d7f97e48eab1524562c3ee99ab85415e9cb8de7 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/swin_transformer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/vgg.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/vgg.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1f9e3b9d37b04fa73b2fd6a748d457bf86b0cb0f Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/vgg.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/__pycache__/vision_transformer.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/__pycache__/vision_transformer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..93e16c4ba4c5f531025793262c28cd6ad0c8e695 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/__pycache__/vision_transformer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/_api.py b/lib/python3.12/site-packages/torchvision/models/_api.py new file mode 100644 index 0000000000000000000000000000000000000000..0999bf7ba6beba91acc0b3374c2307cc8137847c --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/_api.py @@ -0,0 +1,277 @@ +import fnmatch +import importlib +import inspect +import sys +from dataclasses import dataclass +from enum import Enum +from functools import partial +from inspect import signature +from types import ModuleType +from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Set, Type, TypeVar, Union + +from torch import nn + +from .._internally_replaced_utils import load_state_dict_from_url + + +__all__ = ["WeightsEnum", "Weights", "get_model", "get_model_builder", "get_model_weights", "get_weight", "list_models"] + + +@dataclass +class Weights: + """ + This class is used to group important attributes associated with the pre-trained weights. + + Args: + url (str): The location where we find the weights. + transforms (Callable): A callable that constructs the preprocessing method (or validation preset transforms) + needed to use the model. The reason we attach a constructor method rather than an already constructed + object is because the specific object might have memory and thus we want to delay initialization until + needed. + meta (Dict[str, Any]): Stores meta-data related to the weights of the model and its configuration. These can be + informative attributes (for example the number of parameters/flops, recipe link/methods used in training + etc), configuration parameters (for example the `num_classes`) needed to construct the model or important + meta-data (for example the `classes` of a classification model) needed to use the model. + """ + + url: str + transforms: Callable + meta: Dict[str, Any] + + def __eq__(self, other: Any) -> bool: + # We need this custom implementation for correct deep-copy and deserialization behavior. + # TL;DR: After the definition of an enum, creating a new instance, i.e. by deep-copying or deserializing it, + # involves an equality check against the defined members. Unfortunately, the `transforms` attribute is often + # defined with `functools.partial` and `fn = partial(...); assert deepcopy(fn) != fn`. Without custom handling + # for it, the check against the defined members would fail and effectively prevent the weights from being + # deep-copied or deserialized. + # See https://github.com/pytorch/vision/pull/7107 for details. + if not isinstance(other, Weights): + return NotImplemented + + if self.url != other.url: + return False + + if self.meta != other.meta: + return False + + if isinstance(self.transforms, partial) and isinstance(other.transforms, partial): + return ( + self.transforms.func == other.transforms.func + and self.transforms.args == other.transforms.args + and self.transforms.keywords == other.transforms.keywords + ) + else: + return self.transforms == other.transforms + + +class WeightsEnum(Enum): + """ + This class is the parent class of all model weights. Each model building method receives an optional `weights` + parameter with its associated pre-trained weights. It inherits from `Enum` and its values should be of type + `Weights`. + + Args: + value (Weights): The data class entry with the weight information. + """ + + @classmethod + def verify(cls, obj: Any) -> Any: + if obj is not None: + if type(obj) is str: + obj = cls[obj.replace(cls.__name__ + ".", "")] + elif not isinstance(obj, cls): + raise TypeError( + f"Invalid Weight class provided; expected {cls.__name__} but received {obj.__class__.__name__}." + ) + return obj + + def get_state_dict(self, *args: Any, **kwargs: Any) -> Mapping[str, Any]: + return load_state_dict_from_url(self.url, *args, **kwargs) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}.{self._name_}" + + @property + def url(self): + return self.value.url + + @property + def transforms(self): + return self.value.transforms + + @property + def meta(self): + return self.value.meta + + +def get_weight(name: str) -> WeightsEnum: + """ + Gets the weights enum value by its full name. Example: "ResNet50_Weights.IMAGENET1K_V1" + + Args: + name (str): The name of the weight enum entry. + + Returns: + WeightsEnum: The requested weight enum. + """ + try: + enum_name, value_name = name.split(".") + except ValueError: + raise ValueError(f"Invalid weight name provided: '{name}'.") + + base_module_name = ".".join(sys.modules[__name__].__name__.split(".")[:-1]) + base_module = importlib.import_module(base_module_name) + model_modules = [base_module] + [ + x[1] + for x in inspect.getmembers(base_module, inspect.ismodule) + if x[1].__file__.endswith("__init__.py") # type: ignore[union-attr] + ] + + weights_enum = None + for m in model_modules: + potential_class = m.__dict__.get(enum_name, None) + if potential_class is not None and issubclass(potential_class, WeightsEnum): + weights_enum = potential_class + break + + if weights_enum is None: + raise ValueError(f"The weight enum '{enum_name}' for the specific method couldn't be retrieved.") + + return weights_enum[value_name] + + +def get_model_weights(name: Union[Callable, str]) -> Type[WeightsEnum]: + """ + Returns the weights enum class associated to the given model. + + Args: + name (callable or str): The model builder function or the name under which it is registered. + + Returns: + weights_enum (WeightsEnum): The weights enum class associated with the model. + """ + model = get_model_builder(name) if isinstance(name, str) else name + return _get_enum_from_fn(model) + + +def _get_enum_from_fn(fn: Callable) -> Type[WeightsEnum]: + """ + Internal method that gets the weight enum of a specific model builder method. + + Args: + fn (Callable): The builder method used to create the model. + Returns: + WeightsEnum: The requested weight enum. + """ + sig = signature(fn) + if "weights" not in sig.parameters: + raise ValueError("The method is missing the 'weights' argument.") + + ann = signature(fn).parameters["weights"].annotation + weights_enum = None + if isinstance(ann, type) and issubclass(ann, WeightsEnum): + weights_enum = ann + else: + # handle cases like Union[Optional, T] + # TODO: Replace ann.__args__ with typing.get_args(ann) after python >= 3.8 + for t in ann.__args__: # type: ignore[union-attr] + if isinstance(t, type) and issubclass(t, WeightsEnum): + weights_enum = t + break + + if weights_enum is None: + raise ValueError( + "The WeightsEnum class for the specific method couldn't be retrieved. Make sure the typing info is correct." + ) + + return weights_enum + + +M = TypeVar("M", bound=nn.Module) + +BUILTIN_MODELS = {} + + +def register_model(name: Optional[str] = None) -> Callable[[Callable[..., M]], Callable[..., M]]: + def wrapper(fn: Callable[..., M]) -> Callable[..., M]: + key = name if name is not None else fn.__name__ + if key in BUILTIN_MODELS: + raise ValueError(f"An entry is already registered under the name '{key}'.") + BUILTIN_MODELS[key] = fn + return fn + + return wrapper + + +def list_models( + module: Optional[ModuleType] = None, + include: Union[Iterable[str], str, None] = None, + exclude: Union[Iterable[str], str, None] = None, +) -> List[str]: + """ + Returns a list with the names of registered models. + + Args: + module (ModuleType, optional): The module from which we want to extract the available models. + include (str or Iterable[str], optional): Filter(s) for including the models from the set of all models. + Filters are passed to `fnmatch `__ to match Unix shell-style + wildcards. In case of many filters, the results is the union of individual filters. + exclude (str or Iterable[str], optional): Filter(s) applied after include_filters to remove models. + Filter are passed to `fnmatch `__ to match Unix shell-style + wildcards. In case of many filters, the results is removal of all the models that match any individual filter. + + Returns: + models (list): A list with the names of available models. + """ + all_models = { + k for k, v in BUILTIN_MODELS.items() if module is None or v.__module__.rsplit(".", 1)[0] == module.__name__ + } + if include: + models: Set[str] = set() + if isinstance(include, str): + include = [include] + for include_filter in include: + models = models | set(fnmatch.filter(all_models, include_filter)) + else: + models = all_models + + if exclude: + if isinstance(exclude, str): + exclude = [exclude] + for exclude_filter in exclude: + models = models - set(fnmatch.filter(all_models, exclude_filter)) + return sorted(models) + + +def get_model_builder(name: str) -> Callable[..., nn.Module]: + """ + Gets the model name and returns the model builder method. + + Args: + name (str): The name under which the model is registered. + + Returns: + fn (Callable): The model builder method. + """ + name = name.lower() + try: + fn = BUILTIN_MODELS[name] + except KeyError: + raise ValueError(f"Unknown model {name}") + return fn + + +def get_model(name: str, **config: Any) -> nn.Module: + """ + Gets the model name and configuration and returns an instantiated model. + + Args: + name (str): The name under which the model is registered. + **config (Any): parameters passed to the model builder method. + + Returns: + model (nn.Module): The initialized model. + """ + fn = get_model_builder(name) + return fn(**config) diff --git a/lib/python3.12/site-packages/torchvision/models/_meta.py b/lib/python3.12/site-packages/torchvision/models/_meta.py new file mode 100644 index 0000000000000000000000000000000000000000..e66f411c287e0f456448315ba4fd0bfcce281d2b --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/_meta.py @@ -0,0 +1,1554 @@ +""" +This file is part of the private API. Please do not refer to any variables defined here directly as they will be +removed on future versions without warning. +""" + +# This will eventually be replaced with a call at torchvision.datasets.info("imagenet").categories +_IMAGENET_CATEGORIES = [ + "tench", + "goldfish", + "great white shark", + "tiger shark", + "hammerhead", + "electric ray", + "stingray", + "cock", + "hen", + "ostrich", + "brambling", + "goldfinch", + "house finch", + "junco", + "indigo bunting", + "robin", + "bulbul", + "jay", + "magpie", + "chickadee", + "water ouzel", + "kite", + "bald eagle", + "vulture", + "great grey owl", + "European fire salamander", + "common newt", + "eft", + "spotted salamander", + "axolotl", + "bullfrog", + "tree frog", + "tailed frog", + "loggerhead", + "leatherback turtle", + "mud turtle", + "terrapin", + "box turtle", + "banded gecko", + "common iguana", + "American chameleon", + "whiptail", + "agama", + "frilled lizard", + "alligator lizard", + "Gila monster", + "green lizard", + "African chameleon", + "Komodo dragon", + "African crocodile", + "American alligator", + "triceratops", + "thunder snake", + "ringneck snake", + "hognose snake", + "green snake", + "king snake", + "garter snake", + "water snake", + "vine snake", + "night snake", + "boa constrictor", + "rock python", + "Indian cobra", + "green mamba", + "sea snake", + "horned viper", + "diamondback", + "sidewinder", + "trilobite", + "harvestman", + "scorpion", + "black and gold garden spider", + "barn spider", + "garden spider", + "black widow", + "tarantula", + "wolf spider", + "tick", + "centipede", + "black grouse", + "ptarmigan", + "ruffed grouse", + "prairie chicken", + "peacock", + "quail", + "partridge", + "African grey", + "macaw", + "sulphur-crested cockatoo", + "lorikeet", + "coucal", + "bee eater", + "hornbill", + "hummingbird", + "jacamar", + "toucan", + "drake", + "red-breasted merganser", + "goose", + "black swan", + "tusker", + "echidna", + "platypus", + "wallaby", + "koala", + "wombat", + "jellyfish", + "sea anemone", + "brain coral", + "flatworm", + "nematode", + "conch", + "snail", + "slug", + "sea slug", + "chiton", + "chambered nautilus", + "Dungeness crab", + "rock crab", + "fiddler crab", + "king crab", + "American lobster", + "spiny lobster", + "crayfish", + "hermit crab", + "isopod", + "white stork", + "black stork", + "spoonbill", + "flamingo", + "little blue heron", + "American egret", + "bittern", + "crane bird", + "limpkin", + "European gallinule", + "American coot", + "bustard", + "ruddy turnstone", + "red-backed sandpiper", + "redshank", + "dowitcher", + "oystercatcher", + "pelican", + "king penguin", + "albatross", + "grey whale", + "killer whale", + "dugong", + "sea lion", + "Chihuahua", + "Japanese spaniel", + "Maltese dog", + "Pekinese", + "Shih-Tzu", + "Blenheim spaniel", + "papillon", + "toy terrier", + "Rhodesian ridgeback", + "Afghan hound", + "basset", + "beagle", + "bloodhound", + "bluetick", + "black-and-tan coonhound", + "Walker hound", + "English foxhound", + "redbone", + "borzoi", + "Irish wolfhound", + "Italian greyhound", + "whippet", + "Ibizan hound", + "Norwegian elkhound", + "otterhound", + "Saluki", + "Scottish deerhound", + "Weimaraner", + "Staffordshire bullterrier", + "American Staffordshire terrier", + "Bedlington terrier", + "Border terrier", + "Kerry blue terrier", + "Irish terrier", + "Norfolk terrier", + "Norwich terrier", + "Yorkshire terrier", + "wire-haired fox terrier", + "Lakeland terrier", + "Sealyham terrier", + "Airedale", + "cairn", + "Australian terrier", + "Dandie Dinmont", + "Boston bull", + "miniature schnauzer", + "giant schnauzer", + "standard schnauzer", + "Scotch terrier", + "Tibetan terrier", + "silky terrier", + "soft-coated wheaten terrier", + "West Highland white terrier", + "Lhasa", + "flat-coated retriever", + "curly-coated retriever", + "golden retriever", + "Labrador retriever", + "Chesapeake Bay retriever", + "German short-haired pointer", + "vizsla", + "English setter", + "Irish setter", + "Gordon setter", + "Brittany spaniel", + "clumber", + "English springer", + "Welsh springer spaniel", + "cocker spaniel", + "Sussex spaniel", + "Irish water spaniel", + "kuvasz", + "schipperke", + "groenendael", + "malinois", + "briard", + "kelpie", + "komondor", + "Old English sheepdog", + "Shetland sheepdog", + "collie", + "Border collie", + "Bouvier des Flandres", + "Rottweiler", + "German shepherd", + "Doberman", + "miniature pinscher", + "Greater Swiss Mountain dog", + "Bernese mountain dog", + "Appenzeller", + "EntleBucher", + "boxer", + "bull mastiff", + "Tibetan mastiff", + "French bulldog", + "Great Dane", + "Saint Bernard", + "Eskimo dog", + "malamute", + "Siberian husky", + "dalmatian", + "affenpinscher", + "basenji", + "pug", + "Leonberg", + "Newfoundland", + "Great Pyrenees", + "Samoyed", + "Pomeranian", + "chow", + "keeshond", + "Brabancon griffon", + "Pembroke", + "Cardigan", + "toy poodle", + "miniature poodle", + "standard poodle", + "Mexican hairless", + "timber wolf", + "white wolf", + "red wolf", + "coyote", + "dingo", + "dhole", + "African hunting dog", + "hyena", + "red fox", + "kit fox", + "Arctic fox", + "grey fox", + "tabby", + "tiger cat", + "Persian cat", + "Siamese cat", + "Egyptian cat", + "cougar", + "lynx", + "leopard", + "snow leopard", + "jaguar", + "lion", + "tiger", + "cheetah", + "brown bear", + "American black bear", + "ice bear", + "sloth bear", + "mongoose", + "meerkat", + "tiger beetle", + "ladybug", + "ground beetle", + "long-horned beetle", + "leaf beetle", + "dung beetle", + "rhinoceros beetle", + "weevil", + "fly", + "bee", + "ant", + "grasshopper", + "cricket", + "walking stick", + "cockroach", + "mantis", + "cicada", + "leafhopper", + "lacewing", + "dragonfly", + "damselfly", + "admiral", + "ringlet", + "monarch", + "cabbage butterfly", + "sulphur butterfly", + "lycaenid", + "starfish", + "sea urchin", + "sea cucumber", + "wood rabbit", + "hare", + "Angora", + "hamster", + "porcupine", + "fox squirrel", + "marmot", + "beaver", + "guinea pig", + "sorrel", + "zebra", + "hog", + "wild boar", + "warthog", + "hippopotamus", + "ox", + "water buffalo", + "bison", + "ram", + "bighorn", + "ibex", + "hartebeest", + "impala", + "gazelle", + "Arabian camel", + "llama", + "weasel", + "mink", + "polecat", + "black-footed ferret", + "otter", + "skunk", + "badger", + "armadillo", + "three-toed sloth", + "orangutan", + "gorilla", + "chimpanzee", + "gibbon", + "siamang", + "guenon", + "patas", + "baboon", + "macaque", + "langur", + "colobus", + "proboscis monkey", + "marmoset", + "capuchin", + "howler monkey", + "titi", + "spider monkey", + "squirrel monkey", + "Madagascar cat", + "indri", + "Indian elephant", + "African elephant", + "lesser panda", + "giant panda", + "barracouta", + "eel", + "coho", + "rock beauty", + "anemone fish", + "sturgeon", + "gar", + "lionfish", + "puffer", + "abacus", + "abaya", + "academic gown", + "accordion", + "acoustic guitar", + "aircraft carrier", + "airliner", + "airship", + "altar", + "ambulance", + "amphibian", + "analog clock", + "apiary", + "apron", + "ashcan", + "assault rifle", + "backpack", + "bakery", + "balance beam", + "balloon", + "ballpoint", + "Band Aid", + "banjo", + "bannister", + "barbell", + "barber chair", + "barbershop", + "barn", + "barometer", + "barrel", + "barrow", + "baseball", + "basketball", + "bassinet", + "bassoon", + "bathing cap", + "bath towel", + "bathtub", + "beach wagon", + "beacon", + "beaker", + "bearskin", + "beer bottle", + "beer glass", + "bell cote", + "bib", + "bicycle-built-for-two", + "bikini", + "binder", + "binoculars", + "birdhouse", + "boathouse", + "bobsled", + "bolo tie", + "bonnet", + "bookcase", + "bookshop", + "bottlecap", + "bow", + "bow tie", + "brass", + "brassiere", + "breakwater", + "breastplate", + "broom", + "bucket", + "buckle", + "bulletproof vest", + "bullet train", + "butcher shop", + "cab", + "caldron", + "candle", + "cannon", + "canoe", + "can opener", + "cardigan", + "car mirror", + "carousel", + "carpenter's kit", + "carton", + "car wheel", + "cash machine", + "cassette", + "cassette player", + "castle", + "catamaran", + "CD player", + "cello", + "cellular telephone", + "chain", + "chainlink fence", + "chain mail", + "chain saw", + "chest", + "chiffonier", + "chime", + "china cabinet", + "Christmas stocking", + "church", + "cinema", + "cleaver", + "cliff dwelling", + "cloak", + "clog", + "cocktail shaker", + "coffee mug", + "coffeepot", + "coil", + "combination lock", + "computer keyboard", + "confectionery", + "container ship", + "convertible", + "corkscrew", + "cornet", + "cowboy boot", + "cowboy hat", + "cradle", + "crane", + "crash helmet", + "crate", + "crib", + "Crock Pot", + "croquet ball", + "crutch", + "cuirass", + "dam", + "desk", + "desktop computer", + "dial telephone", + "diaper", + "digital clock", + "digital watch", + "dining table", + "dishrag", + "dishwasher", + "disk brake", + "dock", + "dogsled", + "dome", + "doormat", + "drilling platform", + "drum", + "drumstick", + "dumbbell", + "Dutch oven", + "electric fan", + "electric guitar", + "electric locomotive", + "entertainment center", + "envelope", + "espresso maker", + "face powder", + "feather boa", + "file", + "fireboat", + "fire engine", + "fire screen", + "flagpole", + "flute", + "folding chair", + "football helmet", + "forklift", + "fountain", + "fountain pen", + "four-poster", + "freight car", + "French horn", + "frying pan", + "fur coat", + "garbage truck", + "gasmask", + "gas pump", + "goblet", + "go-kart", + "golf ball", + "golfcart", + "gondola", + "gong", + "gown", + "grand piano", + "greenhouse", + "grille", + "grocery store", + "guillotine", + "hair slide", + "hair spray", + "half track", + "hammer", + "hamper", + "hand blower", + "hand-held computer", + "handkerchief", + "hard disc", + "harmonica", + "harp", + "harvester", + "hatchet", + "holster", + "home theater", + "honeycomb", + "hook", + "hoopskirt", + "horizontal bar", + "horse cart", + "hourglass", + "iPod", + "iron", + "jack-o'-lantern", + "jean", + "jeep", + "jersey", + "jigsaw puzzle", + "jinrikisha", + "joystick", + "kimono", + "knee pad", + "knot", + "lab coat", + "ladle", + "lampshade", + "laptop", + "lawn mower", + "lens cap", + "letter opener", + "library", + "lifeboat", + "lighter", + "limousine", + "liner", + "lipstick", + "Loafer", + "lotion", + "loudspeaker", + "loupe", + "lumbermill", + "magnetic compass", + "mailbag", + "mailbox", + "maillot", + "maillot tank suit", + "manhole cover", + "maraca", + "marimba", + "mask", + "matchstick", + "maypole", + "maze", + "measuring cup", + "medicine chest", + "megalith", + "microphone", + "microwave", + "military uniform", + "milk can", + "minibus", + "miniskirt", + "minivan", + "missile", + "mitten", + "mixing bowl", + "mobile home", + "Model T", + "modem", + "monastery", + "monitor", + "moped", + "mortar", + "mortarboard", + "mosque", + "mosquito net", + "motor scooter", + "mountain bike", + "mountain tent", + "mouse", + "mousetrap", + "moving van", + "muzzle", + "nail", + "neck brace", + "necklace", + "nipple", + "notebook", + "obelisk", + "oboe", + "ocarina", + "odometer", + "oil filter", + "organ", + "oscilloscope", + "overskirt", + "oxcart", + "oxygen mask", + "packet", + "paddle", + "paddlewheel", + "padlock", + "paintbrush", + "pajama", + "palace", + "panpipe", + "paper towel", + "parachute", + "parallel bars", + "park bench", + "parking meter", + "passenger car", + "patio", + "pay-phone", + "pedestal", + "pencil box", + "pencil sharpener", + "perfume", + "Petri dish", + "photocopier", + "pick", + "pickelhaube", + "picket fence", + "pickup", + "pier", + "piggy bank", + "pill bottle", + "pillow", + "ping-pong ball", + "pinwheel", + "pirate", + "pitcher", + "plane", + "planetarium", + "plastic bag", + "plate rack", + "plow", + "plunger", + "Polaroid camera", + "pole", + "police van", + "poncho", + "pool table", + "pop bottle", + "pot", + "potter's wheel", + "power drill", + "prayer rug", + "printer", + "prison", + "projectile", + "projector", + "puck", + "punching bag", + "purse", + "quill", + "quilt", + "racer", + "racket", + "radiator", + "radio", + "radio telescope", + "rain barrel", + "recreational vehicle", + "reel", + "reflex camera", + "refrigerator", + "remote control", + "restaurant", + "revolver", + "rifle", + "rocking chair", + "rotisserie", + "rubber eraser", + "rugby ball", + "rule", + "running shoe", + "safe", + "safety pin", + "saltshaker", + "sandal", + "sarong", + "sax", + "scabbard", + "scale", + "school bus", + "schooner", + "scoreboard", + "screen", + "screw", + "screwdriver", + "seat belt", + "sewing machine", + "shield", + "shoe shop", + "shoji", + "shopping basket", + "shopping cart", + "shovel", + "shower cap", + "shower curtain", + "ski", + "ski mask", + "sleeping bag", + "slide rule", + "sliding door", + "slot", + "snorkel", + "snowmobile", + "snowplow", + "soap dispenser", + "soccer ball", + "sock", + "solar dish", + "sombrero", + "soup bowl", + "space bar", + "space heater", + "space shuttle", + "spatula", + "speedboat", + "spider web", + "spindle", + "sports car", + "spotlight", + "stage", + "steam locomotive", + "steel arch bridge", + "steel drum", + "stethoscope", + "stole", + "stone wall", + "stopwatch", + "stove", + "strainer", + "streetcar", + "stretcher", + "studio couch", + "stupa", + "submarine", + "suit", + "sundial", + "sunglass", + "sunglasses", + "sunscreen", + "suspension bridge", + "swab", + "sweatshirt", + "swimming trunks", + "swing", + "switch", + "syringe", + "table lamp", + "tank", + "tape player", + "teapot", + "teddy", + "television", + "tennis ball", + "thatch", + "theater curtain", + "thimble", + "thresher", + "throne", + "tile roof", + "toaster", + "tobacco shop", + "toilet seat", + "torch", + "totem pole", + "tow truck", + "toyshop", + "tractor", + "trailer truck", + "tray", + "trench coat", + "tricycle", + "trimaran", + "tripod", + "triumphal arch", + "trolleybus", + "trombone", + "tub", + "turnstile", + "typewriter keyboard", + "umbrella", + "unicycle", + "upright", + "vacuum", + "vase", + "vault", + "velvet", + "vending machine", + "vestment", + "viaduct", + "violin", + "volleyball", + "waffle iron", + "wall clock", + "wallet", + "wardrobe", + "warplane", + "washbasin", + "washer", + "water bottle", + "water jug", + "water tower", + "whiskey jug", + "whistle", + "wig", + "window screen", + "window shade", + "Windsor tie", + "wine bottle", + "wing", + "wok", + "wooden spoon", + "wool", + "worm fence", + "wreck", + "yawl", + "yurt", + "web site", + "comic book", + "crossword puzzle", + "street sign", + "traffic light", + "book jacket", + "menu", + "plate", + "guacamole", + "consomme", + "hot pot", + "trifle", + "ice cream", + "ice lolly", + "French loaf", + "bagel", + "pretzel", + "cheeseburger", + "hotdog", + "mashed potato", + "head cabbage", + "broccoli", + "cauliflower", + "zucchini", + "spaghetti squash", + "acorn squash", + "butternut squash", + "cucumber", + "artichoke", + "bell pepper", + "cardoon", + "mushroom", + "Granny Smith", + "strawberry", + "orange", + "lemon", + "fig", + "pineapple", + "banana", + "jackfruit", + "custard apple", + "pomegranate", + "hay", + "carbonara", + "chocolate sauce", + "dough", + "meat loaf", + "pizza", + "potpie", + "burrito", + "red wine", + "espresso", + "cup", + "eggnog", + "alp", + "bubble", + "cliff", + "coral reef", + "geyser", + "lakeside", + "promontory", + "sandbar", + "seashore", + "valley", + "volcano", + "ballplayer", + "groom", + "scuba diver", + "rapeseed", + "daisy", + "yellow lady's slipper", + "corn", + "acorn", + "hip", + "buckeye", + "coral fungus", + "agaric", + "gyromitra", + "stinkhorn", + "earthstar", + "hen-of-the-woods", + "bolete", + "ear", + "toilet tissue", +] + +# To be replaced with torchvision.datasets.info("coco").categories +_COCO_CATEGORIES = [ + "__background__", + "person", + "bicycle", + "car", + "motorcycle", + "airplane", + "bus", + "train", + "truck", + "boat", + "traffic light", + "fire hydrant", + "N/A", + "stop sign", + "parking meter", + "bench", + "bird", + "cat", + "dog", + "horse", + "sheep", + "cow", + "elephant", + "bear", + "zebra", + "giraffe", + "N/A", + "backpack", + "umbrella", + "N/A", + "N/A", + "handbag", + "tie", + "suitcase", + "frisbee", + "skis", + "snowboard", + "sports ball", + "kite", + "baseball bat", + "baseball glove", + "skateboard", + "surfboard", + "tennis racket", + "bottle", + "N/A", + "wine glass", + "cup", + "fork", + "knife", + "spoon", + "bowl", + "banana", + "apple", + "sandwich", + "orange", + "broccoli", + "carrot", + "hot dog", + "pizza", + "donut", + "cake", + "chair", + "couch", + "potted plant", + "bed", + "N/A", + "dining table", + "N/A", + "N/A", + "toilet", + "N/A", + "tv", + "laptop", + "mouse", + "remote", + "keyboard", + "cell phone", + "microwave", + "oven", + "toaster", + "sink", + "refrigerator", + "N/A", + "book", + "clock", + "vase", + "scissors", + "teddy bear", + "hair drier", + "toothbrush", +] + +# To be replaced with torchvision.datasets.info("coco_kp") +_COCO_PERSON_CATEGORIES = ["no person", "person"] +_COCO_PERSON_KEYPOINT_NAMES = [ + "nose", + "left_eye", + "right_eye", + "left_ear", + "right_ear", + "left_shoulder", + "right_shoulder", + "left_elbow", + "right_elbow", + "left_wrist", + "right_wrist", + "left_hip", + "right_hip", + "left_knee", + "right_knee", + "left_ankle", + "right_ankle", +] + +# To be replaced with torchvision.datasets.info("voc").categories +_VOC_CATEGORIES = [ + "__background__", + "aeroplane", + "bicycle", + "bird", + "boat", + "bottle", + "bus", + "car", + "cat", + "chair", + "cow", + "diningtable", + "dog", + "horse", + "motorbike", + "person", + "pottedplant", + "sheep", + "sofa", + "train", + "tvmonitor", +] + +# To be replaced with torchvision.datasets.info("kinetics400").categories +_KINETICS400_CATEGORIES = [ + "abseiling", + "air drumming", + "answering questions", + "applauding", + "applying cream", + "archery", + "arm wrestling", + "arranging flowers", + "assembling computer", + "auctioning", + "baby waking up", + "baking cookies", + "balloon blowing", + "bandaging", + "barbequing", + "bartending", + "beatboxing", + "bee keeping", + "belly dancing", + "bench pressing", + "bending back", + "bending metal", + "biking through snow", + "blasting sand", + "blowing glass", + "blowing leaves", + "blowing nose", + "blowing out candles", + "bobsledding", + "bookbinding", + "bouncing on trampoline", + "bowling", + "braiding hair", + "breading or breadcrumbing", + "breakdancing", + "brush painting", + "brushing hair", + "brushing teeth", + "building cabinet", + "building shed", + "bungee jumping", + "busking", + "canoeing or kayaking", + "capoeira", + "carrying baby", + "cartwheeling", + "carving pumpkin", + "catching fish", + "catching or throwing baseball", + "catching or throwing frisbee", + "catching or throwing softball", + "celebrating", + "changing oil", + "changing wheel", + "checking tires", + "cheerleading", + "chopping wood", + "clapping", + "clay pottery making", + "clean and jerk", + "cleaning floor", + "cleaning gutters", + "cleaning pool", + "cleaning shoes", + "cleaning toilet", + "cleaning windows", + "climbing a rope", + "climbing ladder", + "climbing tree", + "contact juggling", + "cooking chicken", + "cooking egg", + "cooking on campfire", + "cooking sausages", + "counting money", + "country line dancing", + "cracking neck", + "crawling baby", + "crossing river", + "crying", + "curling hair", + "cutting nails", + "cutting pineapple", + "cutting watermelon", + "dancing ballet", + "dancing charleston", + "dancing gangnam style", + "dancing macarena", + "deadlifting", + "decorating the christmas tree", + "digging", + "dining", + "disc golfing", + "diving cliff", + "dodgeball", + "doing aerobics", + "doing laundry", + "doing nails", + "drawing", + "dribbling basketball", + "drinking", + "drinking beer", + "drinking shots", + "driving car", + "driving tractor", + "drop kicking", + "drumming fingers", + "dunking basketball", + "dying hair", + "eating burger", + "eating cake", + "eating carrots", + "eating chips", + "eating doughnuts", + "eating hotdog", + "eating ice cream", + "eating spaghetti", + "eating watermelon", + "egg hunting", + "exercising arm", + "exercising with an exercise ball", + "extinguishing fire", + "faceplanting", + "feeding birds", + "feeding fish", + "feeding goats", + "filling eyebrows", + "finger snapping", + "fixing hair", + "flipping pancake", + "flying kite", + "folding clothes", + "folding napkins", + "folding paper", + "front raises", + "frying vegetables", + "garbage collecting", + "gargling", + "getting a haircut", + "getting a tattoo", + "giving or receiving award", + "golf chipping", + "golf driving", + "golf putting", + "grinding meat", + "grooming dog", + "grooming horse", + "gymnastics tumbling", + "hammer throw", + "headbanging", + "headbutting", + "high jump", + "high kick", + "hitting baseball", + "hockey stop", + "holding snake", + "hopscotch", + "hoverboarding", + "hugging", + "hula hooping", + "hurdling", + "hurling (sport)", + "ice climbing", + "ice fishing", + "ice skating", + "ironing", + "javelin throw", + "jetskiing", + "jogging", + "juggling balls", + "juggling fire", + "juggling soccer ball", + "jumping into pool", + "jumpstyle dancing", + "kicking field goal", + "kicking soccer ball", + "kissing", + "kitesurfing", + "knitting", + "krumping", + "laughing", + "laying bricks", + "long jump", + "lunge", + "making a cake", + "making a sandwich", + "making bed", + "making jewelry", + "making pizza", + "making snowman", + "making sushi", + "making tea", + "marching", + "massaging back", + "massaging feet", + "massaging legs", + "massaging person's head", + "milking cow", + "mopping floor", + "motorcycling", + "moving furniture", + "mowing lawn", + "news anchoring", + "opening bottle", + "opening present", + "paragliding", + "parasailing", + "parkour", + "passing American football (in game)", + "passing American football (not in game)", + "peeling apples", + "peeling potatoes", + "petting animal (not cat)", + "petting cat", + "picking fruit", + "planting trees", + "plastering", + "playing accordion", + "playing badminton", + "playing bagpipes", + "playing basketball", + "playing bass guitar", + "playing cards", + "playing cello", + "playing chess", + "playing clarinet", + "playing controller", + "playing cricket", + "playing cymbals", + "playing didgeridoo", + "playing drums", + "playing flute", + "playing guitar", + "playing harmonica", + "playing harp", + "playing ice hockey", + "playing keyboard", + "playing kickball", + "playing monopoly", + "playing organ", + "playing paintball", + "playing piano", + "playing poker", + "playing recorder", + "playing saxophone", + "playing squash or racquetball", + "playing tennis", + "playing trombone", + "playing trumpet", + "playing ukulele", + "playing violin", + "playing volleyball", + "playing xylophone", + "pole vault", + "presenting weather forecast", + "pull ups", + "pumping fist", + "pumping gas", + "punching bag", + "punching person (boxing)", + "push up", + "pushing car", + "pushing cart", + "pushing wheelchair", + "reading book", + "reading newspaper", + "recording music", + "riding a bike", + "riding camel", + "riding elephant", + "riding mechanical bull", + "riding mountain bike", + "riding mule", + "riding or walking with horse", + "riding scooter", + "riding unicycle", + "ripping paper", + "robot dancing", + "rock climbing", + "rock scissors paper", + "roller skating", + "running on treadmill", + "sailing", + "salsa dancing", + "sanding floor", + "scrambling eggs", + "scuba diving", + "setting table", + "shaking hands", + "shaking head", + "sharpening knives", + "sharpening pencil", + "shaving head", + "shaving legs", + "shearing sheep", + "shining shoes", + "shooting basketball", + "shooting goal (soccer)", + "shot put", + "shoveling snow", + "shredding paper", + "shuffling cards", + "side kick", + "sign language interpreting", + "singing", + "situp", + "skateboarding", + "ski jumping", + "skiing (not slalom or crosscountry)", + "skiing crosscountry", + "skiing slalom", + "skipping rope", + "skydiving", + "slacklining", + "slapping", + "sled dog racing", + "smoking", + "smoking hookah", + "snatch weight lifting", + "sneezing", + "sniffing", + "snorkeling", + "snowboarding", + "snowkiting", + "snowmobiling", + "somersaulting", + "spinning poi", + "spray painting", + "spraying", + "springboard diving", + "squat", + "sticking tongue out", + "stomping grapes", + "stretching arm", + "stretching leg", + "strumming guitar", + "surfing crowd", + "surfing water", + "sweeping floor", + "swimming backstroke", + "swimming breast stroke", + "swimming butterfly stroke", + "swing dancing", + "swinging legs", + "swinging on something", + "sword fighting", + "tai chi", + "taking a shower", + "tango dancing", + "tap dancing", + "tapping guitar", + "tapping pen", + "tasting beer", + "tasting food", + "testifying", + "texting", + "throwing axe", + "throwing ball", + "throwing discus", + "tickling", + "tobogganing", + "tossing coin", + "tossing salad", + "training dog", + "trapezing", + "trimming or shaving beard", + "trimming trees", + "triple jump", + "tying bow tie", + "tying knot (not on a tie)", + "tying tie", + "unboxing", + "unloading truck", + "using computer", + "using remote controller (not gaming)", + "using segway", + "vault", + "waiting in line", + "walking the dog", + "washing dishes", + "washing feet", + "washing hair", + "washing hands", + "water skiing", + "water sliding", + "watering plants", + "waxing back", + "waxing chest", + "waxing eyebrows", + "waxing legs", + "weaving basket", + "welding", + "whistling", + "windsurfing", + "wrapping present", + "wrestling", + "writing", + "yawning", + "yoga", + "zumba", +] diff --git a/lib/python3.12/site-packages/torchvision/models/_utils.py b/lib/python3.12/site-packages/torchvision/models/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d59a6220b91f0f2fe85b6bf01948ee5a506cd82a --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/_utils.py @@ -0,0 +1,256 @@ +import functools +import inspect +import warnings +from collections import OrderedDict +from typing import Any, Callable, Dict, Optional, Tuple, TypeVar, Union + +from torch import nn + +from .._utils import sequence_to_str +from ._api import WeightsEnum + + +class IntermediateLayerGetter(nn.ModuleDict): + """ + Module wrapper that returns intermediate layers from a model + + It has a strong assumption that the modules have been registered + into the model in the same order as they are used. + This means that one should **not** reuse the same nn.Module + twice in the forward if you want this to work. + + Additionally, it is only able to query submodules that are directly + assigned to the model. So if `model` is passed, `model.feature1` can + be returned, but not `model.feature1.layer2`. + + Args: + model (nn.Module): model on which we will extract the features + return_layers (Dict[name, new_name]): a dict containing the names + of the modules for which the activations will be returned as + the key of the dict, and the value of the dict is the name + of the returned activation (which the user can specify). + + Examples:: + + >>> m = torchvision.models.resnet18(weights=ResNet18_Weights.DEFAULT) + >>> # extract layer1 and layer3, giving as names `feat1` and feat2` + >>> new_m = torchvision.models._utils.IntermediateLayerGetter(m, + >>> {'layer1': 'feat1', 'layer3': 'feat2'}) + >>> out = new_m(torch.rand(1, 3, 224, 224)) + >>> print([(k, v.shape) for k, v in out.items()]) + >>> [('feat1', torch.Size([1, 64, 56, 56])), + >>> ('feat2', torch.Size([1, 256, 14, 14]))] + """ + + _version = 2 + __annotations__ = { + "return_layers": Dict[str, str], + } + + def __init__(self, model: nn.Module, return_layers: Dict[str, str]) -> None: + if not set(return_layers).issubset([name for name, _ in model.named_children()]): + raise ValueError("return_layers are not present in model") + orig_return_layers = return_layers + return_layers = {str(k): str(v) for k, v in return_layers.items()} + layers = OrderedDict() + for name, module in model.named_children(): + layers[name] = module + if name in return_layers: + del return_layers[name] + if not return_layers: + break + + super().__init__(layers) + self.return_layers = orig_return_layers + + def forward(self, x): + out = OrderedDict() + for name, module in self.items(): + x = module(x) + if name in self.return_layers: + out_name = self.return_layers[name] + out[out_name] = x + return out + + +def _make_divisible(v: float, divisor: int, min_value: Optional[int] = None) -> int: + """ + This function is taken from the original tf repo. + It ensures that all layers have a channel number that is divisible by 8 + It can be seen here: + https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py + """ + if min_value is None: + min_value = divisor + new_v = max(min_value, int(v + divisor / 2) // divisor * divisor) + # Make sure that round down does not go down by more than 10%. + if new_v < 0.9 * v: + new_v += divisor + return new_v + + +D = TypeVar("D") + + +def kwonly_to_pos_or_kw(fn: Callable[..., D]) -> Callable[..., D]: + """Decorates a function that uses keyword only parameters to also allow them being passed as positionals. + + For example, consider the use case of changing the signature of ``old_fn`` into the one from ``new_fn``: + + .. code:: + + def old_fn(foo, bar, baz=None): + ... + + def new_fn(foo, *, bar, baz=None): + ... + + Calling ``old_fn("foo", "bar, "baz")`` was valid, but the same call is no longer valid with ``new_fn``. To keep BC + and at the same time warn the user of the deprecation, this decorator can be used: + + .. code:: + + @kwonly_to_pos_or_kw + def new_fn(foo, *, bar, baz=None): + ... + + new_fn("foo", "bar, "baz") + """ + params = inspect.signature(fn).parameters + + try: + keyword_only_start_idx = next( + idx for idx, param in enumerate(params.values()) if param.kind == param.KEYWORD_ONLY + ) + except StopIteration: + raise TypeError(f"Found no keyword-only parameter on function '{fn.__name__}'") from None + + keyword_only_params = tuple(inspect.signature(fn).parameters)[keyword_only_start_idx:] + + @functools.wraps(fn) + def wrapper(*args: Any, **kwargs: Any) -> D: + args, keyword_only_args = args[:keyword_only_start_idx], args[keyword_only_start_idx:] + if keyword_only_args: + keyword_only_kwargs = dict(zip(keyword_only_params, keyword_only_args)) + warnings.warn( + f"Using {sequence_to_str(tuple(keyword_only_kwargs.keys()), separate_last='and ')} as positional " + f"parameter(s) is deprecated since 0.13 and may be removed in the future. Please use keyword parameter(s) " + f"instead." + ) + kwargs.update(keyword_only_kwargs) + + return fn(*args, **kwargs) + + return wrapper + + +W = TypeVar("W", bound=WeightsEnum) +M = TypeVar("M", bound=nn.Module) +V = TypeVar("V") + + +def handle_legacy_interface(**weights: Tuple[str, Union[Optional[W], Callable[[Dict[str, Any]], Optional[W]]]]): + """Decorates a model builder with the new interface to make it compatible with the old. + + In particular this handles two things: + + 1. Allows positional parameters again, but emits a deprecation warning in case they are used. See + :func:`torchvision.prototype.utils._internal.kwonly_to_pos_or_kw` for details. + 2. Handles the default value change from ``pretrained=False`` to ``weights=None`` and ``pretrained=True`` to + ``weights=Weights`` and emits a deprecation warning with instructions for the new interface. + + Args: + **weights (Tuple[str, Union[Optional[W], Callable[[Dict[str, Any]], Optional[W]]]]): Deprecated parameter + name and default value for the legacy ``pretrained=True``. The default value can be a callable in which + case it will be called with a dictionary of the keyword arguments. The only key that is guaranteed to be in + the dictionary is the deprecated parameter name passed as first element in the tuple. All other parameters + should be accessed with :meth:`~dict.get`. + """ + + def outer_wrapper(builder: Callable[..., M]) -> Callable[..., M]: + @kwonly_to_pos_or_kw + @functools.wraps(builder) + def inner_wrapper(*args: Any, **kwargs: Any) -> M: + for weights_param, (pretrained_param, default) in weights.items(): # type: ignore[union-attr] + # If neither the weights nor the pretrained parameter as passed, or the weights argument already use + # the new style arguments, there is nothing to do. Note that we cannot use `None` as sentinel for the + # weight argument, since it is a valid value. + sentinel = object() + weights_arg = kwargs.get(weights_param, sentinel) + if ( + (weights_param not in kwargs and pretrained_param not in kwargs) + or isinstance(weights_arg, WeightsEnum) + or (isinstance(weights_arg, str) and weights_arg != "legacy") + or weights_arg is None + ): + continue + + # If the pretrained parameter was passed as positional argument, it is now mapped to + # `kwargs[weights_param]`. This happens because the @kwonly_to_pos_or_kw decorator uses the current + # signature to infer the names of positionally passed arguments and thus has no knowledge that there + # used to be a pretrained parameter. + pretrained_positional = weights_arg is not sentinel + if pretrained_positional: + # We put the pretrained argument under its legacy name in the keyword argument dictionary to have + # unified access to the value if the default value is a callable. + kwargs[pretrained_param] = pretrained_arg = kwargs.pop(weights_param) + else: + pretrained_arg = kwargs[pretrained_param] + + if pretrained_arg: + default_weights_arg = default(kwargs) if callable(default) else default + if not isinstance(default_weights_arg, WeightsEnum): + raise ValueError(f"No weights available for model {builder.__name__}") + else: + default_weights_arg = None + + if not pretrained_positional: + warnings.warn( + f"The parameter '{pretrained_param}' is deprecated since 0.13 and may be removed in the future, " + f"please use '{weights_param}' instead." + ) + + msg = ( + f"Arguments other than a weight enum or `None` for '{weights_param}' are deprecated since 0.13 and " + f"may be removed in the future. " + f"The current behavior is equivalent to passing `{weights_param}={default_weights_arg}`." + ) + if pretrained_arg: + msg = ( + f"{msg} You can also use `{weights_param}={type(default_weights_arg).__name__}.DEFAULT` " + f"to get the most up-to-date weights." + ) + warnings.warn(msg) + + del kwargs[pretrained_param] + kwargs[weights_param] = default_weights_arg + + return builder(*args, **kwargs) + + return inner_wrapper + + return outer_wrapper + + +def _ovewrite_named_param(kwargs: Dict[str, Any], param: str, new_value: V) -> None: + if param in kwargs: + if kwargs[param] != new_value: + raise ValueError(f"The parameter '{param}' expected value {new_value} but got {kwargs[param]} instead.") + else: + kwargs[param] = new_value + + +def _ovewrite_value_param(param: str, actual: Optional[V], expected: V) -> V: + if actual is not None: + if actual != expected: + raise ValueError(f"The parameter '{param}' expected value {expected} but got {actual} instead.") + return expected + + +class _ModelURLs(dict): + def __getitem__(self, item): + warnings.warn( + "Accessing the model URLs via the internal dictionary of the module is deprecated since 0.13 and may " + "be removed in the future. Please access them via the appropriate Weights Enum instead." + ) + return super().__getitem__(item) diff --git a/lib/python3.12/site-packages/torchvision/models/alexnet.py b/lib/python3.12/site-packages/torchvision/models/alexnet.py new file mode 100644 index 0000000000000000000000000000000000000000..f85acbeb2148d2aa8f289808e61aa61e2d68e2f9 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/alexnet.py @@ -0,0 +1,119 @@ +from functools import partial +from typing import Any, Optional + +import torch +import torch.nn as nn + +from ..transforms._presets import ImageClassification +from ..utils import _log_api_usage_once +from ._api import register_model, Weights, WeightsEnum +from ._meta import _IMAGENET_CATEGORIES +from ._utils import _ovewrite_named_param, handle_legacy_interface + + +__all__ = ["AlexNet", "AlexNet_Weights", "alexnet"] + + +class AlexNet(nn.Module): + def __init__(self, num_classes: int = 1000, dropout: float = 0.5) -> None: + super().__init__() + _log_api_usage_once(self) + self.features = nn.Sequential( + nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2), + nn.ReLU(inplace=True), + nn.MaxPool2d(kernel_size=3, stride=2), + nn.Conv2d(64, 192, kernel_size=5, padding=2), + nn.ReLU(inplace=True), + nn.MaxPool2d(kernel_size=3, stride=2), + nn.Conv2d(192, 384, kernel_size=3, padding=1), + nn.ReLU(inplace=True), + nn.Conv2d(384, 256, kernel_size=3, padding=1), + nn.ReLU(inplace=True), + nn.Conv2d(256, 256, kernel_size=3, padding=1), + nn.ReLU(inplace=True), + nn.MaxPool2d(kernel_size=3, stride=2), + ) + self.avgpool = nn.AdaptiveAvgPool2d((6, 6)) + self.classifier = nn.Sequential( + nn.Dropout(p=dropout), + nn.Linear(256 * 6 * 6, 4096), + nn.ReLU(inplace=True), + nn.Dropout(p=dropout), + nn.Linear(4096, 4096), + nn.ReLU(inplace=True), + nn.Linear(4096, num_classes), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.features(x) + x = self.avgpool(x) + x = torch.flatten(x, 1) + x = self.classifier(x) + return x + + +class AlexNet_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/alexnet-owt-7be5be79.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + "num_params": 61100840, + "min_size": (63, 63), + "categories": _IMAGENET_CATEGORIES, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#alexnet-and-vgg", + "_metrics": { + "ImageNet-1K": { + "acc@1": 56.522, + "acc@5": 79.066, + } + }, + "_ops": 0.714, + "_file_size": 233.087, + "_docs": """ + These weights reproduce closely the results of the paper using a simplified training recipe. + """, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", AlexNet_Weights.IMAGENET1K_V1)) +def alexnet(*, weights: Optional[AlexNet_Weights] = None, progress: bool = True, **kwargs: Any) -> AlexNet: + """AlexNet model architecture from `One weird trick for parallelizing convolutional neural networks `__. + + .. note:: + AlexNet was originally introduced in the `ImageNet Classification with + Deep Convolutional Neural Networks + `__ + paper. Our implementation is based instead on the "One weird trick" + paper above. + + Args: + weights (:class:`~torchvision.models.AlexNet_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.AlexNet_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.squeezenet.AlexNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.AlexNet_Weights + :members: + """ + + weights = AlexNet_Weights.verify(weights) + + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + + model = AlexNet(**kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model diff --git a/lib/python3.12/site-packages/torchvision/models/convnext.py b/lib/python3.12/site-packages/torchvision/models/convnext.py new file mode 100644 index 0000000000000000000000000000000000000000..444ef3c219efa09d1c720f29db44a8acc8714bbc --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/convnext.py @@ -0,0 +1,414 @@ +from functools import partial +from typing import Any, Callable, List, Optional, Sequence + +import torch +from torch import nn, Tensor +from torch.nn import functional as F + +from ..ops.misc import Conv2dNormActivation, Permute +from ..ops.stochastic_depth import StochasticDepth +from ..transforms._presets import ImageClassification +from ..utils import _log_api_usage_once +from ._api import register_model, Weights, WeightsEnum +from ._meta import _IMAGENET_CATEGORIES +from ._utils import _ovewrite_named_param, handle_legacy_interface + + +__all__ = [ + "ConvNeXt", + "ConvNeXt_Tiny_Weights", + "ConvNeXt_Small_Weights", + "ConvNeXt_Base_Weights", + "ConvNeXt_Large_Weights", + "convnext_tiny", + "convnext_small", + "convnext_base", + "convnext_large", +] + + +class LayerNorm2d(nn.LayerNorm): + def forward(self, x: Tensor) -> Tensor: + x = x.permute(0, 2, 3, 1) + x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + x = x.permute(0, 3, 1, 2) + return x + + +class CNBlock(nn.Module): + def __init__( + self, + dim, + layer_scale: float, + stochastic_depth_prob: float, + norm_layer: Optional[Callable[..., nn.Module]] = None, + ) -> None: + super().__init__() + if norm_layer is None: + norm_layer = partial(nn.LayerNorm, eps=1e-6) + + self.block = nn.Sequential( + nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim, bias=True), + Permute([0, 2, 3, 1]), + norm_layer(dim), + nn.Linear(in_features=dim, out_features=4 * dim, bias=True), + nn.GELU(), + nn.Linear(in_features=4 * dim, out_features=dim, bias=True), + Permute([0, 3, 1, 2]), + ) + self.layer_scale = nn.Parameter(torch.ones(dim, 1, 1) * layer_scale) + self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row") + + def forward(self, input: Tensor) -> Tensor: + result = self.layer_scale * self.block(input) + result = self.stochastic_depth(result) + result += input + return result + + +class CNBlockConfig: + # Stores information listed at Section 3 of the ConvNeXt paper + def __init__( + self, + input_channels: int, + out_channels: Optional[int], + num_layers: int, + ) -> None: + self.input_channels = input_channels + self.out_channels = out_channels + self.num_layers = num_layers + + def __repr__(self) -> str: + s = self.__class__.__name__ + "(" + s += "input_channels={input_channels}" + s += ", out_channels={out_channels}" + s += ", num_layers={num_layers}" + s += ")" + return s.format(**self.__dict__) + + +class ConvNeXt(nn.Module): + def __init__( + self, + block_setting: List[CNBlockConfig], + stochastic_depth_prob: float = 0.0, + layer_scale: float = 1e-6, + num_classes: int = 1000, + block: Optional[Callable[..., nn.Module]] = None, + norm_layer: Optional[Callable[..., nn.Module]] = None, + **kwargs: Any, + ) -> None: + super().__init__() + _log_api_usage_once(self) + + if not block_setting: + raise ValueError("The block_setting should not be empty") + elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])): + raise TypeError("The block_setting should be List[CNBlockConfig]") + + if block is None: + block = CNBlock + + if norm_layer is None: + norm_layer = partial(LayerNorm2d, eps=1e-6) + + layers: List[nn.Module] = [] + + # Stem + firstconv_output_channels = block_setting[0].input_channels + layers.append( + Conv2dNormActivation( + 3, + firstconv_output_channels, + kernel_size=4, + stride=4, + padding=0, + norm_layer=norm_layer, + activation_layer=None, + bias=True, + ) + ) + + total_stage_blocks = sum(cnf.num_layers for cnf in block_setting) + stage_block_id = 0 + for cnf in block_setting: + # Bottlenecks + stage: List[nn.Module] = [] + for _ in range(cnf.num_layers): + # adjust stochastic depth probability based on the depth of the stage block + sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0) + stage.append(block(cnf.input_channels, layer_scale, sd_prob)) + stage_block_id += 1 + layers.append(nn.Sequential(*stage)) + if cnf.out_channels is not None: + # Downsampling + layers.append( + nn.Sequential( + norm_layer(cnf.input_channels), + nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2), + ) + ) + + self.features = nn.Sequential(*layers) + self.avgpool = nn.AdaptiveAvgPool2d(1) + + lastblock = block_setting[-1] + lastconv_output_channels = ( + lastblock.out_channels if lastblock.out_channels is not None else lastblock.input_channels + ) + self.classifier = nn.Sequential( + norm_layer(lastconv_output_channels), nn.Flatten(1), nn.Linear(lastconv_output_channels, num_classes) + ) + + for m in self.modules(): + if isinstance(m, (nn.Conv2d, nn.Linear)): + nn.init.trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.zeros_(m.bias) + + def _forward_impl(self, x: Tensor) -> Tensor: + x = self.features(x) + x = self.avgpool(x) + x = self.classifier(x) + return x + + def forward(self, x: Tensor) -> Tensor: + return self._forward_impl(x) + + +def _convnext( + block_setting: List[CNBlockConfig], + stochastic_depth_prob: float, + weights: Optional[WeightsEnum], + progress: bool, + **kwargs: Any, +) -> ConvNeXt: + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + + model = ConvNeXt(block_setting, stochastic_depth_prob=stochastic_depth_prob, **kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +_COMMON_META = { + "min_size": (32, 32), + "categories": _IMAGENET_CATEGORIES, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#convnext", + "_docs": """ + These weights improve upon the results of the original paper by using a modified version of TorchVision's + `new training recipe + `_. + """, +} + + +class ConvNeXt_Tiny_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/convnext_tiny-983f1562.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=236), + meta={ + **_COMMON_META, + "num_params": 28589128, + "_metrics": { + "ImageNet-1K": { + "acc@1": 82.520, + "acc@5": 96.146, + } + }, + "_ops": 4.456, + "_file_size": 109.119, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class ConvNeXt_Small_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/convnext_small-0c510722.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=230), + meta={ + **_COMMON_META, + "num_params": 50223688, + "_metrics": { + "ImageNet-1K": { + "acc@1": 83.616, + "acc@5": 96.650, + } + }, + "_ops": 8.684, + "_file_size": 191.703, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class ConvNeXt_Base_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/convnext_base-6075fbad.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 88591464, + "_metrics": { + "ImageNet-1K": { + "acc@1": 84.062, + "acc@5": 96.870, + } + }, + "_ops": 15.355, + "_file_size": 338.064, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class ConvNeXt_Large_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/convnext_large-ea097f82.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 197767336, + "_metrics": { + "ImageNet-1K": { + "acc@1": 84.414, + "acc@5": 96.976, + } + }, + "_ops": 34.361, + "_file_size": 754.537, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ConvNeXt_Tiny_Weights.IMAGENET1K_V1)) +def convnext_tiny(*, weights: Optional[ConvNeXt_Tiny_Weights] = None, progress: bool = True, **kwargs: Any) -> ConvNeXt: + """ConvNeXt Tiny model architecture from the + `A ConvNet for the 2020s `_ paper. + + Args: + weights (:class:`~torchvision.models.convnext.ConvNeXt_Tiny_Weights`, optional): The pretrained + weights to use. See :class:`~torchvision.models.convnext.ConvNeXt_Tiny_Weights` + below for more details and possible values. By default, no pre-trained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.convnext.ConvNext`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.ConvNeXt_Tiny_Weights + :members: + """ + weights = ConvNeXt_Tiny_Weights.verify(weights) + + block_setting = [ + CNBlockConfig(96, 192, 3), + CNBlockConfig(192, 384, 3), + CNBlockConfig(384, 768, 9), + CNBlockConfig(768, None, 3), + ] + stochastic_depth_prob = kwargs.pop("stochastic_depth_prob", 0.1) + return _convnext(block_setting, stochastic_depth_prob, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ConvNeXt_Small_Weights.IMAGENET1K_V1)) +def convnext_small( + *, weights: Optional[ConvNeXt_Small_Weights] = None, progress: bool = True, **kwargs: Any +) -> ConvNeXt: + """ConvNeXt Small model architecture from the + `A ConvNet for the 2020s `_ paper. + + Args: + weights (:class:`~torchvision.models.convnext.ConvNeXt_Small_Weights`, optional): The pretrained + weights to use. See :class:`~torchvision.models.convnext.ConvNeXt_Small_Weights` + below for more details and possible values. By default, no pre-trained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.convnext.ConvNext`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.ConvNeXt_Small_Weights + :members: + """ + weights = ConvNeXt_Small_Weights.verify(weights) + + block_setting = [ + CNBlockConfig(96, 192, 3), + CNBlockConfig(192, 384, 3), + CNBlockConfig(384, 768, 27), + CNBlockConfig(768, None, 3), + ] + stochastic_depth_prob = kwargs.pop("stochastic_depth_prob", 0.4) + return _convnext(block_setting, stochastic_depth_prob, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ConvNeXt_Base_Weights.IMAGENET1K_V1)) +def convnext_base(*, weights: Optional[ConvNeXt_Base_Weights] = None, progress: bool = True, **kwargs: Any) -> ConvNeXt: + """ConvNeXt Base model architecture from the + `A ConvNet for the 2020s `_ paper. + + Args: + weights (:class:`~torchvision.models.convnext.ConvNeXt_Base_Weights`, optional): The pretrained + weights to use. See :class:`~torchvision.models.convnext.ConvNeXt_Base_Weights` + below for more details and possible values. By default, no pre-trained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.convnext.ConvNext`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.ConvNeXt_Base_Weights + :members: + """ + weights = ConvNeXt_Base_Weights.verify(weights) + + block_setting = [ + CNBlockConfig(128, 256, 3), + CNBlockConfig(256, 512, 3), + CNBlockConfig(512, 1024, 27), + CNBlockConfig(1024, None, 3), + ] + stochastic_depth_prob = kwargs.pop("stochastic_depth_prob", 0.5) + return _convnext(block_setting, stochastic_depth_prob, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ConvNeXt_Large_Weights.IMAGENET1K_V1)) +def convnext_large( + *, weights: Optional[ConvNeXt_Large_Weights] = None, progress: bool = True, **kwargs: Any +) -> ConvNeXt: + """ConvNeXt Large model architecture from the + `A ConvNet for the 2020s `_ paper. + + Args: + weights (:class:`~torchvision.models.convnext.ConvNeXt_Large_Weights`, optional): The pretrained + weights to use. See :class:`~torchvision.models.convnext.ConvNeXt_Large_Weights` + below for more details and possible values. By default, no pre-trained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.convnext.ConvNext`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.ConvNeXt_Large_Weights + :members: + """ + weights = ConvNeXt_Large_Weights.verify(weights) + + block_setting = [ + CNBlockConfig(192, 384, 3), + CNBlockConfig(384, 768, 3), + CNBlockConfig(768, 1536, 27), + CNBlockConfig(1536, None, 3), + ] + stochastic_depth_prob = kwargs.pop("stochastic_depth_prob", 0.5) + return _convnext(block_setting, stochastic_depth_prob, weights, progress, **kwargs) diff --git a/lib/python3.12/site-packages/torchvision/models/densenet.py b/lib/python3.12/site-packages/torchvision/models/densenet.py new file mode 100644 index 0000000000000000000000000000000000000000..c4c1bc9525a43fbba565dc3f03c6b7231e913fe5 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/densenet.py @@ -0,0 +1,448 @@ +import re +from collections import OrderedDict +from functools import partial +from typing import Any, List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as cp +from torch import Tensor + +from ..transforms._presets import ImageClassification +from ..utils import _log_api_usage_once +from ._api import register_model, Weights, WeightsEnum +from ._meta import _IMAGENET_CATEGORIES +from ._utils import _ovewrite_named_param, handle_legacy_interface + +__all__ = [ + "DenseNet", + "DenseNet121_Weights", + "DenseNet161_Weights", + "DenseNet169_Weights", + "DenseNet201_Weights", + "densenet121", + "densenet161", + "densenet169", + "densenet201", +] + + +class _DenseLayer(nn.Module): + def __init__( + self, num_input_features: int, growth_rate: int, bn_size: int, drop_rate: float, memory_efficient: bool = False + ) -> None: + super().__init__() + self.norm1 = nn.BatchNorm2d(num_input_features) + self.relu1 = nn.ReLU(inplace=True) + self.conv1 = nn.Conv2d(num_input_features, bn_size * growth_rate, kernel_size=1, stride=1, bias=False) + + self.norm2 = nn.BatchNorm2d(bn_size * growth_rate) + self.relu2 = nn.ReLU(inplace=True) + self.conv2 = nn.Conv2d(bn_size * growth_rate, growth_rate, kernel_size=3, stride=1, padding=1, bias=False) + + self.drop_rate = float(drop_rate) + self.memory_efficient = memory_efficient + + def bn_function(self, inputs: List[Tensor]) -> Tensor: + concated_features = torch.cat(inputs, 1) + bottleneck_output = self.conv1(self.relu1(self.norm1(concated_features))) # noqa: T484 + return bottleneck_output + + # todo: rewrite when torchscript supports any + def any_requires_grad(self, input: List[Tensor]) -> bool: + for tensor in input: + if tensor.requires_grad: + return True + return False + + @torch.jit.unused # noqa: T484 + def call_checkpoint_bottleneck(self, input: List[Tensor]) -> Tensor: + def closure(*inputs): + return self.bn_function(inputs) + + return cp.checkpoint(closure, *input, use_reentrant=False) + + @torch.jit._overload_method # noqa: F811 + def forward(self, input: List[Tensor]) -> Tensor: # noqa: F811 + pass + + @torch.jit._overload_method # noqa: F811 + def forward(self, input: Tensor) -> Tensor: # noqa: F811 + pass + + # torchscript does not yet support *args, so we overload method + # allowing it to take either a List[Tensor] or single Tensor + def forward(self, input: Tensor) -> Tensor: # noqa: F811 + if isinstance(input, Tensor): + prev_features = [input] + else: + prev_features = input + + if self.memory_efficient and self.any_requires_grad(prev_features): + if torch.jit.is_scripting(): + raise Exception("Memory Efficient not supported in JIT") + + bottleneck_output = self.call_checkpoint_bottleneck(prev_features) + else: + bottleneck_output = self.bn_function(prev_features) + + new_features = self.conv2(self.relu2(self.norm2(bottleneck_output))) + if self.drop_rate > 0: + new_features = F.dropout(new_features, p=self.drop_rate, training=self.training) + return new_features + + +class _DenseBlock(nn.ModuleDict): + _version = 2 + + def __init__( + self, + num_layers: int, + num_input_features: int, + bn_size: int, + growth_rate: int, + drop_rate: float, + memory_efficient: bool = False, + ) -> None: + super().__init__() + for i in range(num_layers): + layer = _DenseLayer( + num_input_features + i * growth_rate, + growth_rate=growth_rate, + bn_size=bn_size, + drop_rate=drop_rate, + memory_efficient=memory_efficient, + ) + self.add_module("denselayer%d" % (i + 1), layer) + + def forward(self, init_features: Tensor) -> Tensor: + features = [init_features] + for name, layer in self.items(): + new_features = layer(features) + features.append(new_features) + return torch.cat(features, 1) + + +class _Transition(nn.Sequential): + def __init__(self, num_input_features: int, num_output_features: int) -> None: + super().__init__() + self.norm = nn.BatchNorm2d(num_input_features) + self.relu = nn.ReLU(inplace=True) + self.conv = nn.Conv2d(num_input_features, num_output_features, kernel_size=1, stride=1, bias=False) + self.pool = nn.AvgPool2d(kernel_size=2, stride=2) + + +class DenseNet(nn.Module): + r"""Densenet-BC model class, based on + `"Densely Connected Convolutional Networks" `_. + + Args: + growth_rate (int) - how many filters to add each layer (`k` in paper) + block_config (list of 4 ints) - how many layers in each pooling block + num_init_features (int) - the number of filters to learn in the first convolution layer + bn_size (int) - multiplicative factor for number of bottle neck layers + (i.e. bn_size * k features in the bottleneck layer) + drop_rate (float) - dropout rate after each dense layer + num_classes (int) - number of classification classes + memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient, + but slower. Default: *False*. See `"paper" `_. + """ + + def __init__( + self, + growth_rate: int = 32, + block_config: Tuple[int, int, int, int] = (6, 12, 24, 16), + num_init_features: int = 64, + bn_size: int = 4, + drop_rate: float = 0, + num_classes: int = 1000, + memory_efficient: bool = False, + ) -> None: + + super().__init__() + _log_api_usage_once(self) + + # First convolution + self.features = nn.Sequential( + OrderedDict( + [ + ("conv0", nn.Conv2d(3, num_init_features, kernel_size=7, stride=2, padding=3, bias=False)), + ("norm0", nn.BatchNorm2d(num_init_features)), + ("relu0", nn.ReLU(inplace=True)), + ("pool0", nn.MaxPool2d(kernel_size=3, stride=2, padding=1)), + ] + ) + ) + + # Each denseblock + num_features = num_init_features + for i, num_layers in enumerate(block_config): + block = _DenseBlock( + num_layers=num_layers, + num_input_features=num_features, + bn_size=bn_size, + growth_rate=growth_rate, + drop_rate=drop_rate, + memory_efficient=memory_efficient, + ) + self.features.add_module("denseblock%d" % (i + 1), block) + num_features = num_features + num_layers * growth_rate + if i != len(block_config) - 1: + trans = _Transition(num_input_features=num_features, num_output_features=num_features // 2) + self.features.add_module("transition%d" % (i + 1), trans) + num_features = num_features // 2 + + # Final batch norm + self.features.add_module("norm5", nn.BatchNorm2d(num_features)) + + # Linear layer + self.classifier = nn.Linear(num_features, num_classes) + + # Official init from torch repo. + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight) + elif isinstance(m, nn.BatchNorm2d): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Linear): + nn.init.constant_(m.bias, 0) + + def forward(self, x: Tensor) -> Tensor: + features = self.features(x) + out = F.relu(features, inplace=True) + out = F.adaptive_avg_pool2d(out, (1, 1)) + out = torch.flatten(out, 1) + out = self.classifier(out) + return out + + +def _load_state_dict(model: nn.Module, weights: WeightsEnum, progress: bool) -> None: + # '.'s are no longer allowed in module names, but previous _DenseLayer + # has keys 'norm.1', 'relu.1', 'conv.1', 'norm.2', 'relu.2', 'conv.2'. + # They are also in the checkpoints in model_urls. This pattern is used + # to find such keys. + pattern = re.compile( + r"^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$" + ) + + state_dict = weights.get_state_dict(progress=progress, check_hash=True) + for key in list(state_dict.keys()): + res = pattern.match(key) + if res: + new_key = res.group(1) + res.group(2) + state_dict[new_key] = state_dict[key] + del state_dict[key] + model.load_state_dict(state_dict) + + +def _densenet( + growth_rate: int, + block_config: Tuple[int, int, int, int], + num_init_features: int, + weights: Optional[WeightsEnum], + progress: bool, + **kwargs: Any, +) -> DenseNet: + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + + model = DenseNet(growth_rate, block_config, num_init_features, **kwargs) + + if weights is not None: + _load_state_dict(model=model, weights=weights, progress=progress) + + return model + + +_COMMON_META = { + "min_size": (29, 29), + "categories": _IMAGENET_CATEGORIES, + "recipe": "https://github.com/pytorch/vision/pull/116", + "_docs": """These weights are ported from LuaTorch.""", +} + + +class DenseNet121_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/densenet121-a639ec97.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 7978856, + "_metrics": { + "ImageNet-1K": { + "acc@1": 74.434, + "acc@5": 91.972, + } + }, + "_ops": 2.834, + "_file_size": 30.845, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class DenseNet161_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/densenet161-8d451a50.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 28681000, + "_metrics": { + "ImageNet-1K": { + "acc@1": 77.138, + "acc@5": 93.560, + } + }, + "_ops": 7.728, + "_file_size": 110.369, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class DenseNet169_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/densenet169-b2777c0a.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 14149480, + "_metrics": { + "ImageNet-1K": { + "acc@1": 75.600, + "acc@5": 92.806, + } + }, + "_ops": 3.36, + "_file_size": 54.708, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class DenseNet201_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/densenet201-c1103571.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 20013928, + "_metrics": { + "ImageNet-1K": { + "acc@1": 76.896, + "acc@5": 93.370, + } + }, + "_ops": 4.291, + "_file_size": 77.373, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", DenseNet121_Weights.IMAGENET1K_V1)) +def densenet121(*, weights: Optional[DenseNet121_Weights] = None, progress: bool = True, **kwargs: Any) -> DenseNet: + r"""Densenet-121 model from + `Densely Connected Convolutional Networks `_. + + Args: + weights (:class:`~torchvision.models.DenseNet121_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.DenseNet121_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.densenet.DenseNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.DenseNet121_Weights + :members: + """ + weights = DenseNet121_Weights.verify(weights) + + return _densenet(32, (6, 12, 24, 16), 64, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", DenseNet161_Weights.IMAGENET1K_V1)) +def densenet161(*, weights: Optional[DenseNet161_Weights] = None, progress: bool = True, **kwargs: Any) -> DenseNet: + r"""Densenet-161 model from + `Densely Connected Convolutional Networks `_. + + Args: + weights (:class:`~torchvision.models.DenseNet161_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.DenseNet161_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.densenet.DenseNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.DenseNet161_Weights + :members: + """ + weights = DenseNet161_Weights.verify(weights) + + return _densenet(48, (6, 12, 36, 24), 96, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", DenseNet169_Weights.IMAGENET1K_V1)) +def densenet169(*, weights: Optional[DenseNet169_Weights] = None, progress: bool = True, **kwargs: Any) -> DenseNet: + r"""Densenet-169 model from + `Densely Connected Convolutional Networks `_. + + Args: + weights (:class:`~torchvision.models.DenseNet169_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.DenseNet169_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.densenet.DenseNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.DenseNet169_Weights + :members: + """ + weights = DenseNet169_Weights.verify(weights) + + return _densenet(32, (6, 12, 32, 32), 64, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", DenseNet201_Weights.IMAGENET1K_V1)) +def densenet201(*, weights: Optional[DenseNet201_Weights] = None, progress: bool = True, **kwargs: Any) -> DenseNet: + r"""Densenet-201 model from + `Densely Connected Convolutional Networks `_. + + Args: + weights (:class:`~torchvision.models.DenseNet201_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.DenseNet201_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.densenet.DenseNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.DenseNet201_Weights + :members: + """ + weights = DenseNet201_Weights.verify(weights) + + return _densenet(32, (6, 12, 48, 32), 64, weights, progress, **kwargs) diff --git a/lib/python3.12/site-packages/torchvision/models/detection/__init__.py b/lib/python3.12/site-packages/torchvision/models/detection/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4146651c737971cc5a883b6750f2ded3051bc8ea --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/detection/__init__.py @@ -0,0 +1,7 @@ +from .faster_rcnn import * +from .fcos import * +from .keypoint_rcnn import * +from .mask_rcnn import * +from .retinanet import * +from .ssd import * +from .ssdlite import * diff --git a/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4cd05a719ccd41a0d08d743804a97a9a73f14a02 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1d77026bab5c85f56139ab25361b3831b6b48cc0 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/anchor_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/anchor_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..14b4d292b86fe227669368d138279543c6b88abd Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/anchor_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/backbone_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/backbone_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..afb365e64e0e0c7f4a7517709d34e24147fc634e Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/backbone_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/faster_rcnn.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/faster_rcnn.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e64946380e94b9414c296e69912b00b12d0d3c1 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/faster_rcnn.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/fcos.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/fcos.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1d4db8bda7a3abe132b0cb5366a7a36b78812211 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/fcos.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/generalized_rcnn.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/generalized_rcnn.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dca8a2682a78fe85d81f5425c2869b42b6555191 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/generalized_rcnn.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/image_list.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/image_list.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..230e4e58e47240788780b706a802df788752ca82 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/image_list.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/keypoint_rcnn.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/keypoint_rcnn.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dd7c616229b39cefca135ffe623d235a62b711fe Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/keypoint_rcnn.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/mask_rcnn.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/mask_rcnn.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aa100665327ea0e943a0bdd9eeb27f0ac80e9d16 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/mask_rcnn.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/retinanet.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/retinanet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b17e424ec68dc4012e2e8efe963a49b4ed808933 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/retinanet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/roi_heads.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/roi_heads.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2dad624354c10834cd47fa9df4fbaf4b5359d064 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/roi_heads.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/rpn.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/rpn.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bca6e5aee29789bdbd7cb6f679641e171adb0934 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/rpn.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/ssd.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/ssd.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df09a6aa82bbb74dcb4e99349be89db7f3e739c4 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/ssd.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/ssdlite.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/ssdlite.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ff6d569c2a9ddc5309a38eefc49f3785c967de92 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/ssdlite.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/transform.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/transform.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..670e0cd5cffcec8238f0d17724573acf5a2292c3 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/detection/__pycache__/transform.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/detection/_utils.py b/lib/python3.12/site-packages/torchvision/models/detection/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..559db858ac32f3b9f157aff3c22da83abece2a73 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/detection/_utils.py @@ -0,0 +1,540 @@ +import math +from collections import OrderedDict +from typing import Dict, List, Optional, Tuple + +import torch +from torch import nn, Tensor +from torch.nn import functional as F +from torchvision.ops import complete_box_iou_loss, distance_box_iou_loss, FrozenBatchNorm2d, generalized_box_iou_loss + + +class BalancedPositiveNegativeSampler: + """ + This class samples batches, ensuring that they contain a fixed proportion of positives + """ + + def __init__(self, batch_size_per_image: int, positive_fraction: float) -> None: + """ + Args: + batch_size_per_image (int): number of elements to be selected per image + positive_fraction (float): percentage of positive elements per batch + """ + self.batch_size_per_image = batch_size_per_image + self.positive_fraction = positive_fraction + + def __call__(self, matched_idxs: List[Tensor]) -> Tuple[List[Tensor], List[Tensor]]: + """ + Args: + matched_idxs: list of tensors containing -1, 0 or positive values. + Each tensor corresponds to a specific image. + -1 values are ignored, 0 are considered as negatives and > 0 as + positives. + + Returns: + pos_idx (list[tensor]) + neg_idx (list[tensor]) + + Returns two lists of binary masks for each image. + The first list contains the positive elements that were selected, + and the second list the negative example. + """ + pos_idx = [] + neg_idx = [] + for matched_idxs_per_image in matched_idxs: + positive = torch.where(matched_idxs_per_image >= 1)[0] + negative = torch.where(matched_idxs_per_image == 0)[0] + + num_pos = int(self.batch_size_per_image * self.positive_fraction) + # protect against not enough positive examples + num_pos = min(positive.numel(), num_pos) + num_neg = self.batch_size_per_image - num_pos + # protect against not enough negative examples + num_neg = min(negative.numel(), num_neg) + + # randomly select positive and negative examples + perm1 = torch.randperm(positive.numel(), device=positive.device)[:num_pos] + perm2 = torch.randperm(negative.numel(), device=negative.device)[:num_neg] + + pos_idx_per_image = positive[perm1] + neg_idx_per_image = negative[perm2] + + # create binary mask from indices + pos_idx_per_image_mask = torch.zeros_like(matched_idxs_per_image, dtype=torch.uint8) + neg_idx_per_image_mask = torch.zeros_like(matched_idxs_per_image, dtype=torch.uint8) + + pos_idx_per_image_mask[pos_idx_per_image] = 1 + neg_idx_per_image_mask[neg_idx_per_image] = 1 + + pos_idx.append(pos_idx_per_image_mask) + neg_idx.append(neg_idx_per_image_mask) + + return pos_idx, neg_idx + + +@torch.jit._script_if_tracing +def encode_boxes(reference_boxes: Tensor, proposals: Tensor, weights: Tensor) -> Tensor: + """ + Encode a set of proposals with respect to some + reference boxes + + Args: + reference_boxes (Tensor): reference boxes + proposals (Tensor): boxes to be encoded + weights (Tensor[4]): the weights for ``(x, y, w, h)`` + """ + + # perform some unpacking to make it JIT-fusion friendly + wx = weights[0] + wy = weights[1] + ww = weights[2] + wh = weights[3] + + proposals_x1 = proposals[:, 0].unsqueeze(1) + proposals_y1 = proposals[:, 1].unsqueeze(1) + proposals_x2 = proposals[:, 2].unsqueeze(1) + proposals_y2 = proposals[:, 3].unsqueeze(1) + + reference_boxes_x1 = reference_boxes[:, 0].unsqueeze(1) + reference_boxes_y1 = reference_boxes[:, 1].unsqueeze(1) + reference_boxes_x2 = reference_boxes[:, 2].unsqueeze(1) + reference_boxes_y2 = reference_boxes[:, 3].unsqueeze(1) + + # implementation starts here + ex_widths = proposals_x2 - proposals_x1 + ex_heights = proposals_y2 - proposals_y1 + ex_ctr_x = proposals_x1 + 0.5 * ex_widths + ex_ctr_y = proposals_y1 + 0.5 * ex_heights + + gt_widths = reference_boxes_x2 - reference_boxes_x1 + gt_heights = reference_boxes_y2 - reference_boxes_y1 + gt_ctr_x = reference_boxes_x1 + 0.5 * gt_widths + gt_ctr_y = reference_boxes_y1 + 0.5 * gt_heights + + targets_dx = wx * (gt_ctr_x - ex_ctr_x) / ex_widths + targets_dy = wy * (gt_ctr_y - ex_ctr_y) / ex_heights + targets_dw = ww * torch.log(gt_widths / ex_widths) + targets_dh = wh * torch.log(gt_heights / ex_heights) + + targets = torch.cat((targets_dx, targets_dy, targets_dw, targets_dh), dim=1) + return targets + + +class BoxCoder: + """ + This class encodes and decodes a set of bounding boxes into + the representation used for training the regressors. + """ + + def __init__( + self, weights: Tuple[float, float, float, float], bbox_xform_clip: float = math.log(1000.0 / 16) + ) -> None: + """ + Args: + weights (4-element tuple) + bbox_xform_clip (float) + """ + self.weights = weights + self.bbox_xform_clip = bbox_xform_clip + + def encode(self, reference_boxes: List[Tensor], proposals: List[Tensor]) -> List[Tensor]: + boxes_per_image = [len(b) for b in reference_boxes] + reference_boxes = torch.cat(reference_boxes, dim=0) + proposals = torch.cat(proposals, dim=0) + targets = self.encode_single(reference_boxes, proposals) + return targets.split(boxes_per_image, 0) + + def encode_single(self, reference_boxes: Tensor, proposals: Tensor) -> Tensor: + """ + Encode a set of proposals with respect to some + reference boxes + + Args: + reference_boxes (Tensor): reference boxes + proposals (Tensor): boxes to be encoded + """ + dtype = reference_boxes.dtype + device = reference_boxes.device + weights = torch.as_tensor(self.weights, dtype=dtype, device=device) + targets = encode_boxes(reference_boxes, proposals, weights) + + return targets + + def decode(self, rel_codes: Tensor, boxes: List[Tensor]) -> Tensor: + torch._assert( + isinstance(boxes, (list, tuple)), + "This function expects boxes of type list or tuple.", + ) + torch._assert( + isinstance(rel_codes, torch.Tensor), + "This function expects rel_codes of type torch.Tensor.", + ) + boxes_per_image = [b.size(0) for b in boxes] + concat_boxes = torch.cat(boxes, dim=0) + box_sum = 0 + for val in boxes_per_image: + box_sum += val + if box_sum > 0: + rel_codes = rel_codes.reshape(box_sum, -1) + pred_boxes = self.decode_single(rel_codes, concat_boxes) + if box_sum > 0: + pred_boxes = pred_boxes.reshape(box_sum, -1, 4) + return pred_boxes + + def decode_single(self, rel_codes: Tensor, boxes: Tensor) -> Tensor: + """ + From a set of original boxes and encoded relative box offsets, + get the decoded boxes. + + Args: + rel_codes (Tensor): encoded boxes + boxes (Tensor): reference boxes. + """ + + boxes = boxes.to(rel_codes.dtype) + + widths = boxes[:, 2] - boxes[:, 0] + heights = boxes[:, 3] - boxes[:, 1] + ctr_x = boxes[:, 0] + 0.5 * widths + ctr_y = boxes[:, 1] + 0.5 * heights + + wx, wy, ww, wh = self.weights + dx = rel_codes[:, 0::4] / wx + dy = rel_codes[:, 1::4] / wy + dw = rel_codes[:, 2::4] / ww + dh = rel_codes[:, 3::4] / wh + + # Prevent sending too large values into torch.exp() + dw = torch.clamp(dw, max=self.bbox_xform_clip) + dh = torch.clamp(dh, max=self.bbox_xform_clip) + + pred_ctr_x = dx * widths[:, None] + ctr_x[:, None] + pred_ctr_y = dy * heights[:, None] + ctr_y[:, None] + pred_w = torch.exp(dw) * widths[:, None] + pred_h = torch.exp(dh) * heights[:, None] + + # Distance from center to box's corner. + c_to_c_h = torch.tensor(0.5, dtype=pred_ctr_y.dtype, device=pred_h.device) * pred_h + c_to_c_w = torch.tensor(0.5, dtype=pred_ctr_x.dtype, device=pred_w.device) * pred_w + + pred_boxes1 = pred_ctr_x - c_to_c_w + pred_boxes2 = pred_ctr_y - c_to_c_h + pred_boxes3 = pred_ctr_x + c_to_c_w + pred_boxes4 = pred_ctr_y + c_to_c_h + pred_boxes = torch.stack((pred_boxes1, pred_boxes2, pred_boxes3, pred_boxes4), dim=2).flatten(1) + return pred_boxes + + +class BoxLinearCoder: + """ + The linear box-to-box transform defined in FCOS. The transformation is parameterized + by the distance from the center of (square) src box to 4 edges of the target box. + """ + + def __init__(self, normalize_by_size: bool = True) -> None: + """ + Args: + normalize_by_size (bool): normalize deltas by the size of src (anchor) boxes. + """ + self.normalize_by_size = normalize_by_size + + def encode(self, reference_boxes: Tensor, proposals: Tensor) -> Tensor: + """ + Encode a set of proposals with respect to some reference boxes + + Args: + reference_boxes (Tensor): reference boxes + proposals (Tensor): boxes to be encoded + + Returns: + Tensor: the encoded relative box offsets that can be used to + decode the boxes. + + """ + + # get the center of reference_boxes + reference_boxes_ctr_x = 0.5 * (reference_boxes[..., 0] + reference_boxes[..., 2]) + reference_boxes_ctr_y = 0.5 * (reference_boxes[..., 1] + reference_boxes[..., 3]) + + # get box regression transformation deltas + target_l = reference_boxes_ctr_x - proposals[..., 0] + target_t = reference_boxes_ctr_y - proposals[..., 1] + target_r = proposals[..., 2] - reference_boxes_ctr_x + target_b = proposals[..., 3] - reference_boxes_ctr_y + + targets = torch.stack((target_l, target_t, target_r, target_b), dim=-1) + + if self.normalize_by_size: + reference_boxes_w = reference_boxes[..., 2] - reference_boxes[..., 0] + reference_boxes_h = reference_boxes[..., 3] - reference_boxes[..., 1] + reference_boxes_size = torch.stack( + (reference_boxes_w, reference_boxes_h, reference_boxes_w, reference_boxes_h), dim=-1 + ) + targets = targets / reference_boxes_size + return targets + + def decode(self, rel_codes: Tensor, boxes: Tensor) -> Tensor: + + """ + From a set of original boxes and encoded relative box offsets, + get the decoded boxes. + + Args: + rel_codes (Tensor): encoded boxes + boxes (Tensor): reference boxes. + + Returns: + Tensor: the predicted boxes with the encoded relative box offsets. + + .. note:: + This method assumes that ``rel_codes`` and ``boxes`` have same size for 0th dimension. i.e. ``len(rel_codes) == len(boxes)``. + + """ + + boxes = boxes.to(dtype=rel_codes.dtype) + + ctr_x = 0.5 * (boxes[..., 0] + boxes[..., 2]) + ctr_y = 0.5 * (boxes[..., 1] + boxes[..., 3]) + + if self.normalize_by_size: + boxes_w = boxes[..., 2] - boxes[..., 0] + boxes_h = boxes[..., 3] - boxes[..., 1] + + list_box_size = torch.stack((boxes_w, boxes_h, boxes_w, boxes_h), dim=-1) + rel_codes = rel_codes * list_box_size + + pred_boxes1 = ctr_x - rel_codes[..., 0] + pred_boxes2 = ctr_y - rel_codes[..., 1] + pred_boxes3 = ctr_x + rel_codes[..., 2] + pred_boxes4 = ctr_y + rel_codes[..., 3] + + pred_boxes = torch.stack((pred_boxes1, pred_boxes2, pred_boxes3, pred_boxes4), dim=-1) + return pred_boxes + + +class Matcher: + """ + This class assigns to each predicted "element" (e.g., a box) a ground-truth + element. Each predicted element will have exactly zero or one matches; each + ground-truth element may be assigned to zero or more predicted elements. + + Matching is based on the MxN match_quality_matrix, that characterizes how well + each (ground-truth, predicted)-pair match. For example, if the elements are + boxes, the matrix may contain box IoU overlap values. + + The matcher returns a tensor of size N containing the index of the ground-truth + element m that matches to prediction n. If there is no match, a negative value + is returned. + """ + + BELOW_LOW_THRESHOLD = -1 + BETWEEN_THRESHOLDS = -2 + + __annotations__ = { + "BELOW_LOW_THRESHOLD": int, + "BETWEEN_THRESHOLDS": int, + } + + def __init__(self, high_threshold: float, low_threshold: float, allow_low_quality_matches: bool = False) -> None: + """ + Args: + high_threshold (float): quality values greater than or equal to + this value are candidate matches. + low_threshold (float): a lower quality threshold used to stratify + matches into three levels: + 1) matches >= high_threshold + 2) BETWEEN_THRESHOLDS matches in [low_threshold, high_threshold) + 3) BELOW_LOW_THRESHOLD matches in [0, low_threshold) + allow_low_quality_matches (bool): if True, produce additional matches + for predictions that have only low-quality match candidates. See + set_low_quality_matches_ for more details. + """ + self.BELOW_LOW_THRESHOLD = -1 + self.BETWEEN_THRESHOLDS = -2 + torch._assert(low_threshold <= high_threshold, "low_threshold should be <= high_threshold") + self.high_threshold = high_threshold + self.low_threshold = low_threshold + self.allow_low_quality_matches = allow_low_quality_matches + + def __call__(self, match_quality_matrix: Tensor) -> Tensor: + """ + Args: + match_quality_matrix (Tensor[float]): an MxN tensor, containing the + pairwise quality between M ground-truth elements and N predicted elements. + + Returns: + matches (Tensor[int64]): an N tensor where N[i] is a matched gt in + [0, M - 1] or a negative value indicating that prediction i could not + be matched. + """ + if match_quality_matrix.numel() == 0: + # empty targets or proposals not supported during training + if match_quality_matrix.shape[0] == 0: + raise ValueError("No ground-truth boxes available for one of the images during training") + else: + raise ValueError("No proposal boxes available for one of the images during training") + + # match_quality_matrix is M (gt) x N (predicted) + # Max over gt elements (dim 0) to find best gt candidate for each prediction + matched_vals, matches = match_quality_matrix.max(dim=0) + if self.allow_low_quality_matches: + all_matches = matches.clone() + else: + all_matches = None # type: ignore[assignment] + + # Assign candidate matches with low quality to negative (unassigned) values + below_low_threshold = matched_vals < self.low_threshold + between_thresholds = (matched_vals >= self.low_threshold) & (matched_vals < self.high_threshold) + matches[below_low_threshold] = self.BELOW_LOW_THRESHOLD + matches[between_thresholds] = self.BETWEEN_THRESHOLDS + + if self.allow_low_quality_matches: + if all_matches is None: + torch._assert(False, "all_matches should not be None") + else: + self.set_low_quality_matches_(matches, all_matches, match_quality_matrix) + + return matches + + def set_low_quality_matches_(self, matches: Tensor, all_matches: Tensor, match_quality_matrix: Tensor) -> None: + """ + Produce additional matches for predictions that have only low-quality matches. + Specifically, for each ground-truth find the set of predictions that have + maximum overlap with it (including ties); for each prediction in that set, if + it is unmatched, then match it to the ground-truth with which it has the highest + quality value. + """ + # For each gt, find the prediction with which it has the highest quality + highest_quality_foreach_gt, _ = match_quality_matrix.max(dim=1) + # Find the highest quality match available, even if it is low, including ties + gt_pred_pairs_of_highest_quality = torch.where(match_quality_matrix == highest_quality_foreach_gt[:, None]) + # Example gt_pred_pairs_of_highest_quality: + # (tensor([0, 1, 1, 2, 2, 3, 3, 4, 5, 5]), + # tensor([39796, 32055, 32070, 39190, 40255, 40390, 41455, 45470, 45325, 46390])) + # Each element in the first tensor is a gt index, and each element in second tensor is a prediction index + # Note how gt items 1, 2, 3, and 5 each have two ties + + pred_inds_to_update = gt_pred_pairs_of_highest_quality[1] + matches[pred_inds_to_update] = all_matches[pred_inds_to_update] + + +class SSDMatcher(Matcher): + def __init__(self, threshold: float) -> None: + super().__init__(threshold, threshold, allow_low_quality_matches=False) + + def __call__(self, match_quality_matrix: Tensor) -> Tensor: + matches = super().__call__(match_quality_matrix) + + # For each gt, find the prediction with which it has the highest quality + _, highest_quality_pred_foreach_gt = match_quality_matrix.max(dim=1) + matches[highest_quality_pred_foreach_gt] = torch.arange( + highest_quality_pred_foreach_gt.size(0), dtype=torch.int64, device=highest_quality_pred_foreach_gt.device + ) + + return matches + + +def overwrite_eps(model: nn.Module, eps: float) -> None: + """ + This method overwrites the default eps values of all the + FrozenBatchNorm2d layers of the model with the provided value. + This is necessary to address the BC-breaking change introduced + by the bug-fix at pytorch/vision#2933. The overwrite is applied + only when the pretrained weights are loaded to maintain compatibility + with previous versions. + + Args: + model (nn.Module): The model on which we perform the overwrite. + eps (float): The new value of eps. + """ + for module in model.modules(): + if isinstance(module, FrozenBatchNorm2d): + module.eps = eps + + +def retrieve_out_channels(model: nn.Module, size: Tuple[int, int]) -> List[int]: + """ + This method retrieves the number of output channels of a specific model. + + Args: + model (nn.Module): The model for which we estimate the out_channels. + It should return a single Tensor or an OrderedDict[Tensor]. + size (Tuple[int, int]): The size (wxh) of the input. + + Returns: + out_channels (List[int]): A list of the output channels of the model. + """ + in_training = model.training + model.eval() + + with torch.no_grad(): + # Use dummy data to retrieve the feature map sizes to avoid hard-coding their values + device = next(model.parameters()).device + tmp_img = torch.zeros((1, 3, size[1], size[0]), device=device) + features = model(tmp_img) + if isinstance(features, torch.Tensor): + features = OrderedDict([("0", features)]) + out_channels = [x.size(1) for x in features.values()] + + if in_training: + model.train() + + return out_channels + + +@torch.jit.unused +def _fake_cast_onnx(v: Tensor) -> int: + return v # type: ignore[return-value] + + +def _topk_min(input: Tensor, orig_kval: int, axis: int) -> int: + """ + ONNX spec requires the k-value to be less than or equal to the number of inputs along + provided dim. Certain models use the number of elements along a particular axis instead of K + if K exceeds the number of elements along that axis. Previously, python's min() function was + used to determine whether to use the provided k-value or the specified dim axis value. + + However, in cases where the model is being exported in tracing mode, python min() is + static causing the model to be traced incorrectly and eventually fail at the topk node. + In order to avoid this situation, in tracing mode, torch.min() is used instead. + + Args: + input (Tensor): The original input tensor. + orig_kval (int): The provided k-value. + axis(int): Axis along which we retrieve the input size. + + Returns: + min_kval (int): Appropriately selected k-value. + """ + if not torch.jit.is_tracing(): + return min(orig_kval, input.size(axis)) + axis_dim_val = torch._shape_as_tensor(input)[axis].unsqueeze(0) + min_kval = torch.min(torch.cat((torch.tensor([orig_kval], dtype=axis_dim_val.dtype), axis_dim_val), 0)) + return _fake_cast_onnx(min_kval) + + +def _box_loss( + type: str, + box_coder: BoxCoder, + anchors_per_image: Tensor, + matched_gt_boxes_per_image: Tensor, + bbox_regression_per_image: Tensor, + cnf: Optional[Dict[str, float]] = None, +) -> Tensor: + torch._assert(type in ["l1", "smooth_l1", "ciou", "diou", "giou"], f"Unsupported loss: {type}") + + if type == "l1": + target_regression = box_coder.encode_single(matched_gt_boxes_per_image, anchors_per_image) + return F.l1_loss(bbox_regression_per_image, target_regression, reduction="sum") + elif type == "smooth_l1": + target_regression = box_coder.encode_single(matched_gt_boxes_per_image, anchors_per_image) + beta = cnf["beta"] if cnf is not None and "beta" in cnf else 1.0 + return F.smooth_l1_loss(bbox_regression_per_image, target_regression, reduction="sum", beta=beta) + else: + bbox_per_image = box_coder.decode_single(bbox_regression_per_image, anchors_per_image) + eps = cnf["eps"] if cnf is not None and "eps" in cnf else 1e-7 + if type == "ciou": + return complete_box_iou_loss(bbox_per_image, matched_gt_boxes_per_image, reduction="sum", eps=eps) + if type == "diou": + return distance_box_iou_loss(bbox_per_image, matched_gt_boxes_per_image, reduction="sum", eps=eps) + # otherwise giou + return generalized_box_iou_loss(bbox_per_image, matched_gt_boxes_per_image, reduction="sum", eps=eps) diff --git a/lib/python3.12/site-packages/torchvision/models/detection/anchor_utils.py b/lib/python3.12/site-packages/torchvision/models/detection/anchor_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..253f6502a9b6344f5a3da239f2394179a256424e --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/detection/anchor_utils.py @@ -0,0 +1,268 @@ +import math +from typing import List, Optional + +import torch +from torch import nn, Tensor + +from .image_list import ImageList + + +class AnchorGenerator(nn.Module): + """ + Module that generates anchors for a set of feature maps and + image sizes. + + The module support computing anchors at multiple sizes and aspect ratios + per feature map. This module assumes aspect ratio = height / width for + each anchor. + + sizes and aspect_ratios should have the same number of elements, and it should + correspond to the number of feature maps. + + sizes[i] and aspect_ratios[i] can have an arbitrary number of elements, + and AnchorGenerator will output a set of sizes[i] * aspect_ratios[i] anchors + per spatial location for feature map i. + + Args: + sizes (Tuple[Tuple[int]]): + aspect_ratios (Tuple[Tuple[float]]): + """ + + __annotations__ = { + "cell_anchors": List[torch.Tensor], + } + + def __init__( + self, + sizes=((128, 256, 512),), + aspect_ratios=((0.5, 1.0, 2.0),), + ): + super().__init__() + + if not isinstance(sizes[0], (list, tuple)): + # TODO change this + sizes = tuple((s,) for s in sizes) + if not isinstance(aspect_ratios[0], (list, tuple)): + aspect_ratios = (aspect_ratios,) * len(sizes) + + self.sizes = sizes + self.aspect_ratios = aspect_ratios + self.cell_anchors = [ + self.generate_anchors(size, aspect_ratio) for size, aspect_ratio in zip(sizes, aspect_ratios) + ] + + # TODO: https://github.com/pytorch/pytorch/issues/26792 + # For every (aspect_ratios, scales) combination, output a zero-centered anchor with those values. + # (scales, aspect_ratios) are usually an element of zip(self.scales, self.aspect_ratios) + # This method assumes aspect ratio = height / width for an anchor. + def generate_anchors( + self, + scales: List[int], + aspect_ratios: List[float], + dtype: torch.dtype = torch.float32, + device: torch.device = torch.device("cpu"), + ) -> Tensor: + scales = torch.as_tensor(scales, dtype=dtype, device=device) + aspect_ratios = torch.as_tensor(aspect_ratios, dtype=dtype, device=device) + h_ratios = torch.sqrt(aspect_ratios) + w_ratios = 1 / h_ratios + + ws = (w_ratios[:, None] * scales[None, :]).view(-1) + hs = (h_ratios[:, None] * scales[None, :]).view(-1) + + base_anchors = torch.stack([-ws, -hs, ws, hs], dim=1) / 2 + return base_anchors.round() + + def set_cell_anchors(self, dtype: torch.dtype, device: torch.device): + self.cell_anchors = [cell_anchor.to(dtype=dtype, device=device) for cell_anchor in self.cell_anchors] + + def num_anchors_per_location(self) -> List[int]: + return [len(s) * len(a) for s, a in zip(self.sizes, self.aspect_ratios)] + + # For every combination of (a, (g, s), i) in (self.cell_anchors, zip(grid_sizes, strides), 0:2), + # output g[i] anchors that are s[i] distance apart in direction i, with the same dimensions as a. + def grid_anchors(self, grid_sizes: List[List[int]], strides: List[List[Tensor]]) -> List[Tensor]: + anchors = [] + cell_anchors = self.cell_anchors + torch._assert(cell_anchors is not None, "cell_anchors should not be None") + torch._assert( + len(grid_sizes) == len(strides) == len(cell_anchors), + "Anchors should be Tuple[Tuple[int]] because each feature " + "map could potentially have different sizes and aspect ratios. " + "There needs to be a match between the number of " + "feature maps passed and the number of sizes / aspect ratios specified.", + ) + + for size, stride, base_anchors in zip(grid_sizes, strides, cell_anchors): + grid_height, grid_width = size + stride_height, stride_width = stride + device = base_anchors.device + + # For output anchor, compute [x_center, y_center, x_center, y_center] + shifts_x = torch.arange(0, grid_width, dtype=torch.int32, device=device) * stride_width + shifts_y = torch.arange(0, grid_height, dtype=torch.int32, device=device) * stride_height + shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x, indexing="ij") + shift_x = shift_x.reshape(-1) + shift_y = shift_y.reshape(-1) + shifts = torch.stack((shift_x, shift_y, shift_x, shift_y), dim=1) + + # For every (base anchor, output anchor) pair, + # offset each zero-centered base anchor by the center of the output anchor. + anchors.append((shifts.view(-1, 1, 4) + base_anchors.view(1, -1, 4)).reshape(-1, 4)) + + return anchors + + def forward(self, image_list: ImageList, feature_maps: List[Tensor]) -> List[Tensor]: + grid_sizes = [feature_map.shape[-2:] for feature_map in feature_maps] + image_size = image_list.tensors.shape[-2:] + dtype, device = feature_maps[0].dtype, feature_maps[0].device + strides = [ + [ + torch.empty((), dtype=torch.int64, device=device).fill_(image_size[0] // g[0]), + torch.empty((), dtype=torch.int64, device=device).fill_(image_size[1] // g[1]), + ] + for g in grid_sizes + ] + self.set_cell_anchors(dtype, device) + anchors_over_all_feature_maps = self.grid_anchors(grid_sizes, strides) + anchors: List[List[torch.Tensor]] = [] + for _ in range(len(image_list.image_sizes)): + anchors_in_image = [anchors_per_feature_map for anchors_per_feature_map in anchors_over_all_feature_maps] + anchors.append(anchors_in_image) + anchors = [torch.cat(anchors_per_image) for anchors_per_image in anchors] + return anchors + + +class DefaultBoxGenerator(nn.Module): + """ + This module generates the default boxes of SSD for a set of feature maps and image sizes. + + Args: + aspect_ratios (List[List[int]]): A list with all the aspect ratios used in each feature map. + min_ratio (float): The minimum scale :math:`\text{s}_{\text{min}}` of the default boxes used in the estimation + of the scales of each feature map. It is used only if the ``scales`` parameter is not provided. + max_ratio (float): The maximum scale :math:`\text{s}_{\text{max}}` of the default boxes used in the estimation + of the scales of each feature map. It is used only if the ``scales`` parameter is not provided. + scales (List[float]], optional): The scales of the default boxes. If not provided it will be estimated using + the ``min_ratio`` and ``max_ratio`` parameters. + steps (List[int]], optional): It's a hyper-parameter that affects the tiling of default boxes. If not provided + it will be estimated from the data. + clip (bool): Whether the standardized values of default boxes should be clipped between 0 and 1. The clipping + is applied while the boxes are encoded in format ``(cx, cy, w, h)``. + """ + + def __init__( + self, + aspect_ratios: List[List[int]], + min_ratio: float = 0.15, + max_ratio: float = 0.9, + scales: Optional[List[float]] = None, + steps: Optional[List[int]] = None, + clip: bool = True, + ): + super().__init__() + if steps is not None and len(aspect_ratios) != len(steps): + raise ValueError("aspect_ratios and steps should have the same length") + self.aspect_ratios = aspect_ratios + self.steps = steps + self.clip = clip + num_outputs = len(aspect_ratios) + + # Estimation of default boxes scales + if scales is None: + if num_outputs > 1: + range_ratio = max_ratio - min_ratio + self.scales = [min_ratio + range_ratio * k / (num_outputs - 1.0) for k in range(num_outputs)] + self.scales.append(1.0) + else: + self.scales = [min_ratio, max_ratio] + else: + self.scales = scales + + self._wh_pairs = self._generate_wh_pairs(num_outputs) + + def _generate_wh_pairs( + self, num_outputs: int, dtype: torch.dtype = torch.float32, device: torch.device = torch.device("cpu") + ) -> List[Tensor]: + _wh_pairs: List[Tensor] = [] + for k in range(num_outputs): + # Adding the 2 default width-height pairs for aspect ratio 1 and scale s'k + s_k = self.scales[k] + s_prime_k = math.sqrt(self.scales[k] * self.scales[k + 1]) + wh_pairs = [[s_k, s_k], [s_prime_k, s_prime_k]] + + # Adding 2 pairs for each aspect ratio of the feature map k + for ar in self.aspect_ratios[k]: + sq_ar = math.sqrt(ar) + w = self.scales[k] * sq_ar + h = self.scales[k] / sq_ar + wh_pairs.extend([[w, h], [h, w]]) + + _wh_pairs.append(torch.as_tensor(wh_pairs, dtype=dtype, device=device)) + return _wh_pairs + + def num_anchors_per_location(self) -> List[int]: + # Estimate num of anchors based on aspect ratios: 2 default boxes + 2 * ratios of feaure map. + return [2 + 2 * len(r) for r in self.aspect_ratios] + + # Default Boxes calculation based on page 6 of SSD paper + def _grid_default_boxes( + self, grid_sizes: List[List[int]], image_size: List[int], dtype: torch.dtype = torch.float32 + ) -> Tensor: + default_boxes = [] + for k, f_k in enumerate(grid_sizes): + # Now add the default boxes for each width-height pair + if self.steps is not None: + x_f_k = image_size[1] / self.steps[k] + y_f_k = image_size[0] / self.steps[k] + else: + y_f_k, x_f_k = f_k + + shifts_x = ((torch.arange(0, f_k[1]) + 0.5) / x_f_k).to(dtype=dtype) + shifts_y = ((torch.arange(0, f_k[0]) + 0.5) / y_f_k).to(dtype=dtype) + shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x, indexing="ij") + shift_x = shift_x.reshape(-1) + shift_y = shift_y.reshape(-1) + + shifts = torch.stack((shift_x, shift_y) * len(self._wh_pairs[k]), dim=-1).reshape(-1, 2) + # Clipping the default boxes while the boxes are encoded in format (cx, cy, w, h) + _wh_pair = self._wh_pairs[k].clamp(min=0, max=1) if self.clip else self._wh_pairs[k] + wh_pairs = _wh_pair.repeat((f_k[0] * f_k[1]), 1) + + default_box = torch.cat((shifts, wh_pairs), dim=1) + + default_boxes.append(default_box) + + return torch.cat(default_boxes, dim=0) + + def __repr__(self) -> str: + s = ( + f"{self.__class__.__name__}(" + f"aspect_ratios={self.aspect_ratios}" + f", clip={self.clip}" + f", scales={self.scales}" + f", steps={self.steps}" + ")" + ) + return s + + def forward(self, image_list: ImageList, feature_maps: List[Tensor]) -> List[Tensor]: + grid_sizes = [feature_map.shape[-2:] for feature_map in feature_maps] + image_size = image_list.tensors.shape[-2:] + dtype, device = feature_maps[0].dtype, feature_maps[0].device + default_boxes = self._grid_default_boxes(grid_sizes, image_size, dtype=dtype) + default_boxes = default_boxes.to(device) + + dboxes = [] + x_y_size = torch.tensor([image_size[1], image_size[0]], device=default_boxes.device) + for _ in image_list.image_sizes: + dboxes_in_image = default_boxes + dboxes_in_image = torch.cat( + [ + (dboxes_in_image[:, :2] - 0.5 * dboxes_in_image[:, 2:]) * x_y_size, + (dboxes_in_image[:, :2] + 0.5 * dboxes_in_image[:, 2:]) * x_y_size, + ], + -1, + ) + dboxes.append(dboxes_in_image) + return dboxes diff --git a/lib/python3.12/site-packages/torchvision/models/detection/backbone_utils.py b/lib/python3.12/site-packages/torchvision/models/detection/backbone_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..668e6b31696eb949513d07878eada9d468dc99cd --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/detection/backbone_utils.py @@ -0,0 +1,244 @@ +import warnings +from typing import Callable, Dict, List, Optional, Union + +from torch import nn, Tensor +from torchvision.ops import misc as misc_nn_ops +from torchvision.ops.feature_pyramid_network import ExtraFPNBlock, FeaturePyramidNetwork, LastLevelMaxPool + +from .. import mobilenet, resnet +from .._api import _get_enum_from_fn, WeightsEnum +from .._utils import handle_legacy_interface, IntermediateLayerGetter + + +class BackboneWithFPN(nn.Module): + """ + Adds a FPN on top of a model. + Internally, it uses torchvision.models._utils.IntermediateLayerGetter to + extract a submodel that returns the feature maps specified in return_layers. + The same limitations of IntermediateLayerGetter apply here. + Args: + backbone (nn.Module) + return_layers (Dict[name, new_name]): a dict containing the names + of the modules for which the activations will be returned as + the key of the dict, and the value of the dict is the name + of the returned activation (which the user can specify). + in_channels_list (List[int]): number of channels for each feature map + that is returned, in the order they are present in the OrderedDict + out_channels (int): number of channels in the FPN. + norm_layer (callable, optional): Module specifying the normalization layer to use. Default: None + Attributes: + out_channels (int): the number of channels in the FPN + """ + + def __init__( + self, + backbone: nn.Module, + return_layers: Dict[str, str], + in_channels_list: List[int], + out_channels: int, + extra_blocks: Optional[ExtraFPNBlock] = None, + norm_layer: Optional[Callable[..., nn.Module]] = None, + ) -> None: + super().__init__() + + if extra_blocks is None: + extra_blocks = LastLevelMaxPool() + + self.body = IntermediateLayerGetter(backbone, return_layers=return_layers) + self.fpn = FeaturePyramidNetwork( + in_channels_list=in_channels_list, + out_channels=out_channels, + extra_blocks=extra_blocks, + norm_layer=norm_layer, + ) + self.out_channels = out_channels + + def forward(self, x: Tensor) -> Dict[str, Tensor]: + x = self.body(x) + x = self.fpn(x) + return x + + +@handle_legacy_interface( + weights=( + "pretrained", + lambda kwargs: _get_enum_from_fn(resnet.__dict__[kwargs["backbone_name"]])["IMAGENET1K_V1"], + ), +) +def resnet_fpn_backbone( + *, + backbone_name: str, + weights: Optional[WeightsEnum], + norm_layer: Callable[..., nn.Module] = misc_nn_ops.FrozenBatchNorm2d, + trainable_layers: int = 3, + returned_layers: Optional[List[int]] = None, + extra_blocks: Optional[ExtraFPNBlock] = None, +) -> BackboneWithFPN: + """ + Constructs a specified ResNet backbone with FPN on top. Freezes the specified number of layers in the backbone. + + Examples:: + + >>> import torch + >>> from torchvision.models import ResNet50_Weights + >>> from torchvision.models.detection.backbone_utils import resnet_fpn_backbone + >>> backbone = resnet_fpn_backbone(backbone_name='resnet50', weights=ResNet50_Weights.DEFAULT, trainable_layers=3) + >>> # get some dummy image + >>> x = torch.rand(1,3,64,64) + >>> # compute the output + >>> output = backbone(x) + >>> print([(k, v.shape) for k, v in output.items()]) + >>> # returns + >>> [('0', torch.Size([1, 256, 16, 16])), + >>> ('1', torch.Size([1, 256, 8, 8])), + >>> ('2', torch.Size([1, 256, 4, 4])), + >>> ('3', torch.Size([1, 256, 2, 2])), + >>> ('pool', torch.Size([1, 256, 1, 1]))] + + Args: + backbone_name (string): resnet architecture. Possible values are 'resnet18', 'resnet34', 'resnet50', + 'resnet101', 'resnet152', 'resnext50_32x4d', 'resnext101_32x8d', 'wide_resnet50_2', 'wide_resnet101_2' + weights (WeightsEnum, optional): The pretrained weights for the model + norm_layer (callable): it is recommended to use the default value. For details visit: + (https://github.com/facebookresearch/maskrcnn-benchmark/issues/267) + trainable_layers (int): number of trainable (not frozen) layers starting from final block. + Valid values are between 0 and 5, with 5 meaning all backbone layers are trainable. + returned_layers (list of int): The layers of the network to return. Each entry must be in ``[1, 4]``. + By default, all layers are returned. + extra_blocks (ExtraFPNBlock or None): if provided, extra operations will + be performed. It is expected to take the fpn features, the original + features and the names of the original features as input, and returns + a new list of feature maps and their corresponding names. By + default, a ``LastLevelMaxPool`` is used. + """ + backbone = resnet.__dict__[backbone_name](weights=weights, norm_layer=norm_layer) + return _resnet_fpn_extractor(backbone, trainable_layers, returned_layers, extra_blocks) + + +def _resnet_fpn_extractor( + backbone: resnet.ResNet, + trainable_layers: int, + returned_layers: Optional[List[int]] = None, + extra_blocks: Optional[ExtraFPNBlock] = None, + norm_layer: Optional[Callable[..., nn.Module]] = None, +) -> BackboneWithFPN: + + # select layers that won't be frozen + if trainable_layers < 0 or trainable_layers > 5: + raise ValueError(f"Trainable layers should be in the range [0,5], got {trainable_layers}") + layers_to_train = ["layer4", "layer3", "layer2", "layer1", "conv1"][:trainable_layers] + if trainable_layers == 5: + layers_to_train.append("bn1") + for name, parameter in backbone.named_parameters(): + if all([not name.startswith(layer) for layer in layers_to_train]): + parameter.requires_grad_(False) + + if extra_blocks is None: + extra_blocks = LastLevelMaxPool() + + if returned_layers is None: + returned_layers = [1, 2, 3, 4] + if min(returned_layers) <= 0 or max(returned_layers) >= 5: + raise ValueError(f"Each returned layer should be in the range [1,4]. Got {returned_layers}") + return_layers = {f"layer{k}": str(v) for v, k in enumerate(returned_layers)} + + in_channels_stage2 = backbone.inplanes // 8 + in_channels_list = [in_channels_stage2 * 2 ** (i - 1) for i in returned_layers] + out_channels = 256 + return BackboneWithFPN( + backbone, return_layers, in_channels_list, out_channels, extra_blocks=extra_blocks, norm_layer=norm_layer + ) + + +def _validate_trainable_layers( + is_trained: bool, + trainable_backbone_layers: Optional[int], + max_value: int, + default_value: int, +) -> int: + # don't freeze any layers if pretrained model or backbone is not used + if not is_trained: + if trainable_backbone_layers is not None: + warnings.warn( + "Changing trainable_backbone_layers has no effect if " + "neither pretrained nor pretrained_backbone have been set to True, " + f"falling back to trainable_backbone_layers={max_value} so that all layers are trainable" + ) + trainable_backbone_layers = max_value + + # by default freeze first blocks + if trainable_backbone_layers is None: + trainable_backbone_layers = default_value + if trainable_backbone_layers < 0 or trainable_backbone_layers > max_value: + raise ValueError( + f"Trainable backbone layers should be in the range [0,{max_value}], got {trainable_backbone_layers} " + ) + return trainable_backbone_layers + + +@handle_legacy_interface( + weights=( + "pretrained", + lambda kwargs: _get_enum_from_fn(mobilenet.__dict__[kwargs["backbone_name"]])["IMAGENET1K_V1"], + ), +) +def mobilenet_backbone( + *, + backbone_name: str, + weights: Optional[WeightsEnum], + fpn: bool, + norm_layer: Callable[..., nn.Module] = misc_nn_ops.FrozenBatchNorm2d, + trainable_layers: int = 2, + returned_layers: Optional[List[int]] = None, + extra_blocks: Optional[ExtraFPNBlock] = None, +) -> nn.Module: + backbone = mobilenet.__dict__[backbone_name](weights=weights, norm_layer=norm_layer) + return _mobilenet_extractor(backbone, fpn, trainable_layers, returned_layers, extra_blocks) + + +def _mobilenet_extractor( + backbone: Union[mobilenet.MobileNetV2, mobilenet.MobileNetV3], + fpn: bool, + trainable_layers: int, + returned_layers: Optional[List[int]] = None, + extra_blocks: Optional[ExtraFPNBlock] = None, + norm_layer: Optional[Callable[..., nn.Module]] = None, +) -> nn.Module: + backbone = backbone.features + # Gather the indices of blocks which are strided. These are the locations of C1, ..., Cn-1 blocks. + # The first and last blocks are always included because they are the C0 (conv1) and Cn. + stage_indices = [0] + [i for i, b in enumerate(backbone) if getattr(b, "_is_cn", False)] + [len(backbone) - 1] + num_stages = len(stage_indices) + + # find the index of the layer from which we won't freeze + if trainable_layers < 0 or trainable_layers > num_stages: + raise ValueError(f"Trainable layers should be in the range [0,{num_stages}], got {trainable_layers} ") + freeze_before = len(backbone) if trainable_layers == 0 else stage_indices[num_stages - trainable_layers] + + for b in backbone[:freeze_before]: + for parameter in b.parameters(): + parameter.requires_grad_(False) + + out_channels = 256 + if fpn: + if extra_blocks is None: + extra_blocks = LastLevelMaxPool() + + if returned_layers is None: + returned_layers = [num_stages - 2, num_stages - 1] + if min(returned_layers) < 0 or max(returned_layers) >= num_stages: + raise ValueError(f"Each returned layer should be in the range [0,{num_stages - 1}], got {returned_layers} ") + return_layers = {f"{stage_indices[k]}": str(v) for v, k in enumerate(returned_layers)} + + in_channels_list = [backbone[stage_indices[i]].out_channels for i in returned_layers] + return BackboneWithFPN( + backbone, return_layers, in_channels_list, out_channels, extra_blocks=extra_blocks, norm_layer=norm_layer + ) + else: + m = nn.Sequential( + backbone, + # depthwise linear combination of channels to reduce their size + nn.Conv2d(backbone[-1].out_channels, out_channels, 1), + ) + m.out_channels = out_channels # type: ignore[assignment] + return m diff --git a/lib/python3.12/site-packages/torchvision/models/detection/faster_rcnn.py b/lib/python3.12/site-packages/torchvision/models/detection/faster_rcnn.py new file mode 100644 index 0000000000000000000000000000000000000000..18474ee84f4539cfec99d24534acb1e1e74a14b3 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/detection/faster_rcnn.py @@ -0,0 +1,846 @@ +from typing import Any, Callable, List, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +from torch import nn +from torchvision.ops import MultiScaleRoIAlign + +from ...ops import misc as misc_nn_ops +from ...transforms._presets import ObjectDetection +from .._api import register_model, Weights, WeightsEnum +from .._meta import _COCO_CATEGORIES +from .._utils import _ovewrite_value_param, handle_legacy_interface +from ..mobilenetv3 import mobilenet_v3_large, MobileNet_V3_Large_Weights +from ..resnet import resnet50, ResNet50_Weights +from ._utils import overwrite_eps +from .anchor_utils import AnchorGenerator +from .backbone_utils import _mobilenet_extractor, _resnet_fpn_extractor, _validate_trainable_layers +from .generalized_rcnn import GeneralizedRCNN +from .roi_heads import RoIHeads +from .rpn import RegionProposalNetwork, RPNHead +from .transform import GeneralizedRCNNTransform + + +__all__ = [ + "FasterRCNN", + "FasterRCNN_ResNet50_FPN_Weights", + "FasterRCNN_ResNet50_FPN_V2_Weights", + "FasterRCNN_MobileNet_V3_Large_FPN_Weights", + "FasterRCNN_MobileNet_V3_Large_320_FPN_Weights", + "fasterrcnn_resnet50_fpn", + "fasterrcnn_resnet50_fpn_v2", + "fasterrcnn_mobilenet_v3_large_fpn", + "fasterrcnn_mobilenet_v3_large_320_fpn", +] + + +def _default_anchorgen(): + anchor_sizes = ((32,), (64,), (128,), (256,), (512,)) + aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes) + return AnchorGenerator(anchor_sizes, aspect_ratios) + + +class FasterRCNN(GeneralizedRCNN): + """ + Implements Faster R-CNN. + + The input to the model is expected to be a list of tensors, each of shape [C, H, W], one for each + image, and should be in 0-1 range. Different images can have different sizes. + + The behavior of the model changes depending on if it is in training or evaluation mode. + + During training, the model expects both the input tensors and targets (list of dictionary), + containing: + - boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (Int64Tensor[N]): the class label for each ground-truth box + + The model returns a Dict[Tensor] during training, containing the classification and regression + losses for both the RPN and the R-CNN. + + During inference, the model requires only the input tensors, and returns the post-processed + predictions as a List[Dict[Tensor]], one for each input image. The fields of the Dict are as + follows: + - boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (Int64Tensor[N]): the predicted labels for each image + - scores (Tensor[N]): the scores or each prediction + + Args: + backbone (nn.Module): the network used to compute the features for the model. + It should contain an out_channels attribute, which indicates the number of output + channels that each feature map has (and it should be the same for all feature maps). + The backbone should return a single Tensor or and OrderedDict[Tensor]. + num_classes (int): number of output classes of the model (including the background). + If box_predictor is specified, num_classes should be None. + min_size (int): Images are rescaled before feeding them to the backbone: + we attempt to preserve the aspect ratio and scale the shorter edge + to ``min_size``. If the resulting longer edge exceeds ``max_size``, + then downscale so that the longer edge does not exceed ``max_size``. + This may result in the shorter edge beeing lower than ``min_size``. + max_size (int): See ``min_size``. + image_mean (Tuple[float, float, float]): mean values used for input normalization. + They are generally the mean values of the dataset on which the backbone has been trained + on + image_std (Tuple[float, float, float]): std values used for input normalization. + They are generally the std values of the dataset on which the backbone has been trained on + rpn_anchor_generator (AnchorGenerator): module that generates the anchors for a set of feature + maps. + rpn_head (nn.Module): module that computes the objectness and regression deltas from the RPN + rpn_pre_nms_top_n_train (int): number of proposals to keep before applying NMS during training + rpn_pre_nms_top_n_test (int): number of proposals to keep before applying NMS during testing + rpn_post_nms_top_n_train (int): number of proposals to keep after applying NMS during training + rpn_post_nms_top_n_test (int): number of proposals to keep after applying NMS during testing + rpn_nms_thresh (float): NMS threshold used for postprocessing the RPN proposals + rpn_fg_iou_thresh (float): minimum IoU between the anchor and the GT box so that they can be + considered as positive during training of the RPN. + rpn_bg_iou_thresh (float): maximum IoU between the anchor and the GT box so that they can be + considered as negative during training of the RPN. + rpn_batch_size_per_image (int): number of anchors that are sampled during training of the RPN + for computing the loss + rpn_positive_fraction (float): proportion of positive anchors in a mini-batch during training + of the RPN + rpn_score_thresh (float): only return proposals with an objectness score greater than rpn_score_thresh + box_roi_pool (MultiScaleRoIAlign): the module which crops and resizes the feature maps in + the locations indicated by the bounding boxes + box_head (nn.Module): module that takes the cropped feature maps as input + box_predictor (nn.Module): module that takes the output of box_head and returns the + classification logits and box regression deltas. + box_score_thresh (float): during inference, only return proposals with a classification score + greater than box_score_thresh + box_nms_thresh (float): NMS threshold for the prediction head. Used during inference + box_detections_per_img (int): maximum number of detections per image, for all classes. + box_fg_iou_thresh (float): minimum IoU between the proposals and the GT box so that they can be + considered as positive during training of the classification head + box_bg_iou_thresh (float): maximum IoU between the proposals and the GT box so that they can be + considered as negative during training of the classification head + box_batch_size_per_image (int): number of proposals that are sampled during training of the + classification head + box_positive_fraction (float): proportion of positive proposals in a mini-batch during training + of the classification head + bbox_reg_weights (Tuple[float, float, float, float]): weights for the encoding/decoding of the + bounding boxes + + Example:: + + >>> import torch + >>> import torchvision + >>> from torchvision.models.detection import FasterRCNN + >>> from torchvision.models.detection.rpn import AnchorGenerator + >>> # load a pre-trained model for classification and return + >>> # only the features + >>> backbone = torchvision.models.mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT).features + >>> # FasterRCNN needs to know the number of + >>> # output channels in a backbone. For mobilenet_v2, it's 1280, + >>> # so we need to add it here + >>> backbone.out_channels = 1280 + >>> + >>> # let's make the RPN generate 5 x 3 anchors per spatial + >>> # location, with 5 different sizes and 3 different aspect + >>> # ratios. We have a Tuple[Tuple[int]] because each feature + >>> # map could potentially have different sizes and + >>> # aspect ratios + >>> anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),), + >>> aspect_ratios=((0.5, 1.0, 2.0),)) + >>> + >>> # let's define what are the feature maps that we will + >>> # use to perform the region of interest cropping, as well as + >>> # the size of the crop after rescaling. + >>> # if your backbone returns a Tensor, featmap_names is expected to + >>> # be ['0']. More generally, the backbone should return an + >>> # OrderedDict[Tensor], and in featmap_names you can choose which + >>> # feature maps to use. + >>> roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], + >>> output_size=7, + >>> sampling_ratio=2) + >>> + >>> # put the pieces together inside a FasterRCNN model + >>> model = FasterRCNN(backbone, + >>> num_classes=2, + >>> rpn_anchor_generator=anchor_generator, + >>> box_roi_pool=roi_pooler) + >>> model.eval() + >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)] + >>> predictions = model(x) + """ + + def __init__( + self, + backbone, + num_classes=None, + # transform parameters + min_size=800, + max_size=1333, + image_mean=None, + image_std=None, + # RPN parameters + rpn_anchor_generator=None, + rpn_head=None, + rpn_pre_nms_top_n_train=2000, + rpn_pre_nms_top_n_test=1000, + rpn_post_nms_top_n_train=2000, + rpn_post_nms_top_n_test=1000, + rpn_nms_thresh=0.7, + rpn_fg_iou_thresh=0.7, + rpn_bg_iou_thresh=0.3, + rpn_batch_size_per_image=256, + rpn_positive_fraction=0.5, + rpn_score_thresh=0.0, + # Box parameters + box_roi_pool=None, + box_head=None, + box_predictor=None, + box_score_thresh=0.05, + box_nms_thresh=0.5, + box_detections_per_img=100, + box_fg_iou_thresh=0.5, + box_bg_iou_thresh=0.5, + box_batch_size_per_image=512, + box_positive_fraction=0.25, + bbox_reg_weights=None, + **kwargs, + ): + + if not hasattr(backbone, "out_channels"): + raise ValueError( + "backbone should contain an attribute out_channels " + "specifying the number of output channels (assumed to be the " + "same for all the levels)" + ) + + if not isinstance(rpn_anchor_generator, (AnchorGenerator, type(None))): + raise TypeError( + f"rpn_anchor_generator should be of type AnchorGenerator or None instead of {type(rpn_anchor_generator)}" + ) + if not isinstance(box_roi_pool, (MultiScaleRoIAlign, type(None))): + raise TypeError( + f"box_roi_pool should be of type MultiScaleRoIAlign or None instead of {type(box_roi_pool)}" + ) + + if num_classes is not None: + if box_predictor is not None: + raise ValueError("num_classes should be None when box_predictor is specified") + else: + if box_predictor is None: + raise ValueError("num_classes should not be None when box_predictor is not specified") + + out_channels = backbone.out_channels + + if rpn_anchor_generator is None: + rpn_anchor_generator = _default_anchorgen() + if rpn_head is None: + rpn_head = RPNHead(out_channels, rpn_anchor_generator.num_anchors_per_location()[0]) + + rpn_pre_nms_top_n = dict(training=rpn_pre_nms_top_n_train, testing=rpn_pre_nms_top_n_test) + rpn_post_nms_top_n = dict(training=rpn_post_nms_top_n_train, testing=rpn_post_nms_top_n_test) + + rpn = RegionProposalNetwork( + rpn_anchor_generator, + rpn_head, + rpn_fg_iou_thresh, + rpn_bg_iou_thresh, + rpn_batch_size_per_image, + rpn_positive_fraction, + rpn_pre_nms_top_n, + rpn_post_nms_top_n, + rpn_nms_thresh, + score_thresh=rpn_score_thresh, + ) + + if box_roi_pool is None: + box_roi_pool = MultiScaleRoIAlign(featmap_names=["0", "1", "2", "3"], output_size=7, sampling_ratio=2) + + if box_head is None: + resolution = box_roi_pool.output_size[0] + representation_size = 1024 + box_head = TwoMLPHead(out_channels * resolution**2, representation_size) + + if box_predictor is None: + representation_size = 1024 + box_predictor = FastRCNNPredictor(representation_size, num_classes) + + roi_heads = RoIHeads( + # Box + box_roi_pool, + box_head, + box_predictor, + box_fg_iou_thresh, + box_bg_iou_thresh, + box_batch_size_per_image, + box_positive_fraction, + bbox_reg_weights, + box_score_thresh, + box_nms_thresh, + box_detections_per_img, + ) + + if image_mean is None: + image_mean = [0.485, 0.456, 0.406] + if image_std is None: + image_std = [0.229, 0.224, 0.225] + transform = GeneralizedRCNNTransform(min_size, max_size, image_mean, image_std, **kwargs) + + super().__init__(backbone, rpn, roi_heads, transform) + + +class TwoMLPHead(nn.Module): + """ + Standard heads for FPN-based models + + Args: + in_channels (int): number of input channels + representation_size (int): size of the intermediate representation + """ + + def __init__(self, in_channels, representation_size): + super().__init__() + + self.fc6 = nn.Linear(in_channels, representation_size) + self.fc7 = nn.Linear(representation_size, representation_size) + + def forward(self, x): + x = x.flatten(start_dim=1) + + x = F.relu(self.fc6(x)) + x = F.relu(self.fc7(x)) + + return x + + +class FastRCNNConvFCHead(nn.Sequential): + def __init__( + self, + input_size: Tuple[int, int, int], + conv_layers: List[int], + fc_layers: List[int], + norm_layer: Optional[Callable[..., nn.Module]] = None, + ): + """ + Args: + input_size (Tuple[int, int, int]): the input size in CHW format. + conv_layers (list): feature dimensions of each Convolution layer + fc_layers (list): feature dimensions of each FCN layer + norm_layer (callable, optional): Module specifying the normalization layer to use. Default: None + """ + in_channels, in_height, in_width = input_size + + blocks = [] + previous_channels = in_channels + for current_channels in conv_layers: + blocks.append(misc_nn_ops.Conv2dNormActivation(previous_channels, current_channels, norm_layer=norm_layer)) + previous_channels = current_channels + blocks.append(nn.Flatten()) + previous_channels = previous_channels * in_height * in_width + for current_channels in fc_layers: + blocks.append(nn.Linear(previous_channels, current_channels)) + blocks.append(nn.ReLU(inplace=True)) + previous_channels = current_channels + + super().__init__(*blocks) + for layer in self.modules(): + if isinstance(layer, nn.Conv2d): + nn.init.kaiming_normal_(layer.weight, mode="fan_out", nonlinearity="relu") + if layer.bias is not None: + nn.init.zeros_(layer.bias) + + +class FastRCNNPredictor(nn.Module): + """ + Standard classification + bounding box regression layers + for Fast R-CNN. + + Args: + in_channels (int): number of input channels + num_classes (int): number of output classes (including background) + """ + + def __init__(self, in_channels, num_classes): + super().__init__() + self.cls_score = nn.Linear(in_channels, num_classes) + self.bbox_pred = nn.Linear(in_channels, num_classes * 4) + + def forward(self, x): + if x.dim() == 4: + torch._assert( + list(x.shape[2:]) == [1, 1], + f"x has the wrong shape, expecting the last two dimensions to be [1,1] instead of {list(x.shape[2:])}", + ) + x = x.flatten(start_dim=1) + scores = self.cls_score(x) + bbox_deltas = self.bbox_pred(x) + + return scores, bbox_deltas + + +_COMMON_META = { + "categories": _COCO_CATEGORIES, + "min_size": (1, 1), +} + + +class FasterRCNN_ResNet50_FPN_Weights(WeightsEnum): + COCO_V1 = Weights( + url="https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth", + transforms=ObjectDetection, + meta={ + **_COMMON_META, + "num_params": 41755286, + "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#faster-r-cnn-resnet-50-fpn", + "_metrics": { + "COCO-val2017": { + "box_map": 37.0, + } + }, + "_ops": 134.38, + "_file_size": 159.743, + "_docs": """These weights were produced by following a similar training recipe as on the paper.""", + }, + ) + DEFAULT = COCO_V1 + + +class FasterRCNN_ResNet50_FPN_V2_Weights(WeightsEnum): + COCO_V1 = Weights( + url="https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_v2_coco-dd69338a.pth", + transforms=ObjectDetection, + meta={ + **_COMMON_META, + "num_params": 43712278, + "recipe": "https://github.com/pytorch/vision/pull/5763", + "_metrics": { + "COCO-val2017": { + "box_map": 46.7, + } + }, + "_ops": 280.371, + "_file_size": 167.104, + "_docs": """These weights were produced using an enhanced training recipe to boost the model accuracy.""", + }, + ) + DEFAULT = COCO_V1 + + +class FasterRCNN_MobileNet_V3_Large_FPN_Weights(WeightsEnum): + COCO_V1 = Weights( + url="https://download.pytorch.org/models/fasterrcnn_mobilenet_v3_large_fpn-fb6a3cc7.pth", + transforms=ObjectDetection, + meta={ + **_COMMON_META, + "num_params": 19386354, + "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#faster-r-cnn-mobilenetv3-large-fpn", + "_metrics": { + "COCO-val2017": { + "box_map": 32.8, + } + }, + "_ops": 4.494, + "_file_size": 74.239, + "_docs": """These weights were produced by following a similar training recipe as on the paper.""", + }, + ) + DEFAULT = COCO_V1 + + +class FasterRCNN_MobileNet_V3_Large_320_FPN_Weights(WeightsEnum): + COCO_V1 = Weights( + url="https://download.pytorch.org/models/fasterrcnn_mobilenet_v3_large_320_fpn-907ea3f9.pth", + transforms=ObjectDetection, + meta={ + **_COMMON_META, + "num_params": 19386354, + "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#faster-r-cnn-mobilenetv3-large-320-fpn", + "_metrics": { + "COCO-val2017": { + "box_map": 22.8, + } + }, + "_ops": 0.719, + "_file_size": 74.239, + "_docs": """These weights were produced by following a similar training recipe as on the paper.""", + }, + ) + DEFAULT = COCO_V1 + + +@register_model() +@handle_legacy_interface( + weights=("pretrained", FasterRCNN_ResNet50_FPN_Weights.COCO_V1), + weights_backbone=("pretrained_backbone", ResNet50_Weights.IMAGENET1K_V1), +) +def fasterrcnn_resnet50_fpn( + *, + weights: Optional[FasterRCNN_ResNet50_FPN_Weights] = None, + progress: bool = True, + num_classes: Optional[int] = None, + weights_backbone: Optional[ResNet50_Weights] = ResNet50_Weights.IMAGENET1K_V1, + trainable_backbone_layers: Optional[int] = None, + **kwargs: Any, +) -> FasterRCNN: + """ + Faster R-CNN model with a ResNet-50-FPN backbone from the `Faster R-CNN: Towards Real-Time Object + Detection with Region Proposal Networks `__ + paper. + + .. betastatus:: detection module + + The input to the model is expected to be a list of tensors, each of shape ``[C, H, W]``, one for each + image, and should be in ``0-1`` range. Different images can have different sizes. + + The behavior of the model changes depending on if it is in training or evaluation mode. + + During training, the model expects both the input tensors and a targets (list of dictionary), + containing: + + - boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (``Int64Tensor[N]``): the class label for each ground-truth box + + The model returns a ``Dict[Tensor]`` during training, containing the classification and regression + losses for both the RPN and the R-CNN. + + During inference, the model requires only the input tensors, and returns the post-processed + predictions as a ``List[Dict[Tensor]]``, one for each input image. The fields of the ``Dict`` are as + follows, where ``N`` is the number of detections: + + - boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (``Int64Tensor[N]``): the predicted labels for each detection + - scores (``Tensor[N]``): the scores of each detection + + For more details on the output, you may refer to :ref:`instance_seg_output`. + + Faster R-CNN is exportable to ONNX for a fixed batch size with inputs images of fixed size. + + Example:: + + >>> model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=FasterRCNN_ResNet50_FPN_Weights.DEFAULT) + >>> # For training + >>> images, boxes = torch.rand(4, 3, 600, 1200), torch.rand(4, 11, 4) + >>> boxes[:, :, 2:4] = boxes[:, :, 0:2] + boxes[:, :, 2:4] + >>> labels = torch.randint(1, 91, (4, 11)) + >>> images = list(image for image in images) + >>> targets = [] + >>> for i in range(len(images)): + >>> d = {} + >>> d['boxes'] = boxes[i] + >>> d['labels'] = labels[i] + >>> targets.append(d) + >>> output = model(images, targets) + >>> # For inference + >>> model.eval() + >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)] + >>> predictions = model(x) + >>> + >>> # optionally, if you want to export the model to ONNX: + >>> torch.onnx.export(model, x, "faster_rcnn.onnx", opset_version = 11) + + Args: + weights (:class:`~torchvision.models.detection.FasterRCNN_ResNet50_FPN_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.detection.FasterRCNN_ResNet50_FPN_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + num_classes (int, optional): number of output classes of the model (including the background) + weights_backbone (:class:`~torchvision.models.ResNet50_Weights`, optional): The + pretrained weights for the backbone. + trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from + final block. Valid values are between 0 and 5, with 5 meaning all backbone layers are + trainable. If ``None`` is passed (the default) this value is set to 3. + **kwargs: parameters passed to the ``torchvision.models.detection.faster_rcnn.FasterRCNN`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.detection.FasterRCNN_ResNet50_FPN_Weights + :members: + """ + weights = FasterRCNN_ResNet50_FPN_Weights.verify(weights) + weights_backbone = ResNet50_Weights.verify(weights_backbone) + + if weights is not None: + weights_backbone = None + num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"])) + elif num_classes is None: + num_classes = 91 + + is_trained = weights is not None or weights_backbone is not None + trainable_backbone_layers = _validate_trainable_layers(is_trained, trainable_backbone_layers, 5, 3) + norm_layer = misc_nn_ops.FrozenBatchNorm2d if is_trained else nn.BatchNorm2d + + backbone = resnet50(weights=weights_backbone, progress=progress, norm_layer=norm_layer) + backbone = _resnet_fpn_extractor(backbone, trainable_backbone_layers) + model = FasterRCNN(backbone, num_classes=num_classes, **kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + if weights == FasterRCNN_ResNet50_FPN_Weights.COCO_V1: + overwrite_eps(model, 0.0) + + return model + + +@register_model() +@handle_legacy_interface( + weights=("pretrained", FasterRCNN_ResNet50_FPN_V2_Weights.COCO_V1), + weights_backbone=("pretrained_backbone", ResNet50_Weights.IMAGENET1K_V1), +) +def fasterrcnn_resnet50_fpn_v2( + *, + weights: Optional[FasterRCNN_ResNet50_FPN_V2_Weights] = None, + progress: bool = True, + num_classes: Optional[int] = None, + weights_backbone: Optional[ResNet50_Weights] = None, + trainable_backbone_layers: Optional[int] = None, + **kwargs: Any, +) -> FasterRCNN: + """ + Constructs an improved Faster R-CNN model with a ResNet-50-FPN backbone from `Benchmarking Detection + Transfer Learning with Vision Transformers `__ paper. + + .. betastatus:: detection module + + It works similarly to Faster R-CNN with ResNet-50 FPN backbone. See + :func:`~torchvision.models.detection.fasterrcnn_resnet50_fpn` for more + details. + + Args: + weights (:class:`~torchvision.models.detection.FasterRCNN_ResNet50_FPN_V2_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.detection.FasterRCNN_ResNet50_FPN_V2_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + num_classes (int, optional): number of output classes of the model (including the background) + weights_backbone (:class:`~torchvision.models.ResNet50_Weights`, optional): The + pretrained weights for the backbone. + trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from + final block. Valid values are between 0 and 5, with 5 meaning all backbone layers are + trainable. If ``None`` is passed (the default) this value is set to 3. + **kwargs: parameters passed to the ``torchvision.models.detection.faster_rcnn.FasterRCNN`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.detection.FasterRCNN_ResNet50_FPN_V2_Weights + :members: + """ + weights = FasterRCNN_ResNet50_FPN_V2_Weights.verify(weights) + weights_backbone = ResNet50_Weights.verify(weights_backbone) + + if weights is not None: + weights_backbone = None + num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"])) + elif num_classes is None: + num_classes = 91 + + is_trained = weights is not None or weights_backbone is not None + trainable_backbone_layers = _validate_trainable_layers(is_trained, trainable_backbone_layers, 5, 3) + + backbone = resnet50(weights=weights_backbone, progress=progress) + backbone = _resnet_fpn_extractor(backbone, trainable_backbone_layers, norm_layer=nn.BatchNorm2d) + rpn_anchor_generator = _default_anchorgen() + rpn_head = RPNHead(backbone.out_channels, rpn_anchor_generator.num_anchors_per_location()[0], conv_depth=2) + box_head = FastRCNNConvFCHead( + (backbone.out_channels, 7, 7), [256, 256, 256, 256], [1024], norm_layer=nn.BatchNorm2d + ) + model = FasterRCNN( + backbone, + num_classes=num_classes, + rpn_anchor_generator=rpn_anchor_generator, + rpn_head=rpn_head, + box_head=box_head, + **kwargs, + ) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +def _fasterrcnn_mobilenet_v3_large_fpn( + *, + weights: Optional[Union[FasterRCNN_MobileNet_V3_Large_FPN_Weights, FasterRCNN_MobileNet_V3_Large_320_FPN_Weights]], + progress: bool, + num_classes: Optional[int], + weights_backbone: Optional[MobileNet_V3_Large_Weights], + trainable_backbone_layers: Optional[int], + **kwargs: Any, +) -> FasterRCNN: + if weights is not None: + weights_backbone = None + num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"])) + elif num_classes is None: + num_classes = 91 + + is_trained = weights is not None or weights_backbone is not None + trainable_backbone_layers = _validate_trainable_layers(is_trained, trainable_backbone_layers, 6, 3) + norm_layer = misc_nn_ops.FrozenBatchNorm2d if is_trained else nn.BatchNorm2d + + backbone = mobilenet_v3_large(weights=weights_backbone, progress=progress, norm_layer=norm_layer) + backbone = _mobilenet_extractor(backbone, True, trainable_backbone_layers) + anchor_sizes = ( + ( + 32, + 64, + 128, + 256, + 512, + ), + ) * 3 + aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes) + model = FasterRCNN( + backbone, num_classes, rpn_anchor_generator=AnchorGenerator(anchor_sizes, aspect_ratios), **kwargs + ) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +@register_model() +@handle_legacy_interface( + weights=("pretrained", FasterRCNN_MobileNet_V3_Large_320_FPN_Weights.COCO_V1), + weights_backbone=("pretrained_backbone", MobileNet_V3_Large_Weights.IMAGENET1K_V1), +) +def fasterrcnn_mobilenet_v3_large_320_fpn( + *, + weights: Optional[FasterRCNN_MobileNet_V3_Large_320_FPN_Weights] = None, + progress: bool = True, + num_classes: Optional[int] = None, + weights_backbone: Optional[MobileNet_V3_Large_Weights] = MobileNet_V3_Large_Weights.IMAGENET1K_V1, + trainable_backbone_layers: Optional[int] = None, + **kwargs: Any, +) -> FasterRCNN: + """ + Low resolution Faster R-CNN model with a MobileNetV3-Large backbone tuned for mobile use cases. + + .. betastatus:: detection module + + It works similarly to Faster R-CNN with ResNet-50 FPN backbone. See + :func:`~torchvision.models.detection.fasterrcnn_resnet50_fpn` for more + details. + + Example:: + + >>> model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_320_fpn(weights=FasterRCNN_MobileNet_V3_Large_320_FPN_Weights.DEFAULT) + >>> model.eval() + >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)] + >>> predictions = model(x) + + Args: + weights (:class:`~torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_320_FPN_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_320_FPN_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + num_classes (int, optional): number of output classes of the model (including the background) + weights_backbone (:class:`~torchvision.models.MobileNet_V3_Large_Weights`, optional): The + pretrained weights for the backbone. + trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from + final block. Valid values are between 0 and 6, with 6 meaning all backbone layers are + trainable. If ``None`` is passed (the default) this value is set to 3. + **kwargs: parameters passed to the ``torchvision.models.detection.faster_rcnn.FasterRCNN`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_320_FPN_Weights + :members: + """ + weights = FasterRCNN_MobileNet_V3_Large_320_FPN_Weights.verify(weights) + weights_backbone = MobileNet_V3_Large_Weights.verify(weights_backbone) + + defaults = { + "min_size": 320, + "max_size": 640, + "rpn_pre_nms_top_n_test": 150, + "rpn_post_nms_top_n_test": 150, + "rpn_score_thresh": 0.05, + } + + kwargs = {**defaults, **kwargs} + return _fasterrcnn_mobilenet_v3_large_fpn( + weights=weights, + progress=progress, + num_classes=num_classes, + weights_backbone=weights_backbone, + trainable_backbone_layers=trainable_backbone_layers, + **kwargs, + ) + + +@register_model() +@handle_legacy_interface( + weights=("pretrained", FasterRCNN_MobileNet_V3_Large_FPN_Weights.COCO_V1), + weights_backbone=("pretrained_backbone", MobileNet_V3_Large_Weights.IMAGENET1K_V1), +) +def fasterrcnn_mobilenet_v3_large_fpn( + *, + weights: Optional[FasterRCNN_MobileNet_V3_Large_FPN_Weights] = None, + progress: bool = True, + num_classes: Optional[int] = None, + weights_backbone: Optional[MobileNet_V3_Large_Weights] = MobileNet_V3_Large_Weights.IMAGENET1K_V1, + trainable_backbone_layers: Optional[int] = None, + **kwargs: Any, +) -> FasterRCNN: + """ + Constructs a high resolution Faster R-CNN model with a MobileNetV3-Large FPN backbone. + + .. betastatus:: detection module + + It works similarly to Faster R-CNN with ResNet-50 FPN backbone. See + :func:`~torchvision.models.detection.fasterrcnn_resnet50_fpn` for more + details. + + Example:: + + >>> model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_fpn(weights=FasterRCNN_MobileNet_V3_Large_FPN_Weights.DEFAULT) + >>> model.eval() + >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)] + >>> predictions = model(x) + + Args: + weights (:class:`~torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_FPN_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_FPN_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + num_classes (int, optional): number of output classes of the model (including the background) + weights_backbone (:class:`~torchvision.models.MobileNet_V3_Large_Weights`, optional): The + pretrained weights for the backbone. + trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from + final block. Valid values are between 0 and 6, with 6 meaning all backbone layers are + trainable. If ``None`` is passed (the default) this value is set to 3. + **kwargs: parameters passed to the ``torchvision.models.detection.faster_rcnn.FasterRCNN`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_FPN_Weights + :members: + """ + weights = FasterRCNN_MobileNet_V3_Large_FPN_Weights.verify(weights) + weights_backbone = MobileNet_V3_Large_Weights.verify(weights_backbone) + + defaults = { + "rpn_score_thresh": 0.05, + } + + kwargs = {**defaults, **kwargs} + return _fasterrcnn_mobilenet_v3_large_fpn( + weights=weights, + progress=progress, + num_classes=num_classes, + weights_backbone=weights_backbone, + trainable_backbone_layers=trainable_backbone_layers, + **kwargs, + ) diff --git a/lib/python3.12/site-packages/torchvision/models/detection/fcos.py b/lib/python3.12/site-packages/torchvision/models/detection/fcos.py new file mode 100644 index 0000000000000000000000000000000000000000..a86ad2f424c32bd1cf951d474d3ef14bd1bddbb7 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/detection/fcos.py @@ -0,0 +1,775 @@ +import math +import warnings +from collections import OrderedDict +from functools import partial +from typing import Any, Callable, Dict, List, Optional, Tuple + +import torch +from torch import nn, Tensor + +from ...ops import boxes as box_ops, generalized_box_iou_loss, misc as misc_nn_ops, sigmoid_focal_loss +from ...ops.feature_pyramid_network import LastLevelP6P7 +from ...transforms._presets import ObjectDetection +from ...utils import _log_api_usage_once +from .._api import register_model, Weights, WeightsEnum +from .._meta import _COCO_CATEGORIES +from .._utils import _ovewrite_value_param, handle_legacy_interface +from ..resnet import resnet50, ResNet50_Weights +from . import _utils as det_utils +from .anchor_utils import AnchorGenerator +from .backbone_utils import _resnet_fpn_extractor, _validate_trainable_layers +from .transform import GeneralizedRCNNTransform + + +__all__ = [ + "FCOS", + "FCOS_ResNet50_FPN_Weights", + "fcos_resnet50_fpn", +] + + +class FCOSHead(nn.Module): + """ + A regression and classification head for use in FCOS. + + Args: + in_channels (int): number of channels of the input feature + num_anchors (int): number of anchors to be predicted + num_classes (int): number of classes to be predicted + num_convs (Optional[int]): number of conv layer of head. Default: 4. + """ + + __annotations__ = { + "box_coder": det_utils.BoxLinearCoder, + } + + def __init__(self, in_channels: int, num_anchors: int, num_classes: int, num_convs: Optional[int] = 4) -> None: + super().__init__() + self.box_coder = det_utils.BoxLinearCoder(normalize_by_size=True) + self.classification_head = FCOSClassificationHead(in_channels, num_anchors, num_classes, num_convs) + self.regression_head = FCOSRegressionHead(in_channels, num_anchors, num_convs) + + def compute_loss( + self, + targets: List[Dict[str, Tensor]], + head_outputs: Dict[str, Tensor], + anchors: List[Tensor], + matched_idxs: List[Tensor], + ) -> Dict[str, Tensor]: + + cls_logits = head_outputs["cls_logits"] # [N, HWA, C] + bbox_regression = head_outputs["bbox_regression"] # [N, HWA, 4] + bbox_ctrness = head_outputs["bbox_ctrness"] # [N, HWA, 1] + + all_gt_classes_targets = [] + all_gt_boxes_targets = [] + for targets_per_image, matched_idxs_per_image in zip(targets, matched_idxs): + if len(targets_per_image["labels"]) == 0: + gt_classes_targets = targets_per_image["labels"].new_zeros((len(matched_idxs_per_image),)) + gt_boxes_targets = targets_per_image["boxes"].new_zeros((len(matched_idxs_per_image), 4)) + else: + gt_classes_targets = targets_per_image["labels"][matched_idxs_per_image.clip(min=0)] + gt_boxes_targets = targets_per_image["boxes"][matched_idxs_per_image.clip(min=0)] + gt_classes_targets[matched_idxs_per_image < 0] = -1 # background + all_gt_classes_targets.append(gt_classes_targets) + all_gt_boxes_targets.append(gt_boxes_targets) + + # List[Tensor] to Tensor conversion of `all_gt_boxes_target`, `all_gt_classes_targets` and `anchors` + all_gt_boxes_targets, all_gt_classes_targets, anchors = ( + torch.stack(all_gt_boxes_targets), + torch.stack(all_gt_classes_targets), + torch.stack(anchors), + ) + + # compute foregroud + foregroud_mask = all_gt_classes_targets >= 0 + num_foreground = foregroud_mask.sum().item() + + # classification loss + gt_classes_targets = torch.zeros_like(cls_logits) + gt_classes_targets[foregroud_mask, all_gt_classes_targets[foregroud_mask]] = 1.0 + loss_cls = sigmoid_focal_loss(cls_logits, gt_classes_targets, reduction="sum") + + # amp issue: pred_boxes need to convert float + pred_boxes = self.box_coder.decode(bbox_regression, anchors) + + # regression loss: GIoU loss + loss_bbox_reg = generalized_box_iou_loss( + pred_boxes[foregroud_mask], + all_gt_boxes_targets[foregroud_mask], + reduction="sum", + ) + + # ctrness loss + + bbox_reg_targets = self.box_coder.encode(anchors, all_gt_boxes_targets) + + if len(bbox_reg_targets) == 0: + gt_ctrness_targets = bbox_reg_targets.new_zeros(bbox_reg_targets.size()[:-1]) + else: + left_right = bbox_reg_targets[:, :, [0, 2]] + top_bottom = bbox_reg_targets[:, :, [1, 3]] + gt_ctrness_targets = torch.sqrt( + (left_right.min(dim=-1)[0] / left_right.max(dim=-1)[0]) + * (top_bottom.min(dim=-1)[0] / top_bottom.max(dim=-1)[0]) + ) + pred_centerness = bbox_ctrness.squeeze(dim=2) + loss_bbox_ctrness = nn.functional.binary_cross_entropy_with_logits( + pred_centerness[foregroud_mask], gt_ctrness_targets[foregroud_mask], reduction="sum" + ) + + return { + "classification": loss_cls / max(1, num_foreground), + "bbox_regression": loss_bbox_reg / max(1, num_foreground), + "bbox_ctrness": loss_bbox_ctrness / max(1, num_foreground), + } + + def forward(self, x: List[Tensor]) -> Dict[str, Tensor]: + cls_logits = self.classification_head(x) + bbox_regression, bbox_ctrness = self.regression_head(x) + return { + "cls_logits": cls_logits, + "bbox_regression": bbox_regression, + "bbox_ctrness": bbox_ctrness, + } + + +class FCOSClassificationHead(nn.Module): + """ + A classification head for use in FCOS. + + Args: + in_channels (int): number of channels of the input feature. + num_anchors (int): number of anchors to be predicted. + num_classes (int): number of classes to be predicted. + num_convs (Optional[int]): number of conv layer. Default: 4. + prior_probability (Optional[float]): probability of prior. Default: 0.01. + norm_layer: Module specifying the normalization layer to use. + """ + + def __init__( + self, + in_channels: int, + num_anchors: int, + num_classes: int, + num_convs: int = 4, + prior_probability: float = 0.01, + norm_layer: Optional[Callable[..., nn.Module]] = None, + ) -> None: + super().__init__() + + self.num_classes = num_classes + self.num_anchors = num_anchors + + if norm_layer is None: + norm_layer = partial(nn.GroupNorm, 32) + + conv = [] + for _ in range(num_convs): + conv.append(nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)) + conv.append(norm_layer(in_channels)) + conv.append(nn.ReLU()) + self.conv = nn.Sequential(*conv) + + for layer in self.conv.children(): + if isinstance(layer, nn.Conv2d): + torch.nn.init.normal_(layer.weight, std=0.01) + torch.nn.init.constant_(layer.bias, 0) + + self.cls_logits = nn.Conv2d(in_channels, num_anchors * num_classes, kernel_size=3, stride=1, padding=1) + torch.nn.init.normal_(self.cls_logits.weight, std=0.01) + torch.nn.init.constant_(self.cls_logits.bias, -math.log((1 - prior_probability) / prior_probability)) + + def forward(self, x: List[Tensor]) -> Tensor: + all_cls_logits = [] + + for features in x: + cls_logits = self.conv(features) + cls_logits = self.cls_logits(cls_logits) + + # Permute classification output from (N, A * K, H, W) to (N, HWA, K). + N, _, H, W = cls_logits.shape + cls_logits = cls_logits.view(N, -1, self.num_classes, H, W) + cls_logits = cls_logits.permute(0, 3, 4, 1, 2) + cls_logits = cls_logits.reshape(N, -1, self.num_classes) # Size=(N, HWA, 4) + + all_cls_logits.append(cls_logits) + + return torch.cat(all_cls_logits, dim=1) + + +class FCOSRegressionHead(nn.Module): + """ + A regression head for use in FCOS, which combines regression branch and center-ness branch. + This can obtain better performance. + + Reference: `FCOS: A simple and strong anchor-free object detector `_. + + Args: + in_channels (int): number of channels of the input feature + num_anchors (int): number of anchors to be predicted + num_convs (Optional[int]): number of conv layer. Default: 4. + norm_layer: Module specifying the normalization layer to use. + """ + + def __init__( + self, + in_channels: int, + num_anchors: int, + num_convs: int = 4, + norm_layer: Optional[Callable[..., nn.Module]] = None, + ): + super().__init__() + + if norm_layer is None: + norm_layer = partial(nn.GroupNorm, 32) + + conv = [] + for _ in range(num_convs): + conv.append(nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)) + conv.append(norm_layer(in_channels)) + conv.append(nn.ReLU()) + self.conv = nn.Sequential(*conv) + + self.bbox_reg = nn.Conv2d(in_channels, num_anchors * 4, kernel_size=3, stride=1, padding=1) + self.bbox_ctrness = nn.Conv2d(in_channels, num_anchors * 1, kernel_size=3, stride=1, padding=1) + for layer in [self.bbox_reg, self.bbox_ctrness]: + torch.nn.init.normal_(layer.weight, std=0.01) + torch.nn.init.zeros_(layer.bias) + + for layer in self.conv.children(): + if isinstance(layer, nn.Conv2d): + torch.nn.init.normal_(layer.weight, std=0.01) + torch.nn.init.zeros_(layer.bias) + + def forward(self, x: List[Tensor]) -> Tuple[Tensor, Tensor]: + all_bbox_regression = [] + all_bbox_ctrness = [] + + for features in x: + bbox_feature = self.conv(features) + bbox_regression = nn.functional.relu(self.bbox_reg(bbox_feature)) + bbox_ctrness = self.bbox_ctrness(bbox_feature) + + # permute bbox regression output from (N, 4 * A, H, W) to (N, HWA, 4). + N, _, H, W = bbox_regression.shape + bbox_regression = bbox_regression.view(N, -1, 4, H, W) + bbox_regression = bbox_regression.permute(0, 3, 4, 1, 2) + bbox_regression = bbox_regression.reshape(N, -1, 4) # Size=(N, HWA, 4) + all_bbox_regression.append(bbox_regression) + + # permute bbox ctrness output from (N, 1 * A, H, W) to (N, HWA, 1). + bbox_ctrness = bbox_ctrness.view(N, -1, 1, H, W) + bbox_ctrness = bbox_ctrness.permute(0, 3, 4, 1, 2) + bbox_ctrness = bbox_ctrness.reshape(N, -1, 1) + all_bbox_ctrness.append(bbox_ctrness) + + return torch.cat(all_bbox_regression, dim=1), torch.cat(all_bbox_ctrness, dim=1) + + +class FCOS(nn.Module): + """ + Implements FCOS. + + The input to the model is expected to be a list of tensors, each of shape [C, H, W], one for each + image, and should be in 0-1 range. Different images can have different sizes. + + The behavior of the model changes depending on if it is in training or evaluation mode. + + During training, the model expects both the input tensors and targets (list of dictionary), + containing: + - boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (Int64Tensor[N]): the class label for each ground-truth box + + The model returns a Dict[Tensor] during training, containing the classification, regression + and centerness losses. + + During inference, the model requires only the input tensors, and returns the post-processed + predictions as a List[Dict[Tensor]], one for each input image. The fields of the Dict are as + follows: + - boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (Int64Tensor[N]): the predicted labels for each image + - scores (Tensor[N]): the scores for each prediction + + Args: + backbone (nn.Module): the network used to compute the features for the model. + It should contain an out_channels attribute, which indicates the number of output + channels that each feature map has (and it should be the same for all feature maps). + The backbone should return a single Tensor or an OrderedDict[Tensor]. + num_classes (int): number of output classes of the model (including the background). + min_size (int): Images are rescaled before feeding them to the backbone: + we attempt to preserve the aspect ratio and scale the shorter edge + to ``min_size``. If the resulting longer edge exceeds ``max_size``, + then downscale so that the longer edge does not exceed ``max_size``. + This may result in the shorter edge beeing lower than ``min_size``. + max_size (int): See ``min_size``. + image_mean (Tuple[float, float, float]): mean values used for input normalization. + They are generally the mean values of the dataset on which the backbone has been trained + on + image_std (Tuple[float, float, float]): std values used for input normalization. + They are generally the std values of the dataset on which the backbone has been trained on + anchor_generator (AnchorGenerator): module that generates the anchors for a set of feature + maps. For FCOS, only set one anchor for per position of each level, the width and height equal to + the stride of feature map, and set aspect ratio = 1.0, so the center of anchor is equivalent to the point + in FCOS paper. + head (nn.Module): Module run on top of the feature pyramid. + Defaults to a module containing a classification and regression module. + center_sampling_radius (int): radius of the "center" of a groundtruth box, + within which all anchor points are labeled positive. + score_thresh (float): Score threshold used for postprocessing the detections. + nms_thresh (float): NMS threshold used for postprocessing the detections. + detections_per_img (int): Number of best detections to keep after NMS. + topk_candidates (int): Number of best detections to keep before NMS. + + Example: + + >>> import torch + >>> import torchvision + >>> from torchvision.models.detection import FCOS + >>> from torchvision.models.detection.anchor_utils import AnchorGenerator + >>> # load a pre-trained model for classification and return + >>> # only the features + >>> backbone = torchvision.models.mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT).features + >>> # FCOS needs to know the number of + >>> # output channels in a backbone. For mobilenet_v2, it's 1280, + >>> # so we need to add it here + >>> backbone.out_channels = 1280 + >>> + >>> # let's make the network generate 5 x 3 anchors per spatial + >>> # location, with 5 different sizes and 3 different aspect + >>> # ratios. We have a Tuple[Tuple[int]] because each feature + >>> # map could potentially have different sizes and + >>> # aspect ratios + >>> anchor_generator = AnchorGenerator( + >>> sizes=((8,), (16,), (32,), (64,), (128,)), + >>> aspect_ratios=((1.0,),) + >>> ) + >>> + >>> # put the pieces together inside a FCOS model + >>> model = FCOS( + >>> backbone, + >>> num_classes=80, + >>> anchor_generator=anchor_generator, + >>> ) + >>> model.eval() + >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)] + >>> predictions = model(x) + """ + + __annotations__ = { + "box_coder": det_utils.BoxLinearCoder, + } + + def __init__( + self, + backbone: nn.Module, + num_classes: int, + # transform parameters + min_size: int = 800, + max_size: int = 1333, + image_mean: Optional[List[float]] = None, + image_std: Optional[List[float]] = None, + # Anchor parameters + anchor_generator: Optional[AnchorGenerator] = None, + head: Optional[nn.Module] = None, + center_sampling_radius: float = 1.5, + score_thresh: float = 0.2, + nms_thresh: float = 0.6, + detections_per_img: int = 100, + topk_candidates: int = 1000, + **kwargs, + ): + super().__init__() + _log_api_usage_once(self) + + if not hasattr(backbone, "out_channels"): + raise ValueError( + "backbone should contain an attribute out_channels " + "specifying the number of output channels (assumed to be the " + "same for all the levels)" + ) + self.backbone = backbone + + if not isinstance(anchor_generator, (AnchorGenerator, type(None))): + raise TypeError( + f"anchor_generator should be of type AnchorGenerator or None, instead got {type(anchor_generator)}" + ) + + if anchor_generator is None: + anchor_sizes = ((8,), (16,), (32,), (64,), (128,)) # equal to strides of multi-level feature map + aspect_ratios = ((1.0,),) * len(anchor_sizes) # set only one anchor + anchor_generator = AnchorGenerator(anchor_sizes, aspect_ratios) + self.anchor_generator = anchor_generator + if self.anchor_generator.num_anchors_per_location()[0] != 1: + raise ValueError( + f"anchor_generator.num_anchors_per_location()[0] should be 1 instead of {anchor_generator.num_anchors_per_location()[0]}" + ) + + if head is None: + head = FCOSHead(backbone.out_channels, anchor_generator.num_anchors_per_location()[0], num_classes) + self.head = head + + self.box_coder = det_utils.BoxLinearCoder(normalize_by_size=True) + + if image_mean is None: + image_mean = [0.485, 0.456, 0.406] + if image_std is None: + image_std = [0.229, 0.224, 0.225] + self.transform = GeneralizedRCNNTransform(min_size, max_size, image_mean, image_std, **kwargs) + + self.center_sampling_radius = center_sampling_radius + self.score_thresh = score_thresh + self.nms_thresh = nms_thresh + self.detections_per_img = detections_per_img + self.topk_candidates = topk_candidates + + # used only on torchscript mode + self._has_warned = False + + @torch.jit.unused + def eager_outputs( + self, losses: Dict[str, Tensor], detections: List[Dict[str, Tensor]] + ) -> Tuple[Dict[str, Tensor], List[Dict[str, Tensor]]]: + if self.training: + return losses + + return detections + + def compute_loss( + self, + targets: List[Dict[str, Tensor]], + head_outputs: Dict[str, Tensor], + anchors: List[Tensor], + num_anchors_per_level: List[int], + ) -> Dict[str, Tensor]: + matched_idxs = [] + for anchors_per_image, targets_per_image in zip(anchors, targets): + if targets_per_image["boxes"].numel() == 0: + matched_idxs.append( + torch.full((anchors_per_image.size(0),), -1, dtype=torch.int64, device=anchors_per_image.device) + ) + continue + + gt_boxes = targets_per_image["boxes"] + gt_centers = (gt_boxes[:, :2] + gt_boxes[:, 2:]) / 2 # Nx2 + anchor_centers = (anchors_per_image[:, :2] + anchors_per_image[:, 2:]) / 2 # N + anchor_sizes = anchors_per_image[:, 2] - anchors_per_image[:, 0] + # center sampling: anchor point must be close enough to gt center. + pairwise_match = (anchor_centers[:, None, :] - gt_centers[None, :, :]).abs_().max( + dim=2 + ).values < self.center_sampling_radius * anchor_sizes[:, None] + # compute pairwise distance between N points and M boxes + x, y = anchor_centers.unsqueeze(dim=2).unbind(dim=1) # (N, 1) + x0, y0, x1, y1 = gt_boxes.unsqueeze(dim=0).unbind(dim=2) # (1, M) + pairwise_dist = torch.stack([x - x0, y - y0, x1 - x, y1 - y], dim=2) # (N, M) + + # anchor point must be inside gt + pairwise_match &= pairwise_dist.min(dim=2).values > 0 + + # each anchor is only responsible for certain scale range. + lower_bound = anchor_sizes * 4 + lower_bound[: num_anchors_per_level[0]] = 0 + upper_bound = anchor_sizes * 8 + upper_bound[-num_anchors_per_level[-1] :] = float("inf") + pairwise_dist = pairwise_dist.max(dim=2).values + pairwise_match &= (pairwise_dist > lower_bound[:, None]) & (pairwise_dist < upper_bound[:, None]) + + # match the GT box with minimum area, if there are multiple GT matches + gt_areas = (gt_boxes[:, 2] - gt_boxes[:, 0]) * (gt_boxes[:, 3] - gt_boxes[:, 1]) # N + pairwise_match = pairwise_match.to(torch.float32) * (1e8 - gt_areas[None, :]) + min_values, matched_idx = pairwise_match.max(dim=1) # R, per-anchor match + matched_idx[min_values < 1e-5] = -1 # unmatched anchors are assigned -1 + + matched_idxs.append(matched_idx) + + return self.head.compute_loss(targets, head_outputs, anchors, matched_idxs) + + def postprocess_detections( + self, head_outputs: Dict[str, List[Tensor]], anchors: List[List[Tensor]], image_shapes: List[Tuple[int, int]] + ) -> List[Dict[str, Tensor]]: + class_logits = head_outputs["cls_logits"] + box_regression = head_outputs["bbox_regression"] + box_ctrness = head_outputs["bbox_ctrness"] + + num_images = len(image_shapes) + + detections: List[Dict[str, Tensor]] = [] + + for index in range(num_images): + box_regression_per_image = [br[index] for br in box_regression] + logits_per_image = [cl[index] for cl in class_logits] + box_ctrness_per_image = [bc[index] for bc in box_ctrness] + anchors_per_image, image_shape = anchors[index], image_shapes[index] + + image_boxes = [] + image_scores = [] + image_labels = [] + + for box_regression_per_level, logits_per_level, box_ctrness_per_level, anchors_per_level in zip( + box_regression_per_image, logits_per_image, box_ctrness_per_image, anchors_per_image + ): + num_classes = logits_per_level.shape[-1] + + # remove low scoring boxes + scores_per_level = torch.sqrt( + torch.sigmoid(logits_per_level) * torch.sigmoid(box_ctrness_per_level) + ).flatten() + keep_idxs = scores_per_level > self.score_thresh + scores_per_level = scores_per_level[keep_idxs] + topk_idxs = torch.where(keep_idxs)[0] + + # keep only topk scoring predictions + num_topk = det_utils._topk_min(topk_idxs, self.topk_candidates, 0) + scores_per_level, idxs = scores_per_level.topk(num_topk) + topk_idxs = topk_idxs[idxs] + + anchor_idxs = torch.div(topk_idxs, num_classes, rounding_mode="floor") + labels_per_level = topk_idxs % num_classes + + boxes_per_level = self.box_coder.decode( + box_regression_per_level[anchor_idxs], anchors_per_level[anchor_idxs] + ) + boxes_per_level = box_ops.clip_boxes_to_image(boxes_per_level, image_shape) + + image_boxes.append(boxes_per_level) + image_scores.append(scores_per_level) + image_labels.append(labels_per_level) + + image_boxes = torch.cat(image_boxes, dim=0) + image_scores = torch.cat(image_scores, dim=0) + image_labels = torch.cat(image_labels, dim=0) + + # non-maximum suppression + keep = box_ops.batched_nms(image_boxes, image_scores, image_labels, self.nms_thresh) + keep = keep[: self.detections_per_img] + + detections.append( + { + "boxes": image_boxes[keep], + "scores": image_scores[keep], + "labels": image_labels[keep], + } + ) + + return detections + + def forward( + self, + images: List[Tensor], + targets: Optional[List[Dict[str, Tensor]]] = None, + ) -> Tuple[Dict[str, Tensor], List[Dict[str, Tensor]]]: + """ + Args: + images (list[Tensor]): images to be processed + targets (list[Dict[Tensor]]): ground-truth boxes present in the image (optional) + + Returns: + result (list[BoxList] or dict[Tensor]): the output from the model. + During training, it returns a dict[Tensor] which contains the losses. + During testing, it returns list[BoxList] contains additional fields + like `scores`, `labels` and `mask` (for Mask R-CNN models). + """ + if self.training: + + if targets is None: + torch._assert(False, "targets should not be none when in training mode") + else: + for target in targets: + boxes = target["boxes"] + torch._assert(isinstance(boxes, torch.Tensor), "Expected target boxes to be of type Tensor.") + torch._assert( + len(boxes.shape) == 2 and boxes.shape[-1] == 4, + f"Expected target boxes to be a tensor of shape [N, 4], got {boxes.shape}.", + ) + + original_image_sizes: List[Tuple[int, int]] = [] + for img in images: + val = img.shape[-2:] + torch._assert( + len(val) == 2, + f"expecting the last two dimensions of the Tensor to be H and W instead got {img.shape[-2:]}", + ) + original_image_sizes.append((val[0], val[1])) + + # transform the input + images, targets = self.transform(images, targets) + + # Check for degenerate boxes + if targets is not None: + for target_idx, target in enumerate(targets): + boxes = target["boxes"] + degenerate_boxes = boxes[:, 2:] <= boxes[:, :2] + if degenerate_boxes.any(): + # print the first degenerate box + bb_idx = torch.where(degenerate_boxes.any(dim=1))[0][0] + degen_bb: List[float] = boxes[bb_idx].tolist() + torch._assert( + False, + f"All bounding boxes should have positive height and width. Found invalid box {degen_bb} for target at index {target_idx}.", + ) + + # get the features from the backbone + features = self.backbone(images.tensors) + if isinstance(features, torch.Tensor): + features = OrderedDict([("0", features)]) + + features = list(features.values()) + + # compute the fcos heads outputs using the features + head_outputs = self.head(features) + + # create the set of anchors + anchors = self.anchor_generator(images, features) + # recover level sizes + num_anchors_per_level = [x.size(2) * x.size(3) for x in features] + + losses = {} + detections: List[Dict[str, Tensor]] = [] + if self.training: + if targets is None: + torch._assert(False, "targets should not be none when in training mode") + else: + # compute the losses + losses = self.compute_loss(targets, head_outputs, anchors, num_anchors_per_level) + else: + # split outputs per level + split_head_outputs: Dict[str, List[Tensor]] = {} + for k in head_outputs: + split_head_outputs[k] = list(head_outputs[k].split(num_anchors_per_level, dim=1)) + split_anchors = [list(a.split(num_anchors_per_level)) for a in anchors] + + # compute the detections + detections = self.postprocess_detections(split_head_outputs, split_anchors, images.image_sizes) + detections = self.transform.postprocess(detections, images.image_sizes, original_image_sizes) + + if torch.jit.is_scripting(): + if not self._has_warned: + warnings.warn("FCOS always returns a (Losses, Detections) tuple in scripting") + self._has_warned = True + return losses, detections + return self.eager_outputs(losses, detections) + + +class FCOS_ResNet50_FPN_Weights(WeightsEnum): + COCO_V1 = Weights( + url="https://download.pytorch.org/models/fcos_resnet50_fpn_coco-99b0c9b7.pth", + transforms=ObjectDetection, + meta={ + "num_params": 32269600, + "categories": _COCO_CATEGORIES, + "min_size": (1, 1), + "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#fcos-resnet-50-fpn", + "_metrics": { + "COCO-val2017": { + "box_map": 39.2, + } + }, + "_ops": 128.207, + "_file_size": 123.608, + "_docs": """These weights were produced by following a similar training recipe as on the paper.""", + }, + ) + DEFAULT = COCO_V1 + + +@register_model() +@handle_legacy_interface( + weights=("pretrained", FCOS_ResNet50_FPN_Weights.COCO_V1), + weights_backbone=("pretrained_backbone", ResNet50_Weights.IMAGENET1K_V1), +) +def fcos_resnet50_fpn( + *, + weights: Optional[FCOS_ResNet50_FPN_Weights] = None, + progress: bool = True, + num_classes: Optional[int] = None, + weights_backbone: Optional[ResNet50_Weights] = ResNet50_Weights.IMAGENET1K_V1, + trainable_backbone_layers: Optional[int] = None, + **kwargs: Any, +) -> FCOS: + """ + Constructs a FCOS model with a ResNet-50-FPN backbone. + + .. betastatus:: detection module + + Reference: `FCOS: Fully Convolutional One-Stage Object Detection `_. + `FCOS: A simple and strong anchor-free object detector `_. + + The input to the model is expected to be a list of tensors, each of shape ``[C, H, W]``, one for each + image, and should be in ``0-1`` range. Different images can have different sizes. + + The behavior of the model changes depending on if it is in training or evaluation mode. + + During training, the model expects both the input tensors and targets (list of dictionary), + containing: + + - boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (``Int64Tensor[N]``): the class label for each ground-truth box + + The model returns a ``Dict[Tensor]`` during training, containing the classification and regression + losses. + + During inference, the model requires only the input tensors, and returns the post-processed + predictions as a ``List[Dict[Tensor]]``, one for each input image. The fields of the ``Dict`` are as + follows, where ``N`` is the number of detections: + + - boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (``Int64Tensor[N]``): the predicted labels for each detection + - scores (``Tensor[N]``): the scores of each detection + + For more details on the output, you may refer to :ref:`instance_seg_output`. + + Example: + + >>> model = torchvision.models.detection.fcos_resnet50_fpn(weights=FCOS_ResNet50_FPN_Weights.DEFAULT) + >>> model.eval() + >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)] + >>> predictions = model(x) + + Args: + weights (:class:`~torchvision.models.detection.FCOS_ResNet50_FPN_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.detection.FCOS_ResNet50_FPN_Weights` + below for more details, and possible values. By default, no + pre-trained weights are used. + progress (bool): If True, displays a progress bar of the download to stderr + num_classes (int, optional): number of output classes of the model (including the background) + weights_backbone (:class:`~torchvision.models.ResNet50_Weights`, optional): The pretrained weights for + the backbone. + trainable_backbone_layers (int, optional): number of trainable (not frozen) resnet layers starting + from final block. Valid values are between 0 and 5, with 5 meaning all backbone layers are + trainable. If ``None`` is passed (the default) this value is set to 3. Default: None + **kwargs: parameters passed to the ``torchvision.models.detection.FCOS`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.detection.FCOS_ResNet50_FPN_Weights + :members: + """ + weights = FCOS_ResNet50_FPN_Weights.verify(weights) + weights_backbone = ResNet50_Weights.verify(weights_backbone) + + if weights is not None: + weights_backbone = None + num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"])) + elif num_classes is None: + num_classes = 91 + + is_trained = weights is not None or weights_backbone is not None + trainable_backbone_layers = _validate_trainable_layers(is_trained, trainable_backbone_layers, 5, 3) + norm_layer = misc_nn_ops.FrozenBatchNorm2d if is_trained else nn.BatchNorm2d + + backbone = resnet50(weights=weights_backbone, progress=progress, norm_layer=norm_layer) + backbone = _resnet_fpn_extractor( + backbone, trainable_backbone_layers, returned_layers=[2, 3, 4], extra_blocks=LastLevelP6P7(256, 256) + ) + model = FCOS(backbone, num_classes, **kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model diff --git a/lib/python3.12/site-packages/torchvision/models/detection/generalized_rcnn.py b/lib/python3.12/site-packages/torchvision/models/detection/generalized_rcnn.py new file mode 100644 index 0000000000000000000000000000000000000000..b481265077fb5a582402d81aeb3516ffca063653 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/detection/generalized_rcnn.py @@ -0,0 +1,118 @@ +""" +Implements the Generalized R-CNN framework +""" + +import warnings +from collections import OrderedDict +from typing import Dict, List, Optional, Tuple, Union + +import torch +from torch import nn, Tensor + +from ...utils import _log_api_usage_once + + +class GeneralizedRCNN(nn.Module): + """ + Main class for Generalized R-CNN. + + Args: + backbone (nn.Module): + rpn (nn.Module): + roi_heads (nn.Module): takes the features + the proposals from the RPN and computes + detections / masks from it. + transform (nn.Module): performs the data transformation from the inputs to feed into + the model + """ + + def __init__(self, backbone: nn.Module, rpn: nn.Module, roi_heads: nn.Module, transform: nn.Module) -> None: + super().__init__() + _log_api_usage_once(self) + self.transform = transform + self.backbone = backbone + self.rpn = rpn + self.roi_heads = roi_heads + # used only on torchscript mode + self._has_warned = False + + @torch.jit.unused + def eager_outputs(self, losses, detections): + # type: (Dict[str, Tensor], List[Dict[str, Tensor]]) -> Union[Dict[str, Tensor], List[Dict[str, Tensor]]] + if self.training: + return losses + + return detections + + def forward(self, images, targets=None): + # type: (List[Tensor], Optional[List[Dict[str, Tensor]]]) -> Tuple[Dict[str, Tensor], List[Dict[str, Tensor]]] + """ + Args: + images (list[Tensor]): images to be processed + targets (list[Dict[str, Tensor]]): ground-truth boxes present in the image (optional) + + Returns: + result (list[BoxList] or dict[Tensor]): the output from the model. + During training, it returns a dict[Tensor] which contains the losses. + During testing, it returns list[BoxList] contains additional fields + like `scores`, `labels` and `mask` (for Mask R-CNN models). + + """ + if self.training: + if targets is None: + torch._assert(False, "targets should not be none when in training mode") + else: + for target in targets: + boxes = target["boxes"] + if isinstance(boxes, torch.Tensor): + torch._assert( + len(boxes.shape) == 2 and boxes.shape[-1] == 4, + f"Expected target boxes to be a tensor of shape [N, 4], got {boxes.shape}.", + ) + else: + torch._assert(False, f"Expected target boxes to be of type Tensor, got {type(boxes)}.") + + original_image_sizes: List[Tuple[int, int]] = [] + for img in images: + val = img.shape[-2:] + torch._assert( + len(val) == 2, + f"expecting the last two dimensions of the Tensor to be H and W instead got {img.shape[-2:]}", + ) + original_image_sizes.append((val[0], val[1])) + + images, targets = self.transform(images, targets) + + # Check for degenerate boxes + # TODO: Move this to a function + if targets is not None: + for target_idx, target in enumerate(targets): + boxes = target["boxes"] + degenerate_boxes = boxes[:, 2:] <= boxes[:, :2] + if degenerate_boxes.any(): + # print the first degenerate box + bb_idx = torch.where(degenerate_boxes.any(dim=1))[0][0] + degen_bb: List[float] = boxes[bb_idx].tolist() + torch._assert( + False, + "All bounding boxes should have positive height and width." + f" Found invalid box {degen_bb} for target at index {target_idx}.", + ) + + features = self.backbone(images.tensors) + if isinstance(features, torch.Tensor): + features = OrderedDict([("0", features)]) + proposals, proposal_losses = self.rpn(images, features, targets) + detections, detector_losses = self.roi_heads(features, proposals, images.image_sizes, targets) + detections = self.transform.postprocess(detections, images.image_sizes, original_image_sizes) # type: ignore[operator] + + losses = {} + losses.update(detector_losses) + losses.update(proposal_losses) + + if torch.jit.is_scripting(): + if not self._has_warned: + warnings.warn("RCNN always returns a (Losses, Detections) tuple in scripting") + self._has_warned = True + return losses, detections + else: + return self.eager_outputs(losses, detections) diff --git a/lib/python3.12/site-packages/torchvision/models/detection/image_list.py b/lib/python3.12/site-packages/torchvision/models/detection/image_list.py new file mode 100644 index 0000000000000000000000000000000000000000..583866557e4c9ec178e7cc268272db3de1698e41 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/detection/image_list.py @@ -0,0 +1,25 @@ +from typing import List, Tuple + +import torch +from torch import Tensor + + +class ImageList: + """ + Structure that holds a list of images (of possibly + varying sizes) as a single tensor. + This works by padding the images to the same size, + and storing in a field the original sizes of each image + + Args: + tensors (tensor): Tensor containing images. + image_sizes (list[tuple[int, int]]): List of Tuples each containing size of images. + """ + + def __init__(self, tensors: Tensor, image_sizes: List[Tuple[int, int]]) -> None: + self.tensors = tensors + self.image_sizes = image_sizes + + def to(self, device: torch.device) -> "ImageList": + cast_tensor = self.tensors.to(device) + return ImageList(cast_tensor, self.image_sizes) diff --git a/lib/python3.12/site-packages/torchvision/models/detection/keypoint_rcnn.py b/lib/python3.12/site-packages/torchvision/models/detection/keypoint_rcnn.py new file mode 100644 index 0000000000000000000000000000000000000000..5d7ff0ea433a681064a11a22c3e276e253997772 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/detection/keypoint_rcnn.py @@ -0,0 +1,474 @@ +from typing import Any, Optional + +import torch +from torch import nn +from torchvision.ops import MultiScaleRoIAlign + +from ...ops import misc as misc_nn_ops +from ...transforms._presets import ObjectDetection +from .._api import register_model, Weights, WeightsEnum +from .._meta import _COCO_PERSON_CATEGORIES, _COCO_PERSON_KEYPOINT_NAMES +from .._utils import _ovewrite_value_param, handle_legacy_interface +from ..resnet import resnet50, ResNet50_Weights +from ._utils import overwrite_eps +from .backbone_utils import _resnet_fpn_extractor, _validate_trainable_layers +from .faster_rcnn import FasterRCNN + + +__all__ = [ + "KeypointRCNN", + "KeypointRCNN_ResNet50_FPN_Weights", + "keypointrcnn_resnet50_fpn", +] + + +class KeypointRCNN(FasterRCNN): + """ + Implements Keypoint R-CNN. + + The input to the model is expected to be a list of tensors, each of shape [C, H, W], one for each + image, and should be in 0-1 range. Different images can have different sizes. + + The behavior of the model changes depending on if it is in training or evaluation mode. + + During training, the model expects both the input tensors and targets (list of dictionary), + containing: + + - boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (Int64Tensor[N]): the class label for each ground-truth box + - keypoints (FloatTensor[N, K, 3]): the K keypoints location for each of the N instances, in the + format [x, y, visibility], where visibility=0 means that the keypoint is not visible. + + The model returns a Dict[Tensor] during training, containing the classification and regression + losses for both the RPN and the R-CNN, and the keypoint loss. + + During inference, the model requires only the input tensors, and returns the post-processed + predictions as a List[Dict[Tensor]], one for each input image. The fields of the Dict are as + follows: + + - boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (Int64Tensor[N]): the predicted labels for each image + - scores (Tensor[N]): the scores or each prediction + - keypoints (FloatTensor[N, K, 3]): the locations of the predicted keypoints, in [x, y, v] format. + + Args: + backbone (nn.Module): the network used to compute the features for the model. + It should contain an out_channels attribute, which indicates the number of output + channels that each feature map has (and it should be the same for all feature maps). + The backbone should return a single Tensor or and OrderedDict[Tensor]. + num_classes (int): number of output classes of the model (including the background). + If box_predictor is specified, num_classes should be None. + min_size (int): Images are rescaled before feeding them to the backbone: + we attempt to preserve the aspect ratio and scale the shorter edge + to ``min_size``. If the resulting longer edge exceeds ``max_size``, + then downscale so that the longer edge does not exceed ``max_size``. + This may result in the shorter edge beeing lower than ``min_size``. + max_size (int): See ``min_size``. + image_mean (Tuple[float, float, float]): mean values used for input normalization. + They are generally the mean values of the dataset on which the backbone has been trained + on + image_std (Tuple[float, float, float]): std values used for input normalization. + They are generally the std values of the dataset on which the backbone has been trained on + rpn_anchor_generator (AnchorGenerator): module that generates the anchors for a set of feature + maps. + rpn_head (nn.Module): module that computes the objectness and regression deltas from the RPN + rpn_pre_nms_top_n_train (int): number of proposals to keep before applying NMS during training + rpn_pre_nms_top_n_test (int): number of proposals to keep before applying NMS during testing + rpn_post_nms_top_n_train (int): number of proposals to keep after applying NMS during training + rpn_post_nms_top_n_test (int): number of proposals to keep after applying NMS during testing + rpn_nms_thresh (float): NMS threshold used for postprocessing the RPN proposals + rpn_fg_iou_thresh (float): minimum IoU between the anchor and the GT box so that they can be + considered as positive during training of the RPN. + rpn_bg_iou_thresh (float): maximum IoU between the anchor and the GT box so that they can be + considered as negative during training of the RPN. + rpn_batch_size_per_image (int): number of anchors that are sampled during training of the RPN + for computing the loss + rpn_positive_fraction (float): proportion of positive anchors in a mini-batch during training + of the RPN + rpn_score_thresh (float): only return proposals with an objectness score greater than rpn_score_thresh + box_roi_pool (MultiScaleRoIAlign): the module which crops and resizes the feature maps in + the locations indicated by the bounding boxes + box_head (nn.Module): module that takes the cropped feature maps as input + box_predictor (nn.Module): module that takes the output of box_head and returns the + classification logits and box regression deltas. + box_score_thresh (float): during inference, only return proposals with a classification score + greater than box_score_thresh + box_nms_thresh (float): NMS threshold for the prediction head. Used during inference + box_detections_per_img (int): maximum number of detections per image, for all classes. + box_fg_iou_thresh (float): minimum IoU between the proposals and the GT box so that they can be + considered as positive during training of the classification head + box_bg_iou_thresh (float): maximum IoU between the proposals and the GT box so that they can be + considered as negative during training of the classification head + box_batch_size_per_image (int): number of proposals that are sampled during training of the + classification head + box_positive_fraction (float): proportion of positive proposals in a mini-batch during training + of the classification head + bbox_reg_weights (Tuple[float, float, float, float]): weights for the encoding/decoding of the + bounding boxes + keypoint_roi_pool (MultiScaleRoIAlign): the module which crops and resizes the feature maps in + the locations indicated by the bounding boxes, which will be used for the keypoint head. + keypoint_head (nn.Module): module that takes the cropped feature maps as input + keypoint_predictor (nn.Module): module that takes the output of the keypoint_head and returns the + heatmap logits + + Example:: + + >>> import torch + >>> import torchvision + >>> from torchvision.models.detection import KeypointRCNN + >>> from torchvision.models.detection.anchor_utils import AnchorGenerator + >>> + >>> # load a pre-trained model for classification and return + >>> # only the features + >>> backbone = torchvision.models.mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT).features + >>> # KeypointRCNN needs to know the number of + >>> # output channels in a backbone. For mobilenet_v2, it's 1280, + >>> # so we need to add it here + >>> backbone.out_channels = 1280 + >>> + >>> # let's make the RPN generate 5 x 3 anchors per spatial + >>> # location, with 5 different sizes and 3 different aspect + >>> # ratios. We have a Tuple[Tuple[int]] because each feature + >>> # map could potentially have different sizes and + >>> # aspect ratios + >>> anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),), + >>> aspect_ratios=((0.5, 1.0, 2.0),)) + >>> + >>> # let's define what are the feature maps that we will + >>> # use to perform the region of interest cropping, as well as + >>> # the size of the crop after rescaling. + >>> # if your backbone returns a Tensor, featmap_names is expected to + >>> # be ['0']. More generally, the backbone should return an + >>> # OrderedDict[Tensor], and in featmap_names you can choose which + >>> # feature maps to use. + >>> roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], + >>> output_size=7, + >>> sampling_ratio=2) + >>> + >>> keypoint_roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], + >>> output_size=14, + >>> sampling_ratio=2) + >>> # put the pieces together inside a KeypointRCNN model + >>> model = KeypointRCNN(backbone, + >>> num_classes=2, + >>> rpn_anchor_generator=anchor_generator, + >>> box_roi_pool=roi_pooler, + >>> keypoint_roi_pool=keypoint_roi_pooler) + >>> model.eval() + >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)] + >>> predictions = model(x) + """ + + def __init__( + self, + backbone, + num_classes=None, + # transform parameters + min_size=None, + max_size=1333, + image_mean=None, + image_std=None, + # RPN parameters + rpn_anchor_generator=None, + rpn_head=None, + rpn_pre_nms_top_n_train=2000, + rpn_pre_nms_top_n_test=1000, + rpn_post_nms_top_n_train=2000, + rpn_post_nms_top_n_test=1000, + rpn_nms_thresh=0.7, + rpn_fg_iou_thresh=0.7, + rpn_bg_iou_thresh=0.3, + rpn_batch_size_per_image=256, + rpn_positive_fraction=0.5, + rpn_score_thresh=0.0, + # Box parameters + box_roi_pool=None, + box_head=None, + box_predictor=None, + box_score_thresh=0.05, + box_nms_thresh=0.5, + box_detections_per_img=100, + box_fg_iou_thresh=0.5, + box_bg_iou_thresh=0.5, + box_batch_size_per_image=512, + box_positive_fraction=0.25, + bbox_reg_weights=None, + # keypoint parameters + keypoint_roi_pool=None, + keypoint_head=None, + keypoint_predictor=None, + num_keypoints=None, + **kwargs, + ): + + if not isinstance(keypoint_roi_pool, (MultiScaleRoIAlign, type(None))): + raise TypeError( + "keypoint_roi_pool should be of type MultiScaleRoIAlign or None instead of {type(keypoint_roi_pool)}" + ) + if min_size is None: + min_size = (640, 672, 704, 736, 768, 800) + + if num_keypoints is not None: + if keypoint_predictor is not None: + raise ValueError("num_keypoints should be None when keypoint_predictor is specified") + else: + num_keypoints = 17 + + out_channels = backbone.out_channels + + if keypoint_roi_pool is None: + keypoint_roi_pool = MultiScaleRoIAlign(featmap_names=["0", "1", "2", "3"], output_size=14, sampling_ratio=2) + + if keypoint_head is None: + keypoint_layers = tuple(512 for _ in range(8)) + keypoint_head = KeypointRCNNHeads(out_channels, keypoint_layers) + + if keypoint_predictor is None: + keypoint_dim_reduced = 512 # == keypoint_layers[-1] + keypoint_predictor = KeypointRCNNPredictor(keypoint_dim_reduced, num_keypoints) + + super().__init__( + backbone, + num_classes, + # transform parameters + min_size, + max_size, + image_mean, + image_std, + # RPN-specific parameters + rpn_anchor_generator, + rpn_head, + rpn_pre_nms_top_n_train, + rpn_pre_nms_top_n_test, + rpn_post_nms_top_n_train, + rpn_post_nms_top_n_test, + rpn_nms_thresh, + rpn_fg_iou_thresh, + rpn_bg_iou_thresh, + rpn_batch_size_per_image, + rpn_positive_fraction, + rpn_score_thresh, + # Box parameters + box_roi_pool, + box_head, + box_predictor, + box_score_thresh, + box_nms_thresh, + box_detections_per_img, + box_fg_iou_thresh, + box_bg_iou_thresh, + box_batch_size_per_image, + box_positive_fraction, + bbox_reg_weights, + **kwargs, + ) + + self.roi_heads.keypoint_roi_pool = keypoint_roi_pool + self.roi_heads.keypoint_head = keypoint_head + self.roi_heads.keypoint_predictor = keypoint_predictor + + +class KeypointRCNNHeads(nn.Sequential): + def __init__(self, in_channels, layers): + d = [] + next_feature = in_channels + for out_channels in layers: + d.append(nn.Conv2d(next_feature, out_channels, 3, stride=1, padding=1)) + d.append(nn.ReLU(inplace=True)) + next_feature = out_channels + super().__init__(*d) + for m in self.children(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") + nn.init.constant_(m.bias, 0) + + +class KeypointRCNNPredictor(nn.Module): + def __init__(self, in_channels, num_keypoints): + super().__init__() + input_features = in_channels + deconv_kernel = 4 + self.kps_score_lowres = nn.ConvTranspose2d( + input_features, + num_keypoints, + deconv_kernel, + stride=2, + padding=deconv_kernel // 2 - 1, + ) + nn.init.kaiming_normal_(self.kps_score_lowres.weight, mode="fan_out", nonlinearity="relu") + nn.init.constant_(self.kps_score_lowres.bias, 0) + self.up_scale = 2 + self.out_channels = num_keypoints + + def forward(self, x): + x = self.kps_score_lowres(x) + return torch.nn.functional.interpolate( + x, scale_factor=float(self.up_scale), mode="bilinear", align_corners=False, recompute_scale_factor=False + ) + + +_COMMON_META = { + "categories": _COCO_PERSON_CATEGORIES, + "keypoint_names": _COCO_PERSON_KEYPOINT_NAMES, + "min_size": (1, 1), +} + + +class KeypointRCNN_ResNet50_FPN_Weights(WeightsEnum): + COCO_LEGACY = Weights( + url="https://download.pytorch.org/models/keypointrcnn_resnet50_fpn_coco-9f466800.pth", + transforms=ObjectDetection, + meta={ + **_COMMON_META, + "num_params": 59137258, + "recipe": "https://github.com/pytorch/vision/issues/1606", + "_metrics": { + "COCO-val2017": { + "box_map": 50.6, + "kp_map": 61.1, + } + }, + "_ops": 133.924, + "_file_size": 226.054, + "_docs": """ + These weights were produced by following a similar training recipe as on the paper but use a checkpoint + from an early epoch. + """, + }, + ) + COCO_V1 = Weights( + url="https://download.pytorch.org/models/keypointrcnn_resnet50_fpn_coco-fc266e95.pth", + transforms=ObjectDetection, + meta={ + **_COMMON_META, + "num_params": 59137258, + "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#keypoint-r-cnn", + "_metrics": { + "COCO-val2017": { + "box_map": 54.6, + "kp_map": 65.0, + } + }, + "_ops": 137.42, + "_file_size": 226.054, + "_docs": """These weights were produced by following a similar training recipe as on the paper.""", + }, + ) + DEFAULT = COCO_V1 + + +@register_model() +@handle_legacy_interface( + weights=( + "pretrained", + lambda kwargs: KeypointRCNN_ResNet50_FPN_Weights.COCO_LEGACY + if kwargs["pretrained"] == "legacy" + else KeypointRCNN_ResNet50_FPN_Weights.COCO_V1, + ), + weights_backbone=("pretrained_backbone", ResNet50_Weights.IMAGENET1K_V1), +) +def keypointrcnn_resnet50_fpn( + *, + weights: Optional[KeypointRCNN_ResNet50_FPN_Weights] = None, + progress: bool = True, + num_classes: Optional[int] = None, + num_keypoints: Optional[int] = None, + weights_backbone: Optional[ResNet50_Weights] = ResNet50_Weights.IMAGENET1K_V1, + trainable_backbone_layers: Optional[int] = None, + **kwargs: Any, +) -> KeypointRCNN: + """ + Constructs a Keypoint R-CNN model with a ResNet-50-FPN backbone. + + .. betastatus:: detection module + + Reference: `Mask R-CNN `__. + + The input to the model is expected to be a list of tensors, each of shape ``[C, H, W]``, one for each + image, and should be in ``0-1`` range. Different images can have different sizes. + + The behavior of the model changes depending on if it is in training or evaluation mode. + + During training, the model expects both the input tensors and targets (list of dictionary), + containing: + + - boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (``Int64Tensor[N]``): the class label for each ground-truth box + - keypoints (``FloatTensor[N, K, 3]``): the ``K`` keypoints location for each of the ``N`` instances, in the + format ``[x, y, visibility]``, where ``visibility=0`` means that the keypoint is not visible. + + The model returns a ``Dict[Tensor]`` during training, containing the classification and regression + losses for both the RPN and the R-CNN, and the keypoint loss. + + During inference, the model requires only the input tensors, and returns the post-processed + predictions as a ``List[Dict[Tensor]]``, one for each input image. The fields of the ``Dict`` are as + follows, where ``N`` is the number of detected instances: + + - boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (``Int64Tensor[N]``): the predicted labels for each instance + - scores (``Tensor[N]``): the scores or each instance + - keypoints (``FloatTensor[N, K, 3]``): the locations of the predicted keypoints, in ``[x, y, v]`` format. + + For more details on the output, you may refer to :ref:`instance_seg_output`. + + Keypoint R-CNN is exportable to ONNX for a fixed batch size with inputs images of fixed size. + + Example:: + + >>> model = torchvision.models.detection.keypointrcnn_resnet50_fpn(weights=KeypointRCNN_ResNet50_FPN_Weights.DEFAULT) + >>> model.eval() + >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)] + >>> predictions = model(x) + >>> + >>> # optionally, if you want to export the model to ONNX: + >>> torch.onnx.export(model, x, "keypoint_rcnn.onnx", opset_version = 11) + + Args: + weights (:class:`~torchvision.models.detection.KeypointRCNN_ResNet50_FPN_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.detection.KeypointRCNN_ResNet50_FPN_Weights` + below for more details, and possible values. By default, no + pre-trained weights are used. + progress (bool): If True, displays a progress bar of the download to stderr + num_classes (int, optional): number of output classes of the model (including the background) + num_keypoints (int, optional): number of keypoints + weights_backbone (:class:`~torchvision.models.ResNet50_Weights`, optional): The + pretrained weights for the backbone. + trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from final block. + Valid values are between 0 and 5, with 5 meaning all backbone layers are trainable. If ``None`` is + passed (the default) this value is set to 3. + + .. autoclass:: torchvision.models.detection.KeypointRCNN_ResNet50_FPN_Weights + :members: + """ + weights = KeypointRCNN_ResNet50_FPN_Weights.verify(weights) + weights_backbone = ResNet50_Weights.verify(weights_backbone) + + if weights is not None: + weights_backbone = None + num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"])) + num_keypoints = _ovewrite_value_param("num_keypoints", num_keypoints, len(weights.meta["keypoint_names"])) + else: + if num_classes is None: + num_classes = 2 + if num_keypoints is None: + num_keypoints = 17 + + is_trained = weights is not None or weights_backbone is not None + trainable_backbone_layers = _validate_trainable_layers(is_trained, trainable_backbone_layers, 5, 3) + norm_layer = misc_nn_ops.FrozenBatchNorm2d if is_trained else nn.BatchNorm2d + + backbone = resnet50(weights=weights_backbone, progress=progress, norm_layer=norm_layer) + backbone = _resnet_fpn_extractor(backbone, trainable_backbone_layers) + model = KeypointRCNN(backbone, num_classes, num_keypoints=num_keypoints, **kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + if weights == KeypointRCNN_ResNet50_FPN_Weights.COCO_V1: + overwrite_eps(model, 0.0) + + return model diff --git a/lib/python3.12/site-packages/torchvision/models/detection/mask_rcnn.py b/lib/python3.12/site-packages/torchvision/models/detection/mask_rcnn.py new file mode 100644 index 0000000000000000000000000000000000000000..cdabbfd26ca8bbefaefdb6fb8b098afac217b595 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/detection/mask_rcnn.py @@ -0,0 +1,590 @@ +from collections import OrderedDict +from typing import Any, Callable, Optional + +from torch import nn +from torchvision.ops import MultiScaleRoIAlign + +from ...ops import misc as misc_nn_ops +from ...transforms._presets import ObjectDetection +from .._api import register_model, Weights, WeightsEnum +from .._meta import _COCO_CATEGORIES +from .._utils import _ovewrite_value_param, handle_legacy_interface +from ..resnet import resnet50, ResNet50_Weights +from ._utils import overwrite_eps +from .backbone_utils import _resnet_fpn_extractor, _validate_trainable_layers +from .faster_rcnn import _default_anchorgen, FasterRCNN, FastRCNNConvFCHead, RPNHead + + +__all__ = [ + "MaskRCNN", + "MaskRCNN_ResNet50_FPN_Weights", + "MaskRCNN_ResNet50_FPN_V2_Weights", + "maskrcnn_resnet50_fpn", + "maskrcnn_resnet50_fpn_v2", +] + + +class MaskRCNN(FasterRCNN): + """ + Implements Mask R-CNN. + + The input to the model is expected to be a list of tensors, each of shape [C, H, W], one for each + image, and should be in 0-1 range. Different images can have different sizes. + + The behavior of the model changes depending on if it is in training or evaluation mode. + + During training, the model expects both the input tensors and targets (list of dictionary), + containing: + - boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (Int64Tensor[N]): the class label for each ground-truth box + - masks (UInt8Tensor[N, H, W]): the segmentation binary masks for each instance + + The model returns a Dict[Tensor] during training, containing the classification and regression + losses for both the RPN and the R-CNN, and the mask loss. + + During inference, the model requires only the input tensors, and returns the post-processed + predictions as a List[Dict[Tensor]], one for each input image. The fields of the Dict are as + follows: + - boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (Int64Tensor[N]): the predicted labels for each image + - scores (Tensor[N]): the scores or each prediction + - masks (UInt8Tensor[N, 1, H, W]): the predicted masks for each instance, in 0-1 range. In order to + obtain the final segmentation masks, the soft masks can be thresholded, generally + with a value of 0.5 (mask >= 0.5) + + Args: + backbone (nn.Module): the network used to compute the features for the model. + It should contain an out_channels attribute, which indicates the number of output + channels that each feature map has (and it should be the same for all feature maps). + The backbone should return a single Tensor or and OrderedDict[Tensor]. + num_classes (int): number of output classes of the model (including the background). + If box_predictor is specified, num_classes should be None. + min_size (int): Images are rescaled before feeding them to the backbone: + we attempt to preserve the aspect ratio and scale the shorter edge + to ``min_size``. If the resulting longer edge exceeds ``max_size``, + then downscale so that the longer edge does not exceed ``max_size``. + This may result in the shorter edge beeing lower than ``min_size``. + max_size (int): See ``min_size``. + image_mean (Tuple[float, float, float]): mean values used for input normalization. + They are generally the mean values of the dataset on which the backbone has been trained + on + image_std (Tuple[float, float, float]): std values used for input normalization. + They are generally the std values of the dataset on which the backbone has been trained on + rpn_anchor_generator (AnchorGenerator): module that generates the anchors for a set of feature + maps. + rpn_head (nn.Module): module that computes the objectness and regression deltas from the RPN + rpn_pre_nms_top_n_train (int): number of proposals to keep before applying NMS during training + rpn_pre_nms_top_n_test (int): number of proposals to keep before applying NMS during testing + rpn_post_nms_top_n_train (int): number of proposals to keep after applying NMS during training + rpn_post_nms_top_n_test (int): number of proposals to keep after applying NMS during testing + rpn_nms_thresh (float): NMS threshold used for postprocessing the RPN proposals + rpn_fg_iou_thresh (float): minimum IoU between the anchor and the GT box so that they can be + considered as positive during training of the RPN. + rpn_bg_iou_thresh (float): maximum IoU between the anchor and the GT box so that they can be + considered as negative during training of the RPN. + rpn_batch_size_per_image (int): number of anchors that are sampled during training of the RPN + for computing the loss + rpn_positive_fraction (float): proportion of positive anchors in a mini-batch during training + of the RPN + rpn_score_thresh (float): only return proposals with an objectness score greater than rpn_score_thresh + box_roi_pool (MultiScaleRoIAlign): the module which crops and resizes the feature maps in + the locations indicated by the bounding boxes + box_head (nn.Module): module that takes the cropped feature maps as input + box_predictor (nn.Module): module that takes the output of box_head and returns the + classification logits and box regression deltas. + box_score_thresh (float): during inference, only return proposals with a classification score + greater than box_score_thresh + box_nms_thresh (float): NMS threshold for the prediction head. Used during inference + box_detections_per_img (int): maximum number of detections per image, for all classes. + box_fg_iou_thresh (float): minimum IoU between the proposals and the GT box so that they can be + considered as positive during training of the classification head + box_bg_iou_thresh (float): maximum IoU between the proposals and the GT box so that they can be + considered as negative during training of the classification head + box_batch_size_per_image (int): number of proposals that are sampled during training of the + classification head + box_positive_fraction (float): proportion of positive proposals in a mini-batch during training + of the classification head + bbox_reg_weights (Tuple[float, float, float, float]): weights for the encoding/decoding of the + bounding boxes + mask_roi_pool (MultiScaleRoIAlign): the module which crops and resizes the feature maps in + the locations indicated by the bounding boxes, which will be used for the mask head. + mask_head (nn.Module): module that takes the cropped feature maps as input + mask_predictor (nn.Module): module that takes the output of the mask_head and returns the + segmentation mask logits + + Example:: + + >>> import torch + >>> import torchvision + >>> from torchvision.models.detection import MaskRCNN + >>> from torchvision.models.detection.anchor_utils import AnchorGenerator + >>> + >>> # load a pre-trained model for classification and return + >>> # only the features + >>> backbone = torchvision.models.mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT).features + >>> # MaskRCNN needs to know the number of + >>> # output channels in a backbone. For mobilenet_v2, it's 1280 + >>> # so we need to add it here, + >>> backbone.out_channels = 1280 + >>> + >>> # let's make the RPN generate 5 x 3 anchors per spatial + >>> # location, with 5 different sizes and 3 different aspect + >>> # ratios. We have a Tuple[Tuple[int]] because each feature + >>> # map could potentially have different sizes and + >>> # aspect ratios + >>> anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),), + >>> aspect_ratios=((0.5, 1.0, 2.0),)) + >>> + >>> # let's define what are the feature maps that we will + >>> # use to perform the region of interest cropping, as well as + >>> # the size of the crop after rescaling. + >>> # if your backbone returns a Tensor, featmap_names is expected to + >>> # be ['0']. More generally, the backbone should return an + >>> # OrderedDict[Tensor], and in featmap_names you can choose which + >>> # feature maps to use. + >>> roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], + >>> output_size=7, + >>> sampling_ratio=2) + >>> + >>> mask_roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], + >>> output_size=14, + >>> sampling_ratio=2) + >>> # put the pieces together inside a MaskRCNN model + >>> model = MaskRCNN(backbone, + >>> num_classes=2, + >>> rpn_anchor_generator=anchor_generator, + >>> box_roi_pool=roi_pooler, + >>> mask_roi_pool=mask_roi_pooler) + >>> model.eval() + >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)] + >>> predictions = model(x) + """ + + def __init__( + self, + backbone, + num_classes=None, + # transform parameters + min_size=800, + max_size=1333, + image_mean=None, + image_std=None, + # RPN parameters + rpn_anchor_generator=None, + rpn_head=None, + rpn_pre_nms_top_n_train=2000, + rpn_pre_nms_top_n_test=1000, + rpn_post_nms_top_n_train=2000, + rpn_post_nms_top_n_test=1000, + rpn_nms_thresh=0.7, + rpn_fg_iou_thresh=0.7, + rpn_bg_iou_thresh=0.3, + rpn_batch_size_per_image=256, + rpn_positive_fraction=0.5, + rpn_score_thresh=0.0, + # Box parameters + box_roi_pool=None, + box_head=None, + box_predictor=None, + box_score_thresh=0.05, + box_nms_thresh=0.5, + box_detections_per_img=100, + box_fg_iou_thresh=0.5, + box_bg_iou_thresh=0.5, + box_batch_size_per_image=512, + box_positive_fraction=0.25, + bbox_reg_weights=None, + # Mask parameters + mask_roi_pool=None, + mask_head=None, + mask_predictor=None, + **kwargs, + ): + + if not isinstance(mask_roi_pool, (MultiScaleRoIAlign, type(None))): + raise TypeError( + f"mask_roi_pool should be of type MultiScaleRoIAlign or None instead of {type(mask_roi_pool)}" + ) + + if num_classes is not None: + if mask_predictor is not None: + raise ValueError("num_classes should be None when mask_predictor is specified") + + out_channels = backbone.out_channels + + if mask_roi_pool is None: + mask_roi_pool = MultiScaleRoIAlign(featmap_names=["0", "1", "2", "3"], output_size=14, sampling_ratio=2) + + if mask_head is None: + mask_layers = (256, 256, 256, 256) + mask_dilation = 1 + mask_head = MaskRCNNHeads(out_channels, mask_layers, mask_dilation) + + if mask_predictor is None: + mask_predictor_in_channels = 256 # == mask_layers[-1] + mask_dim_reduced = 256 + mask_predictor = MaskRCNNPredictor(mask_predictor_in_channels, mask_dim_reduced, num_classes) + + super().__init__( + backbone, + num_classes, + # transform parameters + min_size, + max_size, + image_mean, + image_std, + # RPN-specific parameters + rpn_anchor_generator, + rpn_head, + rpn_pre_nms_top_n_train, + rpn_pre_nms_top_n_test, + rpn_post_nms_top_n_train, + rpn_post_nms_top_n_test, + rpn_nms_thresh, + rpn_fg_iou_thresh, + rpn_bg_iou_thresh, + rpn_batch_size_per_image, + rpn_positive_fraction, + rpn_score_thresh, + # Box parameters + box_roi_pool, + box_head, + box_predictor, + box_score_thresh, + box_nms_thresh, + box_detections_per_img, + box_fg_iou_thresh, + box_bg_iou_thresh, + box_batch_size_per_image, + box_positive_fraction, + bbox_reg_weights, + **kwargs, + ) + + self.roi_heads.mask_roi_pool = mask_roi_pool + self.roi_heads.mask_head = mask_head + self.roi_heads.mask_predictor = mask_predictor + + +class MaskRCNNHeads(nn.Sequential): + _version = 2 + + def __init__(self, in_channels, layers, dilation, norm_layer: Optional[Callable[..., nn.Module]] = None): + """ + Args: + in_channels (int): number of input channels + layers (list): feature dimensions of each FCN layer + dilation (int): dilation rate of kernel + norm_layer (callable, optional): Module specifying the normalization layer to use. Default: None + """ + blocks = [] + next_feature = in_channels + for layer_features in layers: + blocks.append( + misc_nn_ops.Conv2dNormActivation( + next_feature, + layer_features, + kernel_size=3, + stride=1, + padding=dilation, + dilation=dilation, + norm_layer=norm_layer, + ) + ) + next_feature = layer_features + + super().__init__(*blocks) + for layer in self.modules(): + if isinstance(layer, nn.Conv2d): + nn.init.kaiming_normal_(layer.weight, mode="fan_out", nonlinearity="relu") + if layer.bias is not None: + nn.init.zeros_(layer.bias) + + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + version = local_metadata.get("version", None) + + if version is None or version < 2: + num_blocks = len(self) + for i in range(num_blocks): + for type in ["weight", "bias"]: + old_key = f"{prefix}mask_fcn{i+1}.{type}" + new_key = f"{prefix}{i}.0.{type}" + if old_key in state_dict: + state_dict[new_key] = state_dict.pop(old_key) + + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + +class MaskRCNNPredictor(nn.Sequential): + def __init__(self, in_channels, dim_reduced, num_classes): + super().__init__( + OrderedDict( + [ + ("conv5_mask", nn.ConvTranspose2d(in_channels, dim_reduced, 2, 2, 0)), + ("relu", nn.ReLU(inplace=True)), + ("mask_fcn_logits", nn.Conv2d(dim_reduced, num_classes, 1, 1, 0)), + ] + ) + ) + + for name, param in self.named_parameters(): + if "weight" in name: + nn.init.kaiming_normal_(param, mode="fan_out", nonlinearity="relu") + # elif "bias" in name: + # nn.init.constant_(param, 0) + + +_COMMON_META = { + "categories": _COCO_CATEGORIES, + "min_size": (1, 1), +} + + +class MaskRCNN_ResNet50_FPN_Weights(WeightsEnum): + COCO_V1 = Weights( + url="https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth", + transforms=ObjectDetection, + meta={ + **_COMMON_META, + "num_params": 44401393, + "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#mask-r-cnn", + "_metrics": { + "COCO-val2017": { + "box_map": 37.9, + "mask_map": 34.6, + } + }, + "_ops": 134.38, + "_file_size": 169.84, + "_docs": """These weights were produced by following a similar training recipe as on the paper.""", + }, + ) + DEFAULT = COCO_V1 + + +class MaskRCNN_ResNet50_FPN_V2_Weights(WeightsEnum): + COCO_V1 = Weights( + url="https://download.pytorch.org/models/maskrcnn_resnet50_fpn_v2_coco-73cbd019.pth", + transforms=ObjectDetection, + meta={ + **_COMMON_META, + "num_params": 46359409, + "recipe": "https://github.com/pytorch/vision/pull/5773", + "_metrics": { + "COCO-val2017": { + "box_map": 47.4, + "mask_map": 41.8, + } + }, + "_ops": 333.577, + "_file_size": 177.219, + "_docs": """These weights were produced using an enhanced training recipe to boost the model accuracy.""", + }, + ) + DEFAULT = COCO_V1 + + +@register_model() +@handle_legacy_interface( + weights=("pretrained", MaskRCNN_ResNet50_FPN_Weights.COCO_V1), + weights_backbone=("pretrained_backbone", ResNet50_Weights.IMAGENET1K_V1), +) +def maskrcnn_resnet50_fpn( + *, + weights: Optional[MaskRCNN_ResNet50_FPN_Weights] = None, + progress: bool = True, + num_classes: Optional[int] = None, + weights_backbone: Optional[ResNet50_Weights] = ResNet50_Weights.IMAGENET1K_V1, + trainable_backbone_layers: Optional[int] = None, + **kwargs: Any, +) -> MaskRCNN: + """Mask R-CNN model with a ResNet-50-FPN backbone from the `Mask R-CNN + `_ paper. + + .. betastatus:: detection module + + The input to the model is expected to be a list of tensors, each of shape ``[C, H, W]``, one for each + image, and should be in ``0-1`` range. Different images can have different sizes. + + The behavior of the model changes depending on if it is in training or evaluation mode. + + During training, the model expects both the input tensors and targets (list of dictionary), + containing: + + - boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (``Int64Tensor[N]``): the class label for each ground-truth box + - masks (``UInt8Tensor[N, H, W]``): the segmentation binary masks for each instance + + The model returns a ``Dict[Tensor]`` during training, containing the classification and regression + losses for both the RPN and the R-CNN, and the mask loss. + + During inference, the model requires only the input tensors, and returns the post-processed + predictions as a ``List[Dict[Tensor]]``, one for each input image. The fields of the ``Dict`` are as + follows, where ``N`` is the number of detected instances: + + - boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (``Int64Tensor[N]``): the predicted labels for each instance + - scores (``Tensor[N]``): the scores or each instance + - masks (``UInt8Tensor[N, 1, H, W]``): the predicted masks for each instance, in ``0-1`` range. In order to + obtain the final segmentation masks, the soft masks can be thresholded, generally + with a value of 0.5 (``mask >= 0.5``) + + For more details on the output and on how to plot the masks, you may refer to :ref:`instance_seg_output`. + + Mask R-CNN is exportable to ONNX for a fixed batch size with inputs images of fixed size. + + Example:: + + >>> model = torchvision.models.detection.maskrcnn_resnet50_fpn(weights=MaskRCNN_ResNet50_FPN_Weights.DEFAULT) + >>> model.eval() + >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)] + >>> predictions = model(x) + >>> + >>> # optionally, if you want to export the model to ONNX: + >>> torch.onnx.export(model, x, "mask_rcnn.onnx", opset_version = 11) + + Args: + weights (:class:`~torchvision.models.detection.MaskRCNN_ResNet50_FPN_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.detection.MaskRCNN_ResNet50_FPN_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + num_classes (int, optional): number of output classes of the model (including the background) + weights_backbone (:class:`~torchvision.models.ResNet50_Weights`, optional): The + pretrained weights for the backbone. + trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from + final block. Valid values are between 0 and 5, with 5 meaning all backbone layers are + trainable. If ``None`` is passed (the default) this value is set to 3. + **kwargs: parameters passed to the ``torchvision.models.detection.mask_rcnn.MaskRCNN`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.detection.MaskRCNN_ResNet50_FPN_Weights + :members: + """ + weights = MaskRCNN_ResNet50_FPN_Weights.verify(weights) + weights_backbone = ResNet50_Weights.verify(weights_backbone) + + if weights is not None: + weights_backbone = None + num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"])) + elif num_classes is None: + num_classes = 91 + + is_trained = weights is not None or weights_backbone is not None + trainable_backbone_layers = _validate_trainable_layers(is_trained, trainable_backbone_layers, 5, 3) + norm_layer = misc_nn_ops.FrozenBatchNorm2d if is_trained else nn.BatchNorm2d + + backbone = resnet50(weights=weights_backbone, progress=progress, norm_layer=norm_layer) + backbone = _resnet_fpn_extractor(backbone, trainable_backbone_layers) + model = MaskRCNN(backbone, num_classes=num_classes, **kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + if weights == MaskRCNN_ResNet50_FPN_Weights.COCO_V1: + overwrite_eps(model, 0.0) + + return model + + +@register_model() +@handle_legacy_interface( + weights=("pretrained", MaskRCNN_ResNet50_FPN_V2_Weights.COCO_V1), + weights_backbone=("pretrained_backbone", ResNet50_Weights.IMAGENET1K_V1), +) +def maskrcnn_resnet50_fpn_v2( + *, + weights: Optional[MaskRCNN_ResNet50_FPN_V2_Weights] = None, + progress: bool = True, + num_classes: Optional[int] = None, + weights_backbone: Optional[ResNet50_Weights] = None, + trainable_backbone_layers: Optional[int] = None, + **kwargs: Any, +) -> MaskRCNN: + """Improved Mask R-CNN model with a ResNet-50-FPN backbone from the `Benchmarking Detection Transfer + Learning with Vision Transformers `_ paper. + + .. betastatus:: detection module + + :func:`~torchvision.models.detection.maskrcnn_resnet50_fpn` for more details. + + Args: + weights (:class:`~torchvision.models.detection.MaskRCNN_ResNet50_FPN_V2_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.detection.MaskRCNN_ResNet50_FPN_V2_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + num_classes (int, optional): number of output classes of the model (including the background) + weights_backbone (:class:`~torchvision.models.ResNet50_Weights`, optional): The + pretrained weights for the backbone. + trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from + final block. Valid values are between 0 and 5, with 5 meaning all backbone layers are + trainable. If ``None`` is passed (the default) this value is set to 3. + **kwargs: parameters passed to the ``torchvision.models.detection.mask_rcnn.MaskRCNN`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.detection.MaskRCNN_ResNet50_FPN_V2_Weights + :members: + """ + weights = MaskRCNN_ResNet50_FPN_V2_Weights.verify(weights) + weights_backbone = ResNet50_Weights.verify(weights_backbone) + + if weights is not None: + weights_backbone = None + num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"])) + elif num_classes is None: + num_classes = 91 + + is_trained = weights is not None or weights_backbone is not None + trainable_backbone_layers = _validate_trainable_layers(is_trained, trainable_backbone_layers, 5, 3) + + backbone = resnet50(weights=weights_backbone, progress=progress) + backbone = _resnet_fpn_extractor(backbone, trainable_backbone_layers, norm_layer=nn.BatchNorm2d) + rpn_anchor_generator = _default_anchorgen() + rpn_head = RPNHead(backbone.out_channels, rpn_anchor_generator.num_anchors_per_location()[0], conv_depth=2) + box_head = FastRCNNConvFCHead( + (backbone.out_channels, 7, 7), [256, 256, 256, 256], [1024], norm_layer=nn.BatchNorm2d + ) + mask_head = MaskRCNNHeads(backbone.out_channels, [256, 256, 256, 256], 1, norm_layer=nn.BatchNorm2d) + model = MaskRCNN( + backbone, + num_classes=num_classes, + rpn_anchor_generator=rpn_anchor_generator, + rpn_head=rpn_head, + box_head=box_head, + mask_head=mask_head, + **kwargs, + ) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model diff --git a/lib/python3.12/site-packages/torchvision/models/detection/retinanet.py b/lib/python3.12/site-packages/torchvision/models/detection/retinanet.py new file mode 100644 index 0000000000000000000000000000000000000000..a8cc7755014b6010965108a46c080f71b2d609db --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/detection/retinanet.py @@ -0,0 +1,903 @@ +import math +import warnings +from collections import OrderedDict +from functools import partial +from typing import Any, Callable, Dict, List, Optional, Tuple + +import torch +from torch import nn, Tensor + +from ...ops import boxes as box_ops, misc as misc_nn_ops, sigmoid_focal_loss +from ...ops.feature_pyramid_network import LastLevelP6P7 +from ...transforms._presets import ObjectDetection +from ...utils import _log_api_usage_once +from .._api import register_model, Weights, WeightsEnum +from .._meta import _COCO_CATEGORIES +from .._utils import _ovewrite_value_param, handle_legacy_interface +from ..resnet import resnet50, ResNet50_Weights +from . import _utils as det_utils +from ._utils import _box_loss, overwrite_eps +from .anchor_utils import AnchorGenerator +from .backbone_utils import _resnet_fpn_extractor, _validate_trainable_layers +from .transform import GeneralizedRCNNTransform + + +__all__ = [ + "RetinaNet", + "RetinaNet_ResNet50_FPN_Weights", + "RetinaNet_ResNet50_FPN_V2_Weights", + "retinanet_resnet50_fpn", + "retinanet_resnet50_fpn_v2", +] + + +def _sum(x: List[Tensor]) -> Tensor: + res = x[0] + for i in x[1:]: + res = res + i + return res + + +def _v1_to_v2_weights(state_dict, prefix): + for i in range(4): + for type in ["weight", "bias"]: + old_key = f"{prefix}conv.{2*i}.{type}" + new_key = f"{prefix}conv.{i}.0.{type}" + if old_key in state_dict: + state_dict[new_key] = state_dict.pop(old_key) + + +def _default_anchorgen(): + anchor_sizes = tuple((x, int(x * 2 ** (1.0 / 3)), int(x * 2 ** (2.0 / 3))) for x in [32, 64, 128, 256, 512]) + aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes) + anchor_generator = AnchorGenerator(anchor_sizes, aspect_ratios) + return anchor_generator + + +class RetinaNetHead(nn.Module): + """ + A regression and classification head for use in RetinaNet. + + Args: + in_channels (int): number of channels of the input feature + num_anchors (int): number of anchors to be predicted + num_classes (int): number of classes to be predicted + norm_layer (callable, optional): Module specifying the normalization layer to use. Default: None + """ + + def __init__(self, in_channels, num_anchors, num_classes, norm_layer: Optional[Callable[..., nn.Module]] = None): + super().__init__() + self.classification_head = RetinaNetClassificationHead( + in_channels, num_anchors, num_classes, norm_layer=norm_layer + ) + self.regression_head = RetinaNetRegressionHead(in_channels, num_anchors, norm_layer=norm_layer) + + def compute_loss(self, targets, head_outputs, anchors, matched_idxs): + # type: (List[Dict[str, Tensor]], Dict[str, Tensor], List[Tensor], List[Tensor]) -> Dict[str, Tensor] + return { + "classification": self.classification_head.compute_loss(targets, head_outputs, matched_idxs), + "bbox_regression": self.regression_head.compute_loss(targets, head_outputs, anchors, matched_idxs), + } + + def forward(self, x): + # type: (List[Tensor]) -> Dict[str, Tensor] + return {"cls_logits": self.classification_head(x), "bbox_regression": self.regression_head(x)} + + +class RetinaNetClassificationHead(nn.Module): + """ + A classification head for use in RetinaNet. + + Args: + in_channels (int): number of channels of the input feature + num_anchors (int): number of anchors to be predicted + num_classes (int): number of classes to be predicted + norm_layer (callable, optional): Module specifying the normalization layer to use. Default: None + """ + + _version = 2 + + def __init__( + self, + in_channels, + num_anchors, + num_classes, + prior_probability=0.01, + norm_layer: Optional[Callable[..., nn.Module]] = None, + ): + super().__init__() + + conv = [] + for _ in range(4): + conv.append(misc_nn_ops.Conv2dNormActivation(in_channels, in_channels, norm_layer=norm_layer)) + self.conv = nn.Sequential(*conv) + + for layer in self.conv.modules(): + if isinstance(layer, nn.Conv2d): + torch.nn.init.normal_(layer.weight, std=0.01) + if layer.bias is not None: + torch.nn.init.constant_(layer.bias, 0) + + self.cls_logits = nn.Conv2d(in_channels, num_anchors * num_classes, kernel_size=3, stride=1, padding=1) + torch.nn.init.normal_(self.cls_logits.weight, std=0.01) + torch.nn.init.constant_(self.cls_logits.bias, -math.log((1 - prior_probability) / prior_probability)) + + self.num_classes = num_classes + self.num_anchors = num_anchors + + # This is to fix using det_utils.Matcher.BETWEEN_THRESHOLDS in TorchScript. + # TorchScript doesn't support class attributes. + # https://github.com/pytorch/vision/pull/1697#issuecomment-630255584 + self.BETWEEN_THRESHOLDS = det_utils.Matcher.BETWEEN_THRESHOLDS + + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + version = local_metadata.get("version", None) + + if version is None or version < 2: + _v1_to_v2_weights(state_dict, prefix) + + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + def compute_loss(self, targets, head_outputs, matched_idxs): + # type: (List[Dict[str, Tensor]], Dict[str, Tensor], List[Tensor]) -> Tensor + losses = [] + + cls_logits = head_outputs["cls_logits"] + + for targets_per_image, cls_logits_per_image, matched_idxs_per_image in zip(targets, cls_logits, matched_idxs): + # determine only the foreground + foreground_idxs_per_image = matched_idxs_per_image >= 0 + num_foreground = foreground_idxs_per_image.sum() + + # create the target classification + gt_classes_target = torch.zeros_like(cls_logits_per_image) + gt_classes_target[ + foreground_idxs_per_image, + targets_per_image["labels"][matched_idxs_per_image[foreground_idxs_per_image]], + ] = 1.0 + + # find indices for which anchors should be ignored + valid_idxs_per_image = matched_idxs_per_image != self.BETWEEN_THRESHOLDS + + # compute the classification loss + losses.append( + sigmoid_focal_loss( + cls_logits_per_image[valid_idxs_per_image], + gt_classes_target[valid_idxs_per_image], + reduction="sum", + ) + / max(1, num_foreground) + ) + + return _sum(losses) / len(targets) + + def forward(self, x): + # type: (List[Tensor]) -> Tensor + all_cls_logits = [] + + for features in x: + cls_logits = self.conv(features) + cls_logits = self.cls_logits(cls_logits) + + # Permute classification output from (N, A * K, H, W) to (N, HWA, K). + N, _, H, W = cls_logits.shape + cls_logits = cls_logits.view(N, -1, self.num_classes, H, W) + cls_logits = cls_logits.permute(0, 3, 4, 1, 2) + cls_logits = cls_logits.reshape(N, -1, self.num_classes) # Size=(N, HWA, 4) + + all_cls_logits.append(cls_logits) + + return torch.cat(all_cls_logits, dim=1) + + +class RetinaNetRegressionHead(nn.Module): + """ + A regression head for use in RetinaNet. + + Args: + in_channels (int): number of channels of the input feature + num_anchors (int): number of anchors to be predicted + norm_layer (callable, optional): Module specifying the normalization layer to use. Default: None + """ + + _version = 2 + + __annotations__ = { + "box_coder": det_utils.BoxCoder, + } + + def __init__(self, in_channels, num_anchors, norm_layer: Optional[Callable[..., nn.Module]] = None): + super().__init__() + + conv = [] + for _ in range(4): + conv.append(misc_nn_ops.Conv2dNormActivation(in_channels, in_channels, norm_layer=norm_layer)) + self.conv = nn.Sequential(*conv) + + self.bbox_reg = nn.Conv2d(in_channels, num_anchors * 4, kernel_size=3, stride=1, padding=1) + torch.nn.init.normal_(self.bbox_reg.weight, std=0.01) + torch.nn.init.zeros_(self.bbox_reg.bias) + + for layer in self.conv.modules(): + if isinstance(layer, nn.Conv2d): + torch.nn.init.normal_(layer.weight, std=0.01) + if layer.bias is not None: + torch.nn.init.zeros_(layer.bias) + + self.box_coder = det_utils.BoxCoder(weights=(1.0, 1.0, 1.0, 1.0)) + self._loss_type = "l1" + + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + version = local_metadata.get("version", None) + + if version is None or version < 2: + _v1_to_v2_weights(state_dict, prefix) + + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + def compute_loss(self, targets, head_outputs, anchors, matched_idxs): + # type: (List[Dict[str, Tensor]], Dict[str, Tensor], List[Tensor], List[Tensor]) -> Tensor + losses = [] + + bbox_regression = head_outputs["bbox_regression"] + + for targets_per_image, bbox_regression_per_image, anchors_per_image, matched_idxs_per_image in zip( + targets, bbox_regression, anchors, matched_idxs + ): + # determine only the foreground indices, ignore the rest + foreground_idxs_per_image = torch.where(matched_idxs_per_image >= 0)[0] + num_foreground = foreground_idxs_per_image.numel() + + # select only the foreground boxes + matched_gt_boxes_per_image = targets_per_image["boxes"][matched_idxs_per_image[foreground_idxs_per_image]] + bbox_regression_per_image = bbox_regression_per_image[foreground_idxs_per_image, :] + anchors_per_image = anchors_per_image[foreground_idxs_per_image, :] + + # compute the loss + losses.append( + _box_loss( + self._loss_type, + self.box_coder, + anchors_per_image, + matched_gt_boxes_per_image, + bbox_regression_per_image, + ) + / max(1, num_foreground) + ) + + return _sum(losses) / max(1, len(targets)) + + def forward(self, x): + # type: (List[Tensor]) -> Tensor + all_bbox_regression = [] + + for features in x: + bbox_regression = self.conv(features) + bbox_regression = self.bbox_reg(bbox_regression) + + # Permute bbox regression output from (N, 4 * A, H, W) to (N, HWA, 4). + N, _, H, W = bbox_regression.shape + bbox_regression = bbox_regression.view(N, -1, 4, H, W) + bbox_regression = bbox_regression.permute(0, 3, 4, 1, 2) + bbox_regression = bbox_regression.reshape(N, -1, 4) # Size=(N, HWA, 4) + + all_bbox_regression.append(bbox_regression) + + return torch.cat(all_bbox_regression, dim=1) + + +class RetinaNet(nn.Module): + """ + Implements RetinaNet. + + The input to the model is expected to be a list of tensors, each of shape [C, H, W], one for each + image, and should be in 0-1 range. Different images can have different sizes. + + The behavior of the model changes depending on if it is in training or evaluation mode. + + During training, the model expects both the input tensors and targets (list of dictionary), + containing: + - boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (Int64Tensor[N]): the class label for each ground-truth box + + The model returns a Dict[Tensor] during training, containing the classification and regression + losses. + + During inference, the model requires only the input tensors, and returns the post-processed + predictions as a List[Dict[Tensor]], one for each input image. The fields of the Dict are as + follows: + - boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (Int64Tensor[N]): the predicted labels for each image + - scores (Tensor[N]): the scores for each prediction + + Args: + backbone (nn.Module): the network used to compute the features for the model. + It should contain an out_channels attribute, which indicates the number of output + channels that each feature map has (and it should be the same for all feature maps). + The backbone should return a single Tensor or an OrderedDict[Tensor]. + num_classes (int): number of output classes of the model (including the background). + min_size (int): Images are rescaled before feeding them to the backbone: + we attempt to preserve the aspect ratio and scale the shorter edge + to ``min_size``. If the resulting longer edge exceeds ``max_size``, + then downscale so that the longer edge does not exceed ``max_size``. + This may result in the shorter edge beeing lower than ``min_size``. + max_size (int): See ``min_size``. + image_mean (Tuple[float, float, float]): mean values used for input normalization. + They are generally the mean values of the dataset on which the backbone has been trained + on + image_std (Tuple[float, float, float]): std values used for input normalization. + They are generally the std values of the dataset on which the backbone has been trained on + anchor_generator (AnchorGenerator): module that generates the anchors for a set of feature + maps. + head (nn.Module): Module run on top of the feature pyramid. + Defaults to a module containing a classification and regression module. + score_thresh (float): Score threshold used for postprocessing the detections. + nms_thresh (float): NMS threshold used for postprocessing the detections. + detections_per_img (int): Number of best detections to keep after NMS. + fg_iou_thresh (float): minimum IoU between the anchor and the GT box so that they can be + considered as positive during training. + bg_iou_thresh (float): maximum IoU between the anchor and the GT box so that they can be + considered as negative during training. + topk_candidates (int): Number of best detections to keep before NMS. + + Example: + + >>> import torch + >>> import torchvision + >>> from torchvision.models.detection import RetinaNet + >>> from torchvision.models.detection.anchor_utils import AnchorGenerator + >>> # load a pre-trained model for classification and return + >>> # only the features + >>> backbone = torchvision.models.mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT).features + >>> # RetinaNet needs to know the number of + >>> # output channels in a backbone. For mobilenet_v2, it's 1280, + >>> # so we need to add it here + >>> backbone.out_channels = 1280 + >>> + >>> # let's make the network generate 5 x 3 anchors per spatial + >>> # location, with 5 different sizes and 3 different aspect + >>> # ratios. We have a Tuple[Tuple[int]] because each feature + >>> # map could potentially have different sizes and + >>> # aspect ratios + >>> anchor_generator = AnchorGenerator( + >>> sizes=((32, 64, 128, 256, 512),), + >>> aspect_ratios=((0.5, 1.0, 2.0),) + >>> ) + >>> + >>> # put the pieces together inside a RetinaNet model + >>> model = RetinaNet(backbone, + >>> num_classes=2, + >>> anchor_generator=anchor_generator) + >>> model.eval() + >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)] + >>> predictions = model(x) + """ + + __annotations__ = { + "box_coder": det_utils.BoxCoder, + "proposal_matcher": det_utils.Matcher, + } + + def __init__( + self, + backbone, + num_classes, + # transform parameters + min_size=800, + max_size=1333, + image_mean=None, + image_std=None, + # Anchor parameters + anchor_generator=None, + head=None, + proposal_matcher=None, + score_thresh=0.05, + nms_thresh=0.5, + detections_per_img=300, + fg_iou_thresh=0.5, + bg_iou_thresh=0.4, + topk_candidates=1000, + **kwargs, + ): + super().__init__() + _log_api_usage_once(self) + + if not hasattr(backbone, "out_channels"): + raise ValueError( + "backbone should contain an attribute out_channels " + "specifying the number of output channels (assumed to be the " + "same for all the levels)" + ) + self.backbone = backbone + + if not isinstance(anchor_generator, (AnchorGenerator, type(None))): + raise TypeError( + f"anchor_generator should be of type AnchorGenerator or None instead of {type(anchor_generator)}" + ) + + if anchor_generator is None: + anchor_generator = _default_anchorgen() + self.anchor_generator = anchor_generator + + if head is None: + head = RetinaNetHead(backbone.out_channels, anchor_generator.num_anchors_per_location()[0], num_classes) + self.head = head + + if proposal_matcher is None: + proposal_matcher = det_utils.Matcher( + fg_iou_thresh, + bg_iou_thresh, + allow_low_quality_matches=True, + ) + self.proposal_matcher = proposal_matcher + + self.box_coder = det_utils.BoxCoder(weights=(1.0, 1.0, 1.0, 1.0)) + + if image_mean is None: + image_mean = [0.485, 0.456, 0.406] + if image_std is None: + image_std = [0.229, 0.224, 0.225] + self.transform = GeneralizedRCNNTransform(min_size, max_size, image_mean, image_std, **kwargs) + + self.score_thresh = score_thresh + self.nms_thresh = nms_thresh + self.detections_per_img = detections_per_img + self.topk_candidates = topk_candidates + + # used only on torchscript mode + self._has_warned = False + + @torch.jit.unused + def eager_outputs(self, losses, detections): + # type: (Dict[str, Tensor], List[Dict[str, Tensor]]) -> Tuple[Dict[str, Tensor], List[Dict[str, Tensor]]] + if self.training: + return losses + + return detections + + def compute_loss(self, targets, head_outputs, anchors): + # type: (List[Dict[str, Tensor]], Dict[str, Tensor], List[Tensor]) -> Dict[str, Tensor] + matched_idxs = [] + for anchors_per_image, targets_per_image in zip(anchors, targets): + if targets_per_image["boxes"].numel() == 0: + matched_idxs.append( + torch.full((anchors_per_image.size(0),), -1, dtype=torch.int64, device=anchors_per_image.device) + ) + continue + + match_quality_matrix = box_ops.box_iou(targets_per_image["boxes"], anchors_per_image) + matched_idxs.append(self.proposal_matcher(match_quality_matrix)) + + return self.head.compute_loss(targets, head_outputs, anchors, matched_idxs) + + def postprocess_detections(self, head_outputs, anchors, image_shapes): + # type: (Dict[str, List[Tensor]], List[List[Tensor]], List[Tuple[int, int]]) -> List[Dict[str, Tensor]] + class_logits = head_outputs["cls_logits"] + box_regression = head_outputs["bbox_regression"] + + num_images = len(image_shapes) + + detections: List[Dict[str, Tensor]] = [] + + for index in range(num_images): + box_regression_per_image = [br[index] for br in box_regression] + logits_per_image = [cl[index] for cl in class_logits] + anchors_per_image, image_shape = anchors[index], image_shapes[index] + + image_boxes = [] + image_scores = [] + image_labels = [] + + for box_regression_per_level, logits_per_level, anchors_per_level in zip( + box_regression_per_image, logits_per_image, anchors_per_image + ): + num_classes = logits_per_level.shape[-1] + + # remove low scoring boxes + scores_per_level = torch.sigmoid(logits_per_level).flatten() + keep_idxs = scores_per_level > self.score_thresh + scores_per_level = scores_per_level[keep_idxs] + topk_idxs = torch.where(keep_idxs)[0] + + # keep only topk scoring predictions + num_topk = det_utils._topk_min(topk_idxs, self.topk_candidates, 0) + scores_per_level, idxs = scores_per_level.topk(num_topk) + topk_idxs = topk_idxs[idxs] + + anchor_idxs = torch.div(topk_idxs, num_classes, rounding_mode="floor") + labels_per_level = topk_idxs % num_classes + + boxes_per_level = self.box_coder.decode_single( + box_regression_per_level[anchor_idxs], anchors_per_level[anchor_idxs] + ) + boxes_per_level = box_ops.clip_boxes_to_image(boxes_per_level, image_shape) + + image_boxes.append(boxes_per_level) + image_scores.append(scores_per_level) + image_labels.append(labels_per_level) + + image_boxes = torch.cat(image_boxes, dim=0) + image_scores = torch.cat(image_scores, dim=0) + image_labels = torch.cat(image_labels, dim=0) + + # non-maximum suppression + keep = box_ops.batched_nms(image_boxes, image_scores, image_labels, self.nms_thresh) + keep = keep[: self.detections_per_img] + + detections.append( + { + "boxes": image_boxes[keep], + "scores": image_scores[keep], + "labels": image_labels[keep], + } + ) + + return detections + + def forward(self, images, targets=None): + # type: (List[Tensor], Optional[List[Dict[str, Tensor]]]) -> Tuple[Dict[str, Tensor], List[Dict[str, Tensor]]] + """ + Args: + images (list[Tensor]): images to be processed + targets (list[Dict[Tensor]]): ground-truth boxes present in the image (optional) + + Returns: + result (list[BoxList] or dict[Tensor]): the output from the model. + During training, it returns a dict[Tensor] which contains the losses. + During testing, it returns list[BoxList] contains additional fields + like `scores`, `labels` and `mask` (for Mask R-CNN models). + + """ + if self.training: + if targets is None: + torch._assert(False, "targets should not be none when in training mode") + else: + for target in targets: + boxes = target["boxes"] + torch._assert(isinstance(boxes, torch.Tensor), "Expected target boxes to be of type Tensor.") + torch._assert( + len(boxes.shape) == 2 and boxes.shape[-1] == 4, + "Expected target boxes to be a tensor of shape [N, 4].", + ) + + # get the original image sizes + original_image_sizes: List[Tuple[int, int]] = [] + for img in images: + val = img.shape[-2:] + torch._assert( + len(val) == 2, + f"expecting the last two dimensions of the Tensor to be H and W instead got {img.shape[-2:]}", + ) + original_image_sizes.append((val[0], val[1])) + + # transform the input + images, targets = self.transform(images, targets) + + # Check for degenerate boxes + # TODO: Move this to a function + if targets is not None: + for target_idx, target in enumerate(targets): + boxes = target["boxes"] + degenerate_boxes = boxes[:, 2:] <= boxes[:, :2] + if degenerate_boxes.any(): + # print the first degenerate box + bb_idx = torch.where(degenerate_boxes.any(dim=1))[0][0] + degen_bb: List[float] = boxes[bb_idx].tolist() + torch._assert( + False, + "All bounding boxes should have positive height and width." + f" Found invalid box {degen_bb} for target at index {target_idx}.", + ) + + # get the features from the backbone + features = self.backbone(images.tensors) + if isinstance(features, torch.Tensor): + features = OrderedDict([("0", features)]) + + # TODO: Do we want a list or a dict? + features = list(features.values()) + + # compute the retinanet heads outputs using the features + head_outputs = self.head(features) + + # create the set of anchors + anchors = self.anchor_generator(images, features) + + losses = {} + detections: List[Dict[str, Tensor]] = [] + if self.training: + if targets is None: + torch._assert(False, "targets should not be none when in training mode") + else: + # compute the losses + losses = self.compute_loss(targets, head_outputs, anchors) + else: + # recover level sizes + num_anchors_per_level = [x.size(2) * x.size(3) for x in features] + HW = 0 + for v in num_anchors_per_level: + HW += v + HWA = head_outputs["cls_logits"].size(1) + A = HWA // HW + num_anchors_per_level = [hw * A for hw in num_anchors_per_level] + + # split outputs per level + split_head_outputs: Dict[str, List[Tensor]] = {} + for k in head_outputs: + split_head_outputs[k] = list(head_outputs[k].split(num_anchors_per_level, dim=1)) + split_anchors = [list(a.split(num_anchors_per_level)) for a in anchors] + + # compute the detections + detections = self.postprocess_detections(split_head_outputs, split_anchors, images.image_sizes) + detections = self.transform.postprocess(detections, images.image_sizes, original_image_sizes) + + if torch.jit.is_scripting(): + if not self._has_warned: + warnings.warn("RetinaNet always returns a (Losses, Detections) tuple in scripting") + self._has_warned = True + return losses, detections + return self.eager_outputs(losses, detections) + + +_COMMON_META = { + "categories": _COCO_CATEGORIES, + "min_size": (1, 1), +} + + +class RetinaNet_ResNet50_FPN_Weights(WeightsEnum): + COCO_V1 = Weights( + url="https://download.pytorch.org/models/retinanet_resnet50_fpn_coco-eeacb38b.pth", + transforms=ObjectDetection, + meta={ + **_COMMON_META, + "num_params": 34014999, + "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#retinanet", + "_metrics": { + "COCO-val2017": { + "box_map": 36.4, + } + }, + "_ops": 151.54, + "_file_size": 130.267, + "_docs": """These weights were produced by following a similar training recipe as on the paper.""", + }, + ) + DEFAULT = COCO_V1 + + +class RetinaNet_ResNet50_FPN_V2_Weights(WeightsEnum): + COCO_V1 = Weights( + url="https://download.pytorch.org/models/retinanet_resnet50_fpn_v2_coco-5905b1c5.pth", + transforms=ObjectDetection, + meta={ + **_COMMON_META, + "num_params": 38198935, + "recipe": "https://github.com/pytorch/vision/pull/5756", + "_metrics": { + "COCO-val2017": { + "box_map": 41.5, + } + }, + "_ops": 152.238, + "_file_size": 146.037, + "_docs": """These weights were produced using an enhanced training recipe to boost the model accuracy.""", + }, + ) + DEFAULT = COCO_V1 + + +@register_model() +@handle_legacy_interface( + weights=("pretrained", RetinaNet_ResNet50_FPN_Weights.COCO_V1), + weights_backbone=("pretrained_backbone", ResNet50_Weights.IMAGENET1K_V1), +) +def retinanet_resnet50_fpn( + *, + weights: Optional[RetinaNet_ResNet50_FPN_Weights] = None, + progress: bool = True, + num_classes: Optional[int] = None, + weights_backbone: Optional[ResNet50_Weights] = ResNet50_Weights.IMAGENET1K_V1, + trainable_backbone_layers: Optional[int] = None, + **kwargs: Any, +) -> RetinaNet: + """ + Constructs a RetinaNet model with a ResNet-50-FPN backbone. + + .. betastatus:: detection module + + Reference: `Focal Loss for Dense Object Detection `_. + + The input to the model is expected to be a list of tensors, each of shape ``[C, H, W]``, one for each + image, and should be in ``0-1`` range. Different images can have different sizes. + + The behavior of the model changes depending on if it is in training or evaluation mode. + + During training, the model expects both the input tensors and targets (list of dictionary), + containing: + + - boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (``Int64Tensor[N]``): the class label for each ground-truth box + + The model returns a ``Dict[Tensor]`` during training, containing the classification and regression + losses. + + During inference, the model requires only the input tensors, and returns the post-processed + predictions as a ``List[Dict[Tensor]]``, one for each input image. The fields of the ``Dict`` are as + follows, where ``N`` is the number of detections: + + - boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (``Int64Tensor[N]``): the predicted labels for each detection + - scores (``Tensor[N]``): the scores of each detection + + For more details on the output, you may refer to :ref:`instance_seg_output`. + + Example:: + + >>> model = torchvision.models.detection.retinanet_resnet50_fpn(weights=RetinaNet_ResNet50_FPN_Weights.DEFAULT) + >>> model.eval() + >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)] + >>> predictions = model(x) + + Args: + weights (:class:`~torchvision.models.detection.RetinaNet_ResNet50_FPN_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.detection.RetinaNet_ResNet50_FPN_Weights` + below for more details, and possible values. By default, no + pre-trained weights are used. + progress (bool): If True, displays a progress bar of the download to stderr. Default is True. + num_classes (int, optional): number of output classes of the model (including the background) + weights_backbone (:class:`~torchvision.models.ResNet50_Weights`, optional): The pretrained weights for + the backbone. + trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from final block. + Valid values are between 0 and 5, with 5 meaning all backbone layers are trainable. If ``None`` is + passed (the default) this value is set to 3. + **kwargs: parameters passed to the ``torchvision.models.detection.RetinaNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.detection.RetinaNet_ResNet50_FPN_Weights + :members: + """ + weights = RetinaNet_ResNet50_FPN_Weights.verify(weights) + weights_backbone = ResNet50_Weights.verify(weights_backbone) + + if weights is not None: + weights_backbone = None + num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"])) + elif num_classes is None: + num_classes = 91 + + is_trained = weights is not None or weights_backbone is not None + trainable_backbone_layers = _validate_trainable_layers(is_trained, trainable_backbone_layers, 5, 3) + norm_layer = misc_nn_ops.FrozenBatchNorm2d if is_trained else nn.BatchNorm2d + + backbone = resnet50(weights=weights_backbone, progress=progress, norm_layer=norm_layer) + # skip P2 because it generates too many anchors (according to their paper) + backbone = _resnet_fpn_extractor( + backbone, trainable_backbone_layers, returned_layers=[2, 3, 4], extra_blocks=LastLevelP6P7(256, 256) + ) + model = RetinaNet(backbone, num_classes, **kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + if weights == RetinaNet_ResNet50_FPN_Weights.COCO_V1: + overwrite_eps(model, 0.0) + + return model + + +@register_model() +@handle_legacy_interface( + weights=("pretrained", RetinaNet_ResNet50_FPN_V2_Weights.COCO_V1), + weights_backbone=("pretrained_backbone", ResNet50_Weights.IMAGENET1K_V1), +) +def retinanet_resnet50_fpn_v2( + *, + weights: Optional[RetinaNet_ResNet50_FPN_V2_Weights] = None, + progress: bool = True, + num_classes: Optional[int] = None, + weights_backbone: Optional[ResNet50_Weights] = None, + trainable_backbone_layers: Optional[int] = None, + **kwargs: Any, +) -> RetinaNet: + """ + Constructs an improved RetinaNet model with a ResNet-50-FPN backbone. + + .. betastatus:: detection module + + Reference: `Bridging the Gap Between Anchor-based and Anchor-free Detection via Adaptive Training Sample Selection + `_. + + :func:`~torchvision.models.detection.retinanet_resnet50_fpn` for more details. + + Args: + weights (:class:`~torchvision.models.detection.RetinaNet_ResNet50_FPN_V2_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.detection.RetinaNet_ResNet50_FPN_V2_Weights` + below for more details, and possible values. By default, no + pre-trained weights are used. + progress (bool): If True, displays a progress bar of the download to stderr. Default is True. + num_classes (int, optional): number of output classes of the model (including the background) + weights_backbone (:class:`~torchvision.models.ResNet50_Weights`, optional): The pretrained weights for + the backbone. + trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from final block. + Valid values are between 0 and 5, with 5 meaning all backbone layers are trainable. If ``None`` is + passed (the default) this value is set to 3. + **kwargs: parameters passed to the ``torchvision.models.detection.RetinaNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.detection.RetinaNet_ResNet50_FPN_V2_Weights + :members: + """ + weights = RetinaNet_ResNet50_FPN_V2_Weights.verify(weights) + weights_backbone = ResNet50_Weights.verify(weights_backbone) + + if weights is not None: + weights_backbone = None + num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"])) + elif num_classes is None: + num_classes = 91 + + is_trained = weights is not None or weights_backbone is not None + trainable_backbone_layers = _validate_trainable_layers(is_trained, trainable_backbone_layers, 5, 3) + + backbone = resnet50(weights=weights_backbone, progress=progress) + backbone = _resnet_fpn_extractor( + backbone, trainable_backbone_layers, returned_layers=[2, 3, 4], extra_blocks=LastLevelP6P7(2048, 256) + ) + anchor_generator = _default_anchorgen() + head = RetinaNetHead( + backbone.out_channels, + anchor_generator.num_anchors_per_location()[0], + num_classes, + norm_layer=partial(nn.GroupNorm, 32), + ) + head.regression_head._loss_type = "giou" + model = RetinaNet(backbone, num_classes, anchor_generator=anchor_generator, head=head, **kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model diff --git a/lib/python3.12/site-packages/torchvision/models/detection/roi_heads.py b/lib/python3.12/site-packages/torchvision/models/detection/roi_heads.py new file mode 100644 index 0000000000000000000000000000000000000000..51b210cb6f368c1f4914ffe99287efef6057cba4 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/detection/roi_heads.py @@ -0,0 +1,876 @@ +from typing import Dict, List, Optional, Tuple + +import torch +import torch.nn.functional as F +import torchvision +from torch import nn, Tensor +from torchvision.ops import boxes as box_ops, roi_align + +from . import _utils as det_utils + + +def fastrcnn_loss(class_logits, box_regression, labels, regression_targets): + # type: (Tensor, Tensor, List[Tensor], List[Tensor]) -> Tuple[Tensor, Tensor] + """ + Computes the loss for Faster R-CNN. + + Args: + class_logits (Tensor) + box_regression (Tensor) + labels (list[BoxList]) + regression_targets (Tensor) + + Returns: + classification_loss (Tensor) + box_loss (Tensor) + """ + + labels = torch.cat(labels, dim=0) + regression_targets = torch.cat(regression_targets, dim=0) + + classification_loss = F.cross_entropy(class_logits, labels) + + # get indices that correspond to the regression targets for + # the corresponding ground truth labels, to be used with + # advanced indexing + sampled_pos_inds_subset = torch.where(labels > 0)[0] + labels_pos = labels[sampled_pos_inds_subset] + N, num_classes = class_logits.shape + box_regression = box_regression.reshape(N, box_regression.size(-1) // 4, 4) + + box_loss = F.smooth_l1_loss( + box_regression[sampled_pos_inds_subset, labels_pos], + regression_targets[sampled_pos_inds_subset], + beta=1 / 9, + reduction="sum", + ) + box_loss = box_loss / labels.numel() + + return classification_loss, box_loss + + +def maskrcnn_inference(x, labels): + # type: (Tensor, List[Tensor]) -> List[Tensor] + """ + From the results of the CNN, post process the masks + by taking the mask corresponding to the class with max + probability (which are of fixed size and directly output + by the CNN) and return the masks in the mask field of the BoxList. + + Args: + x (Tensor): the mask logits + labels (list[BoxList]): bounding boxes that are used as + reference, one for ech image + + Returns: + results (list[BoxList]): one BoxList for each image, containing + the extra field mask + """ + mask_prob = x.sigmoid() + + # select masks corresponding to the predicted classes + num_masks = x.shape[0] + boxes_per_image = [label.shape[0] for label in labels] + labels = torch.cat(labels) + index = torch.arange(num_masks, device=labels.device) + mask_prob = mask_prob[index, labels][:, None] + mask_prob = mask_prob.split(boxes_per_image, dim=0) + + return mask_prob + + +def project_masks_on_boxes(gt_masks, boxes, matched_idxs, M): + # type: (Tensor, Tensor, Tensor, int) -> Tensor + """ + Given segmentation masks and the bounding boxes corresponding + to the location of the masks in the image, this function + crops and resizes the masks in the position defined by the + boxes. This prepares the masks for them to be fed to the + loss computation as the targets. + """ + matched_idxs = matched_idxs.to(boxes) + rois = torch.cat([matched_idxs[:, None], boxes], dim=1) + gt_masks = gt_masks[:, None].to(rois) + return roi_align(gt_masks, rois, (M, M), 1.0)[:, 0] + + +def maskrcnn_loss(mask_logits, proposals, gt_masks, gt_labels, mask_matched_idxs): + # type: (Tensor, List[Tensor], List[Tensor], List[Tensor], List[Tensor]) -> Tensor + """ + Args: + proposals (list[BoxList]) + mask_logits (Tensor) + targets (list[BoxList]) + + Return: + mask_loss (Tensor): scalar tensor containing the loss + """ + + discretization_size = mask_logits.shape[-1] + labels = [gt_label[idxs] for gt_label, idxs in zip(gt_labels, mask_matched_idxs)] + mask_targets = [ + project_masks_on_boxes(m, p, i, discretization_size) for m, p, i in zip(gt_masks, proposals, mask_matched_idxs) + ] + + labels = torch.cat(labels, dim=0) + mask_targets = torch.cat(mask_targets, dim=0) + + # torch.mean (in binary_cross_entropy_with_logits) doesn't + # accept empty tensors, so handle it separately + if mask_targets.numel() == 0: + return mask_logits.sum() * 0 + + mask_loss = F.binary_cross_entropy_with_logits( + mask_logits[torch.arange(labels.shape[0], device=labels.device), labels], mask_targets + ) + return mask_loss + + +def keypoints_to_heatmap(keypoints, rois, heatmap_size): + # type: (Tensor, Tensor, int) -> Tuple[Tensor, Tensor] + offset_x = rois[:, 0] + offset_y = rois[:, 1] + scale_x = heatmap_size / (rois[:, 2] - rois[:, 0]) + scale_y = heatmap_size / (rois[:, 3] - rois[:, 1]) + + offset_x = offset_x[:, None] + offset_y = offset_y[:, None] + scale_x = scale_x[:, None] + scale_y = scale_y[:, None] + + x = keypoints[..., 0] + y = keypoints[..., 1] + + x_boundary_inds = x == rois[:, 2][:, None] + y_boundary_inds = y == rois[:, 3][:, None] + + x = (x - offset_x) * scale_x + x = x.floor().long() + y = (y - offset_y) * scale_y + y = y.floor().long() + + x[x_boundary_inds] = heatmap_size - 1 + y[y_boundary_inds] = heatmap_size - 1 + + valid_loc = (x >= 0) & (y >= 0) & (x < heatmap_size) & (y < heatmap_size) + vis = keypoints[..., 2] > 0 + valid = (valid_loc & vis).long() + + lin_ind = y * heatmap_size + x + heatmaps = lin_ind * valid + + return heatmaps, valid + + +def _onnx_heatmaps_to_keypoints( + maps, maps_i, roi_map_width, roi_map_height, widths_i, heights_i, offset_x_i, offset_y_i +): + num_keypoints = torch.scalar_tensor(maps.size(1), dtype=torch.int64) + + width_correction = widths_i / roi_map_width + height_correction = heights_i / roi_map_height + + roi_map = F.interpolate( + maps_i[:, None], size=(int(roi_map_height), int(roi_map_width)), mode="bicubic", align_corners=False + )[:, 0] + + w = torch.scalar_tensor(roi_map.size(2), dtype=torch.int64) + pos = roi_map.reshape(num_keypoints, -1).argmax(dim=1) + + x_int = pos % w + y_int = (pos - x_int) // w + + x = (torch.tensor(0.5, dtype=torch.float32) + x_int.to(dtype=torch.float32)) * width_correction.to( + dtype=torch.float32 + ) + y = (torch.tensor(0.5, dtype=torch.float32) + y_int.to(dtype=torch.float32)) * height_correction.to( + dtype=torch.float32 + ) + + xy_preds_i_0 = x + offset_x_i.to(dtype=torch.float32) + xy_preds_i_1 = y + offset_y_i.to(dtype=torch.float32) + xy_preds_i_2 = torch.ones(xy_preds_i_1.shape, dtype=torch.float32) + xy_preds_i = torch.stack( + [ + xy_preds_i_0.to(dtype=torch.float32), + xy_preds_i_1.to(dtype=torch.float32), + xy_preds_i_2.to(dtype=torch.float32), + ], + 0, + ) + + # TODO: simplify when indexing without rank will be supported by ONNX + base = num_keypoints * num_keypoints + num_keypoints + 1 + ind = torch.arange(num_keypoints) + ind = ind.to(dtype=torch.int64) * base + end_scores_i = ( + roi_map.index_select(1, y_int.to(dtype=torch.int64)) + .index_select(2, x_int.to(dtype=torch.int64)) + .view(-1) + .index_select(0, ind.to(dtype=torch.int64)) + ) + + return xy_preds_i, end_scores_i + + +@torch.jit._script_if_tracing +def _onnx_heatmaps_to_keypoints_loop( + maps, rois, widths_ceil, heights_ceil, widths, heights, offset_x, offset_y, num_keypoints +): + xy_preds = torch.zeros((0, 3, int(num_keypoints)), dtype=torch.float32, device=maps.device) + end_scores = torch.zeros((0, int(num_keypoints)), dtype=torch.float32, device=maps.device) + + for i in range(int(rois.size(0))): + xy_preds_i, end_scores_i = _onnx_heatmaps_to_keypoints( + maps, maps[i], widths_ceil[i], heights_ceil[i], widths[i], heights[i], offset_x[i], offset_y[i] + ) + xy_preds = torch.cat((xy_preds.to(dtype=torch.float32), xy_preds_i.unsqueeze(0).to(dtype=torch.float32)), 0) + end_scores = torch.cat( + (end_scores.to(dtype=torch.float32), end_scores_i.to(dtype=torch.float32).unsqueeze(0)), 0 + ) + return xy_preds, end_scores + + +def heatmaps_to_keypoints(maps, rois): + """Extract predicted keypoint locations from heatmaps. Output has shape + (#rois, 4, #keypoints) with the 4 rows corresponding to (x, y, logit, prob) + for each keypoint. + """ + # This function converts a discrete image coordinate in a HEATMAP_SIZE x + # HEATMAP_SIZE image to a continuous keypoint coordinate. We maintain + # consistency with keypoints_to_heatmap_labels by using the conversion from + # Heckbert 1990: c = d + 0.5, where d is a discrete coordinate and c is a + # continuous coordinate. + offset_x = rois[:, 0] + offset_y = rois[:, 1] + + widths = rois[:, 2] - rois[:, 0] + heights = rois[:, 3] - rois[:, 1] + widths = widths.clamp(min=1) + heights = heights.clamp(min=1) + widths_ceil = widths.ceil() + heights_ceil = heights.ceil() + + num_keypoints = maps.shape[1] + + if torchvision._is_tracing(): + xy_preds, end_scores = _onnx_heatmaps_to_keypoints_loop( + maps, + rois, + widths_ceil, + heights_ceil, + widths, + heights, + offset_x, + offset_y, + torch.scalar_tensor(num_keypoints, dtype=torch.int64), + ) + return xy_preds.permute(0, 2, 1), end_scores + + xy_preds = torch.zeros((len(rois), 3, num_keypoints), dtype=torch.float32, device=maps.device) + end_scores = torch.zeros((len(rois), num_keypoints), dtype=torch.float32, device=maps.device) + for i in range(len(rois)): + roi_map_width = int(widths_ceil[i].item()) + roi_map_height = int(heights_ceil[i].item()) + width_correction = widths[i] / roi_map_width + height_correction = heights[i] / roi_map_height + roi_map = F.interpolate( + maps[i][:, None], size=(roi_map_height, roi_map_width), mode="bicubic", align_corners=False + )[:, 0] + # roi_map_probs = scores_to_probs(roi_map.copy()) + w = roi_map.shape[2] + pos = roi_map.reshape(num_keypoints, -1).argmax(dim=1) + + x_int = pos % w + y_int = torch.div(pos - x_int, w, rounding_mode="floor") + # assert (roi_map_probs[k, y_int, x_int] == + # roi_map_probs[k, :, :].max()) + x = (x_int.float() + 0.5) * width_correction + y = (y_int.float() + 0.5) * height_correction + xy_preds[i, 0, :] = x + offset_x[i] + xy_preds[i, 1, :] = y + offset_y[i] + xy_preds[i, 2, :] = 1 + end_scores[i, :] = roi_map[torch.arange(num_keypoints, device=roi_map.device), y_int, x_int] + + return xy_preds.permute(0, 2, 1), end_scores + + +def keypointrcnn_loss(keypoint_logits, proposals, gt_keypoints, keypoint_matched_idxs): + # type: (Tensor, List[Tensor], List[Tensor], List[Tensor]) -> Tensor + N, K, H, W = keypoint_logits.shape + if H != W: + raise ValueError( + f"keypoint_logits height and width (last two elements of shape) should be equal. Instead got H = {H} and W = {W}" + ) + discretization_size = H + heatmaps = [] + valid = [] + for proposals_per_image, gt_kp_in_image, midx in zip(proposals, gt_keypoints, keypoint_matched_idxs): + kp = gt_kp_in_image[midx] + heatmaps_per_image, valid_per_image = keypoints_to_heatmap(kp, proposals_per_image, discretization_size) + heatmaps.append(heatmaps_per_image.view(-1)) + valid.append(valid_per_image.view(-1)) + + keypoint_targets = torch.cat(heatmaps, dim=0) + valid = torch.cat(valid, dim=0).to(dtype=torch.uint8) + valid = torch.where(valid)[0] + + # torch.mean (in binary_cross_entropy_with_logits) doesn't + # accept empty tensors, so handle it sepaartely + if keypoint_targets.numel() == 0 or len(valid) == 0: + return keypoint_logits.sum() * 0 + + keypoint_logits = keypoint_logits.view(N * K, H * W) + + keypoint_loss = F.cross_entropy(keypoint_logits[valid], keypoint_targets[valid]) + return keypoint_loss + + +def keypointrcnn_inference(x, boxes): + # type: (Tensor, List[Tensor]) -> Tuple[List[Tensor], List[Tensor]] + kp_probs = [] + kp_scores = [] + + boxes_per_image = [box.size(0) for box in boxes] + x2 = x.split(boxes_per_image, dim=0) + + for xx, bb in zip(x2, boxes): + kp_prob, scores = heatmaps_to_keypoints(xx, bb) + kp_probs.append(kp_prob) + kp_scores.append(scores) + + return kp_probs, kp_scores + + +def _onnx_expand_boxes(boxes, scale): + # type: (Tensor, float) -> Tensor + w_half = (boxes[:, 2] - boxes[:, 0]) * 0.5 + h_half = (boxes[:, 3] - boxes[:, 1]) * 0.5 + x_c = (boxes[:, 2] + boxes[:, 0]) * 0.5 + y_c = (boxes[:, 3] + boxes[:, 1]) * 0.5 + + w_half = w_half.to(dtype=torch.float32) * scale + h_half = h_half.to(dtype=torch.float32) * scale + + boxes_exp0 = x_c - w_half + boxes_exp1 = y_c - h_half + boxes_exp2 = x_c + w_half + boxes_exp3 = y_c + h_half + boxes_exp = torch.stack((boxes_exp0, boxes_exp1, boxes_exp2, boxes_exp3), 1) + return boxes_exp + + +# the next two functions should be merged inside Masker +# but are kept here for the moment while we need them +# temporarily for paste_mask_in_image +def expand_boxes(boxes, scale): + # type: (Tensor, float) -> Tensor + if torchvision._is_tracing(): + return _onnx_expand_boxes(boxes, scale) + w_half = (boxes[:, 2] - boxes[:, 0]) * 0.5 + h_half = (boxes[:, 3] - boxes[:, 1]) * 0.5 + x_c = (boxes[:, 2] + boxes[:, 0]) * 0.5 + y_c = (boxes[:, 3] + boxes[:, 1]) * 0.5 + + w_half *= scale + h_half *= scale + + boxes_exp = torch.zeros_like(boxes) + boxes_exp[:, 0] = x_c - w_half + boxes_exp[:, 2] = x_c + w_half + boxes_exp[:, 1] = y_c - h_half + boxes_exp[:, 3] = y_c + h_half + return boxes_exp + + +@torch.jit.unused +def expand_masks_tracing_scale(M, padding): + # type: (int, int) -> float + return torch.tensor(M + 2 * padding).to(torch.float32) / torch.tensor(M).to(torch.float32) + + +def expand_masks(mask, padding): + # type: (Tensor, int) -> Tuple[Tensor, float] + M = mask.shape[-1] + if torch._C._get_tracing_state(): # could not import is_tracing(), not sure why + scale = expand_masks_tracing_scale(M, padding) + else: + scale = float(M + 2 * padding) / M + padded_mask = F.pad(mask, (padding,) * 4) + return padded_mask, scale + + +def paste_mask_in_image(mask, box, im_h, im_w): + # type: (Tensor, Tensor, int, int) -> Tensor + TO_REMOVE = 1 + w = int(box[2] - box[0] + TO_REMOVE) + h = int(box[3] - box[1] + TO_REMOVE) + w = max(w, 1) + h = max(h, 1) + + # Set shape to [batchxCxHxW] + mask = mask.expand((1, 1, -1, -1)) + + # Resize mask + mask = F.interpolate(mask, size=(h, w), mode="bilinear", align_corners=False) + mask = mask[0][0] + + im_mask = torch.zeros((im_h, im_w), dtype=mask.dtype, device=mask.device) + x_0 = max(box[0], 0) + x_1 = min(box[2] + 1, im_w) + y_0 = max(box[1], 0) + y_1 = min(box[3] + 1, im_h) + + im_mask[y_0:y_1, x_0:x_1] = mask[(y_0 - box[1]) : (y_1 - box[1]), (x_0 - box[0]) : (x_1 - box[0])] + return im_mask + + +def _onnx_paste_mask_in_image(mask, box, im_h, im_w): + one = torch.ones(1, dtype=torch.int64) + zero = torch.zeros(1, dtype=torch.int64) + + w = box[2] - box[0] + one + h = box[3] - box[1] + one + w = torch.max(torch.cat((w, one))) + h = torch.max(torch.cat((h, one))) + + # Set shape to [batchxCxHxW] + mask = mask.expand((1, 1, mask.size(0), mask.size(1))) + + # Resize mask + mask = F.interpolate(mask, size=(int(h), int(w)), mode="bilinear", align_corners=False) + mask = mask[0][0] + + x_0 = torch.max(torch.cat((box[0].unsqueeze(0), zero))) + x_1 = torch.min(torch.cat((box[2].unsqueeze(0) + one, im_w.unsqueeze(0)))) + y_0 = torch.max(torch.cat((box[1].unsqueeze(0), zero))) + y_1 = torch.min(torch.cat((box[3].unsqueeze(0) + one, im_h.unsqueeze(0)))) + + unpaded_im_mask = mask[(y_0 - box[1]) : (y_1 - box[1]), (x_0 - box[0]) : (x_1 - box[0])] + + # TODO : replace below with a dynamic padding when support is added in ONNX + + # pad y + zeros_y0 = torch.zeros(y_0, unpaded_im_mask.size(1)) + zeros_y1 = torch.zeros(im_h - y_1, unpaded_im_mask.size(1)) + concat_0 = torch.cat((zeros_y0, unpaded_im_mask.to(dtype=torch.float32), zeros_y1), 0)[0:im_h, :] + # pad x + zeros_x0 = torch.zeros(concat_0.size(0), x_0) + zeros_x1 = torch.zeros(concat_0.size(0), im_w - x_1) + im_mask = torch.cat((zeros_x0, concat_0, zeros_x1), 1)[:, :im_w] + return im_mask + + +@torch.jit._script_if_tracing +def _onnx_paste_masks_in_image_loop(masks, boxes, im_h, im_w): + res_append = torch.zeros(0, im_h, im_w) + for i in range(masks.size(0)): + mask_res = _onnx_paste_mask_in_image(masks[i][0], boxes[i], im_h, im_w) + mask_res = mask_res.unsqueeze(0) + res_append = torch.cat((res_append, mask_res)) + return res_append + + +def paste_masks_in_image(masks, boxes, img_shape, padding=1): + # type: (Tensor, Tensor, Tuple[int, int], int) -> Tensor + masks, scale = expand_masks(masks, padding=padding) + boxes = expand_boxes(boxes, scale).to(dtype=torch.int64) + im_h, im_w = img_shape + + if torchvision._is_tracing(): + return _onnx_paste_masks_in_image_loop( + masks, boxes, torch.scalar_tensor(im_h, dtype=torch.int64), torch.scalar_tensor(im_w, dtype=torch.int64) + )[:, None] + res = [paste_mask_in_image(m[0], b, im_h, im_w) for m, b in zip(masks, boxes)] + if len(res) > 0: + ret = torch.stack(res, dim=0)[:, None] + else: + ret = masks.new_empty((0, 1, im_h, im_w)) + return ret + + +class RoIHeads(nn.Module): + __annotations__ = { + "box_coder": det_utils.BoxCoder, + "proposal_matcher": det_utils.Matcher, + "fg_bg_sampler": det_utils.BalancedPositiveNegativeSampler, + } + + def __init__( + self, + box_roi_pool, + box_head, + box_predictor, + # Faster R-CNN training + fg_iou_thresh, + bg_iou_thresh, + batch_size_per_image, + positive_fraction, + bbox_reg_weights, + # Faster R-CNN inference + score_thresh, + nms_thresh, + detections_per_img, + # Mask + mask_roi_pool=None, + mask_head=None, + mask_predictor=None, + keypoint_roi_pool=None, + keypoint_head=None, + keypoint_predictor=None, + ): + super().__init__() + + self.box_similarity = box_ops.box_iou + # assign ground-truth boxes for each proposal + self.proposal_matcher = det_utils.Matcher(fg_iou_thresh, bg_iou_thresh, allow_low_quality_matches=False) + + self.fg_bg_sampler = det_utils.BalancedPositiveNegativeSampler(batch_size_per_image, positive_fraction) + + if bbox_reg_weights is None: + bbox_reg_weights = (10.0, 10.0, 5.0, 5.0) + self.box_coder = det_utils.BoxCoder(bbox_reg_weights) + + self.box_roi_pool = box_roi_pool + self.box_head = box_head + self.box_predictor = box_predictor + + self.score_thresh = score_thresh + self.nms_thresh = nms_thresh + self.detections_per_img = detections_per_img + + self.mask_roi_pool = mask_roi_pool + self.mask_head = mask_head + self.mask_predictor = mask_predictor + + self.keypoint_roi_pool = keypoint_roi_pool + self.keypoint_head = keypoint_head + self.keypoint_predictor = keypoint_predictor + + def has_mask(self): + if self.mask_roi_pool is None: + return False + if self.mask_head is None: + return False + if self.mask_predictor is None: + return False + return True + + def has_keypoint(self): + if self.keypoint_roi_pool is None: + return False + if self.keypoint_head is None: + return False + if self.keypoint_predictor is None: + return False + return True + + def assign_targets_to_proposals(self, proposals, gt_boxes, gt_labels): + # type: (List[Tensor], List[Tensor], List[Tensor]) -> Tuple[List[Tensor], List[Tensor]] + matched_idxs = [] + labels = [] + for proposals_in_image, gt_boxes_in_image, gt_labels_in_image in zip(proposals, gt_boxes, gt_labels): + + if gt_boxes_in_image.numel() == 0: + # Background image + device = proposals_in_image.device + clamped_matched_idxs_in_image = torch.zeros( + (proposals_in_image.shape[0],), dtype=torch.int64, device=device + ) + labels_in_image = torch.zeros((proposals_in_image.shape[0],), dtype=torch.int64, device=device) + else: + # set to self.box_similarity when https://github.com/pytorch/pytorch/issues/27495 lands + match_quality_matrix = box_ops.box_iou(gt_boxes_in_image, proposals_in_image) + matched_idxs_in_image = self.proposal_matcher(match_quality_matrix) + + clamped_matched_idxs_in_image = matched_idxs_in_image.clamp(min=0) + + labels_in_image = gt_labels_in_image[clamped_matched_idxs_in_image] + labels_in_image = labels_in_image.to(dtype=torch.int64) + + # Label background (below the low threshold) + bg_inds = matched_idxs_in_image == self.proposal_matcher.BELOW_LOW_THRESHOLD + labels_in_image[bg_inds] = 0 + + # Label ignore proposals (between low and high thresholds) + ignore_inds = matched_idxs_in_image == self.proposal_matcher.BETWEEN_THRESHOLDS + labels_in_image[ignore_inds] = -1 # -1 is ignored by sampler + + matched_idxs.append(clamped_matched_idxs_in_image) + labels.append(labels_in_image) + return matched_idxs, labels + + def subsample(self, labels): + # type: (List[Tensor]) -> List[Tensor] + sampled_pos_inds, sampled_neg_inds = self.fg_bg_sampler(labels) + sampled_inds = [] + for img_idx, (pos_inds_img, neg_inds_img) in enumerate(zip(sampled_pos_inds, sampled_neg_inds)): + img_sampled_inds = torch.where(pos_inds_img | neg_inds_img)[0] + sampled_inds.append(img_sampled_inds) + return sampled_inds + + def add_gt_proposals(self, proposals, gt_boxes): + # type: (List[Tensor], List[Tensor]) -> List[Tensor] + proposals = [torch.cat((proposal, gt_box)) for proposal, gt_box in zip(proposals, gt_boxes)] + + return proposals + + def check_targets(self, targets): + # type: (Optional[List[Dict[str, Tensor]]]) -> None + if targets is None: + raise ValueError("targets should not be None") + if not all(["boxes" in t for t in targets]): + raise ValueError("Every element of targets should have a boxes key") + if not all(["labels" in t for t in targets]): + raise ValueError("Every element of targets should have a labels key") + if self.has_mask(): + if not all(["masks" in t for t in targets]): + raise ValueError("Every element of targets should have a masks key") + + def select_training_samples( + self, + proposals, # type: List[Tensor] + targets, # type: Optional[List[Dict[str, Tensor]]] + ): + # type: (...) -> Tuple[List[Tensor], List[Tensor], List[Tensor], List[Tensor]] + self.check_targets(targets) + if targets is None: + raise ValueError("targets should not be None") + dtype = proposals[0].dtype + device = proposals[0].device + + gt_boxes = [t["boxes"].to(dtype) for t in targets] + gt_labels = [t["labels"] for t in targets] + + # append ground-truth bboxes to propos + proposals = self.add_gt_proposals(proposals, gt_boxes) + + # get matching gt indices for each proposal + matched_idxs, labels = self.assign_targets_to_proposals(proposals, gt_boxes, gt_labels) + # sample a fixed proportion of positive-negative proposals + sampled_inds = self.subsample(labels) + matched_gt_boxes = [] + num_images = len(proposals) + for img_id in range(num_images): + img_sampled_inds = sampled_inds[img_id] + proposals[img_id] = proposals[img_id][img_sampled_inds] + labels[img_id] = labels[img_id][img_sampled_inds] + matched_idxs[img_id] = matched_idxs[img_id][img_sampled_inds] + + gt_boxes_in_image = gt_boxes[img_id] + if gt_boxes_in_image.numel() == 0: + gt_boxes_in_image = torch.zeros((1, 4), dtype=dtype, device=device) + matched_gt_boxes.append(gt_boxes_in_image[matched_idxs[img_id]]) + + regression_targets = self.box_coder.encode(matched_gt_boxes, proposals) + return proposals, matched_idxs, labels, regression_targets + + def postprocess_detections( + self, + class_logits, # type: Tensor + box_regression, # type: Tensor + proposals, # type: List[Tensor] + image_shapes, # type: List[Tuple[int, int]] + ): + # type: (...) -> Tuple[List[Tensor], List[Tensor], List[Tensor]] + device = class_logits.device + num_classes = class_logits.shape[-1] + + boxes_per_image = [boxes_in_image.shape[0] for boxes_in_image in proposals] + pred_boxes = self.box_coder.decode(box_regression, proposals) + + pred_scores = F.softmax(class_logits, -1) + + pred_boxes_list = pred_boxes.split(boxes_per_image, 0) + pred_scores_list = pred_scores.split(boxes_per_image, 0) + + all_boxes = [] + all_scores = [] + all_labels = [] + for boxes, scores, image_shape in zip(pred_boxes_list, pred_scores_list, image_shapes): + boxes = box_ops.clip_boxes_to_image(boxes, image_shape) + + # create labels for each prediction + labels = torch.arange(num_classes, device=device) + labels = labels.view(1, -1).expand_as(scores) + + # remove predictions with the background label + boxes = boxes[:, 1:] + scores = scores[:, 1:] + labels = labels[:, 1:] + + # batch everything, by making every class prediction be a separate instance + boxes = boxes.reshape(-1, 4) + scores = scores.reshape(-1) + labels = labels.reshape(-1) + + # remove low scoring boxes + inds = torch.where(scores > self.score_thresh)[0] + boxes, scores, labels = boxes[inds], scores[inds], labels[inds] + + # remove empty boxes + keep = box_ops.remove_small_boxes(boxes, min_size=1e-2) + boxes, scores, labels = boxes[keep], scores[keep], labels[keep] + + # non-maximum suppression, independently done per class + keep = box_ops.batched_nms(boxes, scores, labels, self.nms_thresh) + # keep only topk scoring predictions + keep = keep[: self.detections_per_img] + boxes, scores, labels = boxes[keep], scores[keep], labels[keep] + + all_boxes.append(boxes) + all_scores.append(scores) + all_labels.append(labels) + + return all_boxes, all_scores, all_labels + + def forward( + self, + features, # type: Dict[str, Tensor] + proposals, # type: List[Tensor] + image_shapes, # type: List[Tuple[int, int]] + targets=None, # type: Optional[List[Dict[str, Tensor]]] + ): + # type: (...) -> Tuple[List[Dict[str, Tensor]], Dict[str, Tensor]] + """ + Args: + features (List[Tensor]) + proposals (List[Tensor[N, 4]]) + image_shapes (List[Tuple[H, W]]) + targets (List[Dict]) + """ + if targets is not None: + for t in targets: + # TODO: https://github.com/pytorch/pytorch/issues/26731 + floating_point_types = (torch.float, torch.double, torch.half) + if not t["boxes"].dtype in floating_point_types: + raise TypeError(f"target boxes must of float type, instead got {t['boxes'].dtype}") + if not t["labels"].dtype == torch.int64: + raise TypeError(f"target labels must of int64 type, instead got {t['labels'].dtype}") + if self.has_keypoint(): + if not t["keypoints"].dtype == torch.float32: + raise TypeError(f"target keypoints must of float type, instead got {t['keypoints'].dtype}") + + if self.training: + proposals, matched_idxs, labels, regression_targets = self.select_training_samples(proposals, targets) + else: + labels = None + regression_targets = None + matched_idxs = None + + box_features = self.box_roi_pool(features, proposals, image_shapes) + box_features = self.box_head(box_features) + class_logits, box_regression = self.box_predictor(box_features) + + result: List[Dict[str, torch.Tensor]] = [] + losses = {} + if self.training: + if labels is None: + raise ValueError("labels cannot be None") + if regression_targets is None: + raise ValueError("regression_targets cannot be None") + loss_classifier, loss_box_reg = fastrcnn_loss(class_logits, box_regression, labels, regression_targets) + losses = {"loss_classifier": loss_classifier, "loss_box_reg": loss_box_reg} + else: + boxes, scores, labels = self.postprocess_detections(class_logits, box_regression, proposals, image_shapes) + num_images = len(boxes) + for i in range(num_images): + result.append( + { + "boxes": boxes[i], + "labels": labels[i], + "scores": scores[i], + } + ) + + if self.has_mask(): + mask_proposals = [p["boxes"] for p in result] + if self.training: + if matched_idxs is None: + raise ValueError("if in training, matched_idxs should not be None") + + # during training, only focus on positive boxes + num_images = len(proposals) + mask_proposals = [] + pos_matched_idxs = [] + for img_id in range(num_images): + pos = torch.where(labels[img_id] > 0)[0] + mask_proposals.append(proposals[img_id][pos]) + pos_matched_idxs.append(matched_idxs[img_id][pos]) + else: + pos_matched_idxs = None + + if self.mask_roi_pool is not None: + mask_features = self.mask_roi_pool(features, mask_proposals, image_shapes) + mask_features = self.mask_head(mask_features) + mask_logits = self.mask_predictor(mask_features) + else: + raise Exception("Expected mask_roi_pool to be not None") + + loss_mask = {} + if self.training: + if targets is None or pos_matched_idxs is None or mask_logits is None: + raise ValueError("targets, pos_matched_idxs, mask_logits cannot be None when training") + + gt_masks = [t["masks"] for t in targets] + gt_labels = [t["labels"] for t in targets] + rcnn_loss_mask = maskrcnn_loss(mask_logits, mask_proposals, gt_masks, gt_labels, pos_matched_idxs) + loss_mask = {"loss_mask": rcnn_loss_mask} + else: + labels = [r["labels"] for r in result] + masks_probs = maskrcnn_inference(mask_logits, labels) + for mask_prob, r in zip(masks_probs, result): + r["masks"] = mask_prob + + losses.update(loss_mask) + + # keep none checks in if conditional so torchscript will conditionally + # compile each branch + if ( + self.keypoint_roi_pool is not None + and self.keypoint_head is not None + and self.keypoint_predictor is not None + ): + keypoint_proposals = [p["boxes"] for p in result] + if self.training: + # during training, only focus on positive boxes + num_images = len(proposals) + keypoint_proposals = [] + pos_matched_idxs = [] + if matched_idxs is None: + raise ValueError("if in trainning, matched_idxs should not be None") + + for img_id in range(num_images): + pos = torch.where(labels[img_id] > 0)[0] + keypoint_proposals.append(proposals[img_id][pos]) + pos_matched_idxs.append(matched_idxs[img_id][pos]) + else: + pos_matched_idxs = None + + keypoint_features = self.keypoint_roi_pool(features, keypoint_proposals, image_shapes) + keypoint_features = self.keypoint_head(keypoint_features) + keypoint_logits = self.keypoint_predictor(keypoint_features) + + loss_keypoint = {} + if self.training: + if targets is None or pos_matched_idxs is None: + raise ValueError("both targets and pos_matched_idxs should not be None when in training mode") + + gt_keypoints = [t["keypoints"] for t in targets] + rcnn_loss_keypoint = keypointrcnn_loss( + keypoint_logits, keypoint_proposals, gt_keypoints, pos_matched_idxs + ) + loss_keypoint = {"loss_keypoint": rcnn_loss_keypoint} + else: + if keypoint_logits is None or keypoint_proposals is None: + raise ValueError( + "both keypoint_logits and keypoint_proposals should not be None when not in training mode" + ) + + keypoints_probs, kp_scores = keypointrcnn_inference(keypoint_logits, keypoint_proposals) + for keypoint_prob, kps, r in zip(keypoints_probs, kp_scores, result): + r["keypoints"] = keypoint_prob + r["keypoints_scores"] = kps + losses.update(loss_keypoint) + + return result, losses diff --git a/lib/python3.12/site-packages/torchvision/models/detection/rpn.py b/lib/python3.12/site-packages/torchvision/models/detection/rpn.py new file mode 100644 index 0000000000000000000000000000000000000000..f103181e4c6cba48c1a3b4c97583c5fb6785a8c4 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/detection/rpn.py @@ -0,0 +1,388 @@ +from typing import Dict, List, Optional, Tuple + +import torch +from torch import nn, Tensor +from torch.nn import functional as F +from torchvision.ops import boxes as box_ops, Conv2dNormActivation + +from . import _utils as det_utils + +# Import AnchorGenerator to keep compatibility. +from .anchor_utils import AnchorGenerator # noqa: 401 +from .image_list import ImageList + + +class RPNHead(nn.Module): + """ + Adds a simple RPN Head with classification and regression heads + + Args: + in_channels (int): number of channels of the input feature + num_anchors (int): number of anchors to be predicted + conv_depth (int, optional): number of convolutions + """ + + _version = 2 + + def __init__(self, in_channels: int, num_anchors: int, conv_depth=1) -> None: + super().__init__() + convs = [] + for _ in range(conv_depth): + convs.append(Conv2dNormActivation(in_channels, in_channels, kernel_size=3, norm_layer=None)) + self.conv = nn.Sequential(*convs) + self.cls_logits = nn.Conv2d(in_channels, num_anchors, kernel_size=1, stride=1) + self.bbox_pred = nn.Conv2d(in_channels, num_anchors * 4, kernel_size=1, stride=1) + + for layer in self.modules(): + if isinstance(layer, nn.Conv2d): + torch.nn.init.normal_(layer.weight, std=0.01) # type: ignore[arg-type] + if layer.bias is not None: + torch.nn.init.constant_(layer.bias, 0) # type: ignore[arg-type] + + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + version = local_metadata.get("version", None) + + if version is None or version < 2: + for type in ["weight", "bias"]: + old_key = f"{prefix}conv.{type}" + new_key = f"{prefix}conv.0.0.{type}" + if old_key in state_dict: + state_dict[new_key] = state_dict.pop(old_key) + + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + def forward(self, x: List[Tensor]) -> Tuple[List[Tensor], List[Tensor]]: + logits = [] + bbox_reg = [] + for feature in x: + t = self.conv(feature) + logits.append(self.cls_logits(t)) + bbox_reg.append(self.bbox_pred(t)) + return logits, bbox_reg + + +def permute_and_flatten(layer: Tensor, N: int, A: int, C: int, H: int, W: int) -> Tensor: + layer = layer.view(N, -1, C, H, W) + layer = layer.permute(0, 3, 4, 1, 2) + layer = layer.reshape(N, -1, C) + return layer + + +def concat_box_prediction_layers(box_cls: List[Tensor], box_regression: List[Tensor]) -> Tuple[Tensor, Tensor]: + box_cls_flattened = [] + box_regression_flattened = [] + # for each feature level, permute the outputs to make them be in the + # same format as the labels. Note that the labels are computed for + # all feature levels concatenated, so we keep the same representation + # for the objectness and the box_regression + for box_cls_per_level, box_regression_per_level in zip(box_cls, box_regression): + N, AxC, H, W = box_cls_per_level.shape + Ax4 = box_regression_per_level.shape[1] + A = Ax4 // 4 + C = AxC // A + box_cls_per_level = permute_and_flatten(box_cls_per_level, N, A, C, H, W) + box_cls_flattened.append(box_cls_per_level) + + box_regression_per_level = permute_and_flatten(box_regression_per_level, N, A, 4, H, W) + box_regression_flattened.append(box_regression_per_level) + # concatenate on the first dimension (representing the feature levels), to + # take into account the way the labels were generated (with all feature maps + # being concatenated as well) + box_cls = torch.cat(box_cls_flattened, dim=1).flatten(0, -2) + box_regression = torch.cat(box_regression_flattened, dim=1).reshape(-1, 4) + return box_cls, box_regression + + +class RegionProposalNetwork(torch.nn.Module): + """ + Implements Region Proposal Network (RPN). + + Args: + anchor_generator (AnchorGenerator): module that generates the anchors for a set of feature + maps. + head (nn.Module): module that computes the objectness and regression deltas + fg_iou_thresh (float): minimum IoU between the anchor and the GT box so that they can be + considered as positive during training of the RPN. + bg_iou_thresh (float): maximum IoU between the anchor and the GT box so that they can be + considered as negative during training of the RPN. + batch_size_per_image (int): number of anchors that are sampled during training of the RPN + for computing the loss + positive_fraction (float): proportion of positive anchors in a mini-batch during training + of the RPN + pre_nms_top_n (Dict[str, int]): number of proposals to keep before applying NMS. It should + contain two fields: training and testing, to allow for different values depending + on training or evaluation + post_nms_top_n (Dict[str, int]): number of proposals to keep after applying NMS. It should + contain two fields: training and testing, to allow for different values depending + on training or evaluation + nms_thresh (float): NMS threshold used for postprocessing the RPN proposals + score_thresh (float): only return proposals with an objectness score greater than score_thresh + + """ + + __annotations__ = { + "box_coder": det_utils.BoxCoder, + "proposal_matcher": det_utils.Matcher, + "fg_bg_sampler": det_utils.BalancedPositiveNegativeSampler, + } + + def __init__( + self, + anchor_generator: AnchorGenerator, + head: nn.Module, + # Faster-RCNN Training + fg_iou_thresh: float, + bg_iou_thresh: float, + batch_size_per_image: int, + positive_fraction: float, + # Faster-RCNN Inference + pre_nms_top_n: Dict[str, int], + post_nms_top_n: Dict[str, int], + nms_thresh: float, + score_thresh: float = 0.0, + ) -> None: + super().__init__() + self.anchor_generator = anchor_generator + self.head = head + self.box_coder = det_utils.BoxCoder(weights=(1.0, 1.0, 1.0, 1.0)) + + # used during training + self.box_similarity = box_ops.box_iou + + self.proposal_matcher = det_utils.Matcher( + fg_iou_thresh, + bg_iou_thresh, + allow_low_quality_matches=True, + ) + + self.fg_bg_sampler = det_utils.BalancedPositiveNegativeSampler(batch_size_per_image, positive_fraction) + # used during testing + self._pre_nms_top_n = pre_nms_top_n + self._post_nms_top_n = post_nms_top_n + self.nms_thresh = nms_thresh + self.score_thresh = score_thresh + self.min_size = 1e-3 + + def pre_nms_top_n(self) -> int: + if self.training: + return self._pre_nms_top_n["training"] + return self._pre_nms_top_n["testing"] + + def post_nms_top_n(self) -> int: + if self.training: + return self._post_nms_top_n["training"] + return self._post_nms_top_n["testing"] + + def assign_targets_to_anchors( + self, anchors: List[Tensor], targets: List[Dict[str, Tensor]] + ) -> Tuple[List[Tensor], List[Tensor]]: + + labels = [] + matched_gt_boxes = [] + for anchors_per_image, targets_per_image in zip(anchors, targets): + gt_boxes = targets_per_image["boxes"] + + if gt_boxes.numel() == 0: + # Background image (negative example) + device = anchors_per_image.device + matched_gt_boxes_per_image = torch.zeros(anchors_per_image.shape, dtype=torch.float32, device=device) + labels_per_image = torch.zeros((anchors_per_image.shape[0],), dtype=torch.float32, device=device) + else: + match_quality_matrix = self.box_similarity(gt_boxes, anchors_per_image) + matched_idxs = self.proposal_matcher(match_quality_matrix) + # get the targets corresponding GT for each proposal + # NB: need to clamp the indices because we can have a single + # GT in the image, and matched_idxs can be -2, which goes + # out of bounds + matched_gt_boxes_per_image = gt_boxes[matched_idxs.clamp(min=0)] + + labels_per_image = matched_idxs >= 0 + labels_per_image = labels_per_image.to(dtype=torch.float32) + + # Background (negative examples) + bg_indices = matched_idxs == self.proposal_matcher.BELOW_LOW_THRESHOLD + labels_per_image[bg_indices] = 0.0 + + # discard indices that are between thresholds + inds_to_discard = matched_idxs == self.proposal_matcher.BETWEEN_THRESHOLDS + labels_per_image[inds_to_discard] = -1.0 + + labels.append(labels_per_image) + matched_gt_boxes.append(matched_gt_boxes_per_image) + return labels, matched_gt_boxes + + def _get_top_n_idx(self, objectness: Tensor, num_anchors_per_level: List[int]) -> Tensor: + r = [] + offset = 0 + for ob in objectness.split(num_anchors_per_level, 1): + num_anchors = ob.shape[1] + pre_nms_top_n = det_utils._topk_min(ob, self.pre_nms_top_n(), 1) + _, top_n_idx = ob.topk(pre_nms_top_n, dim=1) + r.append(top_n_idx + offset) + offset += num_anchors + return torch.cat(r, dim=1) + + def filter_proposals( + self, + proposals: Tensor, + objectness: Tensor, + image_shapes: List[Tuple[int, int]], + num_anchors_per_level: List[int], + ) -> Tuple[List[Tensor], List[Tensor]]: + + num_images = proposals.shape[0] + device = proposals.device + # do not backprop through objectness + objectness = objectness.detach() + objectness = objectness.reshape(num_images, -1) + + levels = [ + torch.full((n,), idx, dtype=torch.int64, device=device) for idx, n in enumerate(num_anchors_per_level) + ] + levels = torch.cat(levels, 0) + levels = levels.reshape(1, -1).expand_as(objectness) + + # select top_n boxes independently per level before applying nms + top_n_idx = self._get_top_n_idx(objectness, num_anchors_per_level) + + image_range = torch.arange(num_images, device=device) + batch_idx = image_range[:, None] + + objectness = objectness[batch_idx, top_n_idx] + levels = levels[batch_idx, top_n_idx] + proposals = proposals[batch_idx, top_n_idx] + + objectness_prob = torch.sigmoid(objectness) + + final_boxes = [] + final_scores = [] + for boxes, scores, lvl, img_shape in zip(proposals, objectness_prob, levels, image_shapes): + boxes = box_ops.clip_boxes_to_image(boxes, img_shape) + + # remove small boxes + keep = box_ops.remove_small_boxes(boxes, self.min_size) + boxes, scores, lvl = boxes[keep], scores[keep], lvl[keep] + + # remove low scoring boxes + # use >= for Backwards compatibility + keep = torch.where(scores >= self.score_thresh)[0] + boxes, scores, lvl = boxes[keep], scores[keep], lvl[keep] + + # non-maximum suppression, independently done per level + keep = box_ops.batched_nms(boxes, scores, lvl, self.nms_thresh) + + # keep only topk scoring predictions + keep = keep[: self.post_nms_top_n()] + boxes, scores = boxes[keep], scores[keep] + + final_boxes.append(boxes) + final_scores.append(scores) + return final_boxes, final_scores + + def compute_loss( + self, objectness: Tensor, pred_bbox_deltas: Tensor, labels: List[Tensor], regression_targets: List[Tensor] + ) -> Tuple[Tensor, Tensor]: + """ + Args: + objectness (Tensor) + pred_bbox_deltas (Tensor) + labels (List[Tensor]) + regression_targets (List[Tensor]) + + Returns: + objectness_loss (Tensor) + box_loss (Tensor) + """ + + sampled_pos_inds, sampled_neg_inds = self.fg_bg_sampler(labels) + sampled_pos_inds = torch.where(torch.cat(sampled_pos_inds, dim=0))[0] + sampled_neg_inds = torch.where(torch.cat(sampled_neg_inds, dim=0))[0] + + sampled_inds = torch.cat([sampled_pos_inds, sampled_neg_inds], dim=0) + + objectness = objectness.flatten() + + labels = torch.cat(labels, dim=0) + regression_targets = torch.cat(regression_targets, dim=0) + + box_loss = F.smooth_l1_loss( + pred_bbox_deltas[sampled_pos_inds], + regression_targets[sampled_pos_inds], + beta=1 / 9, + reduction="sum", + ) / (sampled_inds.numel()) + + objectness_loss = F.binary_cross_entropy_with_logits(objectness[sampled_inds], labels[sampled_inds]) + + return objectness_loss, box_loss + + def forward( + self, + images: ImageList, + features: Dict[str, Tensor], + targets: Optional[List[Dict[str, Tensor]]] = None, + ) -> Tuple[List[Tensor], Dict[str, Tensor]]: + + """ + Args: + images (ImageList): images for which we want to compute the predictions + features (Dict[str, Tensor]): features computed from the images that are + used for computing the predictions. Each tensor in the list + correspond to different feature levels + targets (List[Dict[str, Tensor]]): ground-truth boxes present in the image (optional). + If provided, each element in the dict should contain a field `boxes`, + with the locations of the ground-truth boxes. + + Returns: + boxes (List[Tensor]): the predicted boxes from the RPN, one Tensor per + image. + losses (Dict[str, Tensor]): the losses for the model during training. During + testing, it is an empty dict. + """ + # RPN uses all feature maps that are available + features = list(features.values()) + objectness, pred_bbox_deltas = self.head(features) + anchors = self.anchor_generator(images, features) + + num_images = len(anchors) + num_anchors_per_level_shape_tensors = [o[0].shape for o in objectness] + num_anchors_per_level = [s[0] * s[1] * s[2] for s in num_anchors_per_level_shape_tensors] + objectness, pred_bbox_deltas = concat_box_prediction_layers(objectness, pred_bbox_deltas) + # apply pred_bbox_deltas to anchors to obtain the decoded proposals + # note that we detach the deltas because Faster R-CNN do not backprop through + # the proposals + proposals = self.box_coder.decode(pred_bbox_deltas.detach(), anchors) + proposals = proposals.view(num_images, -1, 4) + boxes, scores = self.filter_proposals(proposals, objectness, images.image_sizes, num_anchors_per_level) + + losses = {} + if self.training: + if targets is None: + raise ValueError("targets should not be None") + labels, matched_gt_boxes = self.assign_targets_to_anchors(anchors, targets) + regression_targets = self.box_coder.encode(matched_gt_boxes, anchors) + loss_objectness, loss_rpn_box_reg = self.compute_loss( + objectness, pred_bbox_deltas, labels, regression_targets + ) + losses = { + "loss_objectness": loss_objectness, + "loss_rpn_box_reg": loss_rpn_box_reg, + } + return boxes, losses diff --git a/lib/python3.12/site-packages/torchvision/models/detection/ssd.py b/lib/python3.12/site-packages/torchvision/models/detection/ssd.py new file mode 100644 index 0000000000000000000000000000000000000000..87062d2bc88a5bf17625e0530116aba22941c538 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/detection/ssd.py @@ -0,0 +1,682 @@ +import warnings +from collections import OrderedDict +from typing import Any, Dict, List, Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn, Tensor + +from ...ops import boxes as box_ops +from ...transforms._presets import ObjectDetection +from ...utils import _log_api_usage_once +from .._api import register_model, Weights, WeightsEnum +from .._meta import _COCO_CATEGORIES +from .._utils import _ovewrite_value_param, handle_legacy_interface +from ..vgg import VGG, vgg16, VGG16_Weights +from . import _utils as det_utils +from .anchor_utils import DefaultBoxGenerator +from .backbone_utils import _validate_trainable_layers +from .transform import GeneralizedRCNNTransform + + +__all__ = [ + "SSD300_VGG16_Weights", + "ssd300_vgg16", +] + + +class SSD300_VGG16_Weights(WeightsEnum): + COCO_V1 = Weights( + url="https://download.pytorch.org/models/ssd300_vgg16_coco-b556d3b4.pth", + transforms=ObjectDetection, + meta={ + "num_params": 35641826, + "categories": _COCO_CATEGORIES, + "min_size": (1, 1), + "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#ssd300-vgg16", + "_metrics": { + "COCO-val2017": { + "box_map": 25.1, + } + }, + "_ops": 34.858, + "_file_size": 135.988, + "_docs": """These weights were produced by following a similar training recipe as on the paper.""", + }, + ) + DEFAULT = COCO_V1 + + +def _xavier_init(conv: nn.Module): + for layer in conv.modules(): + if isinstance(layer, nn.Conv2d): + torch.nn.init.xavier_uniform_(layer.weight) + if layer.bias is not None: + torch.nn.init.constant_(layer.bias, 0.0) + + +class SSDHead(nn.Module): + def __init__(self, in_channels: List[int], num_anchors: List[int], num_classes: int): + super().__init__() + self.classification_head = SSDClassificationHead(in_channels, num_anchors, num_classes) + self.regression_head = SSDRegressionHead(in_channels, num_anchors) + + def forward(self, x: List[Tensor]) -> Dict[str, Tensor]: + return { + "bbox_regression": self.regression_head(x), + "cls_logits": self.classification_head(x), + } + + +class SSDScoringHead(nn.Module): + def __init__(self, module_list: nn.ModuleList, num_columns: int): + super().__init__() + self.module_list = module_list + self.num_columns = num_columns + + def _get_result_from_module_list(self, x: Tensor, idx: int) -> Tensor: + """ + This is equivalent to self.module_list[idx](x), + but torchscript doesn't support this yet + """ + num_blocks = len(self.module_list) + if idx < 0: + idx += num_blocks + out = x + for i, module in enumerate(self.module_list): + if i == idx: + out = module(x) + return out + + def forward(self, x: List[Tensor]) -> Tensor: + all_results = [] + + for i, features in enumerate(x): + results = self._get_result_from_module_list(features, i) + + # Permute output from (N, A * K, H, W) to (N, HWA, K). + N, _, H, W = results.shape + results = results.view(N, -1, self.num_columns, H, W) + results = results.permute(0, 3, 4, 1, 2) + results = results.reshape(N, -1, self.num_columns) # Size=(N, HWA, K) + + all_results.append(results) + + return torch.cat(all_results, dim=1) + + +class SSDClassificationHead(SSDScoringHead): + def __init__(self, in_channels: List[int], num_anchors: List[int], num_classes: int): + cls_logits = nn.ModuleList() + for channels, anchors in zip(in_channels, num_anchors): + cls_logits.append(nn.Conv2d(channels, num_classes * anchors, kernel_size=3, padding=1)) + _xavier_init(cls_logits) + super().__init__(cls_logits, num_classes) + + +class SSDRegressionHead(SSDScoringHead): + def __init__(self, in_channels: List[int], num_anchors: List[int]): + bbox_reg = nn.ModuleList() + for channels, anchors in zip(in_channels, num_anchors): + bbox_reg.append(nn.Conv2d(channels, 4 * anchors, kernel_size=3, padding=1)) + _xavier_init(bbox_reg) + super().__init__(bbox_reg, 4) + + +class SSD(nn.Module): + """ + Implements SSD architecture from `"SSD: Single Shot MultiBox Detector" `_. + + The input to the model is expected to be a list of tensors, each of shape [C, H, W], one for each + image, and should be in 0-1 range. Different images can have different sizes, but they will be resized + to a fixed size before passing it to the backbone. + + The behavior of the model changes depending on if it is in training or evaluation mode. + + During training, the model expects both the input tensors and targets (list of dictionary), + containing: + - boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (Int64Tensor[N]): the class label for each ground-truth box + + The model returns a Dict[Tensor] during training, containing the classification and regression + losses. + + During inference, the model requires only the input tensors, and returns the post-processed + predictions as a List[Dict[Tensor]], one for each input image. The fields of the Dict are as + follows, where ``N`` is the number of detections: + + - boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (Int64Tensor[N]): the predicted labels for each detection + - scores (Tensor[N]): the scores for each detection + + Args: + backbone (nn.Module): the network used to compute the features for the model. + It should contain an out_channels attribute with the list of the output channels of + each feature map. The backbone should return a single Tensor or an OrderedDict[Tensor]. + anchor_generator (DefaultBoxGenerator): module that generates the default boxes for a + set of feature maps. + size (Tuple[int, int]): the width and height to which images will be rescaled before feeding them + to the backbone. + num_classes (int): number of output classes of the model (including the background). + image_mean (Tuple[float, float, float]): mean values used for input normalization. + They are generally the mean values of the dataset on which the backbone has been trained + on + image_std (Tuple[float, float, float]): std values used for input normalization. + They are generally the std values of the dataset on which the backbone has been trained on + head (nn.Module, optional): Module run on top of the backbone features. Defaults to a module containing + a classification and regression module. + score_thresh (float): Score threshold used for postprocessing the detections. + nms_thresh (float): NMS threshold used for postprocessing the detections. + detections_per_img (int): Number of best detections to keep after NMS. + iou_thresh (float): minimum IoU between the anchor and the GT box so that they can be + considered as positive during training. + topk_candidates (int): Number of best detections to keep before NMS. + positive_fraction (float): a number between 0 and 1 which indicates the proportion of positive + proposals used during the training of the classification head. It is used to estimate the negative to + positive ratio. + """ + + __annotations__ = { + "box_coder": det_utils.BoxCoder, + "proposal_matcher": det_utils.Matcher, + } + + def __init__( + self, + backbone: nn.Module, + anchor_generator: DefaultBoxGenerator, + size: Tuple[int, int], + num_classes: int, + image_mean: Optional[List[float]] = None, + image_std: Optional[List[float]] = None, + head: Optional[nn.Module] = None, + score_thresh: float = 0.01, + nms_thresh: float = 0.45, + detections_per_img: int = 200, + iou_thresh: float = 0.5, + topk_candidates: int = 400, + positive_fraction: float = 0.25, + **kwargs: Any, + ): + super().__init__() + _log_api_usage_once(self) + + self.backbone = backbone + + self.anchor_generator = anchor_generator + + self.box_coder = det_utils.BoxCoder(weights=(10.0, 10.0, 5.0, 5.0)) + + if head is None: + if hasattr(backbone, "out_channels"): + out_channels = backbone.out_channels + else: + out_channels = det_utils.retrieve_out_channels(backbone, size) + + if len(out_channels) != len(anchor_generator.aspect_ratios): + raise ValueError( + f"The length of the output channels from the backbone ({len(out_channels)}) do not match the length of the anchor generator aspect ratios ({len(anchor_generator.aspect_ratios)})" + ) + + num_anchors = self.anchor_generator.num_anchors_per_location() + head = SSDHead(out_channels, num_anchors, num_classes) + self.head = head + + self.proposal_matcher = det_utils.SSDMatcher(iou_thresh) + + if image_mean is None: + image_mean = [0.485, 0.456, 0.406] + if image_std is None: + image_std = [0.229, 0.224, 0.225] + self.transform = GeneralizedRCNNTransform( + min(size), max(size), image_mean, image_std, size_divisible=1, fixed_size=size, **kwargs + ) + + self.score_thresh = score_thresh + self.nms_thresh = nms_thresh + self.detections_per_img = detections_per_img + self.topk_candidates = topk_candidates + self.neg_to_pos_ratio = (1.0 - positive_fraction) / positive_fraction + + # used only on torchscript mode + self._has_warned = False + + @torch.jit.unused + def eager_outputs( + self, losses: Dict[str, Tensor], detections: List[Dict[str, Tensor]] + ) -> Tuple[Dict[str, Tensor], List[Dict[str, Tensor]]]: + if self.training: + return losses + + return detections + + def compute_loss( + self, + targets: List[Dict[str, Tensor]], + head_outputs: Dict[str, Tensor], + anchors: List[Tensor], + matched_idxs: List[Tensor], + ) -> Dict[str, Tensor]: + bbox_regression = head_outputs["bbox_regression"] + cls_logits = head_outputs["cls_logits"] + + # Match original targets with default boxes + num_foreground = 0 + bbox_loss = [] + cls_targets = [] + for ( + targets_per_image, + bbox_regression_per_image, + cls_logits_per_image, + anchors_per_image, + matched_idxs_per_image, + ) in zip(targets, bbox_regression, cls_logits, anchors, matched_idxs): + # produce the matching between boxes and targets + foreground_idxs_per_image = torch.where(matched_idxs_per_image >= 0)[0] + foreground_matched_idxs_per_image = matched_idxs_per_image[foreground_idxs_per_image] + num_foreground += foreground_matched_idxs_per_image.numel() + + # Calculate regression loss + matched_gt_boxes_per_image = targets_per_image["boxes"][foreground_matched_idxs_per_image] + bbox_regression_per_image = bbox_regression_per_image[foreground_idxs_per_image, :] + anchors_per_image = anchors_per_image[foreground_idxs_per_image, :] + target_regression = self.box_coder.encode_single(matched_gt_boxes_per_image, anchors_per_image) + bbox_loss.append( + torch.nn.functional.smooth_l1_loss(bbox_regression_per_image, target_regression, reduction="sum") + ) + + # Estimate ground truth for class targets + gt_classes_target = torch.zeros( + (cls_logits_per_image.size(0),), + dtype=targets_per_image["labels"].dtype, + device=targets_per_image["labels"].device, + ) + gt_classes_target[foreground_idxs_per_image] = targets_per_image["labels"][ + foreground_matched_idxs_per_image + ] + cls_targets.append(gt_classes_target) + + bbox_loss = torch.stack(bbox_loss) + cls_targets = torch.stack(cls_targets) + + # Calculate classification loss + num_classes = cls_logits.size(-1) + cls_loss = F.cross_entropy(cls_logits.view(-1, num_classes), cls_targets.view(-1), reduction="none").view( + cls_targets.size() + ) + + # Hard Negative Sampling + foreground_idxs = cls_targets > 0 + num_negative = self.neg_to_pos_ratio * foreground_idxs.sum(1, keepdim=True) + # num_negative[num_negative < self.neg_to_pos_ratio] = self.neg_to_pos_ratio + negative_loss = cls_loss.clone() + negative_loss[foreground_idxs] = -float("inf") # use -inf to detect positive values that creeped in the sample + values, idx = negative_loss.sort(1, descending=True) + # background_idxs = torch.logical_and(idx.sort(1)[1] < num_negative, torch.isfinite(values)) + background_idxs = idx.sort(1)[1] < num_negative + + N = max(1, num_foreground) + return { + "bbox_regression": bbox_loss.sum() / N, + "classification": (cls_loss[foreground_idxs].sum() + cls_loss[background_idxs].sum()) / N, + } + + def forward( + self, images: List[Tensor], targets: Optional[List[Dict[str, Tensor]]] = None + ) -> Tuple[Dict[str, Tensor], List[Dict[str, Tensor]]]: + if self.training: + if targets is None: + torch._assert(False, "targets should not be none when in training mode") + else: + for target in targets: + boxes = target["boxes"] + if isinstance(boxes, torch.Tensor): + torch._assert( + len(boxes.shape) == 2 and boxes.shape[-1] == 4, + f"Expected target boxes to be a tensor of shape [N, 4], got {boxes.shape}.", + ) + else: + torch._assert(False, f"Expected target boxes to be of type Tensor, got {type(boxes)}.") + + # get the original image sizes + original_image_sizes: List[Tuple[int, int]] = [] + for img in images: + val = img.shape[-2:] + torch._assert( + len(val) == 2, + f"expecting the last two dimensions of the Tensor to be H and W instead got {img.shape[-2:]}", + ) + original_image_sizes.append((val[0], val[1])) + + # transform the input + images, targets = self.transform(images, targets) + + # Check for degenerate boxes + if targets is not None: + for target_idx, target in enumerate(targets): + boxes = target["boxes"] + degenerate_boxes = boxes[:, 2:] <= boxes[:, :2] + if degenerate_boxes.any(): + bb_idx = torch.where(degenerate_boxes.any(dim=1))[0][0] + degen_bb: List[float] = boxes[bb_idx].tolist() + torch._assert( + False, + "All bounding boxes should have positive height and width." + f" Found invalid box {degen_bb} for target at index {target_idx}.", + ) + + # get the features from the backbone + features = self.backbone(images.tensors) + if isinstance(features, torch.Tensor): + features = OrderedDict([("0", features)]) + + features = list(features.values()) + + # compute the ssd heads outputs using the features + head_outputs = self.head(features) + + # create the set of anchors + anchors = self.anchor_generator(images, features) + + losses = {} + detections: List[Dict[str, Tensor]] = [] + if self.training: + matched_idxs = [] + if targets is None: + torch._assert(False, "targets should not be none when in training mode") + else: + for anchors_per_image, targets_per_image in zip(anchors, targets): + if targets_per_image["boxes"].numel() == 0: + matched_idxs.append( + torch.full( + (anchors_per_image.size(0),), -1, dtype=torch.int64, device=anchors_per_image.device + ) + ) + continue + + match_quality_matrix = box_ops.box_iou(targets_per_image["boxes"], anchors_per_image) + matched_idxs.append(self.proposal_matcher(match_quality_matrix)) + + losses = self.compute_loss(targets, head_outputs, anchors, matched_idxs) + else: + detections = self.postprocess_detections(head_outputs, anchors, images.image_sizes) + detections = self.transform.postprocess(detections, images.image_sizes, original_image_sizes) + + if torch.jit.is_scripting(): + if not self._has_warned: + warnings.warn("SSD always returns a (Losses, Detections) tuple in scripting") + self._has_warned = True + return losses, detections + return self.eager_outputs(losses, detections) + + def postprocess_detections( + self, head_outputs: Dict[str, Tensor], image_anchors: List[Tensor], image_shapes: List[Tuple[int, int]] + ) -> List[Dict[str, Tensor]]: + bbox_regression = head_outputs["bbox_regression"] + pred_scores = F.softmax(head_outputs["cls_logits"], dim=-1) + + num_classes = pred_scores.size(-1) + device = pred_scores.device + + detections: List[Dict[str, Tensor]] = [] + + for boxes, scores, anchors, image_shape in zip(bbox_regression, pred_scores, image_anchors, image_shapes): + boxes = self.box_coder.decode_single(boxes, anchors) + boxes = box_ops.clip_boxes_to_image(boxes, image_shape) + + image_boxes = [] + image_scores = [] + image_labels = [] + for label in range(1, num_classes): + score = scores[:, label] + + keep_idxs = score > self.score_thresh + score = score[keep_idxs] + box = boxes[keep_idxs] + + # keep only topk scoring predictions + num_topk = det_utils._topk_min(score, self.topk_candidates, 0) + score, idxs = score.topk(num_topk) + box = box[idxs] + + image_boxes.append(box) + image_scores.append(score) + image_labels.append(torch.full_like(score, fill_value=label, dtype=torch.int64, device=device)) + + image_boxes = torch.cat(image_boxes, dim=0) + image_scores = torch.cat(image_scores, dim=0) + image_labels = torch.cat(image_labels, dim=0) + + # non-maximum suppression + keep = box_ops.batched_nms(image_boxes, image_scores, image_labels, self.nms_thresh) + keep = keep[: self.detections_per_img] + + detections.append( + { + "boxes": image_boxes[keep], + "scores": image_scores[keep], + "labels": image_labels[keep], + } + ) + return detections + + +class SSDFeatureExtractorVGG(nn.Module): + def __init__(self, backbone: nn.Module, highres: bool): + super().__init__() + + _, _, maxpool3_pos, maxpool4_pos, _ = (i for i, layer in enumerate(backbone) if isinstance(layer, nn.MaxPool2d)) + + # Patch ceil_mode for maxpool3 to get the same WxH output sizes as the paper + backbone[maxpool3_pos].ceil_mode = True + + # parameters used for L2 regularization + rescaling + self.scale_weight = nn.Parameter(torch.ones(512) * 20) + + # Multiple Feature maps - page 4, Fig 2 of SSD paper + self.features = nn.Sequential(*backbone[:maxpool4_pos]) # until conv4_3 + + # SSD300 case - page 4, Fig 2 of SSD paper + extra = nn.ModuleList( + [ + nn.Sequential( + nn.Conv2d(1024, 256, kernel_size=1), + nn.ReLU(inplace=True), + nn.Conv2d(256, 512, kernel_size=3, padding=1, stride=2), # conv8_2 + nn.ReLU(inplace=True), + ), + nn.Sequential( + nn.Conv2d(512, 128, kernel_size=1), + nn.ReLU(inplace=True), + nn.Conv2d(128, 256, kernel_size=3, padding=1, stride=2), # conv9_2 + nn.ReLU(inplace=True), + ), + nn.Sequential( + nn.Conv2d(256, 128, kernel_size=1), + nn.ReLU(inplace=True), + nn.Conv2d(128, 256, kernel_size=3), # conv10_2 + nn.ReLU(inplace=True), + ), + nn.Sequential( + nn.Conv2d(256, 128, kernel_size=1), + nn.ReLU(inplace=True), + nn.Conv2d(128, 256, kernel_size=3), # conv11_2 + nn.ReLU(inplace=True), + ), + ] + ) + if highres: + # Additional layers for the SSD512 case. See page 11, footernote 5. + extra.append( + nn.Sequential( + nn.Conv2d(256, 128, kernel_size=1), + nn.ReLU(inplace=True), + nn.Conv2d(128, 256, kernel_size=4), # conv12_2 + nn.ReLU(inplace=True), + ) + ) + _xavier_init(extra) + + fc = nn.Sequential( + nn.MaxPool2d(kernel_size=3, stride=1, padding=1, ceil_mode=False), # add modified maxpool5 + nn.Conv2d(in_channels=512, out_channels=1024, kernel_size=3, padding=6, dilation=6), # FC6 with atrous + nn.ReLU(inplace=True), + nn.Conv2d(in_channels=1024, out_channels=1024, kernel_size=1), # FC7 + nn.ReLU(inplace=True), + ) + _xavier_init(fc) + extra.insert( + 0, + nn.Sequential( + *backbone[maxpool4_pos:-1], # until conv5_3, skip maxpool5 + fc, + ), + ) + self.extra = extra + + def forward(self, x: Tensor) -> Dict[str, Tensor]: + # L2 regularization + Rescaling of 1st block's feature map + x = self.features(x) + rescaled = self.scale_weight.view(1, -1, 1, 1) * F.normalize(x) + output = [rescaled] + + # Calculating Feature maps for the rest blocks + for block in self.extra: + x = block(x) + output.append(x) + + return OrderedDict([(str(i), v) for i, v in enumerate(output)]) + + +def _vgg_extractor(backbone: VGG, highres: bool, trainable_layers: int): + backbone = backbone.features + # Gather the indices of maxpools. These are the locations of output blocks. + stage_indices = [0] + [i for i, b in enumerate(backbone) if isinstance(b, nn.MaxPool2d)][:-1] + num_stages = len(stage_indices) + + # find the index of the layer from which we won't freeze + torch._assert( + 0 <= trainable_layers <= num_stages, + f"trainable_layers should be in the range [0, {num_stages}]. Instead got {trainable_layers}", + ) + freeze_before = len(backbone) if trainable_layers == 0 else stage_indices[num_stages - trainable_layers] + + for b in backbone[:freeze_before]: + for parameter in b.parameters(): + parameter.requires_grad_(False) + + return SSDFeatureExtractorVGG(backbone, highres) + + +@register_model() +@handle_legacy_interface( + weights=("pretrained", SSD300_VGG16_Weights.COCO_V1), + weights_backbone=("pretrained_backbone", VGG16_Weights.IMAGENET1K_FEATURES), +) +def ssd300_vgg16( + *, + weights: Optional[SSD300_VGG16_Weights] = None, + progress: bool = True, + num_classes: Optional[int] = None, + weights_backbone: Optional[VGG16_Weights] = VGG16_Weights.IMAGENET1K_FEATURES, + trainable_backbone_layers: Optional[int] = None, + **kwargs: Any, +) -> SSD: + """The SSD300 model is based on the `SSD: Single Shot MultiBox Detector + `_ paper. + + .. betastatus:: detection module + + The input to the model is expected to be a list of tensors, each of shape [C, H, W], one for each + image, and should be in 0-1 range. Different images can have different sizes, but they will be resized + to a fixed size before passing it to the backbone. + + The behavior of the model changes depending on if it is in training or evaluation mode. + + During training, the model expects both the input tensors and targets (list of dictionary), + containing: + + - boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (Int64Tensor[N]): the class label for each ground-truth box + + The model returns a Dict[Tensor] during training, containing the classification and regression + losses. + + During inference, the model requires only the input tensors, and returns the post-processed + predictions as a List[Dict[Tensor]], one for each input image. The fields of the Dict are as + follows, where ``N`` is the number of detections: + + - boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with + ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``. + - labels (Int64Tensor[N]): the predicted labels for each detection + - scores (Tensor[N]): the scores for each detection + + Example: + + >>> model = torchvision.models.detection.ssd300_vgg16(weights=SSD300_VGG16_Weights.DEFAULT) + >>> model.eval() + >>> x = [torch.rand(3, 300, 300), torch.rand(3, 500, 400)] + >>> predictions = model(x) + + Args: + weights (:class:`~torchvision.models.detection.SSD300_VGG16_Weights`, optional): The pretrained + weights to use. See + :class:`~torchvision.models.detection.SSD300_VGG16_Weights` + below for more details, and possible values. By default, no + pre-trained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr + Default is True. + num_classes (int, optional): number of output classes of the model (including the background) + weights_backbone (:class:`~torchvision.models.VGG16_Weights`, optional): The pretrained weights for the + backbone + trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from final block. + Valid values are between 0 and 5, with 5 meaning all backbone layers are trainable. If ``None`` is + passed (the default) this value is set to 4. + **kwargs: parameters passed to the ``torchvision.models.detection.SSD`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.detection.SSD300_VGG16_Weights + :members: + """ + weights = SSD300_VGG16_Weights.verify(weights) + weights_backbone = VGG16_Weights.verify(weights_backbone) + + if "size" in kwargs: + warnings.warn("The size of the model is already fixed; ignoring the parameter.") + + if weights is not None: + weights_backbone = None + num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"])) + elif num_classes is None: + num_classes = 91 + + trainable_backbone_layers = _validate_trainable_layers( + weights is not None or weights_backbone is not None, trainable_backbone_layers, 5, 4 + ) + + # Use custom backbones more appropriate for SSD + backbone = vgg16(weights=weights_backbone, progress=progress) + backbone = _vgg_extractor(backbone, False, trainable_backbone_layers) + anchor_generator = DefaultBoxGenerator( + [[2], [2, 3], [2, 3], [2, 3], [2], [2]], + scales=[0.07, 0.15, 0.33, 0.51, 0.69, 0.87, 1.05], + steps=[8, 16, 32, 64, 100, 300], + ) + + defaults = { + # Rescale the input in a way compatible to the backbone + "image_mean": [0.48235, 0.45882, 0.40784], + "image_std": [1.0 / 255.0, 1.0 / 255.0, 1.0 / 255.0], # undo the 0-1 scaling of toTensor + } + kwargs: Any = {**defaults, **kwargs} + model = SSD(backbone, anchor_generator, (300, 300), num_classes, **kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model diff --git a/lib/python3.12/site-packages/torchvision/models/detection/ssdlite.py b/lib/python3.12/site-packages/torchvision/models/detection/ssdlite.py new file mode 100644 index 0000000000000000000000000000000000000000..eda21bf941ef0d4a9051312ebdba6911c6760e8d --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/detection/ssdlite.py @@ -0,0 +1,331 @@ +import warnings +from collections import OrderedDict +from functools import partial +from typing import Any, Callable, Dict, List, Optional, Union + +import torch +from torch import nn, Tensor + +from ...ops.misc import Conv2dNormActivation +from ...transforms._presets import ObjectDetection +from ...utils import _log_api_usage_once +from .. import mobilenet +from .._api import register_model, Weights, WeightsEnum +from .._meta import _COCO_CATEGORIES +from .._utils import _ovewrite_value_param, handle_legacy_interface +from ..mobilenetv3 import mobilenet_v3_large, MobileNet_V3_Large_Weights +from . import _utils as det_utils +from .anchor_utils import DefaultBoxGenerator +from .backbone_utils import _validate_trainable_layers +from .ssd import SSD, SSDScoringHead + + +__all__ = [ + "SSDLite320_MobileNet_V3_Large_Weights", + "ssdlite320_mobilenet_v3_large", +] + + +# Building blocks of SSDlite as described in section 6.2 of MobileNetV2 paper +def _prediction_block( + in_channels: int, out_channels: int, kernel_size: int, norm_layer: Callable[..., nn.Module] +) -> nn.Sequential: + return nn.Sequential( + # 3x3 depthwise with stride 1 and padding 1 + Conv2dNormActivation( + in_channels, + in_channels, + kernel_size=kernel_size, + groups=in_channels, + norm_layer=norm_layer, + activation_layer=nn.ReLU6, + ), + # 1x1 projetion to output channels + nn.Conv2d(in_channels, out_channels, 1), + ) + + +def _extra_block(in_channels: int, out_channels: int, norm_layer: Callable[..., nn.Module]) -> nn.Sequential: + activation = nn.ReLU6 + intermediate_channels = out_channels // 2 + return nn.Sequential( + # 1x1 projection to half output channels + Conv2dNormActivation( + in_channels, intermediate_channels, kernel_size=1, norm_layer=norm_layer, activation_layer=activation + ), + # 3x3 depthwise with stride 2 and padding 1 + Conv2dNormActivation( + intermediate_channels, + intermediate_channels, + kernel_size=3, + stride=2, + groups=intermediate_channels, + norm_layer=norm_layer, + activation_layer=activation, + ), + # 1x1 projetion to output channels + Conv2dNormActivation( + intermediate_channels, out_channels, kernel_size=1, norm_layer=norm_layer, activation_layer=activation + ), + ) + + +def _normal_init(conv: nn.Module): + for layer in conv.modules(): + if isinstance(layer, nn.Conv2d): + torch.nn.init.normal_(layer.weight, mean=0.0, std=0.03) + if layer.bias is not None: + torch.nn.init.constant_(layer.bias, 0.0) + + +class SSDLiteHead(nn.Module): + def __init__( + self, in_channels: List[int], num_anchors: List[int], num_classes: int, norm_layer: Callable[..., nn.Module] + ): + super().__init__() + self.classification_head = SSDLiteClassificationHead(in_channels, num_anchors, num_classes, norm_layer) + self.regression_head = SSDLiteRegressionHead(in_channels, num_anchors, norm_layer) + + def forward(self, x: List[Tensor]) -> Dict[str, Tensor]: + return { + "bbox_regression": self.regression_head(x), + "cls_logits": self.classification_head(x), + } + + +class SSDLiteClassificationHead(SSDScoringHead): + def __init__( + self, in_channels: List[int], num_anchors: List[int], num_classes: int, norm_layer: Callable[..., nn.Module] + ): + cls_logits = nn.ModuleList() + for channels, anchors in zip(in_channels, num_anchors): + cls_logits.append(_prediction_block(channels, num_classes * anchors, 3, norm_layer)) + _normal_init(cls_logits) + super().__init__(cls_logits, num_classes) + + +class SSDLiteRegressionHead(SSDScoringHead): + def __init__(self, in_channels: List[int], num_anchors: List[int], norm_layer: Callable[..., nn.Module]): + bbox_reg = nn.ModuleList() + for channels, anchors in zip(in_channels, num_anchors): + bbox_reg.append(_prediction_block(channels, 4 * anchors, 3, norm_layer)) + _normal_init(bbox_reg) + super().__init__(bbox_reg, 4) + + +class SSDLiteFeatureExtractorMobileNet(nn.Module): + def __init__( + self, + backbone: nn.Module, + c4_pos: int, + norm_layer: Callable[..., nn.Module], + width_mult: float = 1.0, + min_depth: int = 16, + ): + super().__init__() + _log_api_usage_once(self) + + if backbone[c4_pos].use_res_connect: + raise ValueError("backbone[c4_pos].use_res_connect should be False") + + self.features = nn.Sequential( + # As described in section 6.3 of MobileNetV3 paper + nn.Sequential(*backbone[:c4_pos], backbone[c4_pos].block[0]), # from start until C4 expansion layer + nn.Sequential(backbone[c4_pos].block[1:], *backbone[c4_pos + 1 :]), # from C4 depthwise until end + ) + + get_depth = lambda d: max(min_depth, int(d * width_mult)) # noqa: E731 + extra = nn.ModuleList( + [ + _extra_block(backbone[-1].out_channels, get_depth(512), norm_layer), + _extra_block(get_depth(512), get_depth(256), norm_layer), + _extra_block(get_depth(256), get_depth(256), norm_layer), + _extra_block(get_depth(256), get_depth(128), norm_layer), + ] + ) + _normal_init(extra) + + self.extra = extra + + def forward(self, x: Tensor) -> Dict[str, Tensor]: + # Get feature maps from backbone and extra. Can't be refactored due to JIT limitations. + output = [] + for block in self.features: + x = block(x) + output.append(x) + + for block in self.extra: + x = block(x) + output.append(x) + + return OrderedDict([(str(i), v) for i, v in enumerate(output)]) + + +def _mobilenet_extractor( + backbone: Union[mobilenet.MobileNetV2, mobilenet.MobileNetV3], + trainable_layers: int, + norm_layer: Callable[..., nn.Module], +): + backbone = backbone.features + # Gather the indices of blocks which are strided. These are the locations of C1, ..., Cn-1 blocks. + # The first and last blocks are always included because they are the C0 (conv1) and Cn. + stage_indices = [0] + [i for i, b in enumerate(backbone) if getattr(b, "_is_cn", False)] + [len(backbone) - 1] + num_stages = len(stage_indices) + + # find the index of the layer from which we won't freeze + if not 0 <= trainable_layers <= num_stages: + raise ValueError("trainable_layers should be in the range [0, {num_stages}], instead got {trainable_layers}") + freeze_before = len(backbone) if trainable_layers == 0 else stage_indices[num_stages - trainable_layers] + + for b in backbone[:freeze_before]: + for parameter in b.parameters(): + parameter.requires_grad_(False) + + return SSDLiteFeatureExtractorMobileNet(backbone, stage_indices[-2], norm_layer) + + +class SSDLite320_MobileNet_V3_Large_Weights(WeightsEnum): + COCO_V1 = Weights( + url="https://download.pytorch.org/models/ssdlite320_mobilenet_v3_large_coco-a79551df.pth", + transforms=ObjectDetection, + meta={ + "num_params": 3440060, + "categories": _COCO_CATEGORIES, + "min_size": (1, 1), + "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#ssdlite320-mobilenetv3-large", + "_metrics": { + "COCO-val2017": { + "box_map": 21.3, + } + }, + "_ops": 0.583, + "_file_size": 13.418, + "_docs": """These weights were produced by following a similar training recipe as on the paper.""", + }, + ) + DEFAULT = COCO_V1 + + +@register_model() +@handle_legacy_interface( + weights=("pretrained", SSDLite320_MobileNet_V3_Large_Weights.COCO_V1), + weights_backbone=("pretrained_backbone", MobileNet_V3_Large_Weights.IMAGENET1K_V1), +) +def ssdlite320_mobilenet_v3_large( + *, + weights: Optional[SSDLite320_MobileNet_V3_Large_Weights] = None, + progress: bool = True, + num_classes: Optional[int] = None, + weights_backbone: Optional[MobileNet_V3_Large_Weights] = MobileNet_V3_Large_Weights.IMAGENET1K_V1, + trainable_backbone_layers: Optional[int] = None, + norm_layer: Optional[Callable[..., nn.Module]] = None, + **kwargs: Any, +) -> SSD: + """SSDlite model architecture with input size 320x320 and a MobileNetV3 Large backbone, as + described at `Searching for MobileNetV3 `__ and + `MobileNetV2: Inverted Residuals and Linear Bottlenecks `__. + + .. betastatus:: detection module + + See :func:`~torchvision.models.detection.ssd300_vgg16` for more details. + + Example: + + >>> model = torchvision.models.detection.ssdlite320_mobilenet_v3_large(weights=SSDLite320_MobileNet_V3_Large_Weights.DEFAULT) + >>> model.eval() + >>> x = [torch.rand(3, 320, 320), torch.rand(3, 500, 400)] + >>> predictions = model(x) + + Args: + weights (:class:`~torchvision.models.detection.SSDLite320_MobileNet_V3_Large_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.detection.SSDLite320_MobileNet_V3_Large_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + num_classes (int, optional): number of output classes of the model + (including the background). + weights_backbone (:class:`~torchvision.models.MobileNet_V3_Large_Weights`, optional): The pretrained + weights for the backbone. + trainable_backbone_layers (int, optional): number of trainable (not frozen) layers + starting from final block. Valid values are between 0 and 6, with 6 meaning all + backbone layers are trainable. If ``None`` is passed (the default) this value is + set to 6. + norm_layer (callable, optional): Module specifying the normalization layer to use. + **kwargs: parameters passed to the ``torchvision.models.detection.ssd.SSD`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.detection.SSDLite320_MobileNet_V3_Large_Weights + :members: + """ + + weights = SSDLite320_MobileNet_V3_Large_Weights.verify(weights) + weights_backbone = MobileNet_V3_Large_Weights.verify(weights_backbone) + + if "size" in kwargs: + warnings.warn("The size of the model is already fixed; ignoring the parameter.") + + if weights is not None: + weights_backbone = None + num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"])) + elif num_classes is None: + num_classes = 91 + + trainable_backbone_layers = _validate_trainable_layers( + weights is not None or weights_backbone is not None, trainable_backbone_layers, 6, 6 + ) + + # Enable reduced tail if no pretrained backbone is selected. See Table 6 of MobileNetV3 paper. + reduce_tail = weights_backbone is None + + if norm_layer is None: + norm_layer = partial(nn.BatchNorm2d, eps=0.001, momentum=0.03) + + backbone = mobilenet_v3_large( + weights=weights_backbone, progress=progress, norm_layer=norm_layer, reduced_tail=reduce_tail, **kwargs + ) + if weights_backbone is None: + # Change the default initialization scheme if not pretrained + _normal_init(backbone) + backbone = _mobilenet_extractor( + backbone, + trainable_backbone_layers, + norm_layer, + ) + + size = (320, 320) + anchor_generator = DefaultBoxGenerator([[2, 3] for _ in range(6)], min_ratio=0.2, max_ratio=0.95) + out_channels = det_utils.retrieve_out_channels(backbone, size) + num_anchors = anchor_generator.num_anchors_per_location() + if len(out_channels) != len(anchor_generator.aspect_ratios): + raise ValueError( + f"The length of the output channels from the backbone {len(out_channels)} do not match the length of the anchor generator aspect ratios {len(anchor_generator.aspect_ratios)}" + ) + + defaults = { + "score_thresh": 0.001, + "nms_thresh": 0.55, + "detections_per_img": 300, + "topk_candidates": 300, + # Rescale the input in a way compatible to the backbone: + # The following mean/std rescale the data from [0, 1] to [-1, 1] + "image_mean": [0.5, 0.5, 0.5], + "image_std": [0.5, 0.5, 0.5], + } + kwargs: Any = {**defaults, **kwargs} + model = SSD( + backbone, + anchor_generator, + size, + num_classes, + head=SSDLiteHead(out_channels, num_anchors, num_classes, norm_layer), + **kwargs, + ) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model diff --git a/lib/python3.12/site-packages/torchvision/models/detection/transform.py b/lib/python3.12/site-packages/torchvision/models/detection/transform.py new file mode 100644 index 0000000000000000000000000000000000000000..9c569b0aafb0c5464815654c0f343d7fb927dc6c --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/detection/transform.py @@ -0,0 +1,319 @@ +import math +from typing import Any, Dict, List, Optional, Tuple + +import torch +import torchvision +from torch import nn, Tensor + +from .image_list import ImageList +from .roi_heads import paste_masks_in_image + + +@torch.jit.unused +def _get_shape_onnx(image: Tensor) -> Tensor: + from torch.onnx import operators + + return operators.shape_as_tensor(image)[-2:] + + +@torch.jit.unused +def _fake_cast_onnx(v: Tensor) -> float: + # ONNX requires a tensor but here we fake its type for JIT. + return v + + +def _resize_image_and_masks( + image: Tensor, + self_min_size: int, + self_max_size: int, + target: Optional[Dict[str, Tensor]] = None, + fixed_size: Optional[Tuple[int, int]] = None, +) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]: + if torchvision._is_tracing(): + im_shape = _get_shape_onnx(image) + elif torch.jit.is_scripting(): + im_shape = torch.tensor(image.shape[-2:]) + else: + im_shape = image.shape[-2:] + + size: Optional[List[int]] = None + scale_factor: Optional[float] = None + recompute_scale_factor: Optional[bool] = None + if fixed_size is not None: + size = [fixed_size[1], fixed_size[0]] + else: + if torch.jit.is_scripting() or torchvision._is_tracing(): + min_size = torch.min(im_shape).to(dtype=torch.float32) + max_size = torch.max(im_shape).to(dtype=torch.float32) + self_min_size_f = float(self_min_size) + self_max_size_f = float(self_max_size) + scale = torch.min(self_min_size_f / min_size, self_max_size_f / max_size) + + if torchvision._is_tracing(): + scale_factor = _fake_cast_onnx(scale) + else: + scale_factor = scale.item() + + else: + # Do it the normal way + min_size = min(im_shape) + max_size = max(im_shape) + scale_factor = min(self_min_size / min_size, self_max_size / max_size) + + recompute_scale_factor = True + + image = torch.nn.functional.interpolate( + image[None], + size=size, + scale_factor=scale_factor, + mode="bilinear", + recompute_scale_factor=recompute_scale_factor, + align_corners=False, + )[0] + + if target is None: + return image, target + + if "masks" in target: + mask = target["masks"] + mask = torch.nn.functional.interpolate( + mask[:, None].float(), size=size, scale_factor=scale_factor, recompute_scale_factor=recompute_scale_factor + )[:, 0].byte() + target["masks"] = mask + return image, target + + +class GeneralizedRCNNTransform(nn.Module): + """ + Performs input / target transformation before feeding the data to a GeneralizedRCNN + model. + + The transformations it performs are: + - input normalization (mean subtraction and std division) + - input / target resizing to match min_size / max_size + + It returns a ImageList for the inputs, and a List[Dict[Tensor]] for the targets + """ + + def __init__( + self, + min_size: int, + max_size: int, + image_mean: List[float], + image_std: List[float], + size_divisible: int = 32, + fixed_size: Optional[Tuple[int, int]] = None, + **kwargs: Any, + ): + super().__init__() + if not isinstance(min_size, (list, tuple)): + min_size = (min_size,) + self.min_size = min_size + self.max_size = max_size + self.image_mean = image_mean + self.image_std = image_std + self.size_divisible = size_divisible + self.fixed_size = fixed_size + self._skip_resize = kwargs.pop("_skip_resize", False) + + def forward( + self, images: List[Tensor], targets: Optional[List[Dict[str, Tensor]]] = None + ) -> Tuple[ImageList, Optional[List[Dict[str, Tensor]]]]: + images = [img for img in images] + if targets is not None: + # make a copy of targets to avoid modifying it in-place + # once torchscript supports dict comprehension + # this can be simplified as follows + # targets = [{k: v for k,v in t.items()} for t in targets] + targets_copy: List[Dict[str, Tensor]] = [] + for t in targets: + data: Dict[str, Tensor] = {} + for k, v in t.items(): + data[k] = v + targets_copy.append(data) + targets = targets_copy + for i in range(len(images)): + image = images[i] + target_index = targets[i] if targets is not None else None + + if image.dim() != 3: + raise ValueError(f"images is expected to be a list of 3d tensors of shape [C, H, W], got {image.shape}") + image = self.normalize(image) + image, target_index = self.resize(image, target_index) + images[i] = image + if targets is not None and target_index is not None: + targets[i] = target_index + + image_sizes = [img.shape[-2:] for img in images] + images = self.batch_images(images, size_divisible=self.size_divisible) + image_sizes_list: List[Tuple[int, int]] = [] + for image_size in image_sizes: + torch._assert( + len(image_size) == 2, + f"Input tensors expected to have in the last two elements H and W, instead got {image_size}", + ) + image_sizes_list.append((image_size[0], image_size[1])) + + image_list = ImageList(images, image_sizes_list) + return image_list, targets + + def normalize(self, image: Tensor) -> Tensor: + if not image.is_floating_point(): + raise TypeError( + f"Expected input images to be of floating type (in range [0, 1]), " + f"but found type {image.dtype} instead" + ) + dtype, device = image.dtype, image.device + mean = torch.as_tensor(self.image_mean, dtype=dtype, device=device) + std = torch.as_tensor(self.image_std, dtype=dtype, device=device) + return (image - mean[:, None, None]) / std[:, None, None] + + def torch_choice(self, k: List[int]) -> int: + """ + Implements `random.choice` via torch ops, so it can be compiled with + TorchScript and we use PyTorch's RNG (not native RNG) + """ + index = int(torch.empty(1).uniform_(0.0, float(len(k))).item()) + return k[index] + + def resize( + self, + image: Tensor, + target: Optional[Dict[str, Tensor]] = None, + ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]: + h, w = image.shape[-2:] + if self.training: + if self._skip_resize: + return image, target + size = self.torch_choice(self.min_size) + else: + size = self.min_size[-1] + image, target = _resize_image_and_masks(image, size, self.max_size, target, self.fixed_size) + + if target is None: + return image, target + + bbox = target["boxes"] + bbox = resize_boxes(bbox, (h, w), image.shape[-2:]) + target["boxes"] = bbox + + if "keypoints" in target: + keypoints = target["keypoints"] + keypoints = resize_keypoints(keypoints, (h, w), image.shape[-2:]) + target["keypoints"] = keypoints + return image, target + + # _onnx_batch_images() is an implementation of + # batch_images() that is supported by ONNX tracing. + @torch.jit.unused + def _onnx_batch_images(self, images: List[Tensor], size_divisible: int = 32) -> Tensor: + max_size = [] + for i in range(images[0].dim()): + max_size_i = torch.max(torch.stack([img.shape[i] for img in images]).to(torch.float32)).to(torch.int64) + max_size.append(max_size_i) + stride = size_divisible + max_size[1] = (torch.ceil((max_size[1].to(torch.float32)) / stride) * stride).to(torch.int64) + max_size[2] = (torch.ceil((max_size[2].to(torch.float32)) / stride) * stride).to(torch.int64) + max_size = tuple(max_size) + + # work around for + # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img) + # which is not yet supported in onnx + padded_imgs = [] + for img in images: + padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))] + padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0])) + padded_imgs.append(padded_img) + + return torch.stack(padded_imgs) + + def max_by_axis(self, the_list: List[List[int]]) -> List[int]: + maxes = the_list[0] + for sublist in the_list[1:]: + for index, item in enumerate(sublist): + maxes[index] = max(maxes[index], item) + return maxes + + def batch_images(self, images: List[Tensor], size_divisible: int = 32) -> Tensor: + if torchvision._is_tracing(): + # batch_images() does not export well to ONNX + # call _onnx_batch_images() instead + return self._onnx_batch_images(images, size_divisible) + + max_size = self.max_by_axis([list(img.shape) for img in images]) + stride = float(size_divisible) + max_size = list(max_size) + max_size[1] = int(math.ceil(float(max_size[1]) / stride) * stride) + max_size[2] = int(math.ceil(float(max_size[2]) / stride) * stride) + + batch_shape = [len(images)] + max_size + batched_imgs = images[0].new_full(batch_shape, 0) + for i in range(batched_imgs.shape[0]): + img = images[i] + batched_imgs[i, : img.shape[0], : img.shape[1], : img.shape[2]].copy_(img) + + return batched_imgs + + def postprocess( + self, + result: List[Dict[str, Tensor]], + image_shapes: List[Tuple[int, int]], + original_image_sizes: List[Tuple[int, int]], + ) -> List[Dict[str, Tensor]]: + if self.training: + return result + for i, (pred, im_s, o_im_s) in enumerate(zip(result, image_shapes, original_image_sizes)): + boxes = pred["boxes"] + boxes = resize_boxes(boxes, im_s, o_im_s) + result[i]["boxes"] = boxes + if "masks" in pred: + masks = pred["masks"] + masks = paste_masks_in_image(masks, boxes, o_im_s) + result[i]["masks"] = masks + if "keypoints" in pred: + keypoints = pred["keypoints"] + keypoints = resize_keypoints(keypoints, im_s, o_im_s) + result[i]["keypoints"] = keypoints + return result + + def __repr__(self) -> str: + format_string = f"{self.__class__.__name__}(" + _indent = "\n " + format_string += f"{_indent}Normalize(mean={self.image_mean}, std={self.image_std})" + format_string += f"{_indent}Resize(min_size={self.min_size}, max_size={self.max_size}, mode='bilinear')" + format_string += "\n)" + return format_string + + +def resize_keypoints(keypoints: Tensor, original_size: List[int], new_size: List[int]) -> Tensor: + ratios = [ + torch.tensor(s, dtype=torch.float32, device=keypoints.device) + / torch.tensor(s_orig, dtype=torch.float32, device=keypoints.device) + for s, s_orig in zip(new_size, original_size) + ] + ratio_h, ratio_w = ratios + resized_data = keypoints.clone() + if torch._C._get_tracing_state(): + resized_data_0 = resized_data[:, :, 0] * ratio_w + resized_data_1 = resized_data[:, :, 1] * ratio_h + resized_data = torch.stack((resized_data_0, resized_data_1, resized_data[:, :, 2]), dim=2) + else: + resized_data[..., 0] *= ratio_w + resized_data[..., 1] *= ratio_h + return resized_data + + +def resize_boxes(boxes: Tensor, original_size: List[int], new_size: List[int]) -> Tensor: + ratios = [ + torch.tensor(s, dtype=torch.float32, device=boxes.device) + / torch.tensor(s_orig, dtype=torch.float32, device=boxes.device) + for s, s_orig in zip(new_size, original_size) + ] + ratio_height, ratio_width = ratios + xmin, ymin, xmax, ymax = boxes.unbind(1) + + xmin = xmin * ratio_width + xmax = xmax * ratio_width + ymin = ymin * ratio_height + ymax = ymax * ratio_height + return torch.stack((xmin, ymin, xmax, ymax), dim=1) diff --git a/lib/python3.12/site-packages/torchvision/models/efficientnet.py b/lib/python3.12/site-packages/torchvision/models/efficientnet.py new file mode 100644 index 0000000000000000000000000000000000000000..65f0b2fef44702f5c0db9c545c2ecb751a683a72 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/efficientnet.py @@ -0,0 +1,1131 @@ +import copy +import math +from dataclasses import dataclass +from functools import partial +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union + +import torch +from torch import nn, Tensor +from torchvision.ops import StochasticDepth + +from ..ops.misc import Conv2dNormActivation, SqueezeExcitation +from ..transforms._presets import ImageClassification, InterpolationMode +from ..utils import _log_api_usage_once +from ._api import register_model, Weights, WeightsEnum +from ._meta import _IMAGENET_CATEGORIES +from ._utils import _make_divisible, _ovewrite_named_param, handle_legacy_interface + + +__all__ = [ + "EfficientNet", + "EfficientNet_B0_Weights", + "EfficientNet_B1_Weights", + "EfficientNet_B2_Weights", + "EfficientNet_B3_Weights", + "EfficientNet_B4_Weights", + "EfficientNet_B5_Weights", + "EfficientNet_B6_Weights", + "EfficientNet_B7_Weights", + "EfficientNet_V2_S_Weights", + "EfficientNet_V2_M_Weights", + "EfficientNet_V2_L_Weights", + "efficientnet_b0", + "efficientnet_b1", + "efficientnet_b2", + "efficientnet_b3", + "efficientnet_b4", + "efficientnet_b5", + "efficientnet_b6", + "efficientnet_b7", + "efficientnet_v2_s", + "efficientnet_v2_m", + "efficientnet_v2_l", +] + + +@dataclass +class _MBConvConfig: + expand_ratio: float + kernel: int + stride: int + input_channels: int + out_channels: int + num_layers: int + block: Callable[..., nn.Module] + + @staticmethod + def adjust_channels(channels: int, width_mult: float, min_value: Optional[int] = None) -> int: + return _make_divisible(channels * width_mult, 8, min_value) + + +class MBConvConfig(_MBConvConfig): + # Stores information listed at Table 1 of the EfficientNet paper & Table 4 of the EfficientNetV2 paper + def __init__( + self, + expand_ratio: float, + kernel: int, + stride: int, + input_channels: int, + out_channels: int, + num_layers: int, + width_mult: float = 1.0, + depth_mult: float = 1.0, + block: Optional[Callable[..., nn.Module]] = None, + ) -> None: + input_channels = self.adjust_channels(input_channels, width_mult) + out_channels = self.adjust_channels(out_channels, width_mult) + num_layers = self.adjust_depth(num_layers, depth_mult) + if block is None: + block = MBConv + super().__init__(expand_ratio, kernel, stride, input_channels, out_channels, num_layers, block) + + @staticmethod + def adjust_depth(num_layers: int, depth_mult: float): + return int(math.ceil(num_layers * depth_mult)) + + +class FusedMBConvConfig(_MBConvConfig): + # Stores information listed at Table 4 of the EfficientNetV2 paper + def __init__( + self, + expand_ratio: float, + kernel: int, + stride: int, + input_channels: int, + out_channels: int, + num_layers: int, + block: Optional[Callable[..., nn.Module]] = None, + ) -> None: + if block is None: + block = FusedMBConv + super().__init__(expand_ratio, kernel, stride, input_channels, out_channels, num_layers, block) + + +class MBConv(nn.Module): + def __init__( + self, + cnf: MBConvConfig, + stochastic_depth_prob: float, + norm_layer: Callable[..., nn.Module], + se_layer: Callable[..., nn.Module] = SqueezeExcitation, + ) -> None: + super().__init__() + + if not (1 <= cnf.stride <= 2): + raise ValueError("illegal stride value") + + self.use_res_connect = cnf.stride == 1 and cnf.input_channels == cnf.out_channels + + layers: List[nn.Module] = [] + activation_layer = nn.SiLU + + # expand + expanded_channels = cnf.adjust_channels(cnf.input_channels, cnf.expand_ratio) + if expanded_channels != cnf.input_channels: + layers.append( + Conv2dNormActivation( + cnf.input_channels, + expanded_channels, + kernel_size=1, + norm_layer=norm_layer, + activation_layer=activation_layer, + ) + ) + + # depthwise + layers.append( + Conv2dNormActivation( + expanded_channels, + expanded_channels, + kernel_size=cnf.kernel, + stride=cnf.stride, + groups=expanded_channels, + norm_layer=norm_layer, + activation_layer=activation_layer, + ) + ) + + # squeeze and excitation + squeeze_channels = max(1, cnf.input_channels // 4) + layers.append(se_layer(expanded_channels, squeeze_channels, activation=partial(nn.SiLU, inplace=True))) + + # project + layers.append( + Conv2dNormActivation( + expanded_channels, cnf.out_channels, kernel_size=1, norm_layer=norm_layer, activation_layer=None + ) + ) + + self.block = nn.Sequential(*layers) + self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row") + self.out_channels = cnf.out_channels + + def forward(self, input: Tensor) -> Tensor: + result = self.block(input) + if self.use_res_connect: + result = self.stochastic_depth(result) + result += input + return result + + +class FusedMBConv(nn.Module): + def __init__( + self, + cnf: FusedMBConvConfig, + stochastic_depth_prob: float, + norm_layer: Callable[..., nn.Module], + ) -> None: + super().__init__() + + if not (1 <= cnf.stride <= 2): + raise ValueError("illegal stride value") + + self.use_res_connect = cnf.stride == 1 and cnf.input_channels == cnf.out_channels + + layers: List[nn.Module] = [] + activation_layer = nn.SiLU + + expanded_channels = cnf.adjust_channels(cnf.input_channels, cnf.expand_ratio) + if expanded_channels != cnf.input_channels: + # fused expand + layers.append( + Conv2dNormActivation( + cnf.input_channels, + expanded_channels, + kernel_size=cnf.kernel, + stride=cnf.stride, + norm_layer=norm_layer, + activation_layer=activation_layer, + ) + ) + + # project + layers.append( + Conv2dNormActivation( + expanded_channels, cnf.out_channels, kernel_size=1, norm_layer=norm_layer, activation_layer=None + ) + ) + else: + layers.append( + Conv2dNormActivation( + cnf.input_channels, + cnf.out_channels, + kernel_size=cnf.kernel, + stride=cnf.stride, + norm_layer=norm_layer, + activation_layer=activation_layer, + ) + ) + + self.block = nn.Sequential(*layers) + self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row") + self.out_channels = cnf.out_channels + + def forward(self, input: Tensor) -> Tensor: + result = self.block(input) + if self.use_res_connect: + result = self.stochastic_depth(result) + result += input + return result + + +class EfficientNet(nn.Module): + def __init__( + self, + inverted_residual_setting: Sequence[Union[MBConvConfig, FusedMBConvConfig]], + dropout: float, + stochastic_depth_prob: float = 0.2, + num_classes: int = 1000, + norm_layer: Optional[Callable[..., nn.Module]] = None, + last_channel: Optional[int] = None, + ) -> None: + """ + EfficientNet V1 and V2 main class + + Args: + inverted_residual_setting (Sequence[Union[MBConvConfig, FusedMBConvConfig]]): Network structure + dropout (float): The droupout probability + stochastic_depth_prob (float): The stochastic depth probability + num_classes (int): Number of classes + norm_layer (Optional[Callable[..., nn.Module]]): Module specifying the normalization layer to use + last_channel (int): The number of channels on the penultimate layer + """ + super().__init__() + _log_api_usage_once(self) + + if not inverted_residual_setting: + raise ValueError("The inverted_residual_setting should not be empty") + elif not ( + isinstance(inverted_residual_setting, Sequence) + and all([isinstance(s, _MBConvConfig) for s in inverted_residual_setting]) + ): + raise TypeError("The inverted_residual_setting should be List[MBConvConfig]") + + if norm_layer is None: + norm_layer = nn.BatchNorm2d + + layers: List[nn.Module] = [] + + # building first layer + firstconv_output_channels = inverted_residual_setting[0].input_channels + layers.append( + Conv2dNormActivation( + 3, firstconv_output_channels, kernel_size=3, stride=2, norm_layer=norm_layer, activation_layer=nn.SiLU + ) + ) + + # building inverted residual blocks + total_stage_blocks = sum(cnf.num_layers for cnf in inverted_residual_setting) + stage_block_id = 0 + for cnf in inverted_residual_setting: + stage: List[nn.Module] = [] + for _ in range(cnf.num_layers): + # copy to avoid modifications. shallow copy is enough + block_cnf = copy.copy(cnf) + + # overwrite info if not the first conv in the stage + if stage: + block_cnf.input_channels = block_cnf.out_channels + block_cnf.stride = 1 + + # adjust stochastic depth probability based on the depth of the stage block + sd_prob = stochastic_depth_prob * float(stage_block_id) / total_stage_blocks + + stage.append(block_cnf.block(block_cnf, sd_prob, norm_layer)) + stage_block_id += 1 + + layers.append(nn.Sequential(*stage)) + + # building last several layers + lastconv_input_channels = inverted_residual_setting[-1].out_channels + lastconv_output_channels = last_channel if last_channel is not None else 4 * lastconv_input_channels + layers.append( + Conv2dNormActivation( + lastconv_input_channels, + lastconv_output_channels, + kernel_size=1, + norm_layer=norm_layer, + activation_layer=nn.SiLU, + ) + ) + + self.features = nn.Sequential(*layers) + self.avgpool = nn.AdaptiveAvgPool2d(1) + self.classifier = nn.Sequential( + nn.Dropout(p=dropout, inplace=True), + nn.Linear(lastconv_output_channels, num_classes), + ) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode="fan_out") + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): + nn.init.ones_(m.weight) + nn.init.zeros_(m.bias) + elif isinstance(m, nn.Linear): + init_range = 1.0 / math.sqrt(m.out_features) + nn.init.uniform_(m.weight, -init_range, init_range) + nn.init.zeros_(m.bias) + + def _forward_impl(self, x: Tensor) -> Tensor: + x = self.features(x) + + x = self.avgpool(x) + x = torch.flatten(x, 1) + + x = self.classifier(x) + + return x + + def forward(self, x: Tensor) -> Tensor: + return self._forward_impl(x) + + +def _efficientnet( + inverted_residual_setting: Sequence[Union[MBConvConfig, FusedMBConvConfig]], + dropout: float, + last_channel: Optional[int], + weights: Optional[WeightsEnum], + progress: bool, + **kwargs: Any, +) -> EfficientNet: + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + + model = EfficientNet(inverted_residual_setting, dropout, last_channel=last_channel, **kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +def _efficientnet_conf( + arch: str, + **kwargs: Any, +) -> Tuple[Sequence[Union[MBConvConfig, FusedMBConvConfig]], Optional[int]]: + inverted_residual_setting: Sequence[Union[MBConvConfig, FusedMBConvConfig]] + if arch.startswith("efficientnet_b"): + bneck_conf = partial(MBConvConfig, width_mult=kwargs.pop("width_mult"), depth_mult=kwargs.pop("depth_mult")) + inverted_residual_setting = [ + bneck_conf(1, 3, 1, 32, 16, 1), + bneck_conf(6, 3, 2, 16, 24, 2), + bneck_conf(6, 5, 2, 24, 40, 2), + bneck_conf(6, 3, 2, 40, 80, 3), + bneck_conf(6, 5, 1, 80, 112, 3), + bneck_conf(6, 5, 2, 112, 192, 4), + bneck_conf(6, 3, 1, 192, 320, 1), + ] + last_channel = None + elif arch.startswith("efficientnet_v2_s"): + inverted_residual_setting = [ + FusedMBConvConfig(1, 3, 1, 24, 24, 2), + FusedMBConvConfig(4, 3, 2, 24, 48, 4), + FusedMBConvConfig(4, 3, 2, 48, 64, 4), + MBConvConfig(4, 3, 2, 64, 128, 6), + MBConvConfig(6, 3, 1, 128, 160, 9), + MBConvConfig(6, 3, 2, 160, 256, 15), + ] + last_channel = 1280 + elif arch.startswith("efficientnet_v2_m"): + inverted_residual_setting = [ + FusedMBConvConfig(1, 3, 1, 24, 24, 3), + FusedMBConvConfig(4, 3, 2, 24, 48, 5), + FusedMBConvConfig(4, 3, 2, 48, 80, 5), + MBConvConfig(4, 3, 2, 80, 160, 7), + MBConvConfig(6, 3, 1, 160, 176, 14), + MBConvConfig(6, 3, 2, 176, 304, 18), + MBConvConfig(6, 3, 1, 304, 512, 5), + ] + last_channel = 1280 + elif arch.startswith("efficientnet_v2_l"): + inverted_residual_setting = [ + FusedMBConvConfig(1, 3, 1, 32, 32, 4), + FusedMBConvConfig(4, 3, 2, 32, 64, 7), + FusedMBConvConfig(4, 3, 2, 64, 96, 7), + MBConvConfig(4, 3, 2, 96, 192, 10), + MBConvConfig(6, 3, 1, 192, 224, 19), + MBConvConfig(6, 3, 2, 224, 384, 25), + MBConvConfig(6, 3, 1, 384, 640, 7), + ] + last_channel = 1280 + else: + raise ValueError(f"Unsupported model type {arch}") + + return inverted_residual_setting, last_channel + + +_COMMON_META: Dict[str, Any] = { + "categories": _IMAGENET_CATEGORIES, +} + + +_COMMON_META_V1 = { + **_COMMON_META, + "min_size": (1, 1), + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#efficientnet-v1", +} + + +_COMMON_META_V2 = { + **_COMMON_META, + "min_size": (33, 33), + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#efficientnet-v2", +} + + +class EfficientNet_B0_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + # Weights ported from https://github.com/rwightman/pytorch-image-models/ + url="https://download.pytorch.org/models/efficientnet_b0_rwightman-7f5810bc.pth", + transforms=partial( + ImageClassification, crop_size=224, resize_size=256, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_META_V1, + "num_params": 5288548, + "_metrics": { + "ImageNet-1K": { + "acc@1": 77.692, + "acc@5": 93.532, + } + }, + "_ops": 0.386, + "_file_size": 20.451, + "_docs": """These weights are ported from the original paper.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class EfficientNet_B1_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + # Weights ported from https://github.com/rwightman/pytorch-image-models/ + url="https://download.pytorch.org/models/efficientnet_b1_rwightman-bac287d4.pth", + transforms=partial( + ImageClassification, crop_size=240, resize_size=256, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_META_V1, + "num_params": 7794184, + "_metrics": { + "ImageNet-1K": { + "acc@1": 78.642, + "acc@5": 94.186, + } + }, + "_ops": 0.687, + "_file_size": 30.134, + "_docs": """These weights are ported from the original paper.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/efficientnet_b1-c27df63c.pth", + transforms=partial( + ImageClassification, crop_size=240, resize_size=255, interpolation=InterpolationMode.BILINEAR + ), + meta={ + **_COMMON_META_V1, + "num_params": 7794184, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe-with-lr-wd-crop-tuning", + "_metrics": { + "ImageNet-1K": { + "acc@1": 79.838, + "acc@5": 94.934, + } + }, + "_ops": 0.687, + "_file_size": 30.136, + "_docs": """ + These weights improve upon the results of the original paper by using a modified version of TorchVision's + `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class EfficientNet_B2_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + # Weights ported from https://github.com/rwightman/pytorch-image-models/ + url="https://download.pytorch.org/models/efficientnet_b2_rwightman-c35c1473.pth", + transforms=partial( + ImageClassification, crop_size=288, resize_size=288, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_META_V1, + "num_params": 9109994, + "_metrics": { + "ImageNet-1K": { + "acc@1": 80.608, + "acc@5": 95.310, + } + }, + "_ops": 1.088, + "_file_size": 35.174, + "_docs": """These weights are ported from the original paper.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class EfficientNet_B3_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + # Weights ported from https://github.com/rwightman/pytorch-image-models/ + url="https://download.pytorch.org/models/efficientnet_b3_rwightman-b3899882.pth", + transforms=partial( + ImageClassification, crop_size=300, resize_size=320, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_META_V1, + "num_params": 12233232, + "_metrics": { + "ImageNet-1K": { + "acc@1": 82.008, + "acc@5": 96.054, + } + }, + "_ops": 1.827, + "_file_size": 47.184, + "_docs": """These weights are ported from the original paper.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class EfficientNet_B4_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + # Weights ported from https://github.com/rwightman/pytorch-image-models/ + url="https://download.pytorch.org/models/efficientnet_b4_rwightman-23ab8bcd.pth", + transforms=partial( + ImageClassification, crop_size=380, resize_size=384, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_META_V1, + "num_params": 19341616, + "_metrics": { + "ImageNet-1K": { + "acc@1": 83.384, + "acc@5": 96.594, + } + }, + "_ops": 4.394, + "_file_size": 74.489, + "_docs": """These weights are ported from the original paper.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class EfficientNet_B5_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + # Weights ported from https://github.com/lukemelas/EfficientNet-PyTorch/ + url="https://download.pytorch.org/models/efficientnet_b5_lukemelas-1a07897c.pth", + transforms=partial( + ImageClassification, crop_size=456, resize_size=456, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_META_V1, + "num_params": 30389784, + "_metrics": { + "ImageNet-1K": { + "acc@1": 83.444, + "acc@5": 96.628, + } + }, + "_ops": 10.266, + "_file_size": 116.864, + "_docs": """These weights are ported from the original paper.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class EfficientNet_B6_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + # Weights ported from https://github.com/lukemelas/EfficientNet-PyTorch/ + url="https://download.pytorch.org/models/efficientnet_b6_lukemelas-24a108a5.pth", + transforms=partial( + ImageClassification, crop_size=528, resize_size=528, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_META_V1, + "num_params": 43040704, + "_metrics": { + "ImageNet-1K": { + "acc@1": 84.008, + "acc@5": 96.916, + } + }, + "_ops": 19.068, + "_file_size": 165.362, + "_docs": """These weights are ported from the original paper.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class EfficientNet_B7_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + # Weights ported from https://github.com/lukemelas/EfficientNet-PyTorch/ + url="https://download.pytorch.org/models/efficientnet_b7_lukemelas-c5b4e57e.pth", + transforms=partial( + ImageClassification, crop_size=600, resize_size=600, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_META_V1, + "num_params": 66347960, + "_metrics": { + "ImageNet-1K": { + "acc@1": 84.122, + "acc@5": 96.908, + } + }, + "_ops": 37.746, + "_file_size": 254.675, + "_docs": """These weights are ported from the original paper.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class EfficientNet_V2_S_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/efficientnet_v2_s-dd5fe13b.pth", + transforms=partial( + ImageClassification, + crop_size=384, + resize_size=384, + interpolation=InterpolationMode.BILINEAR, + ), + meta={ + **_COMMON_META_V2, + "num_params": 21458488, + "_metrics": { + "ImageNet-1K": { + "acc@1": 84.228, + "acc@5": 96.878, + } + }, + "_ops": 8.366, + "_file_size": 82.704, + "_docs": """ + These weights improve upon the results of the original paper by using a modified version of TorchVision's + `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class EfficientNet_V2_M_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/efficientnet_v2_m-dc08266a.pth", + transforms=partial( + ImageClassification, + crop_size=480, + resize_size=480, + interpolation=InterpolationMode.BILINEAR, + ), + meta={ + **_COMMON_META_V2, + "num_params": 54139356, + "_metrics": { + "ImageNet-1K": { + "acc@1": 85.112, + "acc@5": 97.156, + } + }, + "_ops": 24.582, + "_file_size": 208.01, + "_docs": """ + These weights improve upon the results of the original paper by using a modified version of TorchVision's + `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class EfficientNet_V2_L_Weights(WeightsEnum): + # Weights ported from https://github.com/google/automl/tree/master/efficientnetv2 + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/efficientnet_v2_l-59c71312.pth", + transforms=partial( + ImageClassification, + crop_size=480, + resize_size=480, + interpolation=InterpolationMode.BICUBIC, + mean=(0.5, 0.5, 0.5), + std=(0.5, 0.5, 0.5), + ), + meta={ + **_COMMON_META_V2, + "num_params": 118515272, + "_metrics": { + "ImageNet-1K": { + "acc@1": 85.808, + "acc@5": 97.788, + } + }, + "_ops": 56.08, + "_file_size": 454.573, + "_docs": """These weights are ported from the original paper.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", EfficientNet_B0_Weights.IMAGENET1K_V1)) +def efficientnet_b0( + *, weights: Optional[EfficientNet_B0_Weights] = None, progress: bool = True, **kwargs: Any +) -> EfficientNet: + """EfficientNet B0 model architecture from the `EfficientNet: Rethinking Model Scaling for Convolutional + Neural Networks `_ paper. + + Args: + weights (:class:`~torchvision.models.EfficientNet_B0_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.EfficientNet_B0_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.efficientnet.EfficientNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + .. autoclass:: torchvision.models.EfficientNet_B0_Weights + :members: + """ + weights = EfficientNet_B0_Weights.verify(weights) + + inverted_residual_setting, last_channel = _efficientnet_conf("efficientnet_b0", width_mult=1.0, depth_mult=1.0) + return _efficientnet( + inverted_residual_setting, kwargs.pop("dropout", 0.2), last_channel, weights, progress, **kwargs + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", EfficientNet_B1_Weights.IMAGENET1K_V1)) +def efficientnet_b1( + *, weights: Optional[EfficientNet_B1_Weights] = None, progress: bool = True, **kwargs: Any +) -> EfficientNet: + """EfficientNet B1 model architecture from the `EfficientNet: Rethinking Model Scaling for Convolutional + Neural Networks `_ paper. + + Args: + weights (:class:`~torchvision.models.EfficientNet_B1_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.EfficientNet_B1_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.efficientnet.EfficientNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + .. autoclass:: torchvision.models.EfficientNet_B1_Weights + :members: + """ + weights = EfficientNet_B1_Weights.verify(weights) + + inverted_residual_setting, last_channel = _efficientnet_conf("efficientnet_b1", width_mult=1.0, depth_mult=1.1) + return _efficientnet( + inverted_residual_setting, kwargs.pop("dropout", 0.2), last_channel, weights, progress, **kwargs + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", EfficientNet_B2_Weights.IMAGENET1K_V1)) +def efficientnet_b2( + *, weights: Optional[EfficientNet_B2_Weights] = None, progress: bool = True, **kwargs: Any +) -> EfficientNet: + """EfficientNet B2 model architecture from the `EfficientNet: Rethinking Model Scaling for Convolutional + Neural Networks `_ paper. + + Args: + weights (:class:`~torchvision.models.EfficientNet_B2_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.EfficientNet_B2_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.efficientnet.EfficientNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + .. autoclass:: torchvision.models.EfficientNet_B2_Weights + :members: + """ + weights = EfficientNet_B2_Weights.verify(weights) + + inverted_residual_setting, last_channel = _efficientnet_conf("efficientnet_b2", width_mult=1.1, depth_mult=1.2) + return _efficientnet( + inverted_residual_setting, kwargs.pop("dropout", 0.3), last_channel, weights, progress, **kwargs + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", EfficientNet_B3_Weights.IMAGENET1K_V1)) +def efficientnet_b3( + *, weights: Optional[EfficientNet_B3_Weights] = None, progress: bool = True, **kwargs: Any +) -> EfficientNet: + """EfficientNet B3 model architecture from the `EfficientNet: Rethinking Model Scaling for Convolutional + Neural Networks `_ paper. + + Args: + weights (:class:`~torchvision.models.EfficientNet_B3_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.EfficientNet_B3_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.efficientnet.EfficientNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + .. autoclass:: torchvision.models.EfficientNet_B3_Weights + :members: + """ + weights = EfficientNet_B3_Weights.verify(weights) + + inverted_residual_setting, last_channel = _efficientnet_conf("efficientnet_b3", width_mult=1.2, depth_mult=1.4) + return _efficientnet( + inverted_residual_setting, + kwargs.pop("dropout", 0.3), + last_channel, + weights, + progress, + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", EfficientNet_B4_Weights.IMAGENET1K_V1)) +def efficientnet_b4( + *, weights: Optional[EfficientNet_B4_Weights] = None, progress: bool = True, **kwargs: Any +) -> EfficientNet: + """EfficientNet B4 model architecture from the `EfficientNet: Rethinking Model Scaling for Convolutional + Neural Networks `_ paper. + + Args: + weights (:class:`~torchvision.models.EfficientNet_B4_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.EfficientNet_B4_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.efficientnet.EfficientNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + .. autoclass:: torchvision.models.EfficientNet_B4_Weights + :members: + """ + weights = EfficientNet_B4_Weights.verify(weights) + + inverted_residual_setting, last_channel = _efficientnet_conf("efficientnet_b4", width_mult=1.4, depth_mult=1.8) + return _efficientnet( + inverted_residual_setting, + kwargs.pop("dropout", 0.4), + last_channel, + weights, + progress, + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", EfficientNet_B5_Weights.IMAGENET1K_V1)) +def efficientnet_b5( + *, weights: Optional[EfficientNet_B5_Weights] = None, progress: bool = True, **kwargs: Any +) -> EfficientNet: + """EfficientNet B5 model architecture from the `EfficientNet: Rethinking Model Scaling for Convolutional + Neural Networks `_ paper. + + Args: + weights (:class:`~torchvision.models.EfficientNet_B5_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.EfficientNet_B5_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.efficientnet.EfficientNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + .. autoclass:: torchvision.models.EfficientNet_B5_Weights + :members: + """ + weights = EfficientNet_B5_Weights.verify(weights) + + inverted_residual_setting, last_channel = _efficientnet_conf("efficientnet_b5", width_mult=1.6, depth_mult=2.2) + return _efficientnet( + inverted_residual_setting, + kwargs.pop("dropout", 0.4), + last_channel, + weights, + progress, + norm_layer=partial(nn.BatchNorm2d, eps=0.001, momentum=0.01), + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", EfficientNet_B6_Weights.IMAGENET1K_V1)) +def efficientnet_b6( + *, weights: Optional[EfficientNet_B6_Weights] = None, progress: bool = True, **kwargs: Any +) -> EfficientNet: + """EfficientNet B6 model architecture from the `EfficientNet: Rethinking Model Scaling for Convolutional + Neural Networks `_ paper. + + Args: + weights (:class:`~torchvision.models.EfficientNet_B6_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.EfficientNet_B6_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.efficientnet.EfficientNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + .. autoclass:: torchvision.models.EfficientNet_B6_Weights + :members: + """ + weights = EfficientNet_B6_Weights.verify(weights) + + inverted_residual_setting, last_channel = _efficientnet_conf("efficientnet_b6", width_mult=1.8, depth_mult=2.6) + return _efficientnet( + inverted_residual_setting, + kwargs.pop("dropout", 0.5), + last_channel, + weights, + progress, + norm_layer=partial(nn.BatchNorm2d, eps=0.001, momentum=0.01), + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", EfficientNet_B7_Weights.IMAGENET1K_V1)) +def efficientnet_b7( + *, weights: Optional[EfficientNet_B7_Weights] = None, progress: bool = True, **kwargs: Any +) -> EfficientNet: + """EfficientNet B7 model architecture from the `EfficientNet: Rethinking Model Scaling for Convolutional + Neural Networks `_ paper. + + Args: + weights (:class:`~torchvision.models.EfficientNet_B7_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.EfficientNet_B7_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.efficientnet.EfficientNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + .. autoclass:: torchvision.models.EfficientNet_B7_Weights + :members: + """ + weights = EfficientNet_B7_Weights.verify(weights) + + inverted_residual_setting, last_channel = _efficientnet_conf("efficientnet_b7", width_mult=2.0, depth_mult=3.1) + return _efficientnet( + inverted_residual_setting, + kwargs.pop("dropout", 0.5), + last_channel, + weights, + progress, + norm_layer=partial(nn.BatchNorm2d, eps=0.001, momentum=0.01), + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", EfficientNet_V2_S_Weights.IMAGENET1K_V1)) +def efficientnet_v2_s( + *, weights: Optional[EfficientNet_V2_S_Weights] = None, progress: bool = True, **kwargs: Any +) -> EfficientNet: + """ + Constructs an EfficientNetV2-S architecture from + `EfficientNetV2: Smaller Models and Faster Training `_. + + Args: + weights (:class:`~torchvision.models.EfficientNet_V2_S_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.EfficientNet_V2_S_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.efficientnet.EfficientNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + .. autoclass:: torchvision.models.EfficientNet_V2_S_Weights + :members: + """ + weights = EfficientNet_V2_S_Weights.verify(weights) + + inverted_residual_setting, last_channel = _efficientnet_conf("efficientnet_v2_s") + return _efficientnet( + inverted_residual_setting, + kwargs.pop("dropout", 0.2), + last_channel, + weights, + progress, + norm_layer=partial(nn.BatchNorm2d, eps=1e-03), + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", EfficientNet_V2_M_Weights.IMAGENET1K_V1)) +def efficientnet_v2_m( + *, weights: Optional[EfficientNet_V2_M_Weights] = None, progress: bool = True, **kwargs: Any +) -> EfficientNet: + """ + Constructs an EfficientNetV2-M architecture from + `EfficientNetV2: Smaller Models and Faster Training `_. + + Args: + weights (:class:`~torchvision.models.EfficientNet_V2_M_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.EfficientNet_V2_M_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.efficientnet.EfficientNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + .. autoclass:: torchvision.models.EfficientNet_V2_M_Weights + :members: + """ + weights = EfficientNet_V2_M_Weights.verify(weights) + + inverted_residual_setting, last_channel = _efficientnet_conf("efficientnet_v2_m") + return _efficientnet( + inverted_residual_setting, + kwargs.pop("dropout", 0.3), + last_channel, + weights, + progress, + norm_layer=partial(nn.BatchNorm2d, eps=1e-03), + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", EfficientNet_V2_L_Weights.IMAGENET1K_V1)) +def efficientnet_v2_l( + *, weights: Optional[EfficientNet_V2_L_Weights] = None, progress: bool = True, **kwargs: Any +) -> EfficientNet: + """ + Constructs an EfficientNetV2-L architecture from + `EfficientNetV2: Smaller Models and Faster Training `_. + + Args: + weights (:class:`~torchvision.models.EfficientNet_V2_L_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.EfficientNet_V2_L_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.efficientnet.EfficientNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + .. autoclass:: torchvision.models.EfficientNet_V2_L_Weights + :members: + """ + weights = EfficientNet_V2_L_Weights.verify(weights) + + inverted_residual_setting, last_channel = _efficientnet_conf("efficientnet_v2_l") + return _efficientnet( + inverted_residual_setting, + kwargs.pop("dropout", 0.4), + last_channel, + weights, + progress, + norm_layer=partial(nn.BatchNorm2d, eps=1e-03), + **kwargs, + ) diff --git a/lib/python3.12/site-packages/torchvision/models/feature_extraction.py b/lib/python3.12/site-packages/torchvision/models/feature_extraction.py new file mode 100644 index 0000000000000000000000000000000000000000..a20387a6b8936272f1eadc55e76451244dba83b5 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/feature_extraction.py @@ -0,0 +1,607 @@ +import copy +import inspect +import math +import re +import warnings +from collections import OrderedDict +from copy import deepcopy +from itertools import chain +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import torch +import torchvision +from torch import fx, nn +from torch.fx.graph_module import _CodeOnlyModule, _copy_attr, _USER_PRESERVED_ATTRIBUTES_KEY + + +__all__ = ["create_feature_extractor", "get_graph_node_names"] + + +class LeafModuleAwareTracer(fx.Tracer): + """ + An fx.Tracer that allows the user to specify a set of leaf modules, i.e. + modules that are not to be traced through. The resulting graph ends up + having single nodes referencing calls to the leaf modules' forward methods. + """ + + def __init__(self, *args, **kwargs): + self.leaf_modules = {} + if "leaf_modules" in kwargs: + leaf_modules = kwargs.pop("leaf_modules") + self.leaf_modules = leaf_modules + super().__init__(*args, **kwargs) + + def is_leaf_module(self, m: nn.Module, module_qualname: str) -> bool: + if isinstance(m, tuple(self.leaf_modules)): + return True + return super().is_leaf_module(m, module_qualname) + + +class NodePathTracer(LeafModuleAwareTracer): + """ + NodePathTracer is an FX tracer that, for each operation, also records the + name of the Node from which the operation originated. A node name here is + a `.` separated path walking the hierarchy from top level module down to + leaf operation or leaf module. The name of the top level module is not + included as part of the node name. For example, if we trace a module whose + forward method applies a ReLU module, the name for that node will simply + be 'relu'. + + Some notes on the specifics: + - Nodes are recorded to `self.node_to_qualname` which is a dictionary + mapping a given Node object to its node name. + - Nodes are recorded in the order which they are executed during + tracing. + - When a duplicate node name is encountered, a suffix of the form + _{int} is added. The counter starts from 1. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Track the qualified name of the Node being traced + self.current_module_qualname = "" + # A map from FX Node to the qualified name\# + # NOTE: This is loosely like the "qualified name" mentioned in the + # torch.fx docs https://pytorch.org/docs/stable/fx.html but adapted + # for the purposes of the torchvision feature extractor + self.node_to_qualname = OrderedDict() + + def call_module(self, m: torch.nn.Module, forward: Callable, args, kwargs): + """ + Override of `fx.Tracer.call_module` + This override: + 1) Stores away the qualified name of the caller for restoration later + 2) Adds the qualified name of the caller to + `current_module_qualname` for retrieval by `create_proxy` + 3) Once a leaf module is reached, calls `create_proxy` + 4) Restores the caller's qualified name into current_module_qualname + """ + old_qualname = self.current_module_qualname + try: + module_qualname = self.path_of_module(m) + self.current_module_qualname = module_qualname + if not self.is_leaf_module(m, module_qualname): + out = forward(*args, **kwargs) + return out + return self.create_proxy("call_module", module_qualname, args, kwargs) + finally: + self.current_module_qualname = old_qualname + + def create_proxy( + self, kind: str, target: fx.node.Target, args, kwargs, name=None, type_expr=None, *_ + ) -> fx.proxy.Proxy: + """ + Override of `Tracer.create_proxy`. This override intercepts the recording + of every operation and stores away the current traced module's qualified + name in `node_to_qualname` + """ + proxy = super().create_proxy(kind, target, args, kwargs, name, type_expr) + self.node_to_qualname[proxy.node] = self._get_node_qualname(self.current_module_qualname, proxy.node) + return proxy + + def _get_node_qualname(self, module_qualname: str, node: fx.node.Node) -> str: + node_qualname = module_qualname + + if node.op != "call_module": + # In this case module_qualname from torch.fx doesn't go all the + # way to the leaf function/op, so we need to append it + if len(node_qualname) > 0: + # Only append '.' if we are deeper than the top level module + node_qualname += "." + node_qualname += str(node) + + # Now we need to add an _{index} postfix on any repeated node names + # For modules we do this from scratch + # But for anything else, torch.fx already has a globally scoped + # _{index} postfix. But we want it locally (relative to direct parent) + # scoped. So first we need to undo the torch.fx postfix + if re.match(r".+_[0-9]+$", node_qualname) is not None: + node_qualname = node_qualname.rsplit("_", 1)[0] + + # ... and now we add on our own postfix + for existing_qualname in reversed(self.node_to_qualname.values()): + # Check to see if existing_qualname is of the form + # {node_qualname} or {node_qualname}_{int} + if re.match(rf"{node_qualname}(_[0-9]+)?$", existing_qualname) is not None: + postfix = existing_qualname.replace(node_qualname, "") + if len(postfix): + # existing_qualname is of the form {node_qualname}_{int} + next_index = int(postfix[1:]) + 1 + else: + # existing_qualname is of the form {node_qualname} + next_index = 1 + node_qualname += f"_{next_index}" + break + + return node_qualname + + +def _is_subseq(x, y): + """Check if y is a subsequence of x + https://stackoverflow.com/a/24017747/4391249 + """ + iter_x = iter(x) + return all(any(x_item == y_item for x_item in iter_x) for y_item in y) + + +def _warn_graph_differences(train_tracer: NodePathTracer, eval_tracer: NodePathTracer): + """ + Utility function for warning the user if there are differences between + the train graph nodes and the eval graph nodes. + """ + train_nodes = list(train_tracer.node_to_qualname.values()) + eval_nodes = list(eval_tracer.node_to_qualname.values()) + + if len(train_nodes) == len(eval_nodes) and all(t == e for t, e in zip(train_nodes, eval_nodes)): + return + + suggestion_msg = ( + "When choosing nodes for feature extraction, you may need to specify " + "output nodes for train and eval mode separately." + ) + + if _is_subseq(train_nodes, eval_nodes): + msg = ( + "NOTE: The nodes obtained by tracing the model in eval mode " + "are a subsequence of those obtained in train mode. " + ) + elif _is_subseq(eval_nodes, train_nodes): + msg = ( + "NOTE: The nodes obtained by tracing the model in train mode " + "are a subsequence of those obtained in eval mode. " + ) + else: + msg = "The nodes obtained by tracing the model in train mode are different to those obtained in eval mode. " + warnings.warn(msg + suggestion_msg) + + +def _get_leaf_modules_for_ops() -> List[type]: + members = inspect.getmembers(torchvision.ops) + result = [] + for _, obj in members: + if inspect.isclass(obj) and issubclass(obj, torch.nn.Module): + result.append(obj) + return result + + +def _set_default_tracer_kwargs(original_tr_kwargs: Optional[Dict[str, Any]]) -> Dict[str, Any]: + default_autowrap_modules = (math, torchvision.ops) + default_leaf_modules = _get_leaf_modules_for_ops() + result_tracer_kwargs = {} if original_tr_kwargs is None else original_tr_kwargs + result_tracer_kwargs["autowrap_modules"] = ( + tuple(set(result_tracer_kwargs["autowrap_modules"] + default_autowrap_modules)) + if "autowrap_modules" in result_tracer_kwargs + else default_autowrap_modules + ) + result_tracer_kwargs["leaf_modules"] = ( + list(set(result_tracer_kwargs["leaf_modules"] + default_leaf_modules)) + if "leaf_modules" in result_tracer_kwargs + else default_leaf_modules + ) + return result_tracer_kwargs + + +def get_graph_node_names( + model: nn.Module, + tracer_kwargs: Optional[Dict[str, Any]] = None, + suppress_diff_warning: bool = False, + concrete_args: Optional[Dict[str, Any]] = None, +) -> Tuple[List[str], List[str]]: + """ + Dev utility to return node names in order of execution. See note on node + names under :func:`create_feature_extractor`. Useful for seeing which node + names are available for feature extraction. There are two reasons that + node names can't easily be read directly from the code for a model: + + 1. Not all submodules are traced through. Modules from ``torch.nn`` all + fall within this category. + 2. Nodes representing the repeated application of the same operation + or leaf module get a ``_{counter}`` postfix. + + The model is traced twice: once in train mode, and once in eval mode. Both + sets of node names are returned. + + For more details on the node naming conventions used here, please see the + :ref:`relevant subheading ` in the + `documentation `_. + + Args: + model (nn.Module): model for which we'd like to print node names + tracer_kwargs (dict, optional): a dictionary of keyword arguments for + ``NodePathTracer`` (they are eventually passed onto + `torch.fx.Tracer `_). + By default, it will be set to wrap and make leaf nodes all torchvision ops: + {"autowrap_modules": (math, torchvision.ops,),"leaf_modules": _get_leaf_modules_for_ops(),} + WARNING: In case the user provides tracer_kwargs, above default arguments will be appended to the user + provided dictionary. + suppress_diff_warning (bool, optional): whether to suppress a warning + when there are discrepancies between the train and eval version of + the graph. Defaults to False. + concrete_args (Optional[Dict[str, any]]): Concrete arguments that should + not be treated as Proxies. According to the `Pytorch docs + `_, + this parameter's API may not be guaranteed. + + Returns: + tuple(list, list): a list of node names from tracing the model in + train mode, and another from tracing the model in eval mode. + + Examples:: + + >>> model = torchvision.models.resnet18() + >>> train_nodes, eval_nodes = get_graph_node_names(model) + """ + tracer_kwargs = _set_default_tracer_kwargs(tracer_kwargs) + is_training = model.training + train_tracer = NodePathTracer(**tracer_kwargs) + train_tracer.trace(model.train(), concrete_args=concrete_args) + eval_tracer = NodePathTracer(**tracer_kwargs) + eval_tracer.trace(model.eval(), concrete_args=concrete_args) + train_nodes = list(train_tracer.node_to_qualname.values()) + eval_nodes = list(eval_tracer.node_to_qualname.values()) + if not suppress_diff_warning: + _warn_graph_differences(train_tracer, eval_tracer) + # Restore training state + model.train(is_training) + return train_nodes, eval_nodes + + +class DualGraphModule(fx.GraphModule): + """ + A derivative of `fx.GraphModule`. Differs in the following ways: + - Requires a train and eval version of the underlying graph + - Copies submodules according to the nodes of both train and eval graphs. + - Calling train(mode) switches between train graph and eval graph. + """ + + def __init__( + self, root: torch.nn.Module, train_graph: fx.Graph, eval_graph: fx.Graph, class_name: str = "GraphModule" + ): + """ + Args: + root (nn.Module): module from which the copied module hierarchy is + built + train_graph (fx.Graph): the graph that should be used in train mode + eval_graph (fx.Graph): the graph that should be used in eval mode + """ + super(fx.GraphModule, self).__init__() + + self.__class__.__name__ = class_name + + self.train_graph = train_graph + self.eval_graph = eval_graph + + # Copy all get_attr and call_module ops (indicated by BOTH train and + # eval graphs) + for node in chain(iter(train_graph.nodes), iter(eval_graph.nodes)): + if node.op in ["get_attr", "call_module"]: + if not isinstance(node.target, str): + raise TypeError(f"node.target should be of type str instead of {type(node.target)}") + _copy_attr(root, self, node.target) + + # train mode by default + self.train() + self.graph = train_graph + + # (borrowed from fx.GraphModule): + # Store the Tracer class responsible for creating a Graph separately as part of the + # GraphModule state, except when the Tracer is defined in a local namespace. + # Locally defined Tracers are not pickleable. This is needed because torch.package will + # serialize a GraphModule without retaining the Graph, and needs to use the correct Tracer + # to re-create the Graph during deserialization. + if self.eval_graph._tracer_cls != self.train_graph._tracer_cls: + raise TypeError( + f"Train mode and eval mode should use the same tracer class. Instead got {self.eval_graph._tracer_cls} for eval vs {self.train_graph._tracer_cls} for train" + ) + self._tracer_cls = None + if self.graph._tracer_cls and "" not in self.graph._tracer_cls.__qualname__: + self._tracer_cls = self.graph._tracer_cls + + def train(self, mode=True): + """ + Swap out the graph depending on the selected training mode. + NOTE this should be safe when calling model.eval() because that just + calls this with mode == False. + """ + # NOTE: Only set self.graph if the current graph is not the desired + # one. This saves us from recompiling the graph where not necessary. + if mode and not self.training: + self.graph = self.train_graph + elif not mode and self.training: + self.graph = self.eval_graph + return super().train(mode=mode) + + def _deepcopy_init(self): + # See __deepcopy__ below + return DualGraphModule.__init__ + + def __deepcopy__(self, memo): + # Same as the base class' __deepcopy__ from pytorch, with minor + # modification to account for train_graph and eval_graph + # https://github.com/pytorch/pytorch/blob/f684dbd0026f98f8fa291cab74dbc4d61ba30580/torch/fx/graph_module.py#L875 + # + # This is using a bunch of private stuff from torch, so if that breaks, + # we'll likely have to remove this, along with the associated + # non-regression test. + res = type(self).__new__(type(self)) + memo[id(self)] = res + fake_mod = _CodeOnlyModule(copy.deepcopy(self.__dict__, memo)) + self._deepcopy_init()(res, fake_mod, fake_mod.__dict__["train_graph"], fake_mod.__dict__["eval_graph"]) + + extra_preserved_attrs = [ + "_state_dict_hooks", + "_load_state_dict_pre_hooks", + "_load_state_dict_post_hooks", + "_replace_hook", + "_create_node_hooks", + "_erase_node_hooks", + ] + for attr in extra_preserved_attrs: + if attr in self.__dict__: + setattr(res, attr, copy.deepcopy(self.__dict__[attr], memo)) + res.meta = copy.deepcopy(getattr(self, "meta", {}), memo) + if _USER_PRESERVED_ATTRIBUTES_KEY in res.meta: + for attr_name, attr in res.meta[_USER_PRESERVED_ATTRIBUTES_KEY].items(): + setattr(res, attr_name, attr) + return res + + +def create_feature_extractor( + model: nn.Module, + return_nodes: Optional[Union[List[str], Dict[str, str]]] = None, + train_return_nodes: Optional[Union[List[str], Dict[str, str]]] = None, + eval_return_nodes: Optional[Union[List[str], Dict[str, str]]] = None, + tracer_kwargs: Optional[Dict[str, Any]] = None, + suppress_diff_warning: bool = False, + concrete_args: Optional[Dict[str, Any]] = None, +) -> fx.GraphModule: + """ + Creates a new graph module that returns intermediate nodes from a given + model as dictionary with user specified keys as strings, and the requested + outputs as values. This is achieved by re-writing the computation graph of + the model via FX to return the desired nodes as outputs. All unused nodes + are removed, together with their corresponding parameters. + + Desired output nodes must be specified as a ``.`` separated + path walking the module hierarchy from top level module down to leaf + operation or leaf module. For more details on the node naming conventions + used here, please see the :ref:`relevant subheading ` + in the `documentation `_. + + Not all models will be FX traceable, although with some massaging they can + be made to cooperate. Here's a (not exhaustive) list of tips: + + - If you don't need to trace through a particular, problematic + sub-module, turn it into a "leaf module" by passing a list of + ``leaf_modules`` as one of the ``tracer_kwargs`` (see example below). + It will not be traced through, but rather, the resulting graph will + hold a reference to that module's forward method. + - Likewise, you may turn functions into leaf functions by passing a + list of ``autowrap_functions`` as one of the ``tracer_kwargs`` (see + example below). + - Some inbuilt Python functions can be problematic. For instance, + ``int`` will raise an error during tracing. You may wrap them in your + own function and then pass that in ``autowrap_functions`` as one of + the ``tracer_kwargs``. + + For further information on FX see the + `torch.fx documentation `_. + + Args: + model (nn.Module): model on which we will extract the features + return_nodes (list or dict, optional): either a ``List`` or a ``Dict`` + containing the names (or partial names - see note above) + of the nodes for which the activations will be returned. If it is + a ``Dict``, the keys are the node names, and the values + are the user-specified keys for the graph module's returned + dictionary. If it is a ``List``, it is treated as a ``Dict`` mapping + node specification strings directly to output names. In the case + that ``train_return_nodes`` and ``eval_return_nodes`` are specified, + this should not be specified. + train_return_nodes (list or dict, optional): similar to + ``return_nodes``. This can be used if the return nodes + for train mode are different than those from eval mode. + If this is specified, ``eval_return_nodes`` must also be specified, + and ``return_nodes`` should not be specified. + eval_return_nodes (list or dict, optional): similar to + ``return_nodes``. This can be used if the return nodes + for train mode are different than those from eval mode. + If this is specified, ``train_return_nodes`` must also be specified, + and `return_nodes` should not be specified. + tracer_kwargs (dict, optional): a dictionary of keyword arguments for + ``NodePathTracer`` (which passes them onto it's parent class + `torch.fx.Tracer `_). + By default, it will be set to wrap and make leaf nodes all torchvision ops: + {"autowrap_modules": (math, torchvision.ops,),"leaf_modules": _get_leaf_modules_for_ops(),} + WARNING: In case the user provides tracer_kwargs, above default arguments will be appended to the user + provided dictionary. + suppress_diff_warning (bool, optional): whether to suppress a warning + when there are discrepancies between the train and eval version of + the graph. Defaults to False. + concrete_args (Optional[Dict[str, any]]): Concrete arguments that should + not be treated as Proxies. According to the `Pytorch docs + `_, + this parameter's API may not be guaranteed. + + Examples:: + + >>> # Feature extraction with resnet + >>> model = torchvision.models.resnet18() + >>> # extract layer1 and layer3, giving as names `feat1` and feat2` + >>> model = create_feature_extractor( + >>> model, {'layer1': 'feat1', 'layer3': 'feat2'}) + >>> out = model(torch.rand(1, 3, 224, 224)) + >>> print([(k, v.shape) for k, v in out.items()]) + >>> [('feat1', torch.Size([1, 64, 56, 56])), + >>> ('feat2', torch.Size([1, 256, 14, 14]))] + + >>> # Specifying leaf modules and leaf functions + >>> def leaf_function(x): + >>> # This would raise a TypeError if traced through + >>> return int(x) + >>> + >>> class LeafModule(torch.nn.Module): + >>> def forward(self, x): + >>> # This would raise a TypeError if traced through + >>> int(x.shape[0]) + >>> return torch.nn.functional.relu(x + 4) + >>> + >>> class MyModule(torch.nn.Module): + >>> def __init__(self): + >>> super().__init__() + >>> self.conv = torch.nn.Conv2d(3, 1, 3) + >>> self.leaf_module = LeafModule() + >>> + >>> def forward(self, x): + >>> leaf_function(x.shape[0]) + >>> x = self.conv(x) + >>> return self.leaf_module(x) + >>> + >>> model = create_feature_extractor( + >>> MyModule(), return_nodes=['leaf_module'], + >>> tracer_kwargs={'leaf_modules': [LeafModule], + >>> 'autowrap_functions': [leaf_function]}) + + """ + tracer_kwargs = _set_default_tracer_kwargs(tracer_kwargs) + is_training = model.training + + if all(arg is None for arg in [return_nodes, train_return_nodes, eval_return_nodes]): + + raise ValueError( + "Either `return_nodes` or `train_return_nodes` and `eval_return_nodes` together, should be specified" + ) + + if (train_return_nodes is None) ^ (eval_return_nodes is None): + raise ValueError( + "If any of `train_return_nodes` and `eval_return_nodes` are specified, then both should be specified" + ) + + if not ((return_nodes is None) ^ (train_return_nodes is None)): + raise ValueError("If `train_return_nodes` and `eval_return_nodes` are specified, then both should be specified") + + # Put *_return_nodes into Dict[str, str] format + def to_strdict(n) -> Dict[str, str]: + if isinstance(n, list): + return {str(i): str(i) for i in n} + return {str(k): str(v) for k, v in n.items()} + + if train_return_nodes is None: + return_nodes = to_strdict(return_nodes) + train_return_nodes = deepcopy(return_nodes) + eval_return_nodes = deepcopy(return_nodes) + else: + train_return_nodes = to_strdict(train_return_nodes) + eval_return_nodes = to_strdict(eval_return_nodes) + + # Repeat the tracing and graph rewriting for train and eval mode + tracers = {} + graphs = {} + mode_return_nodes: Dict[str, Dict[str, str]] = {"train": train_return_nodes, "eval": eval_return_nodes} + for mode in ["train", "eval"]: + if mode == "train": + model.train() + elif mode == "eval": + model.eval() + + # Instantiate our NodePathTracer and use that to trace the model + tracer = NodePathTracer(**tracer_kwargs) + graph = tracer.trace(model, concrete_args=concrete_args) + + name = model.__class__.__name__ if isinstance(model, nn.Module) else model.__name__ + graph_module = fx.GraphModule(tracer.root, graph, name) + + available_nodes = list(tracer.node_to_qualname.values()) + # FIXME We don't know if we should expect this to happen + if len(set(available_nodes)) != len(available_nodes): + raise ValueError( + "There are duplicate nodes! Please raise an issue https://github.com/pytorch/vision/issues" + ) + # Check that all outputs in return_nodes are present in the model + for query in mode_return_nodes[mode].keys(): + # To check if a query is available we need to check that at least + # one of the available names starts with it up to a . + if not any([re.match(rf"^{query}(\.|$)", n) is not None for n in available_nodes]): + raise ValueError( + f"node: '{query}' is not present in model. Hint: use " + "`get_graph_node_names` to make sure the " + "`return_nodes` you specified are present. It may even " + "be that you need to specify `train_return_nodes` and " + "`eval_return_nodes` separately." + ) + + # Remove existing output nodes (train mode) + orig_output_nodes = [] + for n in reversed(graph_module.graph.nodes): + if n.op == "output": + orig_output_nodes.append(n) + if not orig_output_nodes: + raise ValueError("No output nodes found in graph_module.graph.nodes") + + for n in orig_output_nodes: + graph_module.graph.erase_node(n) + + # Find nodes corresponding to return_nodes and make them into output_nodes + nodes = [n for n in graph_module.graph.nodes] + output_nodes = OrderedDict() + for n in reversed(nodes): + module_qualname = tracer.node_to_qualname.get(n) + if module_qualname is None: + # NOTE - Know cases where this happens: + # - Node representing creation of a tensor constant - probably + # not interesting as a return node + # - When packing outputs into a named tuple like in InceptionV3 + continue + for query in mode_return_nodes[mode]: + depth = query.count(".") + if ".".join(module_qualname.split(".")[: depth + 1]) == query: + output_nodes[mode_return_nodes[mode][query]] = n + mode_return_nodes[mode].pop(query) + break + output_nodes = OrderedDict(reversed(list(output_nodes.items()))) + + # And add them in the end of the graph + with graph_module.graph.inserting_after(nodes[-1]): + graph_module.graph.output(output_nodes) + + # Remove unused modules / parameters + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + + # Keep track of the tracer and graph, so we can choose the main one + tracers[mode] = tracer + graphs[mode] = graph + + # Warn user if there are any discrepancies between the graphs of the + # train and eval modes + if not suppress_diff_warning: + _warn_graph_differences(tracers["train"], tracers["eval"]) + + # Build the final graph module + graph_module = DualGraphModule(model, graphs["train"], graphs["eval"], class_name=name) + + # Restore original training mode + model.train(is_training) + graph_module.train(is_training) + + return graph_module diff --git a/lib/python3.12/site-packages/torchvision/models/googlenet.py b/lib/python3.12/site-packages/torchvision/models/googlenet.py new file mode 100644 index 0000000000000000000000000000000000000000..1d8ecbcd3b227c19d44f4d4fee532617c3e38ba3 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/googlenet.py @@ -0,0 +1,345 @@ +import warnings +from collections import namedtuple +from functools import partial +from typing import Any, Callable, List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor + +from ..transforms._presets import ImageClassification +from ..utils import _log_api_usage_once +from ._api import register_model, Weights, WeightsEnum +from ._meta import _IMAGENET_CATEGORIES +from ._utils import _ovewrite_named_param, handle_legacy_interface + + +__all__ = ["GoogLeNet", "GoogLeNetOutputs", "_GoogLeNetOutputs", "GoogLeNet_Weights", "googlenet"] + + +GoogLeNetOutputs = namedtuple("GoogLeNetOutputs", ["logits", "aux_logits2", "aux_logits1"]) +GoogLeNetOutputs.__annotations__ = {"logits": Tensor, "aux_logits2": Optional[Tensor], "aux_logits1": Optional[Tensor]} + +# Script annotations failed with _GoogleNetOutputs = namedtuple ... +# _GoogLeNetOutputs set here for backwards compat +_GoogLeNetOutputs = GoogLeNetOutputs + + +class GoogLeNet(nn.Module): + __constants__ = ["aux_logits", "transform_input"] + + def __init__( + self, + num_classes: int = 1000, + aux_logits: bool = True, + transform_input: bool = False, + init_weights: Optional[bool] = None, + blocks: Optional[List[Callable[..., nn.Module]]] = None, + dropout: float = 0.2, + dropout_aux: float = 0.7, + ) -> None: + super().__init__() + _log_api_usage_once(self) + if blocks is None: + blocks = [BasicConv2d, Inception, InceptionAux] + if init_weights is None: + warnings.warn( + "The default weight initialization of GoogleNet will be changed in future releases of " + "torchvision. If you wish to keep the old behavior (which leads to long initialization times" + " due to scipy/scipy#11299), please set init_weights=True.", + FutureWarning, + ) + init_weights = True + if len(blocks) != 3: + raise ValueError(f"blocks length should be 3 instead of {len(blocks)}") + conv_block = blocks[0] + inception_block = blocks[1] + inception_aux_block = blocks[2] + + self.aux_logits = aux_logits + self.transform_input = transform_input + + self.conv1 = conv_block(3, 64, kernel_size=7, stride=2, padding=3) + self.maxpool1 = nn.MaxPool2d(3, stride=2, ceil_mode=True) + self.conv2 = conv_block(64, 64, kernel_size=1) + self.conv3 = conv_block(64, 192, kernel_size=3, padding=1) + self.maxpool2 = nn.MaxPool2d(3, stride=2, ceil_mode=True) + + self.inception3a = inception_block(192, 64, 96, 128, 16, 32, 32) + self.inception3b = inception_block(256, 128, 128, 192, 32, 96, 64) + self.maxpool3 = nn.MaxPool2d(3, stride=2, ceil_mode=True) + + self.inception4a = inception_block(480, 192, 96, 208, 16, 48, 64) + self.inception4b = inception_block(512, 160, 112, 224, 24, 64, 64) + self.inception4c = inception_block(512, 128, 128, 256, 24, 64, 64) + self.inception4d = inception_block(512, 112, 144, 288, 32, 64, 64) + self.inception4e = inception_block(528, 256, 160, 320, 32, 128, 128) + self.maxpool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.inception5a = inception_block(832, 256, 160, 320, 32, 128, 128) + self.inception5b = inception_block(832, 384, 192, 384, 48, 128, 128) + + if aux_logits: + self.aux1 = inception_aux_block(512, num_classes, dropout=dropout_aux) + self.aux2 = inception_aux_block(528, num_classes, dropout=dropout_aux) + else: + self.aux1 = None # type: ignore[assignment] + self.aux2 = None # type: ignore[assignment] + + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.dropout = nn.Dropout(p=dropout) + self.fc = nn.Linear(1024, num_classes) + + if init_weights: + for m in self.modules(): + if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear): + torch.nn.init.trunc_normal_(m.weight, mean=0.0, std=0.01, a=-2, b=2) + elif isinstance(m, nn.BatchNorm2d): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + def _transform_input(self, x: Tensor) -> Tensor: + if self.transform_input: + x_ch0 = torch.unsqueeze(x[:, 0], 1) * (0.229 / 0.5) + (0.485 - 0.5) / 0.5 + x_ch1 = torch.unsqueeze(x[:, 1], 1) * (0.224 / 0.5) + (0.456 - 0.5) / 0.5 + x_ch2 = torch.unsqueeze(x[:, 2], 1) * (0.225 / 0.5) + (0.406 - 0.5) / 0.5 + x = torch.cat((x_ch0, x_ch1, x_ch2), 1) + return x + + def _forward(self, x: Tensor) -> Tuple[Tensor, Optional[Tensor], Optional[Tensor]]: + # N x 3 x 224 x 224 + x = self.conv1(x) + # N x 64 x 112 x 112 + x = self.maxpool1(x) + # N x 64 x 56 x 56 + x = self.conv2(x) + # N x 64 x 56 x 56 + x = self.conv3(x) + # N x 192 x 56 x 56 + x = self.maxpool2(x) + + # N x 192 x 28 x 28 + x = self.inception3a(x) + # N x 256 x 28 x 28 + x = self.inception3b(x) + # N x 480 x 28 x 28 + x = self.maxpool3(x) + # N x 480 x 14 x 14 + x = self.inception4a(x) + # N x 512 x 14 x 14 + aux1: Optional[Tensor] = None + if self.aux1 is not None: + if self.training: + aux1 = self.aux1(x) + + x = self.inception4b(x) + # N x 512 x 14 x 14 + x = self.inception4c(x) + # N x 512 x 14 x 14 + x = self.inception4d(x) + # N x 528 x 14 x 14 + aux2: Optional[Tensor] = None + if self.aux2 is not None: + if self.training: + aux2 = self.aux2(x) + + x = self.inception4e(x) + # N x 832 x 14 x 14 + x = self.maxpool4(x) + # N x 832 x 7 x 7 + x = self.inception5a(x) + # N x 832 x 7 x 7 + x = self.inception5b(x) + # N x 1024 x 7 x 7 + + x = self.avgpool(x) + # N x 1024 x 1 x 1 + x = torch.flatten(x, 1) + # N x 1024 + x = self.dropout(x) + x = self.fc(x) + # N x 1000 (num_classes) + return x, aux2, aux1 + + @torch.jit.unused + def eager_outputs(self, x: Tensor, aux2: Tensor, aux1: Optional[Tensor]) -> GoogLeNetOutputs: + if self.training and self.aux_logits: + return _GoogLeNetOutputs(x, aux2, aux1) + else: + return x # type: ignore[return-value] + + def forward(self, x: Tensor) -> GoogLeNetOutputs: + x = self._transform_input(x) + x, aux2, aux1 = self._forward(x) + aux_defined = self.training and self.aux_logits + if torch.jit.is_scripting(): + if not aux_defined: + warnings.warn("Scripted GoogleNet always returns GoogleNetOutputs Tuple") + return GoogLeNetOutputs(x, aux2, aux1) + else: + return self.eager_outputs(x, aux2, aux1) + + +class Inception(nn.Module): + def __init__( + self, + in_channels: int, + ch1x1: int, + ch3x3red: int, + ch3x3: int, + ch5x5red: int, + ch5x5: int, + pool_proj: int, + conv_block: Optional[Callable[..., nn.Module]] = None, + ) -> None: + super().__init__() + if conv_block is None: + conv_block = BasicConv2d + self.branch1 = conv_block(in_channels, ch1x1, kernel_size=1) + + self.branch2 = nn.Sequential( + conv_block(in_channels, ch3x3red, kernel_size=1), conv_block(ch3x3red, ch3x3, kernel_size=3, padding=1) + ) + + self.branch3 = nn.Sequential( + conv_block(in_channels, ch5x5red, kernel_size=1), + # Here, kernel_size=3 instead of kernel_size=5 is a known bug. + # Please see https://github.com/pytorch/vision/issues/906 for details. + conv_block(ch5x5red, ch5x5, kernel_size=3, padding=1), + ) + + self.branch4 = nn.Sequential( + nn.MaxPool2d(kernel_size=3, stride=1, padding=1, ceil_mode=True), + conv_block(in_channels, pool_proj, kernel_size=1), + ) + + def _forward(self, x: Tensor) -> List[Tensor]: + branch1 = self.branch1(x) + branch2 = self.branch2(x) + branch3 = self.branch3(x) + branch4 = self.branch4(x) + + outputs = [branch1, branch2, branch3, branch4] + return outputs + + def forward(self, x: Tensor) -> Tensor: + outputs = self._forward(x) + return torch.cat(outputs, 1) + + +class InceptionAux(nn.Module): + def __init__( + self, + in_channels: int, + num_classes: int, + conv_block: Optional[Callable[..., nn.Module]] = None, + dropout: float = 0.7, + ) -> None: + super().__init__() + if conv_block is None: + conv_block = BasicConv2d + self.conv = conv_block(in_channels, 128, kernel_size=1) + + self.fc1 = nn.Linear(2048, 1024) + self.fc2 = nn.Linear(1024, num_classes) + self.dropout = nn.Dropout(p=dropout) + + def forward(self, x: Tensor) -> Tensor: + # aux1: N x 512 x 14 x 14, aux2: N x 528 x 14 x 14 + x = F.adaptive_avg_pool2d(x, (4, 4)) + # aux1: N x 512 x 4 x 4, aux2: N x 528 x 4 x 4 + x = self.conv(x) + # N x 128 x 4 x 4 + x = torch.flatten(x, 1) + # N x 2048 + x = F.relu(self.fc1(x), inplace=True) + # N x 1024 + x = self.dropout(x) + # N x 1024 + x = self.fc2(x) + # N x 1000 (num_classes) + + return x + + +class BasicConv2d(nn.Module): + def __init__(self, in_channels: int, out_channels: int, **kwargs: Any) -> None: + super().__init__() + self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs) + self.bn = nn.BatchNorm2d(out_channels, eps=0.001) + + def forward(self, x: Tensor) -> Tensor: + x = self.conv(x) + x = self.bn(x) + return F.relu(x, inplace=True) + + +class GoogLeNet_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/googlenet-1378be20.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + "num_params": 6624904, + "min_size": (15, 15), + "categories": _IMAGENET_CATEGORIES, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#googlenet", + "_metrics": { + "ImageNet-1K": { + "acc@1": 69.778, + "acc@5": 89.530, + } + }, + "_ops": 1.498, + "_file_size": 49.731, + "_docs": """These weights are ported from the original paper.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", GoogLeNet_Weights.IMAGENET1K_V1)) +def googlenet(*, weights: Optional[GoogLeNet_Weights] = None, progress: bool = True, **kwargs: Any) -> GoogLeNet: + """GoogLeNet (Inception v1) model architecture from + `Going Deeper with Convolutions `_. + + Args: + weights (:class:`~torchvision.models.GoogLeNet_Weights`, optional): The + pretrained weights for the model. See + :class:`~torchvision.models.GoogLeNet_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.GoogLeNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + .. autoclass:: torchvision.models.GoogLeNet_Weights + :members: + """ + weights = GoogLeNet_Weights.verify(weights) + + original_aux_logits = kwargs.get("aux_logits", False) + if weights is not None: + if "transform_input" not in kwargs: + _ovewrite_named_param(kwargs, "transform_input", True) + _ovewrite_named_param(kwargs, "aux_logits", True) + _ovewrite_named_param(kwargs, "init_weights", False) + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + + model = GoogLeNet(**kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + if not original_aux_logits: + model.aux_logits = False + model.aux1 = None # type: ignore[assignment] + model.aux2 = None # type: ignore[assignment] + else: + warnings.warn( + "auxiliary heads in the pretrained googlenet model are NOT pretrained, so make sure to train them" + ) + + return model diff --git a/lib/python3.12/site-packages/torchvision/models/inception.py b/lib/python3.12/site-packages/torchvision/models/inception.py new file mode 100644 index 0000000000000000000000000000000000000000..447a7682d62c0e31570aaf1e32c2da55a0d697d1 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/inception.py @@ -0,0 +1,478 @@ +import warnings +from collections import namedtuple +from functools import partial +from typing import Any, Callable, List, Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn, Tensor + +from ..transforms._presets import ImageClassification +from ..utils import _log_api_usage_once +from ._api import register_model, Weights, WeightsEnum +from ._meta import _IMAGENET_CATEGORIES +from ._utils import _ovewrite_named_param, handle_legacy_interface + + +__all__ = ["Inception3", "InceptionOutputs", "_InceptionOutputs", "Inception_V3_Weights", "inception_v3"] + + +InceptionOutputs = namedtuple("InceptionOutputs", ["logits", "aux_logits"]) +InceptionOutputs.__annotations__ = {"logits": Tensor, "aux_logits": Optional[Tensor]} + +# Script annotations failed with _GoogleNetOutputs = namedtuple ... +# _InceptionOutputs set here for backwards compat +_InceptionOutputs = InceptionOutputs + + +class Inception3(nn.Module): + def __init__( + self, + num_classes: int = 1000, + aux_logits: bool = True, + transform_input: bool = False, + inception_blocks: Optional[List[Callable[..., nn.Module]]] = None, + init_weights: Optional[bool] = None, + dropout: float = 0.5, + ) -> None: + super().__init__() + _log_api_usage_once(self) + if inception_blocks is None: + inception_blocks = [BasicConv2d, InceptionA, InceptionB, InceptionC, InceptionD, InceptionE, InceptionAux] + if init_weights is None: + warnings.warn( + "The default weight initialization of inception_v3 will be changed in future releases of " + "torchvision. If you wish to keep the old behavior (which leads to long initialization times" + " due to scipy/scipy#11299), please set init_weights=True.", + FutureWarning, + ) + init_weights = True + if len(inception_blocks) != 7: + raise ValueError(f"length of inception_blocks should be 7 instead of {len(inception_blocks)}") + conv_block = inception_blocks[0] + inception_a = inception_blocks[1] + inception_b = inception_blocks[2] + inception_c = inception_blocks[3] + inception_d = inception_blocks[4] + inception_e = inception_blocks[5] + inception_aux = inception_blocks[6] + + self.aux_logits = aux_logits + self.transform_input = transform_input + self.Conv2d_1a_3x3 = conv_block(3, 32, kernel_size=3, stride=2) + self.Conv2d_2a_3x3 = conv_block(32, 32, kernel_size=3) + self.Conv2d_2b_3x3 = conv_block(32, 64, kernel_size=3, padding=1) + self.maxpool1 = nn.MaxPool2d(kernel_size=3, stride=2) + self.Conv2d_3b_1x1 = conv_block(64, 80, kernel_size=1) + self.Conv2d_4a_3x3 = conv_block(80, 192, kernel_size=3) + self.maxpool2 = nn.MaxPool2d(kernel_size=3, stride=2) + self.Mixed_5b = inception_a(192, pool_features=32) + self.Mixed_5c = inception_a(256, pool_features=64) + self.Mixed_5d = inception_a(288, pool_features=64) + self.Mixed_6a = inception_b(288) + self.Mixed_6b = inception_c(768, channels_7x7=128) + self.Mixed_6c = inception_c(768, channels_7x7=160) + self.Mixed_6d = inception_c(768, channels_7x7=160) + self.Mixed_6e = inception_c(768, channels_7x7=192) + self.AuxLogits: Optional[nn.Module] = None + if aux_logits: + self.AuxLogits = inception_aux(768, num_classes) + self.Mixed_7a = inception_d(768) + self.Mixed_7b = inception_e(1280) + self.Mixed_7c = inception_e(2048) + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.dropout = nn.Dropout(p=dropout) + self.fc = nn.Linear(2048, num_classes) + if init_weights: + for m in self.modules(): + if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear): + stddev = float(m.stddev) if hasattr(m, "stddev") else 0.1 # type: ignore + torch.nn.init.trunc_normal_(m.weight, mean=0.0, std=stddev, a=-2, b=2) + elif isinstance(m, nn.BatchNorm2d): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + def _transform_input(self, x: Tensor) -> Tensor: + if self.transform_input: + x_ch0 = torch.unsqueeze(x[:, 0], 1) * (0.229 / 0.5) + (0.485 - 0.5) / 0.5 + x_ch1 = torch.unsqueeze(x[:, 1], 1) * (0.224 / 0.5) + (0.456 - 0.5) / 0.5 + x_ch2 = torch.unsqueeze(x[:, 2], 1) * (0.225 / 0.5) + (0.406 - 0.5) / 0.5 + x = torch.cat((x_ch0, x_ch1, x_ch2), 1) + return x + + def _forward(self, x: Tensor) -> Tuple[Tensor, Optional[Tensor]]: + # N x 3 x 299 x 299 + x = self.Conv2d_1a_3x3(x) + # N x 32 x 149 x 149 + x = self.Conv2d_2a_3x3(x) + # N x 32 x 147 x 147 + x = self.Conv2d_2b_3x3(x) + # N x 64 x 147 x 147 + x = self.maxpool1(x) + # N x 64 x 73 x 73 + x = self.Conv2d_3b_1x1(x) + # N x 80 x 73 x 73 + x = self.Conv2d_4a_3x3(x) + # N x 192 x 71 x 71 + x = self.maxpool2(x) + # N x 192 x 35 x 35 + x = self.Mixed_5b(x) + # N x 256 x 35 x 35 + x = self.Mixed_5c(x) + # N x 288 x 35 x 35 + x = self.Mixed_5d(x) + # N x 288 x 35 x 35 + x = self.Mixed_6a(x) + # N x 768 x 17 x 17 + x = self.Mixed_6b(x) + # N x 768 x 17 x 17 + x = self.Mixed_6c(x) + # N x 768 x 17 x 17 + x = self.Mixed_6d(x) + # N x 768 x 17 x 17 + x = self.Mixed_6e(x) + # N x 768 x 17 x 17 + aux: Optional[Tensor] = None + if self.AuxLogits is not None: + if self.training: + aux = self.AuxLogits(x) + # N x 768 x 17 x 17 + x = self.Mixed_7a(x) + # N x 1280 x 8 x 8 + x = self.Mixed_7b(x) + # N x 2048 x 8 x 8 + x = self.Mixed_7c(x) + # N x 2048 x 8 x 8 + # Adaptive average pooling + x = self.avgpool(x) + # N x 2048 x 1 x 1 + x = self.dropout(x) + # N x 2048 x 1 x 1 + x = torch.flatten(x, 1) + # N x 2048 + x = self.fc(x) + # N x 1000 (num_classes) + return x, aux + + @torch.jit.unused + def eager_outputs(self, x: Tensor, aux: Optional[Tensor]) -> InceptionOutputs: + if self.training and self.aux_logits: + return InceptionOutputs(x, aux) + else: + return x # type: ignore[return-value] + + def forward(self, x: Tensor) -> InceptionOutputs: + x = self._transform_input(x) + x, aux = self._forward(x) + aux_defined = self.training and self.aux_logits + if torch.jit.is_scripting(): + if not aux_defined: + warnings.warn("Scripted Inception3 always returns Inception3 Tuple") + return InceptionOutputs(x, aux) + else: + return self.eager_outputs(x, aux) + + +class InceptionA(nn.Module): + def __init__( + self, in_channels: int, pool_features: int, conv_block: Optional[Callable[..., nn.Module]] = None + ) -> None: + super().__init__() + if conv_block is None: + conv_block = BasicConv2d + self.branch1x1 = conv_block(in_channels, 64, kernel_size=1) + + self.branch5x5_1 = conv_block(in_channels, 48, kernel_size=1) + self.branch5x5_2 = conv_block(48, 64, kernel_size=5, padding=2) + + self.branch3x3dbl_1 = conv_block(in_channels, 64, kernel_size=1) + self.branch3x3dbl_2 = conv_block(64, 96, kernel_size=3, padding=1) + self.branch3x3dbl_3 = conv_block(96, 96, kernel_size=3, padding=1) + + self.branch_pool = conv_block(in_channels, pool_features, kernel_size=1) + + def _forward(self, x: Tensor) -> List[Tensor]: + branch1x1 = self.branch1x1(x) + + branch5x5 = self.branch5x5_1(x) + branch5x5 = self.branch5x5_2(branch5x5) + + branch3x3dbl = self.branch3x3dbl_1(x) + branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl) + branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl) + + branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1) + branch_pool = self.branch_pool(branch_pool) + + outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool] + return outputs + + def forward(self, x: Tensor) -> Tensor: + outputs = self._forward(x) + return torch.cat(outputs, 1) + + +class InceptionB(nn.Module): + def __init__(self, in_channels: int, conv_block: Optional[Callable[..., nn.Module]] = None) -> None: + super().__init__() + if conv_block is None: + conv_block = BasicConv2d + self.branch3x3 = conv_block(in_channels, 384, kernel_size=3, stride=2) + + self.branch3x3dbl_1 = conv_block(in_channels, 64, kernel_size=1) + self.branch3x3dbl_2 = conv_block(64, 96, kernel_size=3, padding=1) + self.branch3x3dbl_3 = conv_block(96, 96, kernel_size=3, stride=2) + + def _forward(self, x: Tensor) -> List[Tensor]: + branch3x3 = self.branch3x3(x) + + branch3x3dbl = self.branch3x3dbl_1(x) + branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl) + branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl) + + branch_pool = F.max_pool2d(x, kernel_size=3, stride=2) + + outputs = [branch3x3, branch3x3dbl, branch_pool] + return outputs + + def forward(self, x: Tensor) -> Tensor: + outputs = self._forward(x) + return torch.cat(outputs, 1) + + +class InceptionC(nn.Module): + def __init__( + self, in_channels: int, channels_7x7: int, conv_block: Optional[Callable[..., nn.Module]] = None + ) -> None: + super().__init__() + if conv_block is None: + conv_block = BasicConv2d + self.branch1x1 = conv_block(in_channels, 192, kernel_size=1) + + c7 = channels_7x7 + self.branch7x7_1 = conv_block(in_channels, c7, kernel_size=1) + self.branch7x7_2 = conv_block(c7, c7, kernel_size=(1, 7), padding=(0, 3)) + self.branch7x7_3 = conv_block(c7, 192, kernel_size=(7, 1), padding=(3, 0)) + + self.branch7x7dbl_1 = conv_block(in_channels, c7, kernel_size=1) + self.branch7x7dbl_2 = conv_block(c7, c7, kernel_size=(7, 1), padding=(3, 0)) + self.branch7x7dbl_3 = conv_block(c7, c7, kernel_size=(1, 7), padding=(0, 3)) + self.branch7x7dbl_4 = conv_block(c7, c7, kernel_size=(7, 1), padding=(3, 0)) + self.branch7x7dbl_5 = conv_block(c7, 192, kernel_size=(1, 7), padding=(0, 3)) + + self.branch_pool = conv_block(in_channels, 192, kernel_size=1) + + def _forward(self, x: Tensor) -> List[Tensor]: + branch1x1 = self.branch1x1(x) + + branch7x7 = self.branch7x7_1(x) + branch7x7 = self.branch7x7_2(branch7x7) + branch7x7 = self.branch7x7_3(branch7x7) + + branch7x7dbl = self.branch7x7dbl_1(x) + branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl) + branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl) + branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl) + branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl) + + branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1) + branch_pool = self.branch_pool(branch_pool) + + outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool] + return outputs + + def forward(self, x: Tensor) -> Tensor: + outputs = self._forward(x) + return torch.cat(outputs, 1) + + +class InceptionD(nn.Module): + def __init__(self, in_channels: int, conv_block: Optional[Callable[..., nn.Module]] = None) -> None: + super().__init__() + if conv_block is None: + conv_block = BasicConv2d + self.branch3x3_1 = conv_block(in_channels, 192, kernel_size=1) + self.branch3x3_2 = conv_block(192, 320, kernel_size=3, stride=2) + + self.branch7x7x3_1 = conv_block(in_channels, 192, kernel_size=1) + self.branch7x7x3_2 = conv_block(192, 192, kernel_size=(1, 7), padding=(0, 3)) + self.branch7x7x3_3 = conv_block(192, 192, kernel_size=(7, 1), padding=(3, 0)) + self.branch7x7x3_4 = conv_block(192, 192, kernel_size=3, stride=2) + + def _forward(self, x: Tensor) -> List[Tensor]: + branch3x3 = self.branch3x3_1(x) + branch3x3 = self.branch3x3_2(branch3x3) + + branch7x7x3 = self.branch7x7x3_1(x) + branch7x7x3 = self.branch7x7x3_2(branch7x7x3) + branch7x7x3 = self.branch7x7x3_3(branch7x7x3) + branch7x7x3 = self.branch7x7x3_4(branch7x7x3) + + branch_pool = F.max_pool2d(x, kernel_size=3, stride=2) + outputs = [branch3x3, branch7x7x3, branch_pool] + return outputs + + def forward(self, x: Tensor) -> Tensor: + outputs = self._forward(x) + return torch.cat(outputs, 1) + + +class InceptionE(nn.Module): + def __init__(self, in_channels: int, conv_block: Optional[Callable[..., nn.Module]] = None) -> None: + super().__init__() + if conv_block is None: + conv_block = BasicConv2d + self.branch1x1 = conv_block(in_channels, 320, kernel_size=1) + + self.branch3x3_1 = conv_block(in_channels, 384, kernel_size=1) + self.branch3x3_2a = conv_block(384, 384, kernel_size=(1, 3), padding=(0, 1)) + self.branch3x3_2b = conv_block(384, 384, kernel_size=(3, 1), padding=(1, 0)) + + self.branch3x3dbl_1 = conv_block(in_channels, 448, kernel_size=1) + self.branch3x3dbl_2 = conv_block(448, 384, kernel_size=3, padding=1) + self.branch3x3dbl_3a = conv_block(384, 384, kernel_size=(1, 3), padding=(0, 1)) + self.branch3x3dbl_3b = conv_block(384, 384, kernel_size=(3, 1), padding=(1, 0)) + + self.branch_pool = conv_block(in_channels, 192, kernel_size=1) + + def _forward(self, x: Tensor) -> List[Tensor]: + branch1x1 = self.branch1x1(x) + + branch3x3 = self.branch3x3_1(x) + branch3x3 = [ + self.branch3x3_2a(branch3x3), + self.branch3x3_2b(branch3x3), + ] + branch3x3 = torch.cat(branch3x3, 1) + + branch3x3dbl = self.branch3x3dbl_1(x) + branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl) + branch3x3dbl = [ + self.branch3x3dbl_3a(branch3x3dbl), + self.branch3x3dbl_3b(branch3x3dbl), + ] + branch3x3dbl = torch.cat(branch3x3dbl, 1) + + branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1) + branch_pool = self.branch_pool(branch_pool) + + outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool] + return outputs + + def forward(self, x: Tensor) -> Tensor: + outputs = self._forward(x) + return torch.cat(outputs, 1) + + +class InceptionAux(nn.Module): + def __init__( + self, in_channels: int, num_classes: int, conv_block: Optional[Callable[..., nn.Module]] = None + ) -> None: + super().__init__() + if conv_block is None: + conv_block = BasicConv2d + self.conv0 = conv_block(in_channels, 128, kernel_size=1) + self.conv1 = conv_block(128, 768, kernel_size=5) + self.conv1.stddev = 0.01 # type: ignore[assignment] + self.fc = nn.Linear(768, num_classes) + self.fc.stddev = 0.001 # type: ignore[assignment] + + def forward(self, x: Tensor) -> Tensor: + # N x 768 x 17 x 17 + x = F.avg_pool2d(x, kernel_size=5, stride=3) + # N x 768 x 5 x 5 + x = self.conv0(x) + # N x 128 x 5 x 5 + x = self.conv1(x) + # N x 768 x 1 x 1 + # Adaptive average pooling + x = F.adaptive_avg_pool2d(x, (1, 1)) + # N x 768 x 1 x 1 + x = torch.flatten(x, 1) + # N x 768 + x = self.fc(x) + # N x 1000 + return x + + +class BasicConv2d(nn.Module): + def __init__(self, in_channels: int, out_channels: int, **kwargs: Any) -> None: + super().__init__() + self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs) + self.bn = nn.BatchNorm2d(out_channels, eps=0.001) + + def forward(self, x: Tensor) -> Tensor: + x = self.conv(x) + x = self.bn(x) + return F.relu(x, inplace=True) + + +class Inception_V3_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/inception_v3_google-0cc3c7bd.pth", + transforms=partial(ImageClassification, crop_size=299, resize_size=342), + meta={ + "num_params": 27161264, + "min_size": (75, 75), + "categories": _IMAGENET_CATEGORIES, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#inception-v3", + "_metrics": { + "ImageNet-1K": { + "acc@1": 77.294, + "acc@5": 93.450, + } + }, + "_ops": 5.713, + "_file_size": 103.903, + "_docs": """These weights are ported from the original paper.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", Inception_V3_Weights.IMAGENET1K_V1)) +def inception_v3(*, weights: Optional[Inception_V3_Weights] = None, progress: bool = True, **kwargs: Any) -> Inception3: + """ + Inception v3 model architecture from + `Rethinking the Inception Architecture for Computer Vision `_. + + .. note:: + **Important**: In contrast to the other models the inception_v3 expects tensors with a size of + N x 3 x 299 x 299, so ensure your images are sized accordingly. + + Args: + weights (:class:`~torchvision.models.Inception_V3_Weights`, optional): The + pretrained weights for the model. See + :class:`~torchvision.models.Inception_V3_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.Inception3`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.Inception_V3_Weights + :members: + """ + weights = Inception_V3_Weights.verify(weights) + + original_aux_logits = kwargs.get("aux_logits", True) + if weights is not None: + if "transform_input" not in kwargs: + _ovewrite_named_param(kwargs, "transform_input", True) + _ovewrite_named_param(kwargs, "aux_logits", True) + _ovewrite_named_param(kwargs, "init_weights", False) + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + + model = Inception3(**kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + if not original_aux_logits: + model.aux_logits = False + model.AuxLogits = None + + return model diff --git a/lib/python3.12/site-packages/torchvision/models/maxvit.py b/lib/python3.12/site-packages/torchvision/models/maxvit.py new file mode 100644 index 0000000000000000000000000000000000000000..66f49772218d7ee2fe5c76d914be248ac6ae2c9d --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/maxvit.py @@ -0,0 +1,833 @@ +import math +from collections import OrderedDict +from functools import partial +from typing import Any, Callable, List, Optional, Sequence, Tuple + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn, Tensor +from torchvision.models._api import register_model, Weights, WeightsEnum +from torchvision.models._meta import _IMAGENET_CATEGORIES +from torchvision.models._utils import _ovewrite_named_param, handle_legacy_interface +from torchvision.ops.misc import Conv2dNormActivation, SqueezeExcitation +from torchvision.ops.stochastic_depth import StochasticDepth +from torchvision.transforms._presets import ImageClassification, InterpolationMode +from torchvision.utils import _log_api_usage_once + +__all__ = [ + "MaxVit", + "MaxVit_T_Weights", + "maxvit_t", +] + + +def _get_conv_output_shape(input_size: Tuple[int, int], kernel_size: int, stride: int, padding: int) -> Tuple[int, int]: + return ( + (input_size[0] - kernel_size + 2 * padding) // stride + 1, + (input_size[1] - kernel_size + 2 * padding) // stride + 1, + ) + + +def _make_block_input_shapes(input_size: Tuple[int, int], n_blocks: int) -> List[Tuple[int, int]]: + """Util function to check that the input size is correct for a MaxVit configuration.""" + shapes = [] + block_input_shape = _get_conv_output_shape(input_size, 3, 2, 1) + for _ in range(n_blocks): + block_input_shape = _get_conv_output_shape(block_input_shape, 3, 2, 1) + shapes.append(block_input_shape) + return shapes + + +def _get_relative_position_index(height: int, width: int) -> torch.Tensor: + coords = torch.stack(torch.meshgrid([torch.arange(height), torch.arange(width)], indexing="ij")) + coords_flat = torch.flatten(coords, 1) + relative_coords = coords_flat[:, :, None] - coords_flat[:, None, :] + relative_coords = relative_coords.permute(1, 2, 0).contiguous() + relative_coords[:, :, 0] += height - 1 + relative_coords[:, :, 1] += width - 1 + relative_coords[:, :, 0] *= 2 * width - 1 + return relative_coords.sum(-1) + + +class MBConv(nn.Module): + """MBConv: Mobile Inverted Residual Bottleneck. + + Args: + in_channels (int): Number of input channels. + out_channels (int): Number of output channels. + expansion_ratio (float): Expansion ratio in the bottleneck. + squeeze_ratio (float): Squeeze ratio in the SE Layer. + stride (int): Stride of the depthwise convolution. + activation_layer (Callable[..., nn.Module]): Activation function. + norm_layer (Callable[..., nn.Module]): Normalization function. + p_stochastic_dropout (float): Probability of stochastic depth. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + expansion_ratio: float, + squeeze_ratio: float, + stride: int, + activation_layer: Callable[..., nn.Module], + norm_layer: Callable[..., nn.Module], + p_stochastic_dropout: float = 0.0, + ) -> None: + super().__init__() + + proj: Sequence[nn.Module] + self.proj: nn.Module + + should_proj = stride != 1 or in_channels != out_channels + if should_proj: + proj = [nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, bias=True)] + if stride == 2: + proj = [nn.AvgPool2d(kernel_size=3, stride=stride, padding=1)] + proj # type: ignore + self.proj = nn.Sequential(*proj) + else: + self.proj = nn.Identity() # type: ignore + + mid_channels = int(out_channels * expansion_ratio) + sqz_channels = int(out_channels * squeeze_ratio) + + if p_stochastic_dropout: + self.stochastic_depth = StochasticDepth(p_stochastic_dropout, mode="row") # type: ignore + else: + self.stochastic_depth = nn.Identity() # type: ignore + + _layers = OrderedDict() + _layers["pre_norm"] = norm_layer(in_channels) + _layers["conv_a"] = Conv2dNormActivation( + in_channels, + mid_channels, + kernel_size=1, + stride=1, + padding=0, + activation_layer=activation_layer, + norm_layer=norm_layer, + inplace=None, + ) + _layers["conv_b"] = Conv2dNormActivation( + mid_channels, + mid_channels, + kernel_size=3, + stride=stride, + padding=1, + activation_layer=activation_layer, + norm_layer=norm_layer, + groups=mid_channels, + inplace=None, + ) + _layers["squeeze_excitation"] = SqueezeExcitation(mid_channels, sqz_channels, activation=nn.SiLU) + _layers["conv_c"] = nn.Conv2d(in_channels=mid_channels, out_channels=out_channels, kernel_size=1, bias=True) + + self.layers = nn.Sequential(_layers) + + def forward(self, x: Tensor) -> Tensor: + """ + Args: + x (Tensor): Input tensor with expected layout of [B, C, H, W]. + Returns: + Tensor: Output tensor with expected layout of [B, C, H / stride, W / stride]. + """ + res = self.proj(x) + x = self.stochastic_depth(self.layers(x)) + return res + x + + +class RelativePositionalMultiHeadAttention(nn.Module): + """Relative Positional Multi-Head Attention. + + Args: + feat_dim (int): Number of input features. + head_dim (int): Number of features per head. + max_seq_len (int): Maximum sequence length. + """ + + def __init__( + self, + feat_dim: int, + head_dim: int, + max_seq_len: int, + ) -> None: + super().__init__() + + if feat_dim % head_dim != 0: + raise ValueError(f"feat_dim: {feat_dim} must be divisible by head_dim: {head_dim}") + + self.n_heads = feat_dim // head_dim + self.head_dim = head_dim + self.size = int(math.sqrt(max_seq_len)) + self.max_seq_len = max_seq_len + + self.to_qkv = nn.Linear(feat_dim, self.n_heads * self.head_dim * 3) + self.scale_factor = feat_dim**-0.5 + + self.merge = nn.Linear(self.head_dim * self.n_heads, feat_dim) + self.relative_position_bias_table = nn.parameter.Parameter( + torch.empty(((2 * self.size - 1) * (2 * self.size - 1), self.n_heads), dtype=torch.float32), + ) + + self.register_buffer("relative_position_index", _get_relative_position_index(self.size, self.size)) + # initialize with truncated normal the bias + torch.nn.init.trunc_normal_(self.relative_position_bias_table, std=0.02) + + def get_relative_positional_bias(self) -> torch.Tensor: + bias_index = self.relative_position_index.view(-1) # type: ignore + relative_bias = self.relative_position_bias_table[bias_index].view(self.max_seq_len, self.max_seq_len, -1) # type: ignore + relative_bias = relative_bias.permute(2, 0, 1).contiguous() + return relative_bias.unsqueeze(0) + + def forward(self, x: Tensor) -> Tensor: + """ + Args: + x (Tensor): Input tensor with expected layout of [B, G, P, D]. + Returns: + Tensor: Output tensor with expected layout of [B, G, P, D]. + """ + B, G, P, D = x.shape + H, DH = self.n_heads, self.head_dim + + qkv = self.to_qkv(x) + q, k, v = torch.chunk(qkv, 3, dim=-1) + + q = q.reshape(B, G, P, H, DH).permute(0, 1, 3, 2, 4) + k = k.reshape(B, G, P, H, DH).permute(0, 1, 3, 2, 4) + v = v.reshape(B, G, P, H, DH).permute(0, 1, 3, 2, 4) + + k = k * self.scale_factor + dot_prod = torch.einsum("B G H I D, B G H J D -> B G H I J", q, k) + pos_bias = self.get_relative_positional_bias() + + dot_prod = F.softmax(dot_prod + pos_bias, dim=-1) + + out = torch.einsum("B G H I J, B G H J D -> B G H I D", dot_prod, v) + out = out.permute(0, 1, 3, 2, 4).reshape(B, G, P, D) + + out = self.merge(out) + return out + + +class SwapAxes(nn.Module): + """Permute the axes of a tensor.""" + + def __init__(self, a: int, b: int) -> None: + super().__init__() + self.a = a + self.b = b + + def forward(self, x: torch.Tensor) -> torch.Tensor: + res = torch.swapaxes(x, self.a, self.b) + return res + + +class WindowPartition(nn.Module): + """ + Partition the input tensor into non-overlapping windows. + """ + + def __init__(self) -> None: + super().__init__() + + def forward(self, x: Tensor, p: int) -> Tensor: + """ + Args: + x (Tensor): Input tensor with expected layout of [B, C, H, W]. + p (int): Number of partitions. + Returns: + Tensor: Output tensor with expected layout of [B, H/P, W/P, P*P, C]. + """ + B, C, H, W = x.shape + P = p + # chunk up H and W dimensions + x = x.reshape(B, C, H // P, P, W // P, P) + x = x.permute(0, 2, 4, 3, 5, 1) + # colapse P * P dimension + x = x.reshape(B, (H // P) * (W // P), P * P, C) + return x + + +class WindowDepartition(nn.Module): + """ + Departition the input tensor of non-overlapping windows into a feature volume of layout [B, C, H, W]. + """ + + def __init__(self) -> None: + super().__init__() + + def forward(self, x: Tensor, p: int, h_partitions: int, w_partitions: int) -> Tensor: + """ + Args: + x (Tensor): Input tensor with expected layout of [B, (H/P * W/P), P*P, C]. + p (int): Number of partitions. + h_partitions (int): Number of vertical partitions. + w_partitions (int): Number of horizontal partitions. + Returns: + Tensor: Output tensor with expected layout of [B, C, H, W]. + """ + B, G, PP, C = x.shape + P = p + HP, WP = h_partitions, w_partitions + # split P * P dimension into 2 P tile dimensionsa + x = x.reshape(B, HP, WP, P, P, C) + # permute into B, C, HP, P, WP, P + x = x.permute(0, 5, 1, 3, 2, 4) + # reshape into B, C, H, W + x = x.reshape(B, C, HP * P, WP * P) + return x + + +class PartitionAttentionLayer(nn.Module): + """ + Layer for partitioning the input tensor into non-overlapping windows and applying attention to each window. + + Args: + in_channels (int): Number of input channels. + head_dim (int): Dimension of each attention head. + partition_size (int): Size of the partitions. + partition_type (str): Type of partitioning to use. Can be either "grid" or "window". + grid_size (Tuple[int, int]): Size of the grid to partition the input tensor into. + mlp_ratio (int): Ratio of the feature size expansion in the MLP layer. + activation_layer (Callable[..., nn.Module]): Activation function to use. + norm_layer (Callable[..., nn.Module]): Normalization function to use. + attention_dropout (float): Dropout probability for the attention layer. + mlp_dropout (float): Dropout probability for the MLP layer. + p_stochastic_dropout (float): Probability of dropping out a partition. + """ + + def __init__( + self, + in_channels: int, + head_dim: int, + # partitioning parameters + partition_size: int, + partition_type: str, + # grid size needs to be known at initialization time + # because we need to know hamy relative offsets there are in the grid + grid_size: Tuple[int, int], + mlp_ratio: int, + activation_layer: Callable[..., nn.Module], + norm_layer: Callable[..., nn.Module], + attention_dropout: float, + mlp_dropout: float, + p_stochastic_dropout: float, + ) -> None: + super().__init__() + + self.n_heads = in_channels // head_dim + self.head_dim = head_dim + self.n_partitions = grid_size[0] // partition_size + self.partition_type = partition_type + self.grid_size = grid_size + + if partition_type not in ["grid", "window"]: + raise ValueError("partition_type must be either 'grid' or 'window'") + + if partition_type == "window": + self.p, self.g = partition_size, self.n_partitions + else: + self.p, self.g = self.n_partitions, partition_size + + self.partition_op = WindowPartition() + self.departition_op = WindowDepartition() + self.partition_swap = SwapAxes(-2, -3) if partition_type == "grid" else nn.Identity() + self.departition_swap = SwapAxes(-2, -3) if partition_type == "grid" else nn.Identity() + + self.attn_layer = nn.Sequential( + norm_layer(in_channels), + # it's always going to be partition_size ** 2 because + # of the axis swap in the case of grid partitioning + RelativePositionalMultiHeadAttention(in_channels, head_dim, partition_size**2), + nn.Dropout(attention_dropout), + ) + + # pre-normalization similar to transformer layers + self.mlp_layer = nn.Sequential( + nn.LayerNorm(in_channels), + nn.Linear(in_channels, in_channels * mlp_ratio), + activation_layer(), + nn.Linear(in_channels * mlp_ratio, in_channels), + nn.Dropout(mlp_dropout), + ) + + # layer scale factors + self.stochastic_dropout = StochasticDepth(p_stochastic_dropout, mode="row") + + def forward(self, x: Tensor) -> Tensor: + """ + Args: + x (Tensor): Input tensor with expected layout of [B, C, H, W]. + Returns: + Tensor: Output tensor with expected layout of [B, C, H, W]. + """ + + # Undefined behavior if H or W are not divisible by p + # https://github.com/google-research/maxvit/blob/da76cf0d8a6ec668cc31b399c4126186da7da944/maxvit/models/maxvit.py#L766 + gh, gw = self.grid_size[0] // self.p, self.grid_size[1] // self.p + torch._assert( + self.grid_size[0] % self.p == 0 and self.grid_size[1] % self.p == 0, + "Grid size must be divisible by partition size. Got grid size of {} and partition size of {}".format( + self.grid_size, self.p + ), + ) + + x = self.partition_op(x, self.p) + x = self.partition_swap(x) + x = x + self.stochastic_dropout(self.attn_layer(x)) + x = x + self.stochastic_dropout(self.mlp_layer(x)) + x = self.departition_swap(x) + x = self.departition_op(x, self.p, gh, gw) + + return x + + +class MaxVitLayer(nn.Module): + """ + MaxVit layer consisting of a MBConv layer followed by a PartitionAttentionLayer with `window` and a PartitionAttentionLayer with `grid`. + + Args: + in_channels (int): Number of input channels. + out_channels (int): Number of output channels. + expansion_ratio (float): Expansion ratio in the bottleneck. + squeeze_ratio (float): Squeeze ratio in the SE Layer. + stride (int): Stride of the depthwise convolution. + activation_layer (Callable[..., nn.Module]): Activation function. + norm_layer (Callable[..., nn.Module]): Normalization function. + head_dim (int): Dimension of the attention heads. + mlp_ratio (int): Ratio of the MLP layer. + mlp_dropout (float): Dropout probability for the MLP layer. + attention_dropout (float): Dropout probability for the attention layer. + p_stochastic_dropout (float): Probability of stochastic depth. + partition_size (int): Size of the partitions. + grid_size (Tuple[int, int]): Size of the input feature grid. + """ + + def __init__( + self, + # conv parameters + in_channels: int, + out_channels: int, + squeeze_ratio: float, + expansion_ratio: float, + stride: int, + # conv + transformer parameters + norm_layer: Callable[..., nn.Module], + activation_layer: Callable[..., nn.Module], + # transformer parameters + head_dim: int, + mlp_ratio: int, + mlp_dropout: float, + attention_dropout: float, + p_stochastic_dropout: float, + # partitioning parameters + partition_size: int, + grid_size: Tuple[int, int], + ) -> None: + super().__init__() + + layers: OrderedDict = OrderedDict() + + # convolutional layer + layers["MBconv"] = MBConv( + in_channels=in_channels, + out_channels=out_channels, + expansion_ratio=expansion_ratio, + squeeze_ratio=squeeze_ratio, + stride=stride, + activation_layer=activation_layer, + norm_layer=norm_layer, + p_stochastic_dropout=p_stochastic_dropout, + ) + # attention layers, block -> grid + layers["window_attention"] = PartitionAttentionLayer( + in_channels=out_channels, + head_dim=head_dim, + partition_size=partition_size, + partition_type="window", + grid_size=grid_size, + mlp_ratio=mlp_ratio, + activation_layer=activation_layer, + norm_layer=nn.LayerNorm, + attention_dropout=attention_dropout, + mlp_dropout=mlp_dropout, + p_stochastic_dropout=p_stochastic_dropout, + ) + layers["grid_attention"] = PartitionAttentionLayer( + in_channels=out_channels, + head_dim=head_dim, + partition_size=partition_size, + partition_type="grid", + grid_size=grid_size, + mlp_ratio=mlp_ratio, + activation_layer=activation_layer, + norm_layer=nn.LayerNorm, + attention_dropout=attention_dropout, + mlp_dropout=mlp_dropout, + p_stochastic_dropout=p_stochastic_dropout, + ) + self.layers = nn.Sequential(layers) + + def forward(self, x: Tensor) -> Tensor: + """ + Args: + x (Tensor): Input tensor of shape (B, C, H, W). + Returns: + Tensor: Output tensor of shape (B, C, H, W). + """ + x = self.layers(x) + return x + + +class MaxVitBlock(nn.Module): + """ + A MaxVit block consisting of `n_layers` MaxVit layers. + + Args: + in_channels (int): Number of input channels. + out_channels (int): Number of output channels. + expansion_ratio (float): Expansion ratio in the bottleneck. + squeeze_ratio (float): Squeeze ratio in the SE Layer. + activation_layer (Callable[..., nn.Module]): Activation function. + norm_layer (Callable[..., nn.Module]): Normalization function. + head_dim (int): Dimension of the attention heads. + mlp_ratio (int): Ratio of the MLP layer. + mlp_dropout (float): Dropout probability for the MLP layer. + attention_dropout (float): Dropout probability for the attention layer. + p_stochastic_dropout (float): Probability of stochastic depth. + partition_size (int): Size of the partitions. + input_grid_size (Tuple[int, int]): Size of the input feature grid. + n_layers (int): Number of layers in the block. + p_stochastic (List[float]): List of probabilities for stochastic depth for each layer. + """ + + def __init__( + self, + # conv parameters + in_channels: int, + out_channels: int, + squeeze_ratio: float, + expansion_ratio: float, + # conv + transformer parameters + norm_layer: Callable[..., nn.Module], + activation_layer: Callable[..., nn.Module], + # transformer parameters + head_dim: int, + mlp_ratio: int, + mlp_dropout: float, + attention_dropout: float, + # partitioning parameters + partition_size: int, + input_grid_size: Tuple[int, int], + # number of layers + n_layers: int, + p_stochastic: List[float], + ) -> None: + super().__init__() + if not len(p_stochastic) == n_layers: + raise ValueError(f"p_stochastic must have length n_layers={n_layers}, got p_stochastic={p_stochastic}.") + + self.layers = nn.ModuleList() + # account for the first stride of the first layer + self.grid_size = _get_conv_output_shape(input_grid_size, kernel_size=3, stride=2, padding=1) + + for idx, p in enumerate(p_stochastic): + stride = 2 if idx == 0 else 1 + self.layers += [ + MaxVitLayer( + in_channels=in_channels if idx == 0 else out_channels, + out_channels=out_channels, + squeeze_ratio=squeeze_ratio, + expansion_ratio=expansion_ratio, + stride=stride, + norm_layer=norm_layer, + activation_layer=activation_layer, + head_dim=head_dim, + mlp_ratio=mlp_ratio, + mlp_dropout=mlp_dropout, + attention_dropout=attention_dropout, + partition_size=partition_size, + grid_size=self.grid_size, + p_stochastic_dropout=p, + ), + ] + + def forward(self, x: Tensor) -> Tensor: + """ + Args: + x (Tensor): Input tensor of shape (B, C, H, W). + Returns: + Tensor: Output tensor of shape (B, C, H, W). + """ + for layer in self.layers: + x = layer(x) + return x + + +class MaxVit(nn.Module): + """ + Implements MaxVit Transformer from the `MaxViT: Multi-Axis Vision Transformer `_ paper. + Args: + input_size (Tuple[int, int]): Size of the input image. + stem_channels (int): Number of channels in the stem. + partition_size (int): Size of the partitions. + block_channels (List[int]): Number of channels in each block. + block_layers (List[int]): Number of layers in each block. + stochastic_depth_prob (float): Probability of stochastic depth. Expands to a list of probabilities for each layer that scales linearly to the specified value. + squeeze_ratio (float): Squeeze ratio in the SE Layer. Default: 0.25. + expansion_ratio (float): Expansion ratio in the MBConv bottleneck. Default: 4. + norm_layer (Callable[..., nn.Module]): Normalization function. Default: None (setting to None will produce a `BatchNorm2d(eps=1e-3, momentum=0.01)`). + activation_layer (Callable[..., nn.Module]): Activation function Default: nn.GELU. + head_dim (int): Dimension of the attention heads. + mlp_ratio (int): Expansion ratio of the MLP layer. Default: 4. + mlp_dropout (float): Dropout probability for the MLP layer. Default: 0.0. + attention_dropout (float): Dropout probability for the attention layer. Default: 0.0. + num_classes (int): Number of classes. Default: 1000. + """ + + def __init__( + self, + # input size parameters + input_size: Tuple[int, int], + # stem and task parameters + stem_channels: int, + # partitioning parameters + partition_size: int, + # block parameters + block_channels: List[int], + block_layers: List[int], + # attention head dimensions + head_dim: int, + stochastic_depth_prob: float, + # conv + transformer parameters + # norm_layer is applied only to the conv layers + # activation_layer is applied both to conv and transformer layers + norm_layer: Optional[Callable[..., nn.Module]] = None, + activation_layer: Callable[..., nn.Module] = nn.GELU, + # conv parameters + squeeze_ratio: float = 0.25, + expansion_ratio: float = 4, + # transformer parameters + mlp_ratio: int = 4, + mlp_dropout: float = 0.0, + attention_dropout: float = 0.0, + # task parameters + num_classes: int = 1000, + ) -> None: + super().__init__() + _log_api_usage_once(self) + + input_channels = 3 + + # https://github.com/google-research/maxvit/blob/da76cf0d8a6ec668cc31b399c4126186da7da944/maxvit/models/maxvit.py#L1029-L1030 + # for the exact parameters used in batchnorm + if norm_layer is None: + norm_layer = partial(nn.BatchNorm2d, eps=1e-3, momentum=0.01) + + # Make sure input size will be divisible by the partition size in all blocks + # Undefined behavior if H or W are not divisible by p + # https://github.com/google-research/maxvit/blob/da76cf0d8a6ec668cc31b399c4126186da7da944/maxvit/models/maxvit.py#L766 + block_input_sizes = _make_block_input_shapes(input_size, len(block_channels)) + for idx, block_input_size in enumerate(block_input_sizes): + if block_input_size[0] % partition_size != 0 or block_input_size[1] % partition_size != 0: + raise ValueError( + f"Input size {block_input_size} of block {idx} is not divisible by partition size {partition_size}. " + f"Consider changing the partition size or the input size.\n" + f"Current configuration yields the following block input sizes: {block_input_sizes}." + ) + + # stem + self.stem = nn.Sequential( + Conv2dNormActivation( + input_channels, + stem_channels, + 3, + stride=2, + norm_layer=norm_layer, + activation_layer=activation_layer, + bias=False, + inplace=None, + ), + Conv2dNormActivation( + stem_channels, stem_channels, 3, stride=1, norm_layer=None, activation_layer=None, bias=True + ), + ) + + # account for stem stride + input_size = _get_conv_output_shape(input_size, kernel_size=3, stride=2, padding=1) + self.partition_size = partition_size + + # blocks + self.blocks = nn.ModuleList() + in_channels = [stem_channels] + block_channels[:-1] + out_channels = block_channels + + # precompute the stochastich depth probabilities from 0 to stochastic_depth_prob + # since we have N blocks with L layers, we will have N * L probabilities uniformly distributed + # over the range [0, stochastic_depth_prob] + p_stochastic = np.linspace(0, stochastic_depth_prob, sum(block_layers)).tolist() + + p_idx = 0 + for in_channel, out_channel, num_layers in zip(in_channels, out_channels, block_layers): + self.blocks.append( + MaxVitBlock( + in_channels=in_channel, + out_channels=out_channel, + squeeze_ratio=squeeze_ratio, + expansion_ratio=expansion_ratio, + norm_layer=norm_layer, + activation_layer=activation_layer, + head_dim=head_dim, + mlp_ratio=mlp_ratio, + mlp_dropout=mlp_dropout, + attention_dropout=attention_dropout, + partition_size=partition_size, + input_grid_size=input_size, + n_layers=num_layers, + p_stochastic=p_stochastic[p_idx : p_idx + num_layers], + ), + ) + input_size = self.blocks[-1].grid_size # type: ignore[assignment] + p_idx += num_layers + + # see https://github.com/google-research/maxvit/blob/da76cf0d8a6ec668cc31b399c4126186da7da944/maxvit/models/maxvit.py#L1137-L1158 + # for why there is Linear -> Tanh -> Linear + self.classifier = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Flatten(), + nn.LayerNorm(block_channels[-1]), + nn.Linear(block_channels[-1], block_channels[-1]), + nn.Tanh(), + nn.Linear(block_channels[-1], num_classes, bias=False), + ) + + self._init_weights() + + def forward(self, x: Tensor) -> Tensor: + x = self.stem(x) + for block in self.blocks: + x = block(x) + x = self.classifier(x) + return x + + def _init_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, nn.BatchNorm2d): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.zeros_(m.bias) + + +def _maxvit( + # stem parameters + stem_channels: int, + # block parameters + block_channels: List[int], + block_layers: List[int], + stochastic_depth_prob: float, + # partitioning parameters + partition_size: int, + # transformer parameters + head_dim: int, + # Weights API + weights: Optional[WeightsEnum] = None, + progress: bool = False, + # kwargs, + **kwargs: Any, +) -> MaxVit: + + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + assert weights.meta["min_size"][0] == weights.meta["min_size"][1] + _ovewrite_named_param(kwargs, "input_size", weights.meta["min_size"]) + + input_size = kwargs.pop("input_size", (224, 224)) + + model = MaxVit( + stem_channels=stem_channels, + block_channels=block_channels, + block_layers=block_layers, + stochastic_depth_prob=stochastic_depth_prob, + head_dim=head_dim, + partition_size=partition_size, + input_size=input_size, + **kwargs, + ) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +class MaxVit_T_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + # URL empty until official release + url="https://download.pytorch.org/models/maxvit_t-bc5ab103.pth", + transforms=partial( + ImageClassification, crop_size=224, resize_size=224, interpolation=InterpolationMode.BICUBIC + ), + meta={ + "categories": _IMAGENET_CATEGORIES, + "num_params": 30919624, + "min_size": (224, 224), + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#maxvit", + "_metrics": { + "ImageNet-1K": { + "acc@1": 83.700, + "acc@5": 96.722, + } + }, + "_ops": 5.558, + "_file_size": 118.769, + "_docs": """These weights reproduce closely the results of the paper using a similar training recipe. + They were trained with a BatchNorm2D momentum of 0.99 instead of the more correct 0.01.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", MaxVit_T_Weights.IMAGENET1K_V1)) +def maxvit_t(*, weights: Optional[MaxVit_T_Weights] = None, progress: bool = True, **kwargs: Any) -> MaxVit: + """ + Constructs a maxvit_t architecture from + `MaxViT: Multi-Axis Vision Transformer `_. + + Args: + weights (:class:`~torchvision.models.MaxVit_T_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.MaxVit_T_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.maxvit.MaxVit`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.MaxVit_T_Weights + :members: + """ + weights = MaxVit_T_Weights.verify(weights) + + return _maxvit( + stem_channels=64, + block_channels=[64, 128, 256, 512], + block_layers=[2, 2, 5, 2], + head_dim=32, + stochastic_depth_prob=0.2, + partition_size=7, + weights=weights, + progress=progress, + **kwargs, + ) diff --git a/lib/python3.12/site-packages/torchvision/models/mnasnet.py b/lib/python3.12/site-packages/torchvision/models/mnasnet.py new file mode 100644 index 0000000000000000000000000000000000000000..5846111ab1c05b4ebca7ccf9240ac744267b859a --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/mnasnet.py @@ -0,0 +1,434 @@ +import warnings +from functools import partial +from typing import Any, Dict, List, Optional + +import torch +import torch.nn as nn +from torch import Tensor + +from ..transforms._presets import ImageClassification +from ..utils import _log_api_usage_once +from ._api import register_model, Weights, WeightsEnum +from ._meta import _IMAGENET_CATEGORIES +from ._utils import _ovewrite_named_param, handle_legacy_interface + + +__all__ = [ + "MNASNet", + "MNASNet0_5_Weights", + "MNASNet0_75_Weights", + "MNASNet1_0_Weights", + "MNASNet1_3_Weights", + "mnasnet0_5", + "mnasnet0_75", + "mnasnet1_0", + "mnasnet1_3", +] + + +# Paper suggests 0.9997 momentum, for TensorFlow. Equivalent PyTorch momentum is +# 1.0 - tensorflow. +_BN_MOMENTUM = 1 - 0.9997 + + +class _InvertedResidual(nn.Module): + def __init__( + self, in_ch: int, out_ch: int, kernel_size: int, stride: int, expansion_factor: int, bn_momentum: float = 0.1 + ) -> None: + super().__init__() + if stride not in [1, 2]: + raise ValueError(f"stride should be 1 or 2 instead of {stride}") + if kernel_size not in [3, 5]: + raise ValueError(f"kernel_size should be 3 or 5 instead of {kernel_size}") + mid_ch = in_ch * expansion_factor + self.apply_residual = in_ch == out_ch and stride == 1 + self.layers = nn.Sequential( + # Pointwise + nn.Conv2d(in_ch, mid_ch, 1, bias=False), + nn.BatchNorm2d(mid_ch, momentum=bn_momentum), + nn.ReLU(inplace=True), + # Depthwise + nn.Conv2d(mid_ch, mid_ch, kernel_size, padding=kernel_size // 2, stride=stride, groups=mid_ch, bias=False), + nn.BatchNorm2d(mid_ch, momentum=bn_momentum), + nn.ReLU(inplace=True), + # Linear pointwise. Note that there's no activation. + nn.Conv2d(mid_ch, out_ch, 1, bias=False), + nn.BatchNorm2d(out_ch, momentum=bn_momentum), + ) + + def forward(self, input: Tensor) -> Tensor: + if self.apply_residual: + return self.layers(input) + input + else: + return self.layers(input) + + +def _stack( + in_ch: int, out_ch: int, kernel_size: int, stride: int, exp_factor: int, repeats: int, bn_momentum: float +) -> nn.Sequential: + """Creates a stack of inverted residuals.""" + if repeats < 1: + raise ValueError(f"repeats should be >= 1, instead got {repeats}") + # First one has no skip, because feature map size changes. + first = _InvertedResidual(in_ch, out_ch, kernel_size, stride, exp_factor, bn_momentum=bn_momentum) + remaining = [] + for _ in range(1, repeats): + remaining.append(_InvertedResidual(out_ch, out_ch, kernel_size, 1, exp_factor, bn_momentum=bn_momentum)) + return nn.Sequential(first, *remaining) + + +def _round_to_multiple_of(val: float, divisor: int, round_up_bias: float = 0.9) -> int: + """Asymmetric rounding to make `val` divisible by `divisor`. With default + bias, will round up, unless the number is no more than 10% greater than the + smaller divisible value, i.e. (83, 8) -> 80, but (84, 8) -> 88.""" + if not 0.0 < round_up_bias < 1.0: + raise ValueError(f"round_up_bias should be greater than 0.0 and smaller than 1.0 instead of {round_up_bias}") + new_val = max(divisor, int(val + divisor / 2) // divisor * divisor) + return new_val if new_val >= round_up_bias * val else new_val + divisor + + +def _get_depths(alpha: float) -> List[int]: + """Scales tensor depths as in reference MobileNet code, prefers rounding up + rather than down.""" + depths = [32, 16, 24, 40, 80, 96, 192, 320] + return [_round_to_multiple_of(depth * alpha, 8) for depth in depths] + + +class MNASNet(torch.nn.Module): + """MNASNet, as described in https://arxiv.org/abs/1807.11626. This + implements the B1 variant of the model. + >>> model = MNASNet(1.0, num_classes=1000) + >>> x = torch.rand(1, 3, 224, 224) + >>> y = model(x) + >>> y.dim() + 2 + >>> y.nelement() + 1000 + """ + + # Version 2 adds depth scaling in the initial stages of the network. + _version = 2 + + def __init__(self, alpha: float, num_classes: int = 1000, dropout: float = 0.2) -> None: + super().__init__() + _log_api_usage_once(self) + if alpha <= 0.0: + raise ValueError(f"alpha should be greater than 0.0 instead of {alpha}") + self.alpha = alpha + self.num_classes = num_classes + depths = _get_depths(alpha) + layers = [ + # First layer: regular conv. + nn.Conv2d(3, depths[0], 3, padding=1, stride=2, bias=False), + nn.BatchNorm2d(depths[0], momentum=_BN_MOMENTUM), + nn.ReLU(inplace=True), + # Depthwise separable, no skip. + nn.Conv2d(depths[0], depths[0], 3, padding=1, stride=1, groups=depths[0], bias=False), + nn.BatchNorm2d(depths[0], momentum=_BN_MOMENTUM), + nn.ReLU(inplace=True), + nn.Conv2d(depths[0], depths[1], 1, padding=0, stride=1, bias=False), + nn.BatchNorm2d(depths[1], momentum=_BN_MOMENTUM), + # MNASNet blocks: stacks of inverted residuals. + _stack(depths[1], depths[2], 3, 2, 3, 3, _BN_MOMENTUM), + _stack(depths[2], depths[3], 5, 2, 3, 3, _BN_MOMENTUM), + _stack(depths[3], depths[4], 5, 2, 6, 3, _BN_MOMENTUM), + _stack(depths[4], depths[5], 3, 1, 6, 2, _BN_MOMENTUM), + _stack(depths[5], depths[6], 5, 2, 6, 4, _BN_MOMENTUM), + _stack(depths[6], depths[7], 3, 1, 6, 1, _BN_MOMENTUM), + # Final mapping to classifier input. + nn.Conv2d(depths[7], 1280, 1, padding=0, stride=1, bias=False), + nn.BatchNorm2d(1280, momentum=_BN_MOMENTUM), + nn.ReLU(inplace=True), + ] + self.layers = nn.Sequential(*layers) + self.classifier = nn.Sequential(nn.Dropout(p=dropout, inplace=True), nn.Linear(1280, num_classes)) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, nn.BatchNorm2d): + nn.init.ones_(m.weight) + nn.init.zeros_(m.bias) + elif isinstance(m, nn.Linear): + nn.init.kaiming_uniform_(m.weight, mode="fan_out", nonlinearity="sigmoid") + nn.init.zeros_(m.bias) + + def forward(self, x: Tensor) -> Tensor: + x = self.layers(x) + # Equivalent to global avgpool and removing H and W dimensions. + x = x.mean([2, 3]) + return self.classifier(x) + + def _load_from_state_dict( + self, + state_dict: Dict, + prefix: str, + local_metadata: Dict, + strict: bool, + missing_keys: List[str], + unexpected_keys: List[str], + error_msgs: List[str], + ) -> None: + version = local_metadata.get("version", None) + if version not in [1, 2]: + raise ValueError(f"version shluld be set to 1 or 2 instead of {version}") + + if version == 1 and not self.alpha == 1.0: + # In the initial version of the model (v1), stem was fixed-size. + # All other layer configurations were the same. This will patch + # the model so that it's identical to v1. Model with alpha 1.0 is + # unaffected. + depths = _get_depths(self.alpha) + v1_stem = [ + nn.Conv2d(3, 32, 3, padding=1, stride=2, bias=False), + nn.BatchNorm2d(32, momentum=_BN_MOMENTUM), + nn.ReLU(inplace=True), + nn.Conv2d(32, 32, 3, padding=1, stride=1, groups=32, bias=False), + nn.BatchNorm2d(32, momentum=_BN_MOMENTUM), + nn.ReLU(inplace=True), + nn.Conv2d(32, 16, 1, padding=0, stride=1, bias=False), + nn.BatchNorm2d(16, momentum=_BN_MOMENTUM), + _stack(16, depths[2], 3, 2, 3, 3, _BN_MOMENTUM), + ] + for idx, layer in enumerate(v1_stem): + self.layers[idx] = layer + + # The model is now identical to v1, and must be saved as such. + self._version = 1 + warnings.warn( + "A new version of MNASNet model has been implemented. " + "Your checkpoint was saved using the previous version. " + "This checkpoint will load and work as before, but " + "you may want to upgrade by training a newer model or " + "transfer learning from an updated ImageNet checkpoint.", + UserWarning, + ) + + super()._load_from_state_dict( + state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ) + + +_COMMON_META = { + "min_size": (1, 1), + "categories": _IMAGENET_CATEGORIES, + "recipe": "https://github.com/1e100/mnasnet_trainer", +} + + +class MNASNet0_5_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/mnasnet0.5_top1_67.823-3ffadce67e.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 2218512, + "_metrics": { + "ImageNet-1K": { + "acc@1": 67.734, + "acc@5": 87.490, + } + }, + "_ops": 0.104, + "_file_size": 8.591, + "_docs": """These weights reproduce closely the results of the paper.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class MNASNet0_75_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/mnasnet0_75-7090bc5f.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "recipe": "https://github.com/pytorch/vision/pull/6019", + "num_params": 3170208, + "_metrics": { + "ImageNet-1K": { + "acc@1": 71.180, + "acc@5": 90.496, + } + }, + "_ops": 0.215, + "_file_size": 12.303, + "_docs": """ + These weights were trained from scratch by using TorchVision's `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class MNASNet1_0_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/mnasnet1.0_top1_73.512-f206786ef8.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 4383312, + "_metrics": { + "ImageNet-1K": { + "acc@1": 73.456, + "acc@5": 91.510, + } + }, + "_ops": 0.314, + "_file_size": 16.915, + "_docs": """These weights reproduce closely the results of the paper.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class MNASNet1_3_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/mnasnet1_3-a4c69d6f.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "recipe": "https://github.com/pytorch/vision/pull/6019", + "num_params": 6282256, + "_metrics": { + "ImageNet-1K": { + "acc@1": 76.506, + "acc@5": 93.522, + } + }, + "_ops": 0.526, + "_file_size": 24.246, + "_docs": """ + These weights were trained from scratch by using TorchVision's `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +def _mnasnet(alpha: float, weights: Optional[WeightsEnum], progress: bool, **kwargs: Any) -> MNASNet: + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + + model = MNASNet(alpha, **kwargs) + + if weights: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +@register_model() +@handle_legacy_interface(weights=("pretrained", MNASNet0_5_Weights.IMAGENET1K_V1)) +def mnasnet0_5(*, weights: Optional[MNASNet0_5_Weights] = None, progress: bool = True, **kwargs: Any) -> MNASNet: + """MNASNet with depth multiplier of 0.5 from + `MnasNet: Platform-Aware Neural Architecture Search for Mobile + `_ paper. + + Args: + weights (:class:`~torchvision.models.MNASNet0_5_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.MNASNet0_5_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.mnasnet.MNASNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.MNASNet0_5_Weights + :members: + """ + weights = MNASNet0_5_Weights.verify(weights) + + return _mnasnet(0.5, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", MNASNet0_75_Weights.IMAGENET1K_V1)) +def mnasnet0_75(*, weights: Optional[MNASNet0_75_Weights] = None, progress: bool = True, **kwargs: Any) -> MNASNet: + """MNASNet with depth multiplier of 0.75 from + `MnasNet: Platform-Aware Neural Architecture Search for Mobile + `_ paper. + + Args: + weights (:class:`~torchvision.models.MNASNet0_75_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.MNASNet0_75_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.mnasnet.MNASNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.MNASNet0_75_Weights + :members: + """ + weights = MNASNet0_75_Weights.verify(weights) + + return _mnasnet(0.75, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", MNASNet1_0_Weights.IMAGENET1K_V1)) +def mnasnet1_0(*, weights: Optional[MNASNet1_0_Weights] = None, progress: bool = True, **kwargs: Any) -> MNASNet: + """MNASNet with depth multiplier of 1.0 from + `MnasNet: Platform-Aware Neural Architecture Search for Mobile + `_ paper. + + Args: + weights (:class:`~torchvision.models.MNASNet1_0_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.MNASNet1_0_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.mnasnet.MNASNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.MNASNet1_0_Weights + :members: + """ + weights = MNASNet1_0_Weights.verify(weights) + + return _mnasnet(1.0, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", MNASNet1_3_Weights.IMAGENET1K_V1)) +def mnasnet1_3(*, weights: Optional[MNASNet1_3_Weights] = None, progress: bool = True, **kwargs: Any) -> MNASNet: + """MNASNet with depth multiplier of 1.3 from + `MnasNet: Platform-Aware Neural Architecture Search for Mobile + `_ paper. + + Args: + weights (:class:`~torchvision.models.MNASNet1_3_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.MNASNet1_3_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.mnasnet.MNASNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.MNASNet1_3_Weights + :members: + """ + weights = MNASNet1_3_Weights.verify(weights) + + return _mnasnet(1.3, weights, progress, **kwargs) diff --git a/lib/python3.12/site-packages/torchvision/models/mobilenet.py b/lib/python3.12/site-packages/torchvision/models/mobilenet.py new file mode 100644 index 0000000000000000000000000000000000000000..0a270d14d3a4ad9eda62b68c2c01e9fdb710ef38 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/mobilenet.py @@ -0,0 +1,6 @@ +from .mobilenetv2 import * # noqa: F401, F403 +from .mobilenetv3 import * # noqa: F401, F403 +from .mobilenetv2 import __all__ as mv2_all +from .mobilenetv3 import __all__ as mv3_all + +__all__ = mv2_all + mv3_all diff --git a/lib/python3.12/site-packages/torchvision/models/mobilenetv2.py b/lib/python3.12/site-packages/torchvision/models/mobilenetv2.py new file mode 100644 index 0000000000000000000000000000000000000000..fbb6a4981d6dec73c563dfc845329a58eaa3e083 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/mobilenetv2.py @@ -0,0 +1,260 @@ +from functools import partial +from typing import Any, Callable, List, Optional + +import torch +from torch import nn, Tensor + +from ..ops.misc import Conv2dNormActivation +from ..transforms._presets import ImageClassification +from ..utils import _log_api_usage_once +from ._api import register_model, Weights, WeightsEnum +from ._meta import _IMAGENET_CATEGORIES +from ._utils import _make_divisible, _ovewrite_named_param, handle_legacy_interface + + +__all__ = ["MobileNetV2", "MobileNet_V2_Weights", "mobilenet_v2"] + + +# necessary for backwards compatibility +class InvertedResidual(nn.Module): + def __init__( + self, inp: int, oup: int, stride: int, expand_ratio: int, norm_layer: Optional[Callable[..., nn.Module]] = None + ) -> None: + super().__init__() + self.stride = stride + if stride not in [1, 2]: + raise ValueError(f"stride should be 1 or 2 instead of {stride}") + + if norm_layer is None: + norm_layer = nn.BatchNorm2d + + hidden_dim = int(round(inp * expand_ratio)) + self.use_res_connect = self.stride == 1 and inp == oup + + layers: List[nn.Module] = [] + if expand_ratio != 1: + # pw + layers.append( + Conv2dNormActivation(inp, hidden_dim, kernel_size=1, norm_layer=norm_layer, activation_layer=nn.ReLU6) + ) + layers.extend( + [ + # dw + Conv2dNormActivation( + hidden_dim, + hidden_dim, + stride=stride, + groups=hidden_dim, + norm_layer=norm_layer, + activation_layer=nn.ReLU6, + ), + # pw-linear + nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False), + norm_layer(oup), + ] + ) + self.conv = nn.Sequential(*layers) + self.out_channels = oup + self._is_cn = stride > 1 + + def forward(self, x: Tensor) -> Tensor: + if self.use_res_connect: + return x + self.conv(x) + else: + return self.conv(x) + + +class MobileNetV2(nn.Module): + def __init__( + self, + num_classes: int = 1000, + width_mult: float = 1.0, + inverted_residual_setting: Optional[List[List[int]]] = None, + round_nearest: int = 8, + block: Optional[Callable[..., nn.Module]] = None, + norm_layer: Optional[Callable[..., nn.Module]] = None, + dropout: float = 0.2, + ) -> None: + """ + MobileNet V2 main class + + Args: + num_classes (int): Number of classes + width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount + inverted_residual_setting: Network structure + round_nearest (int): Round the number of channels in each layer to be a multiple of this number + Set to 1 to turn off rounding + block: Module specifying inverted residual building block for mobilenet + norm_layer: Module specifying the normalization layer to use + dropout (float): The droupout probability + + """ + super().__init__() + _log_api_usage_once(self) + + if block is None: + block = InvertedResidual + + if norm_layer is None: + norm_layer = nn.BatchNorm2d + + input_channel = 32 + last_channel = 1280 + + if inverted_residual_setting is None: + inverted_residual_setting = [ + # t, c, n, s + [1, 16, 1, 1], + [6, 24, 2, 2], + [6, 32, 3, 2], + [6, 64, 4, 2], + [6, 96, 3, 1], + [6, 160, 3, 2], + [6, 320, 1, 1], + ] + + # only check the first element, assuming user knows t,c,n,s are required + if len(inverted_residual_setting) == 0 or len(inverted_residual_setting[0]) != 4: + raise ValueError( + f"inverted_residual_setting should be non-empty or a 4-element list, got {inverted_residual_setting}" + ) + + # building first layer + input_channel = _make_divisible(input_channel * width_mult, round_nearest) + self.last_channel = _make_divisible(last_channel * max(1.0, width_mult), round_nearest) + features: List[nn.Module] = [ + Conv2dNormActivation(3, input_channel, stride=2, norm_layer=norm_layer, activation_layer=nn.ReLU6) + ] + # building inverted residual blocks + for t, c, n, s in inverted_residual_setting: + output_channel = _make_divisible(c * width_mult, round_nearest) + for i in range(n): + stride = s if i == 0 else 1 + features.append(block(input_channel, output_channel, stride, expand_ratio=t, norm_layer=norm_layer)) + input_channel = output_channel + # building last several layers + features.append( + Conv2dNormActivation( + input_channel, self.last_channel, kernel_size=1, norm_layer=norm_layer, activation_layer=nn.ReLU6 + ) + ) + # make it nn.Sequential + self.features = nn.Sequential(*features) + + # building classifier + self.classifier = nn.Sequential( + nn.Dropout(p=dropout), + nn.Linear(self.last_channel, num_classes), + ) + + # weight initialization + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode="fan_out") + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): + nn.init.ones_(m.weight) + nn.init.zeros_(m.bias) + elif isinstance(m, nn.Linear): + nn.init.normal_(m.weight, 0, 0.01) + nn.init.zeros_(m.bias) + + def _forward_impl(self, x: Tensor) -> Tensor: + # This exists since TorchScript doesn't support inheritance, so the superclass method + # (this one) needs to have a name other than `forward` that can be accessed in a subclass + x = self.features(x) + # Cannot use "squeeze" as batch-size can be 1 + x = nn.functional.adaptive_avg_pool2d(x, (1, 1)) + x = torch.flatten(x, 1) + x = self.classifier(x) + return x + + def forward(self, x: Tensor) -> Tensor: + return self._forward_impl(x) + + +_COMMON_META = { + "num_params": 3504872, + "min_size": (1, 1), + "categories": _IMAGENET_CATEGORIES, +} + + +class MobileNet_V2_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/mobilenet_v2-b0353104.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#mobilenetv2", + "_metrics": { + "ImageNet-1K": { + "acc@1": 71.878, + "acc@5": 90.286, + } + }, + "_ops": 0.301, + "_file_size": 13.555, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/mobilenet_v2-7ebf99e0.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe-with-reg-tuning", + "_metrics": { + "ImageNet-1K": { + "acc@1": 72.154, + "acc@5": 90.822, + } + }, + "_ops": 0.301, + "_file_size": 13.598, + "_docs": """ + These weights improve upon the results of the original paper by using a modified version of TorchVision's + `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", MobileNet_V2_Weights.IMAGENET1K_V1)) +def mobilenet_v2( + *, weights: Optional[MobileNet_V2_Weights] = None, progress: bool = True, **kwargs: Any +) -> MobileNetV2: + """MobileNetV2 architecture from the `MobileNetV2: Inverted Residuals and Linear + Bottlenecks `_ paper. + + Args: + weights (:class:`~torchvision.models.MobileNet_V2_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.MobileNet_V2_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.mobilenetv2.MobileNetV2`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.MobileNet_V2_Weights + :members: + """ + weights = MobileNet_V2_Weights.verify(weights) + + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + + model = MobileNetV2(**kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model diff --git a/lib/python3.12/site-packages/torchvision/models/mobilenetv3.py b/lib/python3.12/site-packages/torchvision/models/mobilenetv3.py new file mode 100644 index 0000000000000000000000000000000000000000..1041d4d149fb1c414991b79156733434a2268e4b --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/mobilenetv3.py @@ -0,0 +1,423 @@ +from functools import partial +from typing import Any, Callable, List, Optional, Sequence + +import torch +from torch import nn, Tensor + +from ..ops.misc import Conv2dNormActivation, SqueezeExcitation as SElayer +from ..transforms._presets import ImageClassification +from ..utils import _log_api_usage_once +from ._api import register_model, Weights, WeightsEnum +from ._meta import _IMAGENET_CATEGORIES +from ._utils import _make_divisible, _ovewrite_named_param, handle_legacy_interface + + +__all__ = [ + "MobileNetV3", + "MobileNet_V3_Large_Weights", + "MobileNet_V3_Small_Weights", + "mobilenet_v3_large", + "mobilenet_v3_small", +] + + +class InvertedResidualConfig: + # Stores information listed at Tables 1 and 2 of the MobileNetV3 paper + def __init__( + self, + input_channels: int, + kernel: int, + expanded_channels: int, + out_channels: int, + use_se: bool, + activation: str, + stride: int, + dilation: int, + width_mult: float, + ): + self.input_channels = self.adjust_channels(input_channels, width_mult) + self.kernel = kernel + self.expanded_channels = self.adjust_channels(expanded_channels, width_mult) + self.out_channels = self.adjust_channels(out_channels, width_mult) + self.use_se = use_se + self.use_hs = activation == "HS" + self.stride = stride + self.dilation = dilation + + @staticmethod + def adjust_channels(channels: int, width_mult: float): + return _make_divisible(channels * width_mult, 8) + + +class InvertedResidual(nn.Module): + # Implemented as described at section 5 of MobileNetV3 paper + def __init__( + self, + cnf: InvertedResidualConfig, + norm_layer: Callable[..., nn.Module], + se_layer: Callable[..., nn.Module] = partial(SElayer, scale_activation=nn.Hardsigmoid), + ): + super().__init__() + if not (1 <= cnf.stride <= 2): + raise ValueError("illegal stride value") + + self.use_res_connect = cnf.stride == 1 and cnf.input_channels == cnf.out_channels + + layers: List[nn.Module] = [] + activation_layer = nn.Hardswish if cnf.use_hs else nn.ReLU + + # expand + if cnf.expanded_channels != cnf.input_channels: + layers.append( + Conv2dNormActivation( + cnf.input_channels, + cnf.expanded_channels, + kernel_size=1, + norm_layer=norm_layer, + activation_layer=activation_layer, + ) + ) + + # depthwise + stride = 1 if cnf.dilation > 1 else cnf.stride + layers.append( + Conv2dNormActivation( + cnf.expanded_channels, + cnf.expanded_channels, + kernel_size=cnf.kernel, + stride=stride, + dilation=cnf.dilation, + groups=cnf.expanded_channels, + norm_layer=norm_layer, + activation_layer=activation_layer, + ) + ) + if cnf.use_se: + squeeze_channels = _make_divisible(cnf.expanded_channels // 4, 8) + layers.append(se_layer(cnf.expanded_channels, squeeze_channels)) + + # project + layers.append( + Conv2dNormActivation( + cnf.expanded_channels, cnf.out_channels, kernel_size=1, norm_layer=norm_layer, activation_layer=None + ) + ) + + self.block = nn.Sequential(*layers) + self.out_channels = cnf.out_channels + self._is_cn = cnf.stride > 1 + + def forward(self, input: Tensor) -> Tensor: + result = self.block(input) + if self.use_res_connect: + result += input + return result + + +class MobileNetV3(nn.Module): + def __init__( + self, + inverted_residual_setting: List[InvertedResidualConfig], + last_channel: int, + num_classes: int = 1000, + block: Optional[Callable[..., nn.Module]] = None, + norm_layer: Optional[Callable[..., nn.Module]] = None, + dropout: float = 0.2, + **kwargs: Any, + ) -> None: + """ + MobileNet V3 main class + + Args: + inverted_residual_setting (List[InvertedResidualConfig]): Network structure + last_channel (int): The number of channels on the penultimate layer + num_classes (int): Number of classes + block (Optional[Callable[..., nn.Module]]): Module specifying inverted residual building block for mobilenet + norm_layer (Optional[Callable[..., nn.Module]]): Module specifying the normalization layer to use + dropout (float): The droupout probability + """ + super().__init__() + _log_api_usage_once(self) + + if not inverted_residual_setting: + raise ValueError("The inverted_residual_setting should not be empty") + elif not ( + isinstance(inverted_residual_setting, Sequence) + and all([isinstance(s, InvertedResidualConfig) for s in inverted_residual_setting]) + ): + raise TypeError("The inverted_residual_setting should be List[InvertedResidualConfig]") + + if block is None: + block = InvertedResidual + + if norm_layer is None: + norm_layer = partial(nn.BatchNorm2d, eps=0.001, momentum=0.01) + + layers: List[nn.Module] = [] + + # building first layer + firstconv_output_channels = inverted_residual_setting[0].input_channels + layers.append( + Conv2dNormActivation( + 3, + firstconv_output_channels, + kernel_size=3, + stride=2, + norm_layer=norm_layer, + activation_layer=nn.Hardswish, + ) + ) + + # building inverted residual blocks + for cnf in inverted_residual_setting: + layers.append(block(cnf, norm_layer)) + + # building last several layers + lastconv_input_channels = inverted_residual_setting[-1].out_channels + lastconv_output_channels = 6 * lastconv_input_channels + layers.append( + Conv2dNormActivation( + lastconv_input_channels, + lastconv_output_channels, + kernel_size=1, + norm_layer=norm_layer, + activation_layer=nn.Hardswish, + ) + ) + + self.features = nn.Sequential(*layers) + self.avgpool = nn.AdaptiveAvgPool2d(1) + self.classifier = nn.Sequential( + nn.Linear(lastconv_output_channels, last_channel), + nn.Hardswish(inplace=True), + nn.Dropout(p=dropout, inplace=True), + nn.Linear(last_channel, num_classes), + ) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode="fan_out") + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): + nn.init.ones_(m.weight) + nn.init.zeros_(m.bias) + elif isinstance(m, nn.Linear): + nn.init.normal_(m.weight, 0, 0.01) + nn.init.zeros_(m.bias) + + def _forward_impl(self, x: Tensor) -> Tensor: + x = self.features(x) + + x = self.avgpool(x) + x = torch.flatten(x, 1) + + x = self.classifier(x) + + return x + + def forward(self, x: Tensor) -> Tensor: + return self._forward_impl(x) + + +def _mobilenet_v3_conf( + arch: str, width_mult: float = 1.0, reduced_tail: bool = False, dilated: bool = False, **kwargs: Any +): + reduce_divider = 2 if reduced_tail else 1 + dilation = 2 if dilated else 1 + + bneck_conf = partial(InvertedResidualConfig, width_mult=width_mult) + adjust_channels = partial(InvertedResidualConfig.adjust_channels, width_mult=width_mult) + + if arch == "mobilenet_v3_large": + inverted_residual_setting = [ + bneck_conf(16, 3, 16, 16, False, "RE", 1, 1), + bneck_conf(16, 3, 64, 24, False, "RE", 2, 1), # C1 + bneck_conf(24, 3, 72, 24, False, "RE", 1, 1), + bneck_conf(24, 5, 72, 40, True, "RE", 2, 1), # C2 + bneck_conf(40, 5, 120, 40, True, "RE", 1, 1), + bneck_conf(40, 5, 120, 40, True, "RE", 1, 1), + bneck_conf(40, 3, 240, 80, False, "HS", 2, 1), # C3 + bneck_conf(80, 3, 200, 80, False, "HS", 1, 1), + bneck_conf(80, 3, 184, 80, False, "HS", 1, 1), + bneck_conf(80, 3, 184, 80, False, "HS", 1, 1), + bneck_conf(80, 3, 480, 112, True, "HS", 1, 1), + bneck_conf(112, 3, 672, 112, True, "HS", 1, 1), + bneck_conf(112, 5, 672, 160 // reduce_divider, True, "HS", 2, dilation), # C4 + bneck_conf(160 // reduce_divider, 5, 960 // reduce_divider, 160 // reduce_divider, True, "HS", 1, dilation), + bneck_conf(160 // reduce_divider, 5, 960 // reduce_divider, 160 // reduce_divider, True, "HS", 1, dilation), + ] + last_channel = adjust_channels(1280 // reduce_divider) # C5 + elif arch == "mobilenet_v3_small": + inverted_residual_setting = [ + bneck_conf(16, 3, 16, 16, True, "RE", 2, 1), # C1 + bneck_conf(16, 3, 72, 24, False, "RE", 2, 1), # C2 + bneck_conf(24, 3, 88, 24, False, "RE", 1, 1), + bneck_conf(24, 5, 96, 40, True, "HS", 2, 1), # C3 + bneck_conf(40, 5, 240, 40, True, "HS", 1, 1), + bneck_conf(40, 5, 240, 40, True, "HS", 1, 1), + bneck_conf(40, 5, 120, 48, True, "HS", 1, 1), + bneck_conf(48, 5, 144, 48, True, "HS", 1, 1), + bneck_conf(48, 5, 288, 96 // reduce_divider, True, "HS", 2, dilation), # C4 + bneck_conf(96 // reduce_divider, 5, 576 // reduce_divider, 96 // reduce_divider, True, "HS", 1, dilation), + bneck_conf(96 // reduce_divider, 5, 576 // reduce_divider, 96 // reduce_divider, True, "HS", 1, dilation), + ] + last_channel = adjust_channels(1024 // reduce_divider) # C5 + else: + raise ValueError(f"Unsupported model type {arch}") + + return inverted_residual_setting, last_channel + + +def _mobilenet_v3( + inverted_residual_setting: List[InvertedResidualConfig], + last_channel: int, + weights: Optional[WeightsEnum], + progress: bool, + **kwargs: Any, +) -> MobileNetV3: + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + + model = MobileNetV3(inverted_residual_setting, last_channel, **kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +_COMMON_META = { + "min_size": (1, 1), + "categories": _IMAGENET_CATEGORIES, +} + + +class MobileNet_V3_Large_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/mobilenet_v3_large-8738ca79.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 5483032, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#mobilenetv3-large--small", + "_metrics": { + "ImageNet-1K": { + "acc@1": 74.042, + "acc@5": 91.340, + } + }, + "_ops": 0.217, + "_file_size": 21.114, + "_docs": """These weights were trained from scratch by using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/mobilenet_v3_large-5c1a4163.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 5483032, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe-with-reg-tuning", + "_metrics": { + "ImageNet-1K": { + "acc@1": 75.274, + "acc@5": 92.566, + } + }, + "_ops": 0.217, + "_file_size": 21.107, + "_docs": """ + These weights improve marginally upon the results of the original paper by using a modified version of + TorchVision's `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class MobileNet_V3_Small_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/mobilenet_v3_small-047dcff4.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 2542856, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#mobilenetv3-large--small", + "_metrics": { + "ImageNet-1K": { + "acc@1": 67.668, + "acc@5": 87.402, + } + }, + "_ops": 0.057, + "_file_size": 9.829, + "_docs": """ + These weights improve upon the results of the original paper by using a simple training recipe. + """, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", MobileNet_V3_Large_Weights.IMAGENET1K_V1)) +def mobilenet_v3_large( + *, weights: Optional[MobileNet_V3_Large_Weights] = None, progress: bool = True, **kwargs: Any +) -> MobileNetV3: + """ + Constructs a large MobileNetV3 architecture from + `Searching for MobileNetV3 `__. + + Args: + weights (:class:`~torchvision.models.MobileNet_V3_Large_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.MobileNet_V3_Large_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.mobilenet.MobileNetV3`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.MobileNet_V3_Large_Weights + :members: + """ + weights = MobileNet_V3_Large_Weights.verify(weights) + + inverted_residual_setting, last_channel = _mobilenet_v3_conf("mobilenet_v3_large", **kwargs) + return _mobilenet_v3(inverted_residual_setting, last_channel, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", MobileNet_V3_Small_Weights.IMAGENET1K_V1)) +def mobilenet_v3_small( + *, weights: Optional[MobileNet_V3_Small_Weights] = None, progress: bool = True, **kwargs: Any +) -> MobileNetV3: + """ + Constructs a small MobileNetV3 architecture from + `Searching for MobileNetV3 `__. + + Args: + weights (:class:`~torchvision.models.MobileNet_V3_Small_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.MobileNet_V3_Small_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.mobilenet.MobileNetV3`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.MobileNet_V3_Small_Weights + :members: + """ + weights = MobileNet_V3_Small_Weights.verify(weights) + + inverted_residual_setting, last_channel = _mobilenet_v3_conf("mobilenet_v3_small", **kwargs) + return _mobilenet_v3(inverted_residual_setting, last_channel, weights, progress, **kwargs) diff --git a/lib/python3.12/site-packages/torchvision/models/optical_flow/__init__.py b/lib/python3.12/site-packages/torchvision/models/optical_flow/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..89d2302f825ff0dbe25d02f6dc7c84d3c0790ad0 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/optical_flow/__init__.py @@ -0,0 +1 @@ +from .raft import * diff --git a/lib/python3.12/site-packages/torchvision/models/optical_flow/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/optical_flow/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6c36c6834cf0ff95611966c2da739d8082cbadc4 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/optical_flow/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/optical_flow/__pycache__/_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/optical_flow/__pycache__/_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c5a3ea38b123209a181abcc40cdf985d8a2f9326 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/optical_flow/__pycache__/_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/optical_flow/__pycache__/raft.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/optical_flow/__pycache__/raft.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e7aeaa4d3eeae7ef9b0064c80bc8ac064ed59629 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/optical_flow/__pycache__/raft.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/optical_flow/_utils.py b/lib/python3.12/site-packages/torchvision/models/optical_flow/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..fa2454a27315d6e560dccb6ea2ce6083da03e256 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/optical_flow/_utils.py @@ -0,0 +1,48 @@ +from typing import Optional + +import torch +import torch.nn.functional as F +from torch import Tensor + + +def grid_sample(img: Tensor, absolute_grid: Tensor, mode: str = "bilinear", align_corners: Optional[bool] = None): + """Same as torch's grid_sample, with absolute pixel coordinates instead of normalized coordinates.""" + h, w = img.shape[-2:] + + xgrid, ygrid = absolute_grid.split([1, 1], dim=-1) + xgrid = 2 * xgrid / (w - 1) - 1 + # Adding condition if h > 1 to enable this function be reused in raft-stereo + if h > 1: + ygrid = 2 * ygrid / (h - 1) - 1 + normalized_grid = torch.cat([xgrid, ygrid], dim=-1) + + return F.grid_sample(img, normalized_grid, mode=mode, align_corners=align_corners) + + +def make_coords_grid(batch_size: int, h: int, w: int, device: str = "cpu"): + device = torch.device(device) + coords = torch.meshgrid(torch.arange(h, device=device), torch.arange(w, device=device), indexing="ij") + coords = torch.stack(coords[::-1], dim=0).float() + return coords[None].repeat(batch_size, 1, 1, 1) + + +def upsample_flow(flow, up_mask: Optional[Tensor] = None, factor: int = 8): + """Upsample flow by the input factor (default 8). + + If up_mask is None we just interpolate. + If up_mask is specified, we upsample using a convex combination of its weights. See paper page 8 and appendix B. + Note that in appendix B the picture assumes a downsample factor of 4 instead of 8. + """ + batch_size, num_channels, h, w = flow.shape + new_h, new_w = h * factor, w * factor + + if up_mask is None: + return factor * F.interpolate(flow, size=(new_h, new_w), mode="bilinear", align_corners=True) + + up_mask = up_mask.view(batch_size, 1, 9, factor, factor, h, w) + up_mask = torch.softmax(up_mask, dim=2) # "convex" == weights sum to 1 + + upsampled_flow = F.unfold(factor * flow, kernel_size=3, padding=1).view(batch_size, num_channels, 9, 1, 1, h, w) + upsampled_flow = torch.sum(up_mask * upsampled_flow, dim=2) + + return upsampled_flow.permute(0, 1, 4, 2, 5, 3).reshape(batch_size, num_channels, new_h, new_w) diff --git a/lib/python3.12/site-packages/torchvision/models/optical_flow/raft.py b/lib/python3.12/site-packages/torchvision/models/optical_flow/raft.py new file mode 100644 index 0000000000000000000000000000000000000000..c294777ee6ffc0a9151f76f13bf2bde018580f9e --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/optical_flow/raft.py @@ -0,0 +1,947 @@ +from typing import List, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from torch.nn.modules.batchnorm import BatchNorm2d +from torch.nn.modules.instancenorm import InstanceNorm2d +from torchvision.ops import Conv2dNormActivation + +from ...transforms._presets import OpticalFlow +from ...utils import _log_api_usage_once +from .._api import register_model, Weights, WeightsEnum +from .._utils import handle_legacy_interface +from ._utils import grid_sample, make_coords_grid, upsample_flow + + +__all__ = ( + "RAFT", + "raft_large", + "raft_small", + "Raft_Large_Weights", + "Raft_Small_Weights", +) + + +class ResidualBlock(nn.Module): + """Slightly modified Residual block with extra relu and biases.""" + + def __init__(self, in_channels, out_channels, *, norm_layer, stride=1, always_project: bool = False): + super().__init__() + + # Note regarding bias=True: + # Usually we can pass bias=False in conv layers followed by a norm layer. + # But in the RAFT training reference, the BatchNorm2d layers are only activated for the first dataset, + # and frozen for the rest of the training process (i.e. set as eval()). The bias term is thus still useful + # for the rest of the datasets. Technically, we could remove the bias for other norm layers like Instance norm + # because these aren't frozen, but we don't bother (also, we wouldn't be able to load the original weights). + self.convnormrelu1 = Conv2dNormActivation( + in_channels, out_channels, norm_layer=norm_layer, kernel_size=3, stride=stride, bias=True + ) + self.convnormrelu2 = Conv2dNormActivation( + out_channels, out_channels, norm_layer=norm_layer, kernel_size=3, bias=True + ) + + # make mypy happy + self.downsample: nn.Module + + if stride == 1 and not always_project: + self.downsample = nn.Identity() + else: + self.downsample = Conv2dNormActivation( + in_channels, + out_channels, + norm_layer=norm_layer, + kernel_size=1, + stride=stride, + bias=True, + activation_layer=None, + ) + + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + y = x + y = self.convnormrelu1(y) + y = self.convnormrelu2(y) + + x = self.downsample(x) + + return self.relu(x + y) + + +class BottleneckBlock(nn.Module): + """Slightly modified BottleNeck block (extra relu and biases)""" + + def __init__(self, in_channels, out_channels, *, norm_layer, stride=1): + super().__init__() + + # See note in ResidualBlock for the reason behind bias=True + self.convnormrelu1 = Conv2dNormActivation( + in_channels, out_channels // 4, norm_layer=norm_layer, kernel_size=1, bias=True + ) + self.convnormrelu2 = Conv2dNormActivation( + out_channels // 4, out_channels // 4, norm_layer=norm_layer, kernel_size=3, stride=stride, bias=True + ) + self.convnormrelu3 = Conv2dNormActivation( + out_channels // 4, out_channels, norm_layer=norm_layer, kernel_size=1, bias=True + ) + self.relu = nn.ReLU(inplace=True) + + if stride == 1: + self.downsample = nn.Identity() + else: + self.downsample = Conv2dNormActivation( + in_channels, + out_channels, + norm_layer=norm_layer, + kernel_size=1, + stride=stride, + bias=True, + activation_layer=None, + ) + + def forward(self, x): + y = x + y = self.convnormrelu1(y) + y = self.convnormrelu2(y) + y = self.convnormrelu3(y) + + x = self.downsample(x) + + return self.relu(x + y) + + +class FeatureEncoder(nn.Module): + """The feature encoder, used both as the actual feature encoder, and as the context encoder. + + It must downsample its input by 8. + """ + + def __init__( + self, *, block=ResidualBlock, layers=(64, 64, 96, 128, 256), strides=(2, 1, 2, 2), norm_layer=nn.BatchNorm2d + ): + super().__init__() + + if len(layers) != 5: + raise ValueError(f"The expected number of layers is 5, instead got {len(layers)}") + + # See note in ResidualBlock for the reason behind bias=True + self.convnormrelu = Conv2dNormActivation( + 3, layers[0], norm_layer=norm_layer, kernel_size=7, stride=strides[0], bias=True + ) + + self.layer1 = self._make_2_blocks(block, layers[0], layers[1], norm_layer=norm_layer, first_stride=strides[1]) + self.layer2 = self._make_2_blocks(block, layers[1], layers[2], norm_layer=norm_layer, first_stride=strides[2]) + self.layer3 = self._make_2_blocks(block, layers[2], layers[3], norm_layer=norm_layer, first_stride=strides[3]) + + self.conv = nn.Conv2d(layers[3], layers[4], kernel_size=1) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") + elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d)): + if m.weight is not None: + nn.init.constant_(m.weight, 1) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + num_downsamples = len(list(filter(lambda s: s == 2, strides))) + self.output_dim = layers[-1] + self.downsample_factor = 2**num_downsamples + + def _make_2_blocks(self, block, in_channels, out_channels, norm_layer, first_stride): + block1 = block(in_channels, out_channels, norm_layer=norm_layer, stride=first_stride) + block2 = block(out_channels, out_channels, norm_layer=norm_layer, stride=1) + return nn.Sequential(block1, block2) + + def forward(self, x): + x = self.convnormrelu(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + + x = self.conv(x) + + return x + + +class MotionEncoder(nn.Module): + """The motion encoder, part of the update block. + + Takes the current predicted flow and the correlation features as input and returns an encoded version of these. + """ + + def __init__(self, *, in_channels_corr, corr_layers=(256, 192), flow_layers=(128, 64), out_channels=128): + super().__init__() + + if len(flow_layers) != 2: + raise ValueError(f"The expected number of flow_layers is 2, instead got {len(flow_layers)}") + if len(corr_layers) not in (1, 2): + raise ValueError(f"The number of corr_layers should be 1 or 2, instead got {len(corr_layers)}") + + self.convcorr1 = Conv2dNormActivation(in_channels_corr, corr_layers[0], norm_layer=None, kernel_size=1) + if len(corr_layers) == 2: + self.convcorr2 = Conv2dNormActivation(corr_layers[0], corr_layers[1], norm_layer=None, kernel_size=3) + else: + self.convcorr2 = nn.Identity() + + self.convflow1 = Conv2dNormActivation(2, flow_layers[0], norm_layer=None, kernel_size=7) + self.convflow2 = Conv2dNormActivation(flow_layers[0], flow_layers[1], norm_layer=None, kernel_size=3) + + # out_channels - 2 because we cat the flow (2 channels) at the end + self.conv = Conv2dNormActivation( + corr_layers[-1] + flow_layers[-1], out_channels - 2, norm_layer=None, kernel_size=3 + ) + + self.out_channels = out_channels + + def forward(self, flow, corr_features): + corr = self.convcorr1(corr_features) + corr = self.convcorr2(corr) + + flow_orig = flow + flow = self.convflow1(flow) + flow = self.convflow2(flow) + + corr_flow = torch.cat([corr, flow], dim=1) + corr_flow = self.conv(corr_flow) + return torch.cat([corr_flow, flow_orig], dim=1) + + +class ConvGRU(nn.Module): + """Convolutional Gru unit.""" + + def __init__(self, *, input_size, hidden_size, kernel_size, padding): + super().__init__() + self.convz = nn.Conv2d(hidden_size + input_size, hidden_size, kernel_size=kernel_size, padding=padding) + self.convr = nn.Conv2d(hidden_size + input_size, hidden_size, kernel_size=kernel_size, padding=padding) + self.convq = nn.Conv2d(hidden_size + input_size, hidden_size, kernel_size=kernel_size, padding=padding) + + def forward(self, h, x): + hx = torch.cat([h, x], dim=1) + z = torch.sigmoid(self.convz(hx)) + r = torch.sigmoid(self.convr(hx)) + q = torch.tanh(self.convq(torch.cat([r * h, x], dim=1))) + h = (1 - z) * h + z * q + return h + + +def _pass_through_h(h, _): + # Declared here for torchscript + return h + + +class RecurrentBlock(nn.Module): + """Recurrent block, part of the update block. + + Takes the current hidden state and the concatenation of (motion encoder output, context) as input. + Returns an updated hidden state. + """ + + def __init__(self, *, input_size, hidden_size, kernel_size=((1, 5), (5, 1)), padding=((0, 2), (2, 0))): + super().__init__() + + if len(kernel_size) != len(padding): + raise ValueError( + f"kernel_size should have the same length as padding, instead got len(kernel_size) = {len(kernel_size)} and len(padding) = {len(padding)}" + ) + if len(kernel_size) not in (1, 2): + raise ValueError(f"kernel_size should either 1 or 2, instead got {len(kernel_size)}") + + self.convgru1 = ConvGRU( + input_size=input_size, hidden_size=hidden_size, kernel_size=kernel_size[0], padding=padding[0] + ) + if len(kernel_size) == 2: + self.convgru2 = ConvGRU( + input_size=input_size, hidden_size=hidden_size, kernel_size=kernel_size[1], padding=padding[1] + ) + else: + self.convgru2 = _pass_through_h + + self.hidden_size = hidden_size + + def forward(self, h, x): + h = self.convgru1(h, x) + h = self.convgru2(h, x) + return h + + +class FlowHead(nn.Module): + """Flow head, part of the update block. + + Takes the hidden state of the recurrent unit as input, and outputs the predicted "delta flow". + """ + + def __init__(self, *, in_channels, hidden_size): + super().__init__() + self.conv1 = nn.Conv2d(in_channels, hidden_size, 3, padding=1) + self.conv2 = nn.Conv2d(hidden_size, 2, 3, padding=1) + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + return self.conv2(self.relu(self.conv1(x))) + + +class UpdateBlock(nn.Module): + """The update block which contains the motion encoder, the recurrent block, and the flow head. + + It must expose a ``hidden_state_size`` attribute which is the hidden state size of its recurrent block. + """ + + def __init__(self, *, motion_encoder, recurrent_block, flow_head): + super().__init__() + self.motion_encoder = motion_encoder + self.recurrent_block = recurrent_block + self.flow_head = flow_head + + self.hidden_state_size = recurrent_block.hidden_size + + def forward(self, hidden_state, context, corr_features, flow): + motion_features = self.motion_encoder(flow, corr_features) + x = torch.cat([context, motion_features], dim=1) + + hidden_state = self.recurrent_block(hidden_state, x) + delta_flow = self.flow_head(hidden_state) + return hidden_state, delta_flow + + +class MaskPredictor(nn.Module): + """Mask predictor to be used when upsampling the predicted flow. + + It takes the hidden state of the recurrent unit as input and outputs the mask. + This is not used in the raft-small model. + """ + + def __init__(self, *, in_channels, hidden_size, multiplier=0.25): + super().__init__() + self.convrelu = Conv2dNormActivation(in_channels, hidden_size, norm_layer=None, kernel_size=3) + # 8 * 8 * 9 because the predicted flow is downsampled by 8, from the downsampling of the initial FeatureEncoder, + # and we interpolate with all 9 surrounding neighbors. See paper and appendix B. + self.conv = nn.Conv2d(hidden_size, 8 * 8 * 9, 1, padding=0) + + # In the original code, they use a factor of 0.25 to "downweight the gradients" of that branch. + # See e.g. https://github.com/princeton-vl/RAFT/issues/119#issuecomment-953950419 + # or https://github.com/princeton-vl/RAFT/issues/24. + # It doesn't seem to affect epe significantly and can likely be set to 1. + self.multiplier = multiplier + + def forward(self, x): + x = self.convrelu(x) + x = self.conv(x) + return self.multiplier * x + + +class CorrBlock(nn.Module): + """The correlation block. + + Creates a correlation pyramid with ``num_levels`` levels from the outputs of the feature encoder, + and then indexes from this pyramid to create correlation features. + The "indexing" of a given centroid pixel x' is done by concatenating its surrounding neighbors that + are within a ``radius``, according to the infinity norm (see paper section 3.2). + Note: typo in the paper, it should be infinity norm, not 1-norm. + """ + + def __init__(self, *, num_levels: int = 4, radius: int = 4): + super().__init__() + self.num_levels = num_levels + self.radius = radius + + self.corr_pyramid: List[Tensor] = [torch.tensor(0)] # useless, but torchscript is otherwise confused :') + + # The neighborhood of a centroid pixel x' is {x' + delta, ||delta||_inf <= radius} + # so it's a square surrounding x', and its sides have a length of 2 * radius + 1 + # The paper claims that it's ||.||_1 instead of ||.||_inf but it's a typo: + # https://github.com/princeton-vl/RAFT/issues/122 + self.out_channels = num_levels * (2 * radius + 1) ** 2 + + def build_pyramid(self, fmap1, fmap2): + """Build the correlation pyramid from two feature maps. + + The correlation volume is first computed as the dot product of each pair (pixel_in_fmap1, pixel_in_fmap2) + The last 2 dimensions of the correlation volume are then pooled num_levels times at different resolutions + to build the correlation pyramid. + """ + + if fmap1.shape != fmap2.shape: + raise ValueError( + f"Input feature maps should have the same shape, instead got {fmap1.shape} (fmap1.shape) != {fmap2.shape} (fmap2.shape)" + ) + + # Explaining min_fmap_size below: the fmaps are down-sampled (num_levels - 1) times by a factor of 2. + # The last corr_volume most have at least 2 values (hence the 2* factor), otherwise grid_sample() would + # produce nans in its output. + min_fmap_size = 2 * (2 ** (self.num_levels - 1)) + if any(fmap_size < min_fmap_size for fmap_size in fmap1.shape[-2:]): + raise ValueError( + "Feature maps are too small to be down-sampled by the correlation pyramid. " + f"H and W of feature maps should be at least {min_fmap_size}; got: {fmap1.shape[-2:]}. " + "Remember that input images to the model are downsampled by 8, so that means their " + f"dimensions should be at least 8 * {min_fmap_size} = {8 * min_fmap_size}." + ) + + corr_volume = self._compute_corr_volume(fmap1, fmap2) + + batch_size, h, w, num_channels, _, _ = corr_volume.shape # _, _ = h, w + corr_volume = corr_volume.reshape(batch_size * h * w, num_channels, h, w) + self.corr_pyramid = [corr_volume] + for _ in range(self.num_levels - 1): + corr_volume = F.avg_pool2d(corr_volume, kernel_size=2, stride=2) + self.corr_pyramid.append(corr_volume) + + def index_pyramid(self, centroids_coords): + """Return correlation features by indexing from the pyramid.""" + neighborhood_side_len = 2 * self.radius + 1 # see note in __init__ about out_channels + di = torch.linspace(-self.radius, self.radius, neighborhood_side_len) + dj = torch.linspace(-self.radius, self.radius, neighborhood_side_len) + delta = torch.stack(torch.meshgrid(di, dj, indexing="ij"), dim=-1).to(centroids_coords.device) + delta = delta.view(1, neighborhood_side_len, neighborhood_side_len, 2) + + batch_size, _, h, w = centroids_coords.shape # _ = 2 + centroids_coords = centroids_coords.permute(0, 2, 3, 1).reshape(batch_size * h * w, 1, 1, 2) + + indexed_pyramid = [] + for corr_volume in self.corr_pyramid: + sampling_coords = centroids_coords + delta # end shape is (batch_size * h * w, side_len, side_len, 2) + indexed_corr_volume = grid_sample(corr_volume, sampling_coords, align_corners=True, mode="bilinear").view( + batch_size, h, w, -1 + ) + indexed_pyramid.append(indexed_corr_volume) + centroids_coords = centroids_coords / 2 + + corr_features = torch.cat(indexed_pyramid, dim=-1).permute(0, 3, 1, 2).contiguous() + + expected_output_shape = (batch_size, self.out_channels, h, w) + if corr_features.shape != expected_output_shape: + raise ValueError( + f"Output shape of index pyramid is incorrect. Should be {expected_output_shape}, got {corr_features.shape}" + ) + + return corr_features + + def _compute_corr_volume(self, fmap1, fmap2): + batch_size, num_channels, h, w = fmap1.shape + fmap1 = fmap1.view(batch_size, num_channels, h * w) + fmap2 = fmap2.view(batch_size, num_channels, h * w) + + corr = torch.matmul(fmap1.transpose(1, 2), fmap2) + corr = corr.view(batch_size, h, w, 1, h, w) + return corr / torch.sqrt(torch.tensor(num_channels)) + + +class RAFT(nn.Module): + def __init__(self, *, feature_encoder, context_encoder, corr_block, update_block, mask_predictor=None): + """RAFT model from + `RAFT: Recurrent All Pairs Field Transforms for Optical Flow `_. + + args: + feature_encoder (nn.Module): The feature encoder. It must downsample the input by 8. + Its input is the concatenation of ``image1`` and ``image2``. + context_encoder (nn.Module): The context encoder. It must downsample the input by 8. + Its input is ``image1``. As in the original implementation, its output will be split into 2 parts: + + - one part will be used as the actual "context", passed to the recurrent unit of the ``update_block`` + - one part will be used to initialize the hidden state of the recurrent unit of + the ``update_block`` + + These 2 parts are split according to the ``hidden_state_size`` of the ``update_block``, so the output + of the ``context_encoder`` must be strictly greater than ``hidden_state_size``. + + corr_block (nn.Module): The correlation block, which creates a correlation pyramid from the output of the + ``feature_encoder``, and then indexes from this pyramid to create correlation features. It must expose + 2 methods: + + - a ``build_pyramid`` method that takes ``feature_map_1`` and ``feature_map_2`` as input (these are the + output of the ``feature_encoder``). + - a ``index_pyramid`` method that takes the coordinates of the centroid pixels as input, and returns + the correlation features. See paper section 3.2. + + It must expose an ``out_channels`` attribute. + + update_block (nn.Module): The update block, which contains the motion encoder, the recurrent unit, and the + flow head. It takes as input the hidden state of its recurrent unit, the context, the correlation + features, and the current predicted flow. It outputs an updated hidden state, and the ``delta_flow`` + prediction (see paper appendix A). It must expose a ``hidden_state_size`` attribute. + mask_predictor (nn.Module, optional): Predicts the mask that will be used to upsample the predicted flow. + The output channel must be 8 * 8 * 9 - see paper section 3.3, and Appendix B. + If ``None`` (default), the flow is upsampled using interpolation. + """ + super().__init__() + _log_api_usage_once(self) + + self.feature_encoder = feature_encoder + self.context_encoder = context_encoder + self.corr_block = corr_block + self.update_block = update_block + + self.mask_predictor = mask_predictor + + if not hasattr(self.update_block, "hidden_state_size"): + raise ValueError("The update_block parameter should expose a 'hidden_state_size' attribute.") + + def forward(self, image1, image2, num_flow_updates: int = 12): + + batch_size, _, h, w = image1.shape + if (h, w) != image2.shape[-2:]: + raise ValueError(f"input images should have the same shape, instead got ({h}, {w}) != {image2.shape[-2:]}") + if not (h % 8 == 0) and (w % 8 == 0): + raise ValueError(f"input image H and W should be divisible by 8, instead got {h} (h) and {w} (w)") + + fmaps = self.feature_encoder(torch.cat([image1, image2], dim=0)) + fmap1, fmap2 = torch.chunk(fmaps, chunks=2, dim=0) + if fmap1.shape[-2:] != (h // 8, w // 8): + raise ValueError("The feature encoder should downsample H and W by 8") + + self.corr_block.build_pyramid(fmap1, fmap2) + + context_out = self.context_encoder(image1) + if context_out.shape[-2:] != (h // 8, w // 8): + raise ValueError("The context encoder should downsample H and W by 8") + + # As in the original paper, the actual output of the context encoder is split in 2 parts: + # - one part is used to initialize the hidden state of the recurent units of the update block + # - the rest is the "actual" context. + hidden_state_size = self.update_block.hidden_state_size + out_channels_context = context_out.shape[1] - hidden_state_size + if out_channels_context <= 0: + raise ValueError( + f"The context encoder outputs {context_out.shape[1]} channels, but it should have at strictly more than hidden_state={hidden_state_size} channels" + ) + hidden_state, context = torch.split(context_out, [hidden_state_size, out_channels_context], dim=1) + hidden_state = torch.tanh(hidden_state) + context = F.relu(context) + + coords0 = make_coords_grid(batch_size, h // 8, w // 8).to(fmap1.device) + coords1 = make_coords_grid(batch_size, h // 8, w // 8).to(fmap1.device) + + flow_predictions = [] + for _ in range(num_flow_updates): + coords1 = coords1.detach() # Don't backpropagate gradients through this branch, see paper + corr_features = self.corr_block.index_pyramid(centroids_coords=coords1) + + flow = coords1 - coords0 + hidden_state, delta_flow = self.update_block(hidden_state, context, corr_features, flow) + + coords1 = coords1 + delta_flow + + up_mask = None if self.mask_predictor is None else self.mask_predictor(hidden_state) + upsampled_flow = upsample_flow(flow=(coords1 - coords0), up_mask=up_mask) + flow_predictions.append(upsampled_flow) + + return flow_predictions + + +_COMMON_META = { + "min_size": (128, 128), +} + + +class Raft_Large_Weights(WeightsEnum): + """The metrics reported here are as follows. + + ``epe`` is the "end-point-error" and indicates how far (in pixels) the + predicted flow is from its true value. This is averaged over all pixels + of all images. ``per_image_epe`` is similar, but the average is different: + the epe is first computed on each image independently, and then averaged + over all images. This corresponds to "Fl-epe" (sometimes written "F1-epe") + in the original paper, and it's only used on Kitti. ``fl-all`` is also a + Kitti-specific metric, defined by the author of the dataset and used for the + Kitti leaderboard. It corresponds to the average of pixels whose epe is + either <3px, or <5% of flow's 2-norm. + """ + + C_T_V1 = Weights( + # Weights ported from https://github.com/princeton-vl/RAFT + url="https://download.pytorch.org/models/raft_large_C_T_V1-22a6c225.pth", + transforms=OpticalFlow, + meta={ + **_COMMON_META, + "num_params": 5257536, + "recipe": "https://github.com/princeton-vl/RAFT", + "_metrics": { + "Sintel-Train-Cleanpass": {"epe": 1.4411}, + "Sintel-Train-Finalpass": {"epe": 2.7894}, + "Kitti-Train": {"per_image_epe": 5.0172, "fl_all": 17.4506}, + }, + "_ops": 211.007, + "_file_size": 20.129, + "_docs": """These weights were ported from the original paper. They + are trained on :class:`~torchvision.datasets.FlyingChairs` + + :class:`~torchvision.datasets.FlyingThings3D`.""", + }, + ) + + C_T_V2 = Weights( + url="https://download.pytorch.org/models/raft_large_C_T_V2-1bb1363a.pth", + transforms=OpticalFlow, + meta={ + **_COMMON_META, + "num_params": 5257536, + "recipe": "https://github.com/pytorch/vision/tree/main/references/optical_flow", + "_metrics": { + "Sintel-Train-Cleanpass": {"epe": 1.3822}, + "Sintel-Train-Finalpass": {"epe": 2.7161}, + "Kitti-Train": {"per_image_epe": 4.5118, "fl_all": 16.0679}, + }, + "_ops": 211.007, + "_file_size": 20.129, + "_docs": """These weights were trained from scratch on + :class:`~torchvision.datasets.FlyingChairs` + + :class:`~torchvision.datasets.FlyingThings3D`.""", + }, + ) + + C_T_SKHT_V1 = Weights( + # Weights ported from https://github.com/princeton-vl/RAFT + url="https://download.pytorch.org/models/raft_large_C_T_SKHT_V1-0b8c9e55.pth", + transforms=OpticalFlow, + meta={ + **_COMMON_META, + "num_params": 5257536, + "recipe": "https://github.com/princeton-vl/RAFT", + "_metrics": { + "Sintel-Test-Cleanpass": {"epe": 1.94}, + "Sintel-Test-Finalpass": {"epe": 3.18}, + }, + "_ops": 211.007, + "_file_size": 20.129, + "_docs": """ + These weights were ported from the original paper. They are + trained on :class:`~torchvision.datasets.FlyingChairs` + + :class:`~torchvision.datasets.FlyingThings3D` and fine-tuned on + Sintel. The Sintel fine-tuning step is a combination of + :class:`~torchvision.datasets.Sintel`, + :class:`~torchvision.datasets.KittiFlow`, + :class:`~torchvision.datasets.HD1K`, and + :class:`~torchvision.datasets.FlyingThings3D` (clean pass). + """, + }, + ) + + C_T_SKHT_V2 = Weights( + url="https://download.pytorch.org/models/raft_large_C_T_SKHT_V2-ff5fadd5.pth", + transforms=OpticalFlow, + meta={ + **_COMMON_META, + "num_params": 5257536, + "recipe": "https://github.com/pytorch/vision/tree/main/references/optical_flow", + "_metrics": { + "Sintel-Test-Cleanpass": {"epe": 1.819}, + "Sintel-Test-Finalpass": {"epe": 3.067}, + }, + "_ops": 211.007, + "_file_size": 20.129, + "_docs": """ + These weights were trained from scratch. They are + pre-trained on :class:`~torchvision.datasets.FlyingChairs` + + :class:`~torchvision.datasets.FlyingThings3D` and then + fine-tuned on Sintel. The Sintel fine-tuning step is a + combination of :class:`~torchvision.datasets.Sintel`, + :class:`~torchvision.datasets.KittiFlow`, + :class:`~torchvision.datasets.HD1K`, and + :class:`~torchvision.datasets.FlyingThings3D` (clean pass). + """, + }, + ) + + C_T_SKHT_K_V1 = Weights( + # Weights ported from https://github.com/princeton-vl/RAFT + url="https://download.pytorch.org/models/raft_large_C_T_SKHT_K_V1-4a6a5039.pth", + transforms=OpticalFlow, + meta={ + **_COMMON_META, + "num_params": 5257536, + "recipe": "https://github.com/princeton-vl/RAFT", + "_metrics": { + "Kitti-Test": {"fl_all": 5.10}, + }, + "_ops": 211.007, + "_file_size": 20.129, + "_docs": """ + These weights were ported from the original paper. They are + pre-trained on :class:`~torchvision.datasets.FlyingChairs` + + :class:`~torchvision.datasets.FlyingThings3D`, + fine-tuned on Sintel, and then fine-tuned on + :class:`~torchvision.datasets.KittiFlow`. The Sintel fine-tuning + step was described above. + """, + }, + ) + + C_T_SKHT_K_V2 = Weights( + url="https://download.pytorch.org/models/raft_large_C_T_SKHT_K_V2-b5c70766.pth", + transforms=OpticalFlow, + meta={ + **_COMMON_META, + "num_params": 5257536, + "recipe": "https://github.com/pytorch/vision/tree/main/references/optical_flow", + "_metrics": { + "Kitti-Test": {"fl_all": 5.19}, + }, + "_ops": 211.007, + "_file_size": 20.129, + "_docs": """ + These weights were trained from scratch. They are + pre-trained on :class:`~torchvision.datasets.FlyingChairs` + + :class:`~torchvision.datasets.FlyingThings3D`, + fine-tuned on Sintel, and then fine-tuned on + :class:`~torchvision.datasets.KittiFlow`. The Sintel fine-tuning + step was described above. + """, + }, + ) + + DEFAULT = C_T_SKHT_V2 + + +class Raft_Small_Weights(WeightsEnum): + """The metrics reported here are as follows. + + ``epe`` is the "end-point-error" and indicates how far (in pixels) the + predicted flow is from its true value. This is averaged over all pixels + of all images. ``per_image_epe`` is similar, but the average is different: + the epe is first computed on each image independently, and then averaged + over all images. This corresponds to "Fl-epe" (sometimes written "F1-epe") + in the original paper, and it's only used on Kitti. ``fl-all`` is also a + Kitti-specific metric, defined by the author of the dataset and used for the + Kitti leaderboard. It corresponds to the average of pixels whose epe is + either <3px, or <5% of flow's 2-norm. + """ + + C_T_V1 = Weights( + # Weights ported from https://github.com/princeton-vl/RAFT + url="https://download.pytorch.org/models/raft_small_C_T_V1-ad48884c.pth", + transforms=OpticalFlow, + meta={ + **_COMMON_META, + "num_params": 990162, + "recipe": "https://github.com/princeton-vl/RAFT", + "_metrics": { + "Sintel-Train-Cleanpass": {"epe": 2.1231}, + "Sintel-Train-Finalpass": {"epe": 3.2790}, + "Kitti-Train": {"per_image_epe": 7.6557, "fl_all": 25.2801}, + }, + "_ops": 47.655, + "_file_size": 3.821, + "_docs": """These weights were ported from the original paper. They + are trained on :class:`~torchvision.datasets.FlyingChairs` + + :class:`~torchvision.datasets.FlyingThings3D`.""", + }, + ) + C_T_V2 = Weights( + url="https://download.pytorch.org/models/raft_small_C_T_V2-01064c6d.pth", + transforms=OpticalFlow, + meta={ + **_COMMON_META, + "num_params": 990162, + "recipe": "https://github.com/pytorch/vision/tree/main/references/optical_flow", + "_metrics": { + "Sintel-Train-Cleanpass": {"epe": 1.9901}, + "Sintel-Train-Finalpass": {"epe": 3.2831}, + "Kitti-Train": {"per_image_epe": 7.5978, "fl_all": 25.2369}, + }, + "_ops": 47.655, + "_file_size": 3.821, + "_docs": """These weights were trained from scratch on + :class:`~torchvision.datasets.FlyingChairs` + + :class:`~torchvision.datasets.FlyingThings3D`.""", + }, + ) + + DEFAULT = C_T_V2 + + +def _raft( + *, + weights=None, + progress=False, + # Feature encoder + feature_encoder_layers, + feature_encoder_block, + feature_encoder_norm_layer, + # Context encoder + context_encoder_layers, + context_encoder_block, + context_encoder_norm_layer, + # Correlation block + corr_block_num_levels, + corr_block_radius, + # Motion encoder + motion_encoder_corr_layers, + motion_encoder_flow_layers, + motion_encoder_out_channels, + # Recurrent block + recurrent_block_hidden_state_size, + recurrent_block_kernel_size, + recurrent_block_padding, + # Flow Head + flow_head_hidden_size, + # Mask predictor + use_mask_predictor, + **kwargs, +): + feature_encoder = kwargs.pop("feature_encoder", None) or FeatureEncoder( + block=feature_encoder_block, layers=feature_encoder_layers, norm_layer=feature_encoder_norm_layer + ) + context_encoder = kwargs.pop("context_encoder", None) or FeatureEncoder( + block=context_encoder_block, layers=context_encoder_layers, norm_layer=context_encoder_norm_layer + ) + + corr_block = kwargs.pop("corr_block", None) or CorrBlock(num_levels=corr_block_num_levels, radius=corr_block_radius) + + update_block = kwargs.pop("update_block", None) + if update_block is None: + motion_encoder = MotionEncoder( + in_channels_corr=corr_block.out_channels, + corr_layers=motion_encoder_corr_layers, + flow_layers=motion_encoder_flow_layers, + out_channels=motion_encoder_out_channels, + ) + + # See comments in forward pass of RAFT class about why we split the output of the context encoder + out_channels_context = context_encoder_layers[-1] - recurrent_block_hidden_state_size + recurrent_block = RecurrentBlock( + input_size=motion_encoder.out_channels + out_channels_context, + hidden_size=recurrent_block_hidden_state_size, + kernel_size=recurrent_block_kernel_size, + padding=recurrent_block_padding, + ) + + flow_head = FlowHead(in_channels=recurrent_block_hidden_state_size, hidden_size=flow_head_hidden_size) + + update_block = UpdateBlock(motion_encoder=motion_encoder, recurrent_block=recurrent_block, flow_head=flow_head) + + mask_predictor = kwargs.pop("mask_predictor", None) + if mask_predictor is None and use_mask_predictor: + mask_predictor = MaskPredictor( + in_channels=recurrent_block_hidden_state_size, + hidden_size=256, + multiplier=0.25, # See comment in MaskPredictor about this + ) + + model = RAFT( + feature_encoder=feature_encoder, + context_encoder=context_encoder, + corr_block=corr_block, + update_block=update_block, + mask_predictor=mask_predictor, + **kwargs, # not really needed, all params should be consumed by now + ) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +@register_model() +@handle_legacy_interface(weights=("pretrained", Raft_Large_Weights.C_T_SKHT_V2)) +def raft_large(*, weights: Optional[Raft_Large_Weights] = None, progress=True, **kwargs) -> RAFT: + """RAFT model from + `RAFT: Recurrent All Pairs Field Transforms for Optical Flow `_. + + Please see the example below for a tutorial on how to use this model. + + Args: + weights(:class:`~torchvision.models.optical_flow.Raft_Large_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.optical_flow.Raft_Large_Weights` + below for more details, and possible values. By default, no + pre-trained weights are used. + progress (bool): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.optical_flow.RAFT`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.optical_flow.Raft_Large_Weights + :members: + """ + + weights = Raft_Large_Weights.verify(weights) + + return _raft( + weights=weights, + progress=progress, + # Feature encoder + feature_encoder_layers=(64, 64, 96, 128, 256), + feature_encoder_block=ResidualBlock, + feature_encoder_norm_layer=InstanceNorm2d, + # Context encoder + context_encoder_layers=(64, 64, 96, 128, 256), + context_encoder_block=ResidualBlock, + context_encoder_norm_layer=BatchNorm2d, + # Correlation block + corr_block_num_levels=4, + corr_block_radius=4, + # Motion encoder + motion_encoder_corr_layers=(256, 192), + motion_encoder_flow_layers=(128, 64), + motion_encoder_out_channels=128, + # Recurrent block + recurrent_block_hidden_state_size=128, + recurrent_block_kernel_size=((1, 5), (5, 1)), + recurrent_block_padding=((0, 2), (2, 0)), + # Flow head + flow_head_hidden_size=256, + # Mask predictor + use_mask_predictor=True, + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", Raft_Small_Weights.C_T_V2)) +def raft_small(*, weights: Optional[Raft_Small_Weights] = None, progress=True, **kwargs) -> RAFT: + """RAFT "small" model from + `RAFT: Recurrent All Pairs Field Transforms for Optical Flow `__. + + Please see the example below for a tutorial on how to use this model. + + Args: + weights(:class:`~torchvision.models.optical_flow.Raft_Small_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.optical_flow.Raft_Small_Weights` + below for more details, and possible values. By default, no + pre-trained weights are used. + progress (bool): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.optical_flow.RAFT`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.optical_flow.Raft_Small_Weights + :members: + """ + weights = Raft_Small_Weights.verify(weights) + + return _raft( + weights=weights, + progress=progress, + # Feature encoder + feature_encoder_layers=(32, 32, 64, 96, 128), + feature_encoder_block=BottleneckBlock, + feature_encoder_norm_layer=InstanceNorm2d, + # Context encoder + context_encoder_layers=(32, 32, 64, 96, 160), + context_encoder_block=BottleneckBlock, + context_encoder_norm_layer=None, + # Correlation block + corr_block_num_levels=4, + corr_block_radius=3, + # Motion encoder + motion_encoder_corr_layers=(96,), + motion_encoder_flow_layers=(64, 32), + motion_encoder_out_channels=82, + # Recurrent block + recurrent_block_hidden_state_size=96, + recurrent_block_kernel_size=(3,), + recurrent_block_padding=(1,), + # Flow head + flow_head_hidden_size=128, + # Mask predictor + use_mask_predictor=False, + **kwargs, + ) diff --git a/lib/python3.12/site-packages/torchvision/models/quantization/__init__.py b/lib/python3.12/site-packages/torchvision/models/quantization/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..da8bbba3567b0b9110429354d89b65ec679a2fd5 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/quantization/__init__.py @@ -0,0 +1,5 @@ +from .googlenet import * +from .inception import * +from .mobilenet import * +from .resnet import * +from .shufflenetv2 import * diff --git a/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd65a3d27ea5e8603371d248a8bf769ece032bdb Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/googlenet.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/googlenet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..483e9f2f9a87a7f06407663e2fc3c86acce36bda Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/googlenet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/inception.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/inception.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0da0f44981d4c879b4e02b101870931536061b0d Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/inception.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/mobilenet.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/mobilenet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bdfa38eb4e5a37bf3a1fc577b98544fba0107dd9 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/mobilenet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/mobilenetv2.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/mobilenetv2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..56bb17d37e501655fcce77b654a5b401a24d0801 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/mobilenetv2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/mobilenetv3.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/mobilenetv3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4d79bdb81ec0d6654b80d8a1ba6ce10491c6b6a6 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/mobilenetv3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/resnet.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/resnet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ce59bb159fa9aa8819e7a8136b9be0df938ff56b Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/resnet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/shufflenetv2.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/shufflenetv2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..11fc151003b6592ab10ce91cc09fa475c7d22833 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/shufflenetv2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..906c5c88056c07f77e19caaa4143c5b1d7eb48cb Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/quantization/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/quantization/googlenet.py b/lib/python3.12/site-packages/torchvision/models/quantization/googlenet.py new file mode 100644 index 0000000000000000000000000000000000000000..30ef3356ba13108b9bdc4c90a9ab4cb7f92e445a --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/quantization/googlenet.py @@ -0,0 +1,210 @@ +import warnings +from functools import partial +from typing import Any, Optional, Union + +import torch +import torch.nn as nn +from torch import Tensor +from torch.nn import functional as F + +from ...transforms._presets import ImageClassification +from .._api import register_model, Weights, WeightsEnum +from .._meta import _IMAGENET_CATEGORIES +from .._utils import _ovewrite_named_param, handle_legacy_interface +from ..googlenet import BasicConv2d, GoogLeNet, GoogLeNet_Weights, GoogLeNetOutputs, Inception, InceptionAux +from .utils import _fuse_modules, _replace_relu, quantize_model + + +__all__ = [ + "QuantizableGoogLeNet", + "GoogLeNet_QuantizedWeights", + "googlenet", +] + + +class QuantizableBasicConv2d(BasicConv2d): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.relu = nn.ReLU() + + def forward(self, x: Tensor) -> Tensor: + x = self.conv(x) + x = self.bn(x) + x = self.relu(x) + return x + + def fuse_model(self, is_qat: Optional[bool] = None) -> None: + _fuse_modules(self, ["conv", "bn", "relu"], is_qat, inplace=True) + + +class QuantizableInception(Inception): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, conv_block=QuantizableBasicConv2d, **kwargs) # type: ignore[misc] + self.cat = nn.quantized.FloatFunctional() + + def forward(self, x: Tensor) -> Tensor: + outputs = self._forward(x) + return self.cat.cat(outputs, 1) + + +class QuantizableInceptionAux(InceptionAux): + # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659 + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, conv_block=QuantizableBasicConv2d, **kwargs) # type: ignore[misc] + self.relu = nn.ReLU() + + def forward(self, x: Tensor) -> Tensor: + # aux1: N x 512 x 14 x 14, aux2: N x 528 x 14 x 14 + x = F.adaptive_avg_pool2d(x, (4, 4)) + # aux1: N x 512 x 4 x 4, aux2: N x 528 x 4 x 4 + x = self.conv(x) + # N x 128 x 4 x 4 + x = torch.flatten(x, 1) + # N x 2048 + x = self.relu(self.fc1(x)) + # N x 1024 + x = self.dropout(x) + # N x 1024 + x = self.fc2(x) + # N x 1000 (num_classes) + + return x + + +class QuantizableGoogLeNet(GoogLeNet): + # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659 + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__( # type: ignore[misc] + *args, blocks=[QuantizableBasicConv2d, QuantizableInception, QuantizableInceptionAux], **kwargs + ) + self.quant = torch.ao.quantization.QuantStub() + self.dequant = torch.ao.quantization.DeQuantStub() + + def forward(self, x: Tensor) -> GoogLeNetOutputs: + x = self._transform_input(x) + x = self.quant(x) + x, aux1, aux2 = self._forward(x) + x = self.dequant(x) + aux_defined = self.training and self.aux_logits + if torch.jit.is_scripting(): + if not aux_defined: + warnings.warn("Scripted QuantizableGoogleNet always returns GoogleNetOutputs Tuple") + return GoogLeNetOutputs(x, aux2, aux1) + else: + return self.eager_outputs(x, aux2, aux1) + + def fuse_model(self, is_qat: Optional[bool] = None) -> None: + r"""Fuse conv/bn/relu modules in googlenet model + + Fuse conv+bn+relu/ conv+relu/conv+bn modules to prepare for quantization. + Model is modified in place. Note that this operation does not change numerics + and the model after modification is in floating point + """ + + for m in self.modules(): + if type(m) is QuantizableBasicConv2d: + m.fuse_model(is_qat) + + +class GoogLeNet_QuantizedWeights(WeightsEnum): + IMAGENET1K_FBGEMM_V1 = Weights( + url="https://download.pytorch.org/models/quantized/googlenet_fbgemm-c81f6644.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + "num_params": 6624904, + "min_size": (15, 15), + "categories": _IMAGENET_CATEGORIES, + "backend": "fbgemm", + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#post-training-quantized-models", + "unquantized": GoogLeNet_Weights.IMAGENET1K_V1, + "_metrics": { + "ImageNet-1K": { + "acc@1": 69.826, + "acc@5": 89.404, + } + }, + "_ops": 1.498, + "_file_size": 12.618, + "_docs": """ + These weights were produced by doing Post Training Quantization (eager mode) on top of the unquantized + weights listed below. + """, + }, + ) + DEFAULT = IMAGENET1K_FBGEMM_V1 + + +@register_model(name="quantized_googlenet") +@handle_legacy_interface( + weights=( + "pretrained", + lambda kwargs: GoogLeNet_QuantizedWeights.IMAGENET1K_FBGEMM_V1 + if kwargs.get("quantize", False) + else GoogLeNet_Weights.IMAGENET1K_V1, + ) +) +def googlenet( + *, + weights: Optional[Union[GoogLeNet_QuantizedWeights, GoogLeNet_Weights]] = None, + progress: bool = True, + quantize: bool = False, + **kwargs: Any, +) -> QuantizableGoogLeNet: + """GoogLeNet (Inception v1) model architecture from `Going Deeper with Convolutions `__. + + .. note:: + Note that ``quantize = True`` returns a quantized model with 8 bit + weights. Quantized models only support inference and run on CPUs. + GPU inference is not yet supported. + + Args: + weights (:class:`~torchvision.models.quantization.GoogLeNet_QuantizedWeights` or :class:`~torchvision.models.GoogLeNet_Weights`, optional): The + pretrained weights for the model. See + :class:`~torchvision.models.quantization.GoogLeNet_QuantizedWeights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + quantize (bool, optional): If True, return a quantized version of the model. Default is False. + **kwargs: parameters passed to the ``torchvision.models.quantization.QuantizableGoogLeNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.quantization.GoogLeNet_QuantizedWeights + :members: + + .. autoclass:: torchvision.models.GoogLeNet_Weights + :members: + :noindex: + """ + weights = (GoogLeNet_QuantizedWeights if quantize else GoogLeNet_Weights).verify(weights) + + original_aux_logits = kwargs.get("aux_logits", False) + if weights is not None: + if "transform_input" not in kwargs: + _ovewrite_named_param(kwargs, "transform_input", True) + _ovewrite_named_param(kwargs, "aux_logits", True) + _ovewrite_named_param(kwargs, "init_weights", False) + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + if "backend" in weights.meta: + _ovewrite_named_param(kwargs, "backend", weights.meta["backend"]) + backend = kwargs.pop("backend", "fbgemm") + + model = QuantizableGoogLeNet(**kwargs) + _replace_relu(model) + if quantize: + quantize_model(model, backend) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + if not original_aux_logits: + model.aux_logits = False + model.aux1 = None # type: ignore[assignment] + model.aux2 = None # type: ignore[assignment] + else: + warnings.warn( + "auxiliary heads in the pretrained googlenet model are NOT pretrained, so make sure to train them" + ) + + return model diff --git a/lib/python3.12/site-packages/torchvision/models/quantization/inception.py b/lib/python3.12/site-packages/torchvision/models/quantization/inception.py new file mode 100644 index 0000000000000000000000000000000000000000..75c126697e99befd6ae7d3c1ee88fb8542e06d31 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/quantization/inception.py @@ -0,0 +1,273 @@ +import warnings +from functools import partial +from typing import Any, List, Optional, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from torchvision.models import inception as inception_module +from torchvision.models.inception import Inception_V3_Weights, InceptionOutputs + +from ...transforms._presets import ImageClassification +from .._api import register_model, Weights, WeightsEnum +from .._meta import _IMAGENET_CATEGORIES +from .._utils import _ovewrite_named_param, handle_legacy_interface +from .utils import _fuse_modules, _replace_relu, quantize_model + + +__all__ = [ + "QuantizableInception3", + "Inception_V3_QuantizedWeights", + "inception_v3", +] + + +class QuantizableBasicConv2d(inception_module.BasicConv2d): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.relu = nn.ReLU() + + def forward(self, x: Tensor) -> Tensor: + x = self.conv(x) + x = self.bn(x) + x = self.relu(x) + return x + + def fuse_model(self, is_qat: Optional[bool] = None) -> None: + _fuse_modules(self, ["conv", "bn", "relu"], is_qat, inplace=True) + + +class QuantizableInceptionA(inception_module.InceptionA): + # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659 + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, conv_block=QuantizableBasicConv2d, **kwargs) # type: ignore[misc] + self.myop = nn.quantized.FloatFunctional() + + def forward(self, x: Tensor) -> Tensor: + outputs = self._forward(x) + return self.myop.cat(outputs, 1) + + +class QuantizableInceptionB(inception_module.InceptionB): + # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659 + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, conv_block=QuantizableBasicConv2d, **kwargs) # type: ignore[misc] + self.myop = nn.quantized.FloatFunctional() + + def forward(self, x: Tensor) -> Tensor: + outputs = self._forward(x) + return self.myop.cat(outputs, 1) + + +class QuantizableInceptionC(inception_module.InceptionC): + # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659 + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, conv_block=QuantizableBasicConv2d, **kwargs) # type: ignore[misc] + self.myop = nn.quantized.FloatFunctional() + + def forward(self, x: Tensor) -> Tensor: + outputs = self._forward(x) + return self.myop.cat(outputs, 1) + + +class QuantizableInceptionD(inception_module.InceptionD): + # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659 + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, conv_block=QuantizableBasicConv2d, **kwargs) # type: ignore[misc] + self.myop = nn.quantized.FloatFunctional() + + def forward(self, x: Tensor) -> Tensor: + outputs = self._forward(x) + return self.myop.cat(outputs, 1) + + +class QuantizableInceptionE(inception_module.InceptionE): + # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659 + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, conv_block=QuantizableBasicConv2d, **kwargs) # type: ignore[misc] + self.myop1 = nn.quantized.FloatFunctional() + self.myop2 = nn.quantized.FloatFunctional() + self.myop3 = nn.quantized.FloatFunctional() + + def _forward(self, x: Tensor) -> List[Tensor]: + branch1x1 = self.branch1x1(x) + + branch3x3 = self.branch3x3_1(x) + branch3x3 = [self.branch3x3_2a(branch3x3), self.branch3x3_2b(branch3x3)] + branch3x3 = self.myop1.cat(branch3x3, 1) + + branch3x3dbl = self.branch3x3dbl_1(x) + branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl) + branch3x3dbl = [ + self.branch3x3dbl_3a(branch3x3dbl), + self.branch3x3dbl_3b(branch3x3dbl), + ] + branch3x3dbl = self.myop2.cat(branch3x3dbl, 1) + + branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1) + branch_pool = self.branch_pool(branch_pool) + + outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool] + return outputs + + def forward(self, x: Tensor) -> Tensor: + outputs = self._forward(x) + return self.myop3.cat(outputs, 1) + + +class QuantizableInceptionAux(inception_module.InceptionAux): + # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659 + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, conv_block=QuantizableBasicConv2d, **kwargs) # type: ignore[misc] + + +class QuantizableInception3(inception_module.Inception3): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__( # type: ignore[misc] + *args, + inception_blocks=[ + QuantizableBasicConv2d, + QuantizableInceptionA, + QuantizableInceptionB, + QuantizableInceptionC, + QuantizableInceptionD, + QuantizableInceptionE, + QuantizableInceptionAux, + ], + **kwargs, + ) + self.quant = torch.ao.quantization.QuantStub() + self.dequant = torch.ao.quantization.DeQuantStub() + + def forward(self, x: Tensor) -> InceptionOutputs: + x = self._transform_input(x) + x = self.quant(x) + x, aux = self._forward(x) + x = self.dequant(x) + aux_defined = self.training and self.aux_logits + if torch.jit.is_scripting(): + if not aux_defined: + warnings.warn("Scripted QuantizableInception3 always returns QuantizableInception3 Tuple") + return InceptionOutputs(x, aux) + else: + return self.eager_outputs(x, aux) + + def fuse_model(self, is_qat: Optional[bool] = None) -> None: + r"""Fuse conv/bn/relu modules in inception model + + Fuse conv+bn+relu/ conv+relu/conv+bn modules to prepare for quantization. + Model is modified in place. Note that this operation does not change numerics + and the model after modification is in floating point + """ + + for m in self.modules(): + if type(m) is QuantizableBasicConv2d: + m.fuse_model(is_qat) + + +class Inception_V3_QuantizedWeights(WeightsEnum): + IMAGENET1K_FBGEMM_V1 = Weights( + url="https://download.pytorch.org/models/quantized/inception_v3_google_fbgemm-a2837893.pth", + transforms=partial(ImageClassification, crop_size=299, resize_size=342), + meta={ + "num_params": 27161264, + "min_size": (75, 75), + "categories": _IMAGENET_CATEGORIES, + "backend": "fbgemm", + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#post-training-quantized-models", + "unquantized": Inception_V3_Weights.IMAGENET1K_V1, + "_metrics": { + "ImageNet-1K": { + "acc@1": 77.176, + "acc@5": 93.354, + } + }, + "_ops": 5.713, + "_file_size": 23.146, + "_docs": """ + These weights were produced by doing Post Training Quantization (eager mode) on top of the unquantized + weights listed below. + """, + }, + ) + DEFAULT = IMAGENET1K_FBGEMM_V1 + + +@register_model(name="quantized_inception_v3") +@handle_legacy_interface( + weights=( + "pretrained", + lambda kwargs: Inception_V3_QuantizedWeights.IMAGENET1K_FBGEMM_V1 + if kwargs.get("quantize", False) + else Inception_V3_Weights.IMAGENET1K_V1, + ) +) +def inception_v3( + *, + weights: Optional[Union[Inception_V3_QuantizedWeights, Inception_V3_Weights]] = None, + progress: bool = True, + quantize: bool = False, + **kwargs: Any, +) -> QuantizableInception3: + r"""Inception v3 model architecture from + `Rethinking the Inception Architecture for Computer Vision `__. + + .. note:: + **Important**: In contrast to the other models the inception_v3 expects tensors with a size of + N x 3 x 299 x 299, so ensure your images are sized accordingly. + + .. note:: + Note that ``quantize = True`` returns a quantized model with 8 bit + weights. Quantized models only support inference and run on CPUs. + GPU inference is not yet supported. + + Args: + weights (:class:`~torchvision.models.quantization.Inception_V3_QuantizedWeights` or :class:`~torchvision.models.Inception_V3_Weights`, optional): The pretrained + weights for the model. See + :class:`~torchvision.models.quantization.Inception_V3_QuantizedWeights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. + Default is True. + quantize (bool, optional): If True, return a quantized version of the model. + Default is False. + **kwargs: parameters passed to the ``torchvision.models.quantization.QuantizableInception3`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.quantization.Inception_V3_QuantizedWeights + :members: + + .. autoclass:: torchvision.models.Inception_V3_Weights + :members: + :noindex: + """ + weights = (Inception_V3_QuantizedWeights if quantize else Inception_V3_Weights).verify(weights) + + original_aux_logits = kwargs.get("aux_logits", False) + if weights is not None: + if "transform_input" not in kwargs: + _ovewrite_named_param(kwargs, "transform_input", True) + _ovewrite_named_param(kwargs, "aux_logits", True) + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + if "backend" in weights.meta: + _ovewrite_named_param(kwargs, "backend", weights.meta["backend"]) + backend = kwargs.pop("backend", "fbgemm") + + model = QuantizableInception3(**kwargs) + _replace_relu(model) + if quantize: + quantize_model(model, backend) + + if weights is not None: + if quantize and not original_aux_logits: + model.aux_logits = False + model.AuxLogits = None + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + if not quantize and not original_aux_logits: + model.aux_logits = False + model.AuxLogits = None + + return model diff --git a/lib/python3.12/site-packages/torchvision/models/quantization/mobilenet.py b/lib/python3.12/site-packages/torchvision/models/quantization/mobilenet.py new file mode 100644 index 0000000000000000000000000000000000000000..0a270d14d3a4ad9eda62b68c2c01e9fdb710ef38 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/quantization/mobilenet.py @@ -0,0 +1,6 @@ +from .mobilenetv2 import * # noqa: F401, F403 +from .mobilenetv3 import * # noqa: F401, F403 +from .mobilenetv2 import __all__ as mv2_all +from .mobilenetv3 import __all__ as mv3_all + +__all__ = mv2_all + mv3_all diff --git a/lib/python3.12/site-packages/torchvision/models/quantization/mobilenetv2.py b/lib/python3.12/site-packages/torchvision/models/quantization/mobilenetv2.py new file mode 100644 index 0000000000000000000000000000000000000000..4700bb4af931072f1aee3403c1e8c461ec33c76d --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/quantization/mobilenetv2.py @@ -0,0 +1,154 @@ +from functools import partial +from typing import Any, Optional, Union + +from torch import nn, Tensor +from torch.ao.quantization import DeQuantStub, QuantStub +from torchvision.models.mobilenetv2 import InvertedResidual, MobileNet_V2_Weights, MobileNetV2 + +from ...ops.misc import Conv2dNormActivation +from ...transforms._presets import ImageClassification +from .._api import register_model, Weights, WeightsEnum +from .._meta import _IMAGENET_CATEGORIES +from .._utils import _ovewrite_named_param, handle_legacy_interface +from .utils import _fuse_modules, _replace_relu, quantize_model + + +__all__ = [ + "QuantizableMobileNetV2", + "MobileNet_V2_QuantizedWeights", + "mobilenet_v2", +] + + +class QuantizableInvertedResidual(InvertedResidual): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.skip_add = nn.quantized.FloatFunctional() + + def forward(self, x: Tensor) -> Tensor: + if self.use_res_connect: + return self.skip_add.add(x, self.conv(x)) + else: + return self.conv(x) + + def fuse_model(self, is_qat: Optional[bool] = None) -> None: + for idx in range(len(self.conv)): + if type(self.conv[idx]) is nn.Conv2d: + _fuse_modules(self.conv, [str(idx), str(idx + 1)], is_qat, inplace=True) + + +class QuantizableMobileNetV2(MobileNetV2): + def __init__(self, *args: Any, **kwargs: Any) -> None: + """ + MobileNet V2 main class + + Args: + Inherits args from floating point MobileNetV2 + """ + super().__init__(*args, **kwargs) + self.quant = QuantStub() + self.dequant = DeQuantStub() + + def forward(self, x: Tensor) -> Tensor: + x = self.quant(x) + x = self._forward_impl(x) + x = self.dequant(x) + return x + + def fuse_model(self, is_qat: Optional[bool] = None) -> None: + for m in self.modules(): + if type(m) is Conv2dNormActivation: + _fuse_modules(m, ["0", "1", "2"], is_qat, inplace=True) + if type(m) is QuantizableInvertedResidual: + m.fuse_model(is_qat) + + +class MobileNet_V2_QuantizedWeights(WeightsEnum): + IMAGENET1K_QNNPACK_V1 = Weights( + url="https://download.pytorch.org/models/quantized/mobilenet_v2_qnnpack_37f702c5.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + "num_params": 3504872, + "min_size": (1, 1), + "categories": _IMAGENET_CATEGORIES, + "backend": "qnnpack", + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#qat-mobilenetv2", + "unquantized": MobileNet_V2_Weights.IMAGENET1K_V1, + "_metrics": { + "ImageNet-1K": { + "acc@1": 71.658, + "acc@5": 90.150, + } + }, + "_ops": 0.301, + "_file_size": 3.423, + "_docs": """ + These weights were produced by doing Quantization Aware Training (eager mode) on top of the unquantized + weights listed below. + """, + }, + ) + DEFAULT = IMAGENET1K_QNNPACK_V1 + + +@register_model(name="quantized_mobilenet_v2") +@handle_legacy_interface( + weights=( + "pretrained", + lambda kwargs: MobileNet_V2_QuantizedWeights.IMAGENET1K_QNNPACK_V1 + if kwargs.get("quantize", False) + else MobileNet_V2_Weights.IMAGENET1K_V1, + ) +) +def mobilenet_v2( + *, + weights: Optional[Union[MobileNet_V2_QuantizedWeights, MobileNet_V2_Weights]] = None, + progress: bool = True, + quantize: bool = False, + **kwargs: Any, +) -> QuantizableMobileNetV2: + """ + Constructs a MobileNetV2 architecture from + `MobileNetV2: Inverted Residuals and Linear Bottlenecks + `_. + + .. note:: + Note that ``quantize = True`` returns a quantized model with 8 bit + weights. Quantized models only support inference and run on CPUs. + GPU inference is not yet supported. + + Args: + weights (:class:`~torchvision.models.quantization.MobileNet_V2_QuantizedWeights` or :class:`~torchvision.models.MobileNet_V2_Weights`, optional): The + pretrained weights for the model. See + :class:`~torchvision.models.quantization.MobileNet_V2_QuantizedWeights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + quantize (bool, optional): If True, returns a quantized version of the model. Default is False. + **kwargs: parameters passed to the ``torchvision.models.quantization.QuantizableMobileNetV2`` + base class. Please refer to the `source code + `_ + for more details about this class. + .. autoclass:: torchvision.models.quantization.MobileNet_V2_QuantizedWeights + :members: + .. autoclass:: torchvision.models.MobileNet_V2_Weights + :members: + :noindex: + """ + weights = (MobileNet_V2_QuantizedWeights if quantize else MobileNet_V2_Weights).verify(weights) + + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + if "backend" in weights.meta: + _ovewrite_named_param(kwargs, "backend", weights.meta["backend"]) + backend = kwargs.pop("backend", "qnnpack") + + model = QuantizableMobileNetV2(block=QuantizableInvertedResidual, **kwargs) + _replace_relu(model) + if quantize: + quantize_model(model, backend) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model diff --git a/lib/python3.12/site-packages/torchvision/models/quantization/mobilenetv3.py b/lib/python3.12/site-packages/torchvision/models/quantization/mobilenetv3.py new file mode 100644 index 0000000000000000000000000000000000000000..f1fdcfec9570d35683efb10344e667d3f4487fce --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/quantization/mobilenetv3.py @@ -0,0 +1,237 @@ +from functools import partial +from typing import Any, List, Optional, Union + +import torch +from torch import nn, Tensor +from torch.ao.quantization import DeQuantStub, QuantStub + +from ...ops.misc import Conv2dNormActivation, SqueezeExcitation +from ...transforms._presets import ImageClassification +from .._api import register_model, Weights, WeightsEnum +from .._meta import _IMAGENET_CATEGORIES +from .._utils import _ovewrite_named_param, handle_legacy_interface +from ..mobilenetv3 import ( + _mobilenet_v3_conf, + InvertedResidual, + InvertedResidualConfig, + MobileNet_V3_Large_Weights, + MobileNetV3, +) +from .utils import _fuse_modules, _replace_relu + + +__all__ = [ + "QuantizableMobileNetV3", + "MobileNet_V3_Large_QuantizedWeights", + "mobilenet_v3_large", +] + + +class QuantizableSqueezeExcitation(SqueezeExcitation): + _version = 2 + + def __init__(self, *args: Any, **kwargs: Any) -> None: + kwargs["scale_activation"] = nn.Hardsigmoid + super().__init__(*args, **kwargs) + self.skip_mul = nn.quantized.FloatFunctional() + + def forward(self, input: Tensor) -> Tensor: + return self.skip_mul.mul(self._scale(input), input) + + def fuse_model(self, is_qat: Optional[bool] = None) -> None: + _fuse_modules(self, ["fc1", "activation"], is_qat, inplace=True) + + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + version = local_metadata.get("version", None) + + if hasattr(self, "qconfig") and (version is None or version < 2): + default_state_dict = { + "scale_activation.activation_post_process.scale": torch.tensor([1.0]), + "scale_activation.activation_post_process.activation_post_process.scale": torch.tensor([1.0]), + "scale_activation.activation_post_process.zero_point": torch.tensor([0], dtype=torch.int32), + "scale_activation.activation_post_process.activation_post_process.zero_point": torch.tensor( + [0], dtype=torch.int32 + ), + "scale_activation.activation_post_process.fake_quant_enabled": torch.tensor([1]), + "scale_activation.activation_post_process.observer_enabled": torch.tensor([1]), + } + for k, v in default_state_dict.items(): + full_key = prefix + k + if full_key not in state_dict: + state_dict[full_key] = v + + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + +class QuantizableInvertedResidual(InvertedResidual): + # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659 + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, se_layer=QuantizableSqueezeExcitation, **kwargs) # type: ignore[misc] + self.skip_add = nn.quantized.FloatFunctional() + + def forward(self, x: Tensor) -> Tensor: + if self.use_res_connect: + return self.skip_add.add(x, self.block(x)) + else: + return self.block(x) + + +class QuantizableMobileNetV3(MobileNetV3): + def __init__(self, *args: Any, **kwargs: Any) -> None: + """ + MobileNet V3 main class + + Args: + Inherits args from floating point MobileNetV3 + """ + super().__init__(*args, **kwargs) + self.quant = QuantStub() + self.dequant = DeQuantStub() + + def forward(self, x: Tensor) -> Tensor: + x = self.quant(x) + x = self._forward_impl(x) + x = self.dequant(x) + return x + + def fuse_model(self, is_qat: Optional[bool] = None) -> None: + for m in self.modules(): + if type(m) is Conv2dNormActivation: + modules_to_fuse = ["0", "1"] + if len(m) == 3 and type(m[2]) is nn.ReLU: + modules_to_fuse.append("2") + _fuse_modules(m, modules_to_fuse, is_qat, inplace=True) + elif type(m) is QuantizableSqueezeExcitation: + m.fuse_model(is_qat) + + +def _mobilenet_v3_model( + inverted_residual_setting: List[InvertedResidualConfig], + last_channel: int, + weights: Optional[WeightsEnum], + progress: bool, + quantize: bool, + **kwargs: Any, +) -> QuantizableMobileNetV3: + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + if "backend" in weights.meta: + _ovewrite_named_param(kwargs, "backend", weights.meta["backend"]) + backend = kwargs.pop("backend", "qnnpack") + + model = QuantizableMobileNetV3(inverted_residual_setting, last_channel, block=QuantizableInvertedResidual, **kwargs) + _replace_relu(model) + + if quantize: + # Instead of quantizing the model and then loading the quantized weights we take a different approach. + # We prepare the QAT model, load the QAT weights from training and then convert it. + # This is done to avoid extremely low accuracies observed on the specific model. This is rather a workaround + # for an unresolved bug on the eager quantization API detailed at: https://github.com/pytorch/vision/issues/5890 + model.fuse_model(is_qat=True) + model.qconfig = torch.ao.quantization.get_default_qat_qconfig(backend) + torch.ao.quantization.prepare_qat(model, inplace=True) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + if quantize: + torch.ao.quantization.convert(model, inplace=True) + model.eval() + + return model + + +class MobileNet_V3_Large_QuantizedWeights(WeightsEnum): + IMAGENET1K_QNNPACK_V1 = Weights( + url="https://download.pytorch.org/models/quantized/mobilenet_v3_large_qnnpack-5bcacf28.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + "num_params": 5483032, + "min_size": (1, 1), + "categories": _IMAGENET_CATEGORIES, + "backend": "qnnpack", + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#qat-mobilenetv3", + "unquantized": MobileNet_V3_Large_Weights.IMAGENET1K_V1, + "_metrics": { + "ImageNet-1K": { + "acc@1": 73.004, + "acc@5": 90.858, + } + }, + "_ops": 0.217, + "_file_size": 21.554, + "_docs": """ + These weights were produced by doing Quantization Aware Training (eager mode) on top of the unquantized + weights listed below. + """, + }, + ) + DEFAULT = IMAGENET1K_QNNPACK_V1 + + +@register_model(name="quantized_mobilenet_v3_large") +@handle_legacy_interface( + weights=( + "pretrained", + lambda kwargs: MobileNet_V3_Large_QuantizedWeights.IMAGENET1K_QNNPACK_V1 + if kwargs.get("quantize", False) + else MobileNet_V3_Large_Weights.IMAGENET1K_V1, + ) +) +def mobilenet_v3_large( + *, + weights: Optional[Union[MobileNet_V3_Large_QuantizedWeights, MobileNet_V3_Large_Weights]] = None, + progress: bool = True, + quantize: bool = False, + **kwargs: Any, +) -> QuantizableMobileNetV3: + """ + MobileNetV3 (Large) model from + `Searching for MobileNetV3 `_. + + .. note:: + Note that ``quantize = True`` returns a quantized model with 8 bit + weights. Quantized models only support inference and run on CPUs. + GPU inference is not yet supported. + + Args: + weights (:class:`~torchvision.models.quantization.MobileNet_V3_Large_QuantizedWeights` or :class:`~torchvision.models.MobileNet_V3_Large_Weights`, optional): The + pretrained weights for the model. See + :class:`~torchvision.models.quantization.MobileNet_V3_Large_QuantizedWeights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool): If True, displays a progress bar of the + download to stderr. Default is True. + quantize (bool): If True, return a quantized version of the model. Default is False. + **kwargs: parameters passed to the ``torchvision.models.quantization.MobileNet_V3_Large_QuantizedWeights`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.quantization.MobileNet_V3_Large_QuantizedWeights + :members: + .. autoclass:: torchvision.models.MobileNet_V3_Large_Weights + :members: + :noindex: + """ + weights = (MobileNet_V3_Large_QuantizedWeights if quantize else MobileNet_V3_Large_Weights).verify(weights) + + inverted_residual_setting, last_channel = _mobilenet_v3_conf("mobilenet_v3_large", **kwargs) + return _mobilenet_v3_model(inverted_residual_setting, last_channel, weights, progress, quantize, **kwargs) diff --git a/lib/python3.12/site-packages/torchvision/models/quantization/resnet.py b/lib/python3.12/site-packages/torchvision/models/quantization/resnet.py new file mode 100644 index 0000000000000000000000000000000000000000..39958a010fbd335709bc77a1aaf26c996584a398 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/quantization/resnet.py @@ -0,0 +1,484 @@ +from functools import partial +from typing import Any, List, Optional, Type, Union + +import torch +import torch.nn as nn +from torch import Tensor +from torchvision.models.resnet import ( + BasicBlock, + Bottleneck, + ResNet, + ResNet18_Weights, + ResNet50_Weights, + ResNeXt101_32X8D_Weights, + ResNeXt101_64X4D_Weights, +) + +from ...transforms._presets import ImageClassification +from .._api import register_model, Weights, WeightsEnum +from .._meta import _IMAGENET_CATEGORIES +from .._utils import _ovewrite_named_param, handle_legacy_interface +from .utils import _fuse_modules, _replace_relu, quantize_model + + +__all__ = [ + "QuantizableResNet", + "ResNet18_QuantizedWeights", + "ResNet50_QuantizedWeights", + "ResNeXt101_32X8D_QuantizedWeights", + "ResNeXt101_64X4D_QuantizedWeights", + "resnet18", + "resnet50", + "resnext101_32x8d", + "resnext101_64x4d", +] + + +class QuantizableBasicBlock(BasicBlock): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.add_relu = torch.nn.quantized.FloatFunctional() + + def forward(self, x: Tensor) -> Tensor: + identity = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + + if self.downsample is not None: + identity = self.downsample(x) + + out = self.add_relu.add_relu(out, identity) + + return out + + def fuse_model(self, is_qat: Optional[bool] = None) -> None: + _fuse_modules(self, [["conv1", "bn1", "relu"], ["conv2", "bn2"]], is_qat, inplace=True) + if self.downsample: + _fuse_modules(self.downsample, ["0", "1"], is_qat, inplace=True) + + +class QuantizableBottleneck(Bottleneck): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.skip_add_relu = nn.quantized.FloatFunctional() + self.relu1 = nn.ReLU(inplace=False) + self.relu2 = nn.ReLU(inplace=False) + + def forward(self, x: Tensor) -> Tensor: + identity = x + out = self.conv1(x) + out = self.bn1(out) + out = self.relu1(out) + out = self.conv2(out) + out = self.bn2(out) + out = self.relu2(out) + + out = self.conv3(out) + out = self.bn3(out) + + if self.downsample is not None: + identity = self.downsample(x) + out = self.skip_add_relu.add_relu(out, identity) + + return out + + def fuse_model(self, is_qat: Optional[bool] = None) -> None: + _fuse_modules( + self, [["conv1", "bn1", "relu1"], ["conv2", "bn2", "relu2"], ["conv3", "bn3"]], is_qat, inplace=True + ) + if self.downsample: + _fuse_modules(self.downsample, ["0", "1"], is_qat, inplace=True) + + +class QuantizableResNet(ResNet): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + + self.quant = torch.ao.quantization.QuantStub() + self.dequant = torch.ao.quantization.DeQuantStub() + + def forward(self, x: Tensor) -> Tensor: + x = self.quant(x) + # Ensure scriptability + # super(QuantizableResNet,self).forward(x) + # is not scriptable + x = self._forward_impl(x) + x = self.dequant(x) + return x + + def fuse_model(self, is_qat: Optional[bool] = None) -> None: + r"""Fuse conv/bn/relu modules in resnet models + + Fuse conv+bn+relu/ Conv+relu/conv+Bn modules to prepare for quantization. + Model is modified in place. Note that this operation does not change numerics + and the model after modification is in floating point + """ + _fuse_modules(self, ["conv1", "bn1", "relu"], is_qat, inplace=True) + for m in self.modules(): + if type(m) is QuantizableBottleneck or type(m) is QuantizableBasicBlock: + m.fuse_model(is_qat) + + +def _resnet( + block: Type[Union[QuantizableBasicBlock, QuantizableBottleneck]], + layers: List[int], + weights: Optional[WeightsEnum], + progress: bool, + quantize: bool, + **kwargs: Any, +) -> QuantizableResNet: + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + if "backend" in weights.meta: + _ovewrite_named_param(kwargs, "backend", weights.meta["backend"]) + backend = kwargs.pop("backend", "fbgemm") + + model = QuantizableResNet(block, layers, **kwargs) + _replace_relu(model) + if quantize: + quantize_model(model, backend) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +_COMMON_META = { + "min_size": (1, 1), + "categories": _IMAGENET_CATEGORIES, + "backend": "fbgemm", + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#post-training-quantized-models", + "_docs": """ + These weights were produced by doing Post Training Quantization (eager mode) on top of the unquantized + weights listed below. + """, +} + + +class ResNet18_QuantizedWeights(WeightsEnum): + IMAGENET1K_FBGEMM_V1 = Weights( + url="https://download.pytorch.org/models/quantized/resnet18_fbgemm_16fa66dd.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 11689512, + "unquantized": ResNet18_Weights.IMAGENET1K_V1, + "_metrics": { + "ImageNet-1K": { + "acc@1": 69.494, + "acc@5": 88.882, + } + }, + "_ops": 1.814, + "_file_size": 11.238, + }, + ) + DEFAULT = IMAGENET1K_FBGEMM_V1 + + +class ResNet50_QuantizedWeights(WeightsEnum): + IMAGENET1K_FBGEMM_V1 = Weights( + url="https://download.pytorch.org/models/quantized/resnet50_fbgemm_bf931d71.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 25557032, + "unquantized": ResNet50_Weights.IMAGENET1K_V1, + "_metrics": { + "ImageNet-1K": { + "acc@1": 75.920, + "acc@5": 92.814, + } + }, + "_ops": 4.089, + "_file_size": 24.759, + }, + ) + IMAGENET1K_FBGEMM_V2 = Weights( + url="https://download.pytorch.org/models/quantized/resnet50_fbgemm-23753f79.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 25557032, + "unquantized": ResNet50_Weights.IMAGENET1K_V2, + "_metrics": { + "ImageNet-1K": { + "acc@1": 80.282, + "acc@5": 94.976, + } + }, + "_ops": 4.089, + "_file_size": 24.953, + }, + ) + DEFAULT = IMAGENET1K_FBGEMM_V2 + + +class ResNeXt101_32X8D_QuantizedWeights(WeightsEnum): + IMAGENET1K_FBGEMM_V1 = Weights( + url="https://download.pytorch.org/models/quantized/resnext101_32x8_fbgemm_09835ccf.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 88791336, + "unquantized": ResNeXt101_32X8D_Weights.IMAGENET1K_V1, + "_metrics": { + "ImageNet-1K": { + "acc@1": 78.986, + "acc@5": 94.480, + } + }, + "_ops": 16.414, + "_file_size": 86.034, + }, + ) + IMAGENET1K_FBGEMM_V2 = Weights( + url="https://download.pytorch.org/models/quantized/resnext101_32x8_fbgemm-ee16d00c.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 88791336, + "unquantized": ResNeXt101_32X8D_Weights.IMAGENET1K_V2, + "_metrics": { + "ImageNet-1K": { + "acc@1": 82.574, + "acc@5": 96.132, + } + }, + "_ops": 16.414, + "_file_size": 86.645, + }, + ) + DEFAULT = IMAGENET1K_FBGEMM_V2 + + +class ResNeXt101_64X4D_QuantizedWeights(WeightsEnum): + IMAGENET1K_FBGEMM_V1 = Weights( + url="https://download.pytorch.org/models/quantized/resnext101_64x4d_fbgemm-605a1cb3.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 83455272, + "recipe": "https://github.com/pytorch/vision/pull/5935", + "unquantized": ResNeXt101_64X4D_Weights.IMAGENET1K_V1, + "_metrics": { + "ImageNet-1K": { + "acc@1": 82.898, + "acc@5": 96.326, + } + }, + "_ops": 15.46, + "_file_size": 81.556, + }, + ) + DEFAULT = IMAGENET1K_FBGEMM_V1 + + +@register_model(name="quantized_resnet18") +@handle_legacy_interface( + weights=( + "pretrained", + lambda kwargs: ResNet18_QuantizedWeights.IMAGENET1K_FBGEMM_V1 + if kwargs.get("quantize", False) + else ResNet18_Weights.IMAGENET1K_V1, + ) +) +def resnet18( + *, + weights: Optional[Union[ResNet18_QuantizedWeights, ResNet18_Weights]] = None, + progress: bool = True, + quantize: bool = False, + **kwargs: Any, +) -> QuantizableResNet: + """ResNet-18 model from + `Deep Residual Learning for Image Recognition `_ + + .. note:: + Note that ``quantize = True`` returns a quantized model with 8 bit + weights. Quantized models only support inference and run on CPUs. + GPU inference is not yet supported. + + Args: + weights (:class:`~torchvision.models.quantization.ResNet18_QuantizedWeights` or :class:`~torchvision.models.ResNet18_Weights`, optional): The + pretrained weights for the model. See + :class:`~torchvision.models.quantization.ResNet18_QuantizedWeights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + quantize (bool, optional): If True, return a quantized version of the model. Default is False. + **kwargs: parameters passed to the ``torchvision.models.quantization.QuantizableResNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.quantization.ResNet18_QuantizedWeights + :members: + + .. autoclass:: torchvision.models.ResNet18_Weights + :members: + :noindex: + """ + weights = (ResNet18_QuantizedWeights if quantize else ResNet18_Weights).verify(weights) + + return _resnet(QuantizableBasicBlock, [2, 2, 2, 2], weights, progress, quantize, **kwargs) + + +@register_model(name="quantized_resnet50") +@handle_legacy_interface( + weights=( + "pretrained", + lambda kwargs: ResNet50_QuantizedWeights.IMAGENET1K_FBGEMM_V1 + if kwargs.get("quantize", False) + else ResNet50_Weights.IMAGENET1K_V1, + ) +) +def resnet50( + *, + weights: Optional[Union[ResNet50_QuantizedWeights, ResNet50_Weights]] = None, + progress: bool = True, + quantize: bool = False, + **kwargs: Any, +) -> QuantizableResNet: + """ResNet-50 model from + `Deep Residual Learning for Image Recognition `_ + + .. note:: + Note that ``quantize = True`` returns a quantized model with 8 bit + weights. Quantized models only support inference and run on CPUs. + GPU inference is not yet supported. + + Args: + weights (:class:`~torchvision.models.quantization.ResNet50_QuantizedWeights` or :class:`~torchvision.models.ResNet50_Weights`, optional): The + pretrained weights for the model. See + :class:`~torchvision.models.quantization.ResNet50_QuantizedWeights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + quantize (bool, optional): If True, return a quantized version of the model. Default is False. + **kwargs: parameters passed to the ``torchvision.models.quantization.QuantizableResNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.quantization.ResNet50_QuantizedWeights + :members: + + .. autoclass:: torchvision.models.ResNet50_Weights + :members: + :noindex: + """ + weights = (ResNet50_QuantizedWeights if quantize else ResNet50_Weights).verify(weights) + + return _resnet(QuantizableBottleneck, [3, 4, 6, 3], weights, progress, quantize, **kwargs) + + +@register_model(name="quantized_resnext101_32x8d") +@handle_legacy_interface( + weights=( + "pretrained", + lambda kwargs: ResNeXt101_32X8D_QuantizedWeights.IMAGENET1K_FBGEMM_V1 + if kwargs.get("quantize", False) + else ResNeXt101_32X8D_Weights.IMAGENET1K_V1, + ) +) +def resnext101_32x8d( + *, + weights: Optional[Union[ResNeXt101_32X8D_QuantizedWeights, ResNeXt101_32X8D_Weights]] = None, + progress: bool = True, + quantize: bool = False, + **kwargs: Any, +) -> QuantizableResNet: + """ResNeXt-101 32x8d model from + `Aggregated Residual Transformation for Deep Neural Networks `_ + + .. note:: + Note that ``quantize = True`` returns a quantized model with 8 bit + weights. Quantized models only support inference and run on CPUs. + GPU inference is not yet supported. + + Args: + weights (:class:`~torchvision.models.quantization.ResNeXt101_32X8D_QuantizedWeights` or :class:`~torchvision.models.ResNeXt101_32X8D_Weights`, optional): The + pretrained weights for the model. See + :class:`~torchvision.models.quantization.ResNet101_32X8D_QuantizedWeights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + quantize (bool, optional): If True, return a quantized version of the model. Default is False. + **kwargs: parameters passed to the ``torchvision.models.quantization.QuantizableResNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.quantization.ResNeXt101_32X8D_QuantizedWeights + :members: + + .. autoclass:: torchvision.models.ResNeXt101_32X8D_Weights + :members: + :noindex: + """ + weights = (ResNeXt101_32X8D_QuantizedWeights if quantize else ResNeXt101_32X8D_Weights).verify(weights) + + _ovewrite_named_param(kwargs, "groups", 32) + _ovewrite_named_param(kwargs, "width_per_group", 8) + return _resnet(QuantizableBottleneck, [3, 4, 23, 3], weights, progress, quantize, **kwargs) + + +@register_model(name="quantized_resnext101_64x4d") +@handle_legacy_interface( + weights=( + "pretrained", + lambda kwargs: ResNeXt101_64X4D_QuantizedWeights.IMAGENET1K_FBGEMM_V1 + if kwargs.get("quantize", False) + else ResNeXt101_64X4D_Weights.IMAGENET1K_V1, + ) +) +def resnext101_64x4d( + *, + weights: Optional[Union[ResNeXt101_64X4D_QuantizedWeights, ResNeXt101_64X4D_Weights]] = None, + progress: bool = True, + quantize: bool = False, + **kwargs: Any, +) -> QuantizableResNet: + """ResNeXt-101 64x4d model from + `Aggregated Residual Transformation for Deep Neural Networks `_ + + .. note:: + Note that ``quantize = True`` returns a quantized model with 8 bit + weights. Quantized models only support inference and run on CPUs. + GPU inference is not yet supported. + + Args: + weights (:class:`~torchvision.models.quantization.ResNeXt101_64X4D_QuantizedWeights` or :class:`~torchvision.models.ResNeXt101_64X4D_Weights`, optional): The + pretrained weights for the model. See + :class:`~torchvision.models.quantization.ResNet101_64X4D_QuantizedWeights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + quantize (bool, optional): If True, return a quantized version of the model. Default is False. + **kwargs: parameters passed to the ``torchvision.models.quantization.QuantizableResNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.quantization.ResNeXt101_64X4D_QuantizedWeights + :members: + + .. autoclass:: torchvision.models.ResNeXt101_64X4D_Weights + :members: + :noindex: + """ + weights = (ResNeXt101_64X4D_QuantizedWeights if quantize else ResNeXt101_64X4D_Weights).verify(weights) + + _ovewrite_named_param(kwargs, "groups", 64) + _ovewrite_named_param(kwargs, "width_per_group", 4) + return _resnet(QuantizableBottleneck, [3, 4, 23, 3], weights, progress, quantize, **kwargs) diff --git a/lib/python3.12/site-packages/torchvision/models/quantization/shufflenetv2.py b/lib/python3.12/site-packages/torchvision/models/quantization/shufflenetv2.py new file mode 100644 index 0000000000000000000000000000000000000000..3e1b01356a74b8e4f16d66811060be698cfed199 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/quantization/shufflenetv2.py @@ -0,0 +1,427 @@ +from functools import partial +from typing import Any, List, Optional, Union + +import torch +import torch.nn as nn +from torch import Tensor +from torchvision.models import shufflenetv2 + +from ...transforms._presets import ImageClassification +from .._api import register_model, Weights, WeightsEnum +from .._meta import _IMAGENET_CATEGORIES +from .._utils import _ovewrite_named_param, handle_legacy_interface +from ..shufflenetv2 import ( + ShuffleNet_V2_X0_5_Weights, + ShuffleNet_V2_X1_0_Weights, + ShuffleNet_V2_X1_5_Weights, + ShuffleNet_V2_X2_0_Weights, +) +from .utils import _fuse_modules, _replace_relu, quantize_model + + +__all__ = [ + "QuantizableShuffleNetV2", + "ShuffleNet_V2_X0_5_QuantizedWeights", + "ShuffleNet_V2_X1_0_QuantizedWeights", + "ShuffleNet_V2_X1_5_QuantizedWeights", + "ShuffleNet_V2_X2_0_QuantizedWeights", + "shufflenet_v2_x0_5", + "shufflenet_v2_x1_0", + "shufflenet_v2_x1_5", + "shufflenet_v2_x2_0", +] + + +class QuantizableInvertedResidual(shufflenetv2.InvertedResidual): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.cat = nn.quantized.FloatFunctional() + + def forward(self, x: Tensor) -> Tensor: + if self.stride == 1: + x1, x2 = x.chunk(2, dim=1) + out = self.cat.cat([x1, self.branch2(x2)], dim=1) + else: + out = self.cat.cat([self.branch1(x), self.branch2(x)], dim=1) + + out = shufflenetv2.channel_shuffle(out, 2) + + return out + + +class QuantizableShuffleNetV2(shufflenetv2.ShuffleNetV2): + # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659 + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, inverted_residual=QuantizableInvertedResidual, **kwargs) # type: ignore[misc] + self.quant = torch.ao.quantization.QuantStub() + self.dequant = torch.ao.quantization.DeQuantStub() + + def forward(self, x: Tensor) -> Tensor: + x = self.quant(x) + x = self._forward_impl(x) + x = self.dequant(x) + return x + + def fuse_model(self, is_qat: Optional[bool] = None) -> None: + r"""Fuse conv/bn/relu modules in shufflenetv2 model + + Fuse conv+bn+relu/ conv+relu/conv+bn modules to prepare for quantization. + Model is modified in place. + + .. note:: + Note that this operation does not change numerics + and the model after modification is in floating point + """ + for name, m in self._modules.items(): + if name in ["conv1", "conv5"] and m is not None: + _fuse_modules(m, [["0", "1", "2"]], is_qat, inplace=True) + for m in self.modules(): + if type(m) is QuantizableInvertedResidual: + if len(m.branch1._modules.items()) > 0: + _fuse_modules(m.branch1, [["0", "1"], ["2", "3", "4"]], is_qat, inplace=True) + _fuse_modules( + m.branch2, + [["0", "1", "2"], ["3", "4"], ["5", "6", "7"]], + is_qat, + inplace=True, + ) + + +def _shufflenetv2( + stages_repeats: List[int], + stages_out_channels: List[int], + *, + weights: Optional[WeightsEnum], + progress: bool, + quantize: bool, + **kwargs: Any, +) -> QuantizableShuffleNetV2: + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + if "backend" in weights.meta: + _ovewrite_named_param(kwargs, "backend", weights.meta["backend"]) + backend = kwargs.pop("backend", "fbgemm") + + model = QuantizableShuffleNetV2(stages_repeats, stages_out_channels, **kwargs) + _replace_relu(model) + if quantize: + quantize_model(model, backend) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +_COMMON_META = { + "min_size": (1, 1), + "categories": _IMAGENET_CATEGORIES, + "backend": "fbgemm", + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#post-training-quantized-models", + "_docs": """ + These weights were produced by doing Post Training Quantization (eager mode) on top of the unquantized + weights listed below. + """, +} + + +class ShuffleNet_V2_X0_5_QuantizedWeights(WeightsEnum): + IMAGENET1K_FBGEMM_V1 = Weights( + url="https://download.pytorch.org/models/quantized/shufflenetv2_x0.5_fbgemm-00845098.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 1366792, + "unquantized": ShuffleNet_V2_X0_5_Weights.IMAGENET1K_V1, + "_metrics": { + "ImageNet-1K": { + "acc@1": 57.972, + "acc@5": 79.780, + } + }, + "_ops": 0.04, + "_file_size": 1.501, + }, + ) + DEFAULT = IMAGENET1K_FBGEMM_V1 + + +class ShuffleNet_V2_X1_0_QuantizedWeights(WeightsEnum): + IMAGENET1K_FBGEMM_V1 = Weights( + url="https://download.pytorch.org/models/quantized/shufflenetv2_x1_fbgemm-1e62bb32.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 2278604, + "unquantized": ShuffleNet_V2_X1_0_Weights.IMAGENET1K_V1, + "_metrics": { + "ImageNet-1K": { + "acc@1": 68.360, + "acc@5": 87.582, + } + }, + "_ops": 0.145, + "_file_size": 2.334, + }, + ) + DEFAULT = IMAGENET1K_FBGEMM_V1 + + +class ShuffleNet_V2_X1_5_QuantizedWeights(WeightsEnum): + IMAGENET1K_FBGEMM_V1 = Weights( + url="https://download.pytorch.org/models/quantized/shufflenetv2_x1_5_fbgemm-d7401f05.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "recipe": "https://github.com/pytorch/vision/pull/5906", + "num_params": 3503624, + "unquantized": ShuffleNet_V2_X1_5_Weights.IMAGENET1K_V1, + "_metrics": { + "ImageNet-1K": { + "acc@1": 72.052, + "acc@5": 90.700, + } + }, + "_ops": 0.296, + "_file_size": 3.672, + }, + ) + DEFAULT = IMAGENET1K_FBGEMM_V1 + + +class ShuffleNet_V2_X2_0_QuantizedWeights(WeightsEnum): + IMAGENET1K_FBGEMM_V1 = Weights( + url="https://download.pytorch.org/models/quantized/shufflenetv2_x2_0_fbgemm-5cac526c.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "recipe": "https://github.com/pytorch/vision/pull/5906", + "num_params": 7393996, + "unquantized": ShuffleNet_V2_X2_0_Weights.IMAGENET1K_V1, + "_metrics": { + "ImageNet-1K": { + "acc@1": 75.354, + "acc@5": 92.488, + } + }, + "_ops": 0.583, + "_file_size": 7.467, + }, + ) + DEFAULT = IMAGENET1K_FBGEMM_V1 + + +@register_model(name="quantized_shufflenet_v2_x0_5") +@handle_legacy_interface( + weights=( + "pretrained", + lambda kwargs: ShuffleNet_V2_X0_5_QuantizedWeights.IMAGENET1K_FBGEMM_V1 + if kwargs.get("quantize", False) + else ShuffleNet_V2_X0_5_Weights.IMAGENET1K_V1, + ) +) +def shufflenet_v2_x0_5( + *, + weights: Optional[Union[ShuffleNet_V2_X0_5_QuantizedWeights, ShuffleNet_V2_X0_5_Weights]] = None, + progress: bool = True, + quantize: bool = False, + **kwargs: Any, +) -> QuantizableShuffleNetV2: + """ + Constructs a ShuffleNetV2 with 0.5x output channels, as described in + `ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design + `__. + + .. note:: + Note that ``quantize = True`` returns a quantized model with 8 bit + weights. Quantized models only support inference and run on CPUs. + GPU inference is not yet supported. + + Args: + weights (:class:`~torchvision.models.quantization.ShuffleNet_V2_X0_5_QuantizedWeights` or :class:`~torchvision.models.ShuffleNet_V2_X0_5_Weights`, optional): The + pretrained weights for the model. See + :class:`~torchvision.models.quantization.ShuffleNet_V2_X0_5_QuantizedWeights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. + Default is True. + quantize (bool, optional): If True, return a quantized version of the model. + Default is False. + **kwargs: parameters passed to the ``torchvision.models.quantization.ShuffleNet_V2_X0_5_QuantizedWeights`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.quantization.ShuffleNet_V2_X0_5_QuantizedWeights + :members: + + .. autoclass:: torchvision.models.ShuffleNet_V2_X0_5_Weights + :members: + :noindex: + """ + weights = (ShuffleNet_V2_X0_5_QuantizedWeights if quantize else ShuffleNet_V2_X0_5_Weights).verify(weights) + return _shufflenetv2( + [4, 8, 4], [24, 48, 96, 192, 1024], weights=weights, progress=progress, quantize=quantize, **kwargs + ) + + +@register_model(name="quantized_shufflenet_v2_x1_0") +@handle_legacy_interface( + weights=( + "pretrained", + lambda kwargs: ShuffleNet_V2_X1_0_QuantizedWeights.IMAGENET1K_FBGEMM_V1 + if kwargs.get("quantize", False) + else ShuffleNet_V2_X1_0_Weights.IMAGENET1K_V1, + ) +) +def shufflenet_v2_x1_0( + *, + weights: Optional[Union[ShuffleNet_V2_X1_0_QuantizedWeights, ShuffleNet_V2_X1_0_Weights]] = None, + progress: bool = True, + quantize: bool = False, + **kwargs: Any, +) -> QuantizableShuffleNetV2: + """ + Constructs a ShuffleNetV2 with 1.0x output channels, as described in + `ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design + `__. + + .. note:: + Note that ``quantize = True`` returns a quantized model with 8 bit + weights. Quantized models only support inference and run on CPUs. + GPU inference is not yet supported. + + Args: + weights (:class:`~torchvision.models.quantization.ShuffleNet_V2_X1_0_QuantizedWeights` or :class:`~torchvision.models.ShuffleNet_V2_X1_0_Weights`, optional): The + pretrained weights for the model. See + :class:`~torchvision.models.quantization.ShuffleNet_V2_X1_0_QuantizedWeights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. + Default is True. + quantize (bool, optional): If True, return a quantized version of the model. + Default is False. + **kwargs: parameters passed to the ``torchvision.models.quantization.ShuffleNet_V2_X1_0_QuantizedWeights`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.quantization.ShuffleNet_V2_X1_0_QuantizedWeights + :members: + + .. autoclass:: torchvision.models.ShuffleNet_V2_X1_0_Weights + :members: + :noindex: + """ + weights = (ShuffleNet_V2_X1_0_QuantizedWeights if quantize else ShuffleNet_V2_X1_0_Weights).verify(weights) + return _shufflenetv2( + [4, 8, 4], [24, 116, 232, 464, 1024], weights=weights, progress=progress, quantize=quantize, **kwargs + ) + + +@register_model(name="quantized_shufflenet_v2_x1_5") +@handle_legacy_interface( + weights=( + "pretrained", + lambda kwargs: ShuffleNet_V2_X1_5_QuantizedWeights.IMAGENET1K_FBGEMM_V1 + if kwargs.get("quantize", False) + else ShuffleNet_V2_X1_5_Weights.IMAGENET1K_V1, + ) +) +def shufflenet_v2_x1_5( + *, + weights: Optional[Union[ShuffleNet_V2_X1_5_QuantizedWeights, ShuffleNet_V2_X1_5_Weights]] = None, + progress: bool = True, + quantize: bool = False, + **kwargs: Any, +) -> QuantizableShuffleNetV2: + """ + Constructs a ShuffleNetV2 with 1.5x output channels, as described in + `ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design + `__. + + .. note:: + Note that ``quantize = True`` returns a quantized model with 8 bit + weights. Quantized models only support inference and run on CPUs. + GPU inference is not yet supported. + + Args: + weights (:class:`~torchvision.models.quantization.ShuffleNet_V2_X1_5_QuantizedWeights` or :class:`~torchvision.models.ShuffleNet_V2_X1_5_Weights`, optional): The + pretrained weights for the model. See + :class:`~torchvision.models.quantization.ShuffleNet_V2_X1_5_QuantizedWeights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. + Default is True. + quantize (bool, optional): If True, return a quantized version of the model. + Default is False. + **kwargs: parameters passed to the ``torchvision.models.quantization.ShuffleNet_V2_X1_5_QuantizedWeights`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.quantization.ShuffleNet_V2_X1_5_QuantizedWeights + :members: + + .. autoclass:: torchvision.models.ShuffleNet_V2_X1_5_Weights + :members: + :noindex: + """ + weights = (ShuffleNet_V2_X1_5_QuantizedWeights if quantize else ShuffleNet_V2_X1_5_Weights).verify(weights) + return _shufflenetv2( + [4, 8, 4], [24, 176, 352, 704, 1024], weights=weights, progress=progress, quantize=quantize, **kwargs + ) + + +@register_model(name="quantized_shufflenet_v2_x2_0") +@handle_legacy_interface( + weights=( + "pretrained", + lambda kwargs: ShuffleNet_V2_X2_0_QuantizedWeights.IMAGENET1K_FBGEMM_V1 + if kwargs.get("quantize", False) + else ShuffleNet_V2_X2_0_Weights.IMAGENET1K_V1, + ) +) +def shufflenet_v2_x2_0( + *, + weights: Optional[Union[ShuffleNet_V2_X2_0_QuantizedWeights, ShuffleNet_V2_X2_0_Weights]] = None, + progress: bool = True, + quantize: bool = False, + **kwargs: Any, +) -> QuantizableShuffleNetV2: + """ + Constructs a ShuffleNetV2 with 2.0x output channels, as described in + `ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design + `__. + + .. note:: + Note that ``quantize = True`` returns a quantized model with 8 bit + weights. Quantized models only support inference and run on CPUs. + GPU inference is not yet supported. + + Args: + weights (:class:`~torchvision.models.quantization.ShuffleNet_V2_X2_0_QuantizedWeights` or :class:`~torchvision.models.ShuffleNet_V2_X2_0_Weights`, optional): The + pretrained weights for the model. See + :class:`~torchvision.models.quantization.ShuffleNet_V2_X2_0_QuantizedWeights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. + Default is True. + quantize (bool, optional): If True, return a quantized version of the model. + Default is False. + **kwargs: parameters passed to the ``torchvision.models.quantization.ShuffleNet_V2_X2_0_QuantizedWeights`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.quantization.ShuffleNet_V2_X2_0_QuantizedWeights + :members: + + .. autoclass:: torchvision.models.ShuffleNet_V2_X2_0_Weights + :members: + :noindex: + """ + weights = (ShuffleNet_V2_X2_0_QuantizedWeights if quantize else ShuffleNet_V2_X2_0_Weights).verify(weights) + return _shufflenetv2( + [4, 8, 4], [24, 244, 488, 976, 2048], weights=weights, progress=progress, quantize=quantize, **kwargs + ) diff --git a/lib/python3.12/site-packages/torchvision/models/quantization/utils.py b/lib/python3.12/site-packages/torchvision/models/quantization/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a21e2af8e016568e79b25e35ec774f39f0595c3a --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/quantization/utils.py @@ -0,0 +1,51 @@ +from typing import Any, List, Optional, Union + +import torch +from torch import nn + + +def _replace_relu(module: nn.Module) -> None: + reassign = {} + for name, mod in module.named_children(): + _replace_relu(mod) + # Checking for explicit type instead of instance + # as we only want to replace modules of the exact type + # not inherited classes + if type(mod) is nn.ReLU or type(mod) is nn.ReLU6: + reassign[name] = nn.ReLU(inplace=False) + + for key, value in reassign.items(): + module._modules[key] = value + + +def quantize_model(model: nn.Module, backend: str) -> None: + _dummy_input_data = torch.rand(1, 3, 299, 299) + if backend not in torch.backends.quantized.supported_engines: + raise RuntimeError("Quantized backend not supported ") + torch.backends.quantized.engine = backend + model.eval() + # Make sure that weight qconfig matches that of the serialized models + if backend == "fbgemm": + model.qconfig = torch.ao.quantization.QConfig( # type: ignore[assignment] + activation=torch.ao.quantization.default_observer, + weight=torch.ao.quantization.default_per_channel_weight_observer, + ) + elif backend == "qnnpack": + model.qconfig = torch.ao.quantization.QConfig( # type: ignore[assignment] + activation=torch.ao.quantization.default_observer, weight=torch.ao.quantization.default_weight_observer + ) + + # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659 + model.fuse_model() # type: ignore[operator] + torch.ao.quantization.prepare(model, inplace=True) + model(_dummy_input_data) + torch.ao.quantization.convert(model, inplace=True) + + +def _fuse_modules( + model: nn.Module, modules_to_fuse: Union[List[str], List[List[str]]], is_qat: Optional[bool], **kwargs: Any +): + if is_qat is None: + is_qat = model.training + method = torch.ao.quantization.fuse_modules_qat if is_qat else torch.ao.quantization.fuse_modules + return method(model, modules_to_fuse, **kwargs) diff --git a/lib/python3.12/site-packages/torchvision/models/regnet.py b/lib/python3.12/site-packages/torchvision/models/regnet.py new file mode 100644 index 0000000000000000000000000000000000000000..f37b2994e48fd202cccd4db6dc63b97f7c0c1ae7 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/regnet.py @@ -0,0 +1,1571 @@ +import math +from collections import OrderedDict +from functools import partial +from typing import Any, Callable, Dict, List, Optional, Tuple + +import torch +from torch import nn, Tensor + +from ..ops.misc import Conv2dNormActivation, SqueezeExcitation +from ..transforms._presets import ImageClassification, InterpolationMode +from ..utils import _log_api_usage_once +from ._api import register_model, Weights, WeightsEnum +from ._meta import _IMAGENET_CATEGORIES +from ._utils import _make_divisible, _ovewrite_named_param, handle_legacy_interface + + +__all__ = [ + "RegNet", + "RegNet_Y_400MF_Weights", + "RegNet_Y_800MF_Weights", + "RegNet_Y_1_6GF_Weights", + "RegNet_Y_3_2GF_Weights", + "RegNet_Y_8GF_Weights", + "RegNet_Y_16GF_Weights", + "RegNet_Y_32GF_Weights", + "RegNet_Y_128GF_Weights", + "RegNet_X_400MF_Weights", + "RegNet_X_800MF_Weights", + "RegNet_X_1_6GF_Weights", + "RegNet_X_3_2GF_Weights", + "RegNet_X_8GF_Weights", + "RegNet_X_16GF_Weights", + "RegNet_X_32GF_Weights", + "regnet_y_400mf", + "regnet_y_800mf", + "regnet_y_1_6gf", + "regnet_y_3_2gf", + "regnet_y_8gf", + "regnet_y_16gf", + "regnet_y_32gf", + "regnet_y_128gf", + "regnet_x_400mf", + "regnet_x_800mf", + "regnet_x_1_6gf", + "regnet_x_3_2gf", + "regnet_x_8gf", + "regnet_x_16gf", + "regnet_x_32gf", +] + + +class SimpleStemIN(Conv2dNormActivation): + """Simple stem for ImageNet: 3x3, BN, ReLU.""" + + def __init__( + self, + width_in: int, + width_out: int, + norm_layer: Callable[..., nn.Module], + activation_layer: Callable[..., nn.Module], + ) -> None: + super().__init__( + width_in, width_out, kernel_size=3, stride=2, norm_layer=norm_layer, activation_layer=activation_layer + ) + + +class BottleneckTransform(nn.Sequential): + """Bottleneck transformation: 1x1, 3x3 [+SE], 1x1.""" + + def __init__( + self, + width_in: int, + width_out: int, + stride: int, + norm_layer: Callable[..., nn.Module], + activation_layer: Callable[..., nn.Module], + group_width: int, + bottleneck_multiplier: float, + se_ratio: Optional[float], + ) -> None: + layers: OrderedDict[str, nn.Module] = OrderedDict() + w_b = int(round(width_out * bottleneck_multiplier)) + g = w_b // group_width + + layers["a"] = Conv2dNormActivation( + width_in, w_b, kernel_size=1, stride=1, norm_layer=norm_layer, activation_layer=activation_layer + ) + layers["b"] = Conv2dNormActivation( + w_b, w_b, kernel_size=3, stride=stride, groups=g, norm_layer=norm_layer, activation_layer=activation_layer + ) + + if se_ratio: + # The SE reduction ratio is defined with respect to the + # beginning of the block + width_se_out = int(round(se_ratio * width_in)) + layers["se"] = SqueezeExcitation( + input_channels=w_b, + squeeze_channels=width_se_out, + activation=activation_layer, + ) + + layers["c"] = Conv2dNormActivation( + w_b, width_out, kernel_size=1, stride=1, norm_layer=norm_layer, activation_layer=None + ) + super().__init__(layers) + + +class ResBottleneckBlock(nn.Module): + """Residual bottleneck block: x + F(x), F = bottleneck transform.""" + + def __init__( + self, + width_in: int, + width_out: int, + stride: int, + norm_layer: Callable[..., nn.Module], + activation_layer: Callable[..., nn.Module], + group_width: int = 1, + bottleneck_multiplier: float = 1.0, + se_ratio: Optional[float] = None, + ) -> None: + super().__init__() + + # Use skip connection with projection if shape changes + self.proj = None + should_proj = (width_in != width_out) or (stride != 1) + if should_proj: + self.proj = Conv2dNormActivation( + width_in, width_out, kernel_size=1, stride=stride, norm_layer=norm_layer, activation_layer=None + ) + self.f = BottleneckTransform( + width_in, + width_out, + stride, + norm_layer, + activation_layer, + group_width, + bottleneck_multiplier, + se_ratio, + ) + self.activation = activation_layer(inplace=True) + + def forward(self, x: Tensor) -> Tensor: + if self.proj is not None: + x = self.proj(x) + self.f(x) + else: + x = x + self.f(x) + return self.activation(x) + + +class AnyStage(nn.Sequential): + """AnyNet stage (sequence of blocks w/ the same output shape).""" + + def __init__( + self, + width_in: int, + width_out: int, + stride: int, + depth: int, + block_constructor: Callable[..., nn.Module], + norm_layer: Callable[..., nn.Module], + activation_layer: Callable[..., nn.Module], + group_width: int, + bottleneck_multiplier: float, + se_ratio: Optional[float] = None, + stage_index: int = 0, + ) -> None: + super().__init__() + + for i in range(depth): + block = block_constructor( + width_in if i == 0 else width_out, + width_out, + stride if i == 0 else 1, + norm_layer, + activation_layer, + group_width, + bottleneck_multiplier, + se_ratio, + ) + + self.add_module(f"block{stage_index}-{i}", block) + + +class BlockParams: + def __init__( + self, + depths: List[int], + widths: List[int], + group_widths: List[int], + bottleneck_multipliers: List[float], + strides: List[int], + se_ratio: Optional[float] = None, + ) -> None: + self.depths = depths + self.widths = widths + self.group_widths = group_widths + self.bottleneck_multipliers = bottleneck_multipliers + self.strides = strides + self.se_ratio = se_ratio + + @classmethod + def from_init_params( + cls, + depth: int, + w_0: int, + w_a: float, + w_m: float, + group_width: int, + bottleneck_multiplier: float = 1.0, + se_ratio: Optional[float] = None, + **kwargs: Any, + ) -> "BlockParams": + """ + Programmatically compute all the per-block settings, + given the RegNet parameters. + + The first step is to compute the quantized linear block parameters, + in log space. Key parameters are: + - `w_a` is the width progression slope + - `w_0` is the initial width + - `w_m` is the width stepping in the log space + + In other terms + `log(block_width) = log(w_0) + w_m * block_capacity`, + with `bock_capacity` ramping up following the w_0 and w_a params. + This block width is finally quantized to multiples of 8. + + The second step is to compute the parameters per stage, + taking into account the skip connection and the final 1x1 convolutions. + We use the fact that the output width is constant within a stage. + """ + + QUANT = 8 + STRIDE = 2 + + if w_a < 0 or w_0 <= 0 or w_m <= 1 or w_0 % 8 != 0: + raise ValueError("Invalid RegNet settings") + # Compute the block widths. Each stage has one unique block width + widths_cont = torch.arange(depth) * w_a + w_0 + block_capacity = torch.round(torch.log(widths_cont / w_0) / math.log(w_m)) + block_widths = (torch.round(torch.divide(w_0 * torch.pow(w_m, block_capacity), QUANT)) * QUANT).int().tolist() + num_stages = len(set(block_widths)) + + # Convert to per stage parameters + split_helper = zip( + block_widths + [0], + [0] + block_widths, + block_widths + [0], + [0] + block_widths, + ) + splits = [w != wp or r != rp for w, wp, r, rp in split_helper] + + stage_widths = [w for w, t in zip(block_widths, splits[:-1]) if t] + stage_depths = torch.diff(torch.tensor([d for d, t in enumerate(splits) if t])).int().tolist() + + strides = [STRIDE] * num_stages + bottleneck_multipliers = [bottleneck_multiplier] * num_stages + group_widths = [group_width] * num_stages + + # Adjust the compatibility of stage widths and group widths + stage_widths, group_widths = cls._adjust_widths_groups_compatibilty( + stage_widths, bottleneck_multipliers, group_widths + ) + + return cls( + depths=stage_depths, + widths=stage_widths, + group_widths=group_widths, + bottleneck_multipliers=bottleneck_multipliers, + strides=strides, + se_ratio=se_ratio, + ) + + def _get_expanded_params(self): + return zip(self.widths, self.strides, self.depths, self.group_widths, self.bottleneck_multipliers) + + @staticmethod + def _adjust_widths_groups_compatibilty( + stage_widths: List[int], bottleneck_ratios: List[float], group_widths: List[int] + ) -> Tuple[List[int], List[int]]: + """ + Adjusts the compatibility of widths and groups, + depending on the bottleneck ratio. + """ + # Compute all widths for the current settings + widths = [int(w * b) for w, b in zip(stage_widths, bottleneck_ratios)] + group_widths_min = [min(g, w_bot) for g, w_bot in zip(group_widths, widths)] + + # Compute the adjusted widths so that stage and group widths fit + ws_bot = [_make_divisible(w_bot, g) for w_bot, g in zip(widths, group_widths_min)] + stage_widths = [int(w_bot / b) for w_bot, b in zip(ws_bot, bottleneck_ratios)] + return stage_widths, group_widths_min + + +class RegNet(nn.Module): + def __init__( + self, + block_params: BlockParams, + num_classes: int = 1000, + stem_width: int = 32, + stem_type: Optional[Callable[..., nn.Module]] = None, + block_type: Optional[Callable[..., nn.Module]] = None, + norm_layer: Optional[Callable[..., nn.Module]] = None, + activation: Optional[Callable[..., nn.Module]] = None, + ) -> None: + super().__init__() + _log_api_usage_once(self) + + if stem_type is None: + stem_type = SimpleStemIN + if norm_layer is None: + norm_layer = nn.BatchNorm2d + if block_type is None: + block_type = ResBottleneckBlock + if activation is None: + activation = nn.ReLU + + # Ad hoc stem + self.stem = stem_type( + 3, # width_in + stem_width, + norm_layer, + activation, + ) + + current_width = stem_width + + blocks = [] + for i, ( + width_out, + stride, + depth, + group_width, + bottleneck_multiplier, + ) in enumerate(block_params._get_expanded_params()): + blocks.append( + ( + f"block{i+1}", + AnyStage( + current_width, + width_out, + stride, + depth, + block_type, + norm_layer, + activation, + group_width, + bottleneck_multiplier, + block_params.se_ratio, + stage_index=i + 1, + ), + ) + ) + + current_width = width_out + + self.trunk_output = nn.Sequential(OrderedDict(blocks)) + + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.fc = nn.Linear(in_features=current_width, out_features=num_classes) + + # Performs ResNet-style weight initialization + for m in self.modules(): + if isinstance(m, nn.Conv2d): + # Note that there is no bias due to BN + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + nn.init.normal_(m.weight, mean=0.0, std=math.sqrt(2.0 / fan_out)) + elif isinstance(m, nn.BatchNorm2d): + nn.init.ones_(m.weight) + nn.init.zeros_(m.bias) + elif isinstance(m, nn.Linear): + nn.init.normal_(m.weight, mean=0.0, std=0.01) + nn.init.zeros_(m.bias) + + def forward(self, x: Tensor) -> Tensor: + x = self.stem(x) + x = self.trunk_output(x) + + x = self.avgpool(x) + x = x.flatten(start_dim=1) + x = self.fc(x) + + return x + + +def _regnet( + block_params: BlockParams, + weights: Optional[WeightsEnum], + progress: bool, + **kwargs: Any, +) -> RegNet: + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + + norm_layer = kwargs.pop("norm_layer", partial(nn.BatchNorm2d, eps=1e-05, momentum=0.1)) + model = RegNet(block_params, norm_layer=norm_layer, **kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +_COMMON_META: Dict[str, Any] = { + "min_size": (1, 1), + "categories": _IMAGENET_CATEGORIES, +} + +_COMMON_SWAG_META = { + **_COMMON_META, + "recipe": "https://github.com/facebookresearch/SWAG", + "license": "https://github.com/facebookresearch/SWAG/blob/main/LICENSE", +} + + +class RegNet_Y_400MF_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/regnet_y_400mf-c65dace8.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 4344144, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#small-models", + "_metrics": { + "ImageNet-1K": { + "acc@1": 74.046, + "acc@5": 91.716, + } + }, + "_ops": 0.402, + "_file_size": 16.806, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/regnet_y_400mf-e6988f5f.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 4344144, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe", + "_metrics": { + "ImageNet-1K": { + "acc@1": 75.804, + "acc@5": 92.742, + } + }, + "_ops": 0.402, + "_file_size": 16.806, + "_docs": """ + These weights improve upon the results of the original paper by using a modified version of TorchVision's + `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class RegNet_Y_800MF_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/regnet_y_800mf-1b27b58c.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 6432512, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#small-models", + "_metrics": { + "ImageNet-1K": { + "acc@1": 76.420, + "acc@5": 93.136, + } + }, + "_ops": 0.834, + "_file_size": 24.774, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/regnet_y_800mf-58fc7688.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 6432512, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe", + "_metrics": { + "ImageNet-1K": { + "acc@1": 78.828, + "acc@5": 94.502, + } + }, + "_ops": 0.834, + "_file_size": 24.774, + "_docs": """ + These weights improve upon the results of the original paper by using a modified version of TorchVision's + `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class RegNet_Y_1_6GF_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/regnet_y_1_6gf-b11a554e.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 11202430, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#small-models", + "_metrics": { + "ImageNet-1K": { + "acc@1": 77.950, + "acc@5": 93.966, + } + }, + "_ops": 1.612, + "_file_size": 43.152, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/regnet_y_1_6gf-0d7bc02a.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 11202430, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe", + "_metrics": { + "ImageNet-1K": { + "acc@1": 80.876, + "acc@5": 95.444, + } + }, + "_ops": 1.612, + "_file_size": 43.152, + "_docs": """ + These weights improve upon the results of the original paper by using a modified version of TorchVision's + `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class RegNet_Y_3_2GF_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/regnet_y_3_2gf-b5a9779c.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 19436338, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#medium-models", + "_metrics": { + "ImageNet-1K": { + "acc@1": 78.948, + "acc@5": 94.576, + } + }, + "_ops": 3.176, + "_file_size": 74.567, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/regnet_y_3_2gf-9180c971.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 19436338, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe", + "_metrics": { + "ImageNet-1K": { + "acc@1": 81.982, + "acc@5": 95.972, + } + }, + "_ops": 3.176, + "_file_size": 74.567, + "_docs": """ + These weights improve upon the results of the original paper by using a modified version of TorchVision's + `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class RegNet_Y_8GF_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/regnet_y_8gf-d0d0e4a8.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 39381472, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#medium-models", + "_metrics": { + "ImageNet-1K": { + "acc@1": 80.032, + "acc@5": 95.048, + } + }, + "_ops": 8.473, + "_file_size": 150.701, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/regnet_y_8gf-dc2b1b54.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 39381472, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe", + "_metrics": { + "ImageNet-1K": { + "acc@1": 82.828, + "acc@5": 96.330, + } + }, + "_ops": 8.473, + "_file_size": 150.701, + "_docs": """ + These weights improve upon the results of the original paper by using a modified version of TorchVision's + `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class RegNet_Y_16GF_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/regnet_y_16gf-9e6ed7dd.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 83590140, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#large-models", + "_metrics": { + "ImageNet-1K": { + "acc@1": 80.424, + "acc@5": 95.240, + } + }, + "_ops": 15.912, + "_file_size": 319.49, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/regnet_y_16gf-3e4a00f9.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 83590140, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe", + "_metrics": { + "ImageNet-1K": { + "acc@1": 82.886, + "acc@5": 96.328, + } + }, + "_ops": 15.912, + "_file_size": 319.49, + "_docs": """ + These weights improve upon the results of the original paper by using a modified version of TorchVision's + `new training recipe + `_. + """, + }, + ) + IMAGENET1K_SWAG_E2E_V1 = Weights( + url="https://download.pytorch.org/models/regnet_y_16gf_swag-43afe44d.pth", + transforms=partial( + ImageClassification, crop_size=384, resize_size=384, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_SWAG_META, + "num_params": 83590140, + "_metrics": { + "ImageNet-1K": { + "acc@1": 86.012, + "acc@5": 98.054, + } + }, + "_ops": 46.735, + "_file_size": 319.49, + "_docs": """ + These weights are learnt via transfer learning by end-to-end fine-tuning the original + `SWAG `_ weights on ImageNet-1K data. + """, + }, + ) + IMAGENET1K_SWAG_LINEAR_V1 = Weights( + url="https://download.pytorch.org/models/regnet_y_16gf_lc_swag-f3ec0043.pth", + transforms=partial( + ImageClassification, crop_size=224, resize_size=224, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_SWAG_META, + "recipe": "https://github.com/pytorch/vision/pull/5793", + "num_params": 83590140, + "_metrics": { + "ImageNet-1K": { + "acc@1": 83.976, + "acc@5": 97.244, + } + }, + "_ops": 15.912, + "_file_size": 319.49, + "_docs": """ + These weights are composed of the original frozen `SWAG `_ trunk + weights and a linear classifier learnt on top of them trained on ImageNet-1K data. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class RegNet_Y_32GF_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/regnet_y_32gf-4dee3f7a.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 145046770, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#large-models", + "_metrics": { + "ImageNet-1K": { + "acc@1": 80.878, + "acc@5": 95.340, + } + }, + "_ops": 32.28, + "_file_size": 554.076, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/regnet_y_32gf-8db6d4b5.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 145046770, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe", + "_metrics": { + "ImageNet-1K": { + "acc@1": 83.368, + "acc@5": 96.498, + } + }, + "_ops": 32.28, + "_file_size": 554.076, + "_docs": """ + These weights improve upon the results of the original paper by using a modified version of TorchVision's + `new training recipe + `_. + """, + }, + ) + IMAGENET1K_SWAG_E2E_V1 = Weights( + url="https://download.pytorch.org/models/regnet_y_32gf_swag-04fdfa75.pth", + transforms=partial( + ImageClassification, crop_size=384, resize_size=384, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_SWAG_META, + "num_params": 145046770, + "_metrics": { + "ImageNet-1K": { + "acc@1": 86.838, + "acc@5": 98.362, + } + }, + "_ops": 94.826, + "_file_size": 554.076, + "_docs": """ + These weights are learnt via transfer learning by end-to-end fine-tuning the original + `SWAG `_ weights on ImageNet-1K data. + """, + }, + ) + IMAGENET1K_SWAG_LINEAR_V1 = Weights( + url="https://download.pytorch.org/models/regnet_y_32gf_lc_swag-e1583746.pth", + transforms=partial( + ImageClassification, crop_size=224, resize_size=224, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_SWAG_META, + "recipe": "https://github.com/pytorch/vision/pull/5793", + "num_params": 145046770, + "_metrics": { + "ImageNet-1K": { + "acc@1": 84.622, + "acc@5": 97.480, + } + }, + "_ops": 32.28, + "_file_size": 554.076, + "_docs": """ + These weights are composed of the original frozen `SWAG `_ trunk + weights and a linear classifier learnt on top of them trained on ImageNet-1K data. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class RegNet_Y_128GF_Weights(WeightsEnum): + IMAGENET1K_SWAG_E2E_V1 = Weights( + url="https://download.pytorch.org/models/regnet_y_128gf_swag-c8ce3e52.pth", + transforms=partial( + ImageClassification, crop_size=384, resize_size=384, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_SWAG_META, + "num_params": 644812894, + "_metrics": { + "ImageNet-1K": { + "acc@1": 88.228, + "acc@5": 98.682, + } + }, + "_ops": 374.57, + "_file_size": 2461.564, + "_docs": """ + These weights are learnt via transfer learning by end-to-end fine-tuning the original + `SWAG `_ weights on ImageNet-1K data. + """, + }, + ) + IMAGENET1K_SWAG_LINEAR_V1 = Weights( + url="https://download.pytorch.org/models/regnet_y_128gf_lc_swag-cbe8ce12.pth", + transforms=partial( + ImageClassification, crop_size=224, resize_size=224, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_SWAG_META, + "recipe": "https://github.com/pytorch/vision/pull/5793", + "num_params": 644812894, + "_metrics": { + "ImageNet-1K": { + "acc@1": 86.068, + "acc@5": 97.844, + } + }, + "_ops": 127.518, + "_file_size": 2461.564, + "_docs": """ + These weights are composed of the original frozen `SWAG `_ trunk + weights and a linear classifier learnt on top of them trained on ImageNet-1K data. + """, + }, + ) + DEFAULT = IMAGENET1K_SWAG_E2E_V1 + + +class RegNet_X_400MF_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/regnet_x_400mf-adf1edd5.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 5495976, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#small-models", + "_metrics": { + "ImageNet-1K": { + "acc@1": 72.834, + "acc@5": 90.950, + } + }, + "_ops": 0.414, + "_file_size": 21.258, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/regnet_x_400mf-62229a5f.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 5495976, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe-with-fixres", + "_metrics": { + "ImageNet-1K": { + "acc@1": 74.864, + "acc@5": 92.322, + } + }, + "_ops": 0.414, + "_file_size": 21.257, + "_docs": """ + These weights improve upon the results of the original paper by using a modified version of TorchVision's + `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class RegNet_X_800MF_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/regnet_x_800mf-ad17e45c.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 7259656, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#small-models", + "_metrics": { + "ImageNet-1K": { + "acc@1": 75.212, + "acc@5": 92.348, + } + }, + "_ops": 0.8, + "_file_size": 27.945, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/regnet_x_800mf-94a99ebd.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 7259656, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe-with-fixres", + "_metrics": { + "ImageNet-1K": { + "acc@1": 77.522, + "acc@5": 93.826, + } + }, + "_ops": 0.8, + "_file_size": 27.945, + "_docs": """ + These weights improve upon the results of the original paper by using a modified version of TorchVision's + `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class RegNet_X_1_6GF_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/regnet_x_1_6gf-e3633e7f.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 9190136, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#small-models", + "_metrics": { + "ImageNet-1K": { + "acc@1": 77.040, + "acc@5": 93.440, + } + }, + "_ops": 1.603, + "_file_size": 35.339, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/regnet_x_1_6gf-a12f2b72.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 9190136, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe-with-fixres", + "_metrics": { + "ImageNet-1K": { + "acc@1": 79.668, + "acc@5": 94.922, + } + }, + "_ops": 1.603, + "_file_size": 35.339, + "_docs": """ + These weights improve upon the results of the original paper by using a modified version of TorchVision's + `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class RegNet_X_3_2GF_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/regnet_x_3_2gf-f342aeae.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 15296552, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#medium-models", + "_metrics": { + "ImageNet-1K": { + "acc@1": 78.364, + "acc@5": 93.992, + } + }, + "_ops": 3.177, + "_file_size": 58.756, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/regnet_x_3_2gf-7071aa85.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 15296552, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe", + "_metrics": { + "ImageNet-1K": { + "acc@1": 81.196, + "acc@5": 95.430, + } + }, + "_ops": 3.177, + "_file_size": 58.756, + "_docs": """ + These weights improve upon the results of the original paper by using a modified version of TorchVision's + `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class RegNet_X_8GF_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/regnet_x_8gf-03ceed89.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 39572648, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#medium-models", + "_metrics": { + "ImageNet-1K": { + "acc@1": 79.344, + "acc@5": 94.686, + } + }, + "_ops": 7.995, + "_file_size": 151.456, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/regnet_x_8gf-2b70d774.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 39572648, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe", + "_metrics": { + "ImageNet-1K": { + "acc@1": 81.682, + "acc@5": 95.678, + } + }, + "_ops": 7.995, + "_file_size": 151.456, + "_docs": """ + These weights improve upon the results of the original paper by using a modified version of TorchVision's + `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class RegNet_X_16GF_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/regnet_x_16gf-2007eb11.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 54278536, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#medium-models", + "_metrics": { + "ImageNet-1K": { + "acc@1": 80.058, + "acc@5": 94.944, + } + }, + "_ops": 15.941, + "_file_size": 207.627, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/regnet_x_16gf-ba3796d7.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 54278536, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe", + "_metrics": { + "ImageNet-1K": { + "acc@1": 82.716, + "acc@5": 96.196, + } + }, + "_ops": 15.941, + "_file_size": 207.627, + "_docs": """ + These weights improve upon the results of the original paper by using a modified version of TorchVision's + `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class RegNet_X_32GF_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/regnet_x_32gf-9d47f8d0.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 107811560, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#large-models", + "_metrics": { + "ImageNet-1K": { + "acc@1": 80.622, + "acc@5": 95.248, + } + }, + "_ops": 31.736, + "_file_size": 412.039, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/regnet_x_32gf-6eb8fdc6.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 107811560, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe", + "_metrics": { + "ImageNet-1K": { + "acc@1": 83.014, + "acc@5": 96.288, + } + }, + "_ops": 31.736, + "_file_size": 412.039, + "_docs": """ + These weights improve upon the results of the original paper by using a modified version of TorchVision's + `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", RegNet_Y_400MF_Weights.IMAGENET1K_V1)) +def regnet_y_400mf(*, weights: Optional[RegNet_Y_400MF_Weights] = None, progress: bool = True, **kwargs: Any) -> RegNet: + """ + Constructs a RegNetY_400MF architecture from + `Designing Network Design Spaces `_. + + Args: + weights (:class:`~torchvision.models.RegNet_Y_400MF_Weights`, optional): The pretrained weights to use. + See :class:`~torchvision.models.RegNet_Y_400MF_Weights` below for more details and possible values. + By default, no pretrained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to either ``torchvision.models.regnet.RegNet`` or + ``torchvision.models.regnet.BlockParams`` class. Please refer to the `source code + `_ + for more detail about the classes. + + .. autoclass:: torchvision.models.RegNet_Y_400MF_Weights + :members: + """ + weights = RegNet_Y_400MF_Weights.verify(weights) + + params = BlockParams.from_init_params(depth=16, w_0=48, w_a=27.89, w_m=2.09, group_width=8, se_ratio=0.25, **kwargs) + return _regnet(params, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", RegNet_Y_800MF_Weights.IMAGENET1K_V1)) +def regnet_y_800mf(*, weights: Optional[RegNet_Y_800MF_Weights] = None, progress: bool = True, **kwargs: Any) -> RegNet: + """ + Constructs a RegNetY_800MF architecture from + `Designing Network Design Spaces `_. + + Args: + weights (:class:`~torchvision.models.RegNet_Y_800MF_Weights`, optional): The pretrained weights to use. + See :class:`~torchvision.models.RegNet_Y_800MF_Weights` below for more details and possible values. + By default, no pretrained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to either ``torchvision.models.regnet.RegNet`` or + ``torchvision.models.regnet.BlockParams`` class. Please refer to the `source code + `_ + for more detail about the classes. + + .. autoclass:: torchvision.models.RegNet_Y_800MF_Weights + :members: + """ + weights = RegNet_Y_800MF_Weights.verify(weights) + + params = BlockParams.from_init_params(depth=14, w_0=56, w_a=38.84, w_m=2.4, group_width=16, se_ratio=0.25, **kwargs) + return _regnet(params, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", RegNet_Y_1_6GF_Weights.IMAGENET1K_V1)) +def regnet_y_1_6gf(*, weights: Optional[RegNet_Y_1_6GF_Weights] = None, progress: bool = True, **kwargs: Any) -> RegNet: + """ + Constructs a RegNetY_1.6GF architecture from + `Designing Network Design Spaces `_. + + Args: + weights (:class:`~torchvision.models.RegNet_Y_1_6GF_Weights`, optional): The pretrained weights to use. + See :class:`~torchvision.models.RegNet_Y_1_6GF_Weights` below for more details and possible values. + By default, no pretrained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to either ``torchvision.models.regnet.RegNet`` or + ``torchvision.models.regnet.BlockParams`` class. Please refer to the `source code + `_ + for more detail about the classes. + + .. autoclass:: torchvision.models.RegNet_Y_1_6GF_Weights + :members: + """ + weights = RegNet_Y_1_6GF_Weights.verify(weights) + + params = BlockParams.from_init_params( + depth=27, w_0=48, w_a=20.71, w_m=2.65, group_width=24, se_ratio=0.25, **kwargs + ) + return _regnet(params, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", RegNet_Y_3_2GF_Weights.IMAGENET1K_V1)) +def regnet_y_3_2gf(*, weights: Optional[RegNet_Y_3_2GF_Weights] = None, progress: bool = True, **kwargs: Any) -> RegNet: + """ + Constructs a RegNetY_3.2GF architecture from + `Designing Network Design Spaces `_. + + Args: + weights (:class:`~torchvision.models.RegNet_Y_3_2GF_Weights`, optional): The pretrained weights to use. + See :class:`~torchvision.models.RegNet_Y_3_2GF_Weights` below for more details and possible values. + By default, no pretrained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to either ``torchvision.models.regnet.RegNet`` or + ``torchvision.models.regnet.BlockParams`` class. Please refer to the `source code + `_ + for more detail about the classes. + + .. autoclass:: torchvision.models.RegNet_Y_3_2GF_Weights + :members: + """ + weights = RegNet_Y_3_2GF_Weights.verify(weights) + + params = BlockParams.from_init_params( + depth=21, w_0=80, w_a=42.63, w_m=2.66, group_width=24, se_ratio=0.25, **kwargs + ) + return _regnet(params, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", RegNet_Y_8GF_Weights.IMAGENET1K_V1)) +def regnet_y_8gf(*, weights: Optional[RegNet_Y_8GF_Weights] = None, progress: bool = True, **kwargs: Any) -> RegNet: + """ + Constructs a RegNetY_8GF architecture from + `Designing Network Design Spaces `_. + + Args: + weights (:class:`~torchvision.models.RegNet_Y_8GF_Weights`, optional): The pretrained weights to use. + See :class:`~torchvision.models.RegNet_Y_8GF_Weights` below for more details and possible values. + By default, no pretrained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to either ``torchvision.models.regnet.RegNet`` or + ``torchvision.models.regnet.BlockParams`` class. Please refer to the `source code + `_ + for more detail about the classes. + + .. autoclass:: torchvision.models.RegNet_Y_8GF_Weights + :members: + """ + weights = RegNet_Y_8GF_Weights.verify(weights) + + params = BlockParams.from_init_params( + depth=17, w_0=192, w_a=76.82, w_m=2.19, group_width=56, se_ratio=0.25, **kwargs + ) + return _regnet(params, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", RegNet_Y_16GF_Weights.IMAGENET1K_V1)) +def regnet_y_16gf(*, weights: Optional[RegNet_Y_16GF_Weights] = None, progress: bool = True, **kwargs: Any) -> RegNet: + """ + Constructs a RegNetY_16GF architecture from + `Designing Network Design Spaces `_. + + Args: + weights (:class:`~torchvision.models.RegNet_Y_16GF_Weights`, optional): The pretrained weights to use. + See :class:`~torchvision.models.RegNet_Y_16GF_Weights` below for more details and possible values. + By default, no pretrained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to either ``torchvision.models.regnet.RegNet`` or + ``torchvision.models.regnet.BlockParams`` class. Please refer to the `source code + `_ + for more detail about the classes. + + .. autoclass:: torchvision.models.RegNet_Y_16GF_Weights + :members: + """ + weights = RegNet_Y_16GF_Weights.verify(weights) + + params = BlockParams.from_init_params( + depth=18, w_0=200, w_a=106.23, w_m=2.48, group_width=112, se_ratio=0.25, **kwargs + ) + return _regnet(params, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", RegNet_Y_32GF_Weights.IMAGENET1K_V1)) +def regnet_y_32gf(*, weights: Optional[RegNet_Y_32GF_Weights] = None, progress: bool = True, **kwargs: Any) -> RegNet: + """ + Constructs a RegNetY_32GF architecture from + `Designing Network Design Spaces `_. + + Args: + weights (:class:`~torchvision.models.RegNet_Y_32GF_Weights`, optional): The pretrained weights to use. + See :class:`~torchvision.models.RegNet_Y_32GF_Weights` below for more details and possible values. + By default, no pretrained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to either ``torchvision.models.regnet.RegNet`` or + ``torchvision.models.regnet.BlockParams`` class. Please refer to the `source code + `_ + for more detail about the classes. + + .. autoclass:: torchvision.models.RegNet_Y_32GF_Weights + :members: + """ + weights = RegNet_Y_32GF_Weights.verify(weights) + + params = BlockParams.from_init_params( + depth=20, w_0=232, w_a=115.89, w_m=2.53, group_width=232, se_ratio=0.25, **kwargs + ) + return _regnet(params, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", None)) +def regnet_y_128gf(*, weights: Optional[RegNet_Y_128GF_Weights] = None, progress: bool = True, **kwargs: Any) -> RegNet: + """ + Constructs a RegNetY_128GF architecture from + `Designing Network Design Spaces `_. + + Args: + weights (:class:`~torchvision.models.RegNet_Y_128GF_Weights`, optional): The pretrained weights to use. + See :class:`~torchvision.models.RegNet_Y_128GF_Weights` below for more details and possible values. + By default, no pretrained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to either ``torchvision.models.regnet.RegNet`` or + ``torchvision.models.regnet.BlockParams`` class. Please refer to the `source code + `_ + for more detail about the classes. + + .. autoclass:: torchvision.models.RegNet_Y_128GF_Weights + :members: + """ + weights = RegNet_Y_128GF_Weights.verify(weights) + + params = BlockParams.from_init_params( + depth=27, w_0=456, w_a=160.83, w_m=2.52, group_width=264, se_ratio=0.25, **kwargs + ) + return _regnet(params, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", RegNet_X_400MF_Weights.IMAGENET1K_V1)) +def regnet_x_400mf(*, weights: Optional[RegNet_X_400MF_Weights] = None, progress: bool = True, **kwargs: Any) -> RegNet: + """ + Constructs a RegNetX_400MF architecture from + `Designing Network Design Spaces `_. + + Args: + weights (:class:`~torchvision.models.RegNet_X_400MF_Weights`, optional): The pretrained weights to use. + See :class:`~torchvision.models.RegNet_X_400MF_Weights` below for more details and possible values. + By default, no pretrained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to either ``torchvision.models.regnet.RegNet`` or + ``torchvision.models.regnet.BlockParams`` class. Please refer to the `source code + `_ + for more detail about the classes. + + .. autoclass:: torchvision.models.RegNet_X_400MF_Weights + :members: + """ + weights = RegNet_X_400MF_Weights.verify(weights) + + params = BlockParams.from_init_params(depth=22, w_0=24, w_a=24.48, w_m=2.54, group_width=16, **kwargs) + return _regnet(params, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", RegNet_X_800MF_Weights.IMAGENET1K_V1)) +def regnet_x_800mf(*, weights: Optional[RegNet_X_800MF_Weights] = None, progress: bool = True, **kwargs: Any) -> RegNet: + """ + Constructs a RegNetX_800MF architecture from + `Designing Network Design Spaces `_. + + Args: + weights (:class:`~torchvision.models.RegNet_X_800MF_Weights`, optional): The pretrained weights to use. + See :class:`~torchvision.models.RegNet_X_800MF_Weights` below for more details and possible values. + By default, no pretrained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to either ``torchvision.models.regnet.RegNet`` or + ``torchvision.models.regnet.BlockParams`` class. Please refer to the `source code + `_ + for more detail about the classes. + + .. autoclass:: torchvision.models.RegNet_X_800MF_Weights + :members: + """ + weights = RegNet_X_800MF_Weights.verify(weights) + + params = BlockParams.from_init_params(depth=16, w_0=56, w_a=35.73, w_m=2.28, group_width=16, **kwargs) + return _regnet(params, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", RegNet_X_1_6GF_Weights.IMAGENET1K_V1)) +def regnet_x_1_6gf(*, weights: Optional[RegNet_X_1_6GF_Weights] = None, progress: bool = True, **kwargs: Any) -> RegNet: + """ + Constructs a RegNetX_1.6GF architecture from + `Designing Network Design Spaces `_. + + Args: + weights (:class:`~torchvision.models.RegNet_X_1_6GF_Weights`, optional): The pretrained weights to use. + See :class:`~torchvision.models.RegNet_X_1_6GF_Weights` below for more details and possible values. + By default, no pretrained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to either ``torchvision.models.regnet.RegNet`` or + ``torchvision.models.regnet.BlockParams`` class. Please refer to the `source code + `_ + for more detail about the classes. + + .. autoclass:: torchvision.models.RegNet_X_1_6GF_Weights + :members: + """ + weights = RegNet_X_1_6GF_Weights.verify(weights) + + params = BlockParams.from_init_params(depth=18, w_0=80, w_a=34.01, w_m=2.25, group_width=24, **kwargs) + return _regnet(params, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", RegNet_X_3_2GF_Weights.IMAGENET1K_V1)) +def regnet_x_3_2gf(*, weights: Optional[RegNet_X_3_2GF_Weights] = None, progress: bool = True, **kwargs: Any) -> RegNet: + """ + Constructs a RegNetX_3.2GF architecture from + `Designing Network Design Spaces `_. + + Args: + weights (:class:`~torchvision.models.RegNet_X_3_2GF_Weights`, optional): The pretrained weights to use. + See :class:`~torchvision.models.RegNet_X_3_2GF_Weights` below for more details and possible values. + By default, no pretrained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to either ``torchvision.models.regnet.RegNet`` or + ``torchvision.models.regnet.BlockParams`` class. Please refer to the `source code + `_ + for more detail about the classes. + + .. autoclass:: torchvision.models.RegNet_X_3_2GF_Weights + :members: + """ + weights = RegNet_X_3_2GF_Weights.verify(weights) + + params = BlockParams.from_init_params(depth=25, w_0=88, w_a=26.31, w_m=2.25, group_width=48, **kwargs) + return _regnet(params, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", RegNet_X_8GF_Weights.IMAGENET1K_V1)) +def regnet_x_8gf(*, weights: Optional[RegNet_X_8GF_Weights] = None, progress: bool = True, **kwargs: Any) -> RegNet: + """ + Constructs a RegNetX_8GF architecture from + `Designing Network Design Spaces `_. + + Args: + weights (:class:`~torchvision.models.RegNet_X_8GF_Weights`, optional): The pretrained weights to use. + See :class:`~torchvision.models.RegNet_X_8GF_Weights` below for more details and possible values. + By default, no pretrained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to either ``torchvision.models.regnet.RegNet`` or + ``torchvision.models.regnet.BlockParams`` class. Please refer to the `source code + `_ + for more detail about the classes. + + .. autoclass:: torchvision.models.RegNet_X_8GF_Weights + :members: + """ + weights = RegNet_X_8GF_Weights.verify(weights) + + params = BlockParams.from_init_params(depth=23, w_0=80, w_a=49.56, w_m=2.88, group_width=120, **kwargs) + return _regnet(params, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", RegNet_X_16GF_Weights.IMAGENET1K_V1)) +def regnet_x_16gf(*, weights: Optional[RegNet_X_16GF_Weights] = None, progress: bool = True, **kwargs: Any) -> RegNet: + """ + Constructs a RegNetX_16GF architecture from + `Designing Network Design Spaces `_. + + Args: + weights (:class:`~torchvision.models.RegNet_X_16GF_Weights`, optional): The pretrained weights to use. + See :class:`~torchvision.models.RegNet_X_16GF_Weights` below for more details and possible values. + By default, no pretrained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to either ``torchvision.models.regnet.RegNet`` or + ``torchvision.models.regnet.BlockParams`` class. Please refer to the `source code + `_ + for more detail about the classes. + + .. autoclass:: torchvision.models.RegNet_X_16GF_Weights + :members: + """ + weights = RegNet_X_16GF_Weights.verify(weights) + + params = BlockParams.from_init_params(depth=22, w_0=216, w_a=55.59, w_m=2.1, group_width=128, **kwargs) + return _regnet(params, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", RegNet_X_32GF_Weights.IMAGENET1K_V1)) +def regnet_x_32gf(*, weights: Optional[RegNet_X_32GF_Weights] = None, progress: bool = True, **kwargs: Any) -> RegNet: + """ + Constructs a RegNetX_32GF architecture from + `Designing Network Design Spaces `_. + + Args: + weights (:class:`~torchvision.models.RegNet_X_32GF_Weights`, optional): The pretrained weights to use. + See :class:`~torchvision.models.RegNet_X_32GF_Weights` below for more details and possible values. + By default, no pretrained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to either ``torchvision.models.regnet.RegNet`` or + ``torchvision.models.regnet.BlockParams`` class. Please refer to the `source code + `_ + for more detail about the classes. + + .. autoclass:: torchvision.models.RegNet_X_32GF_Weights + :members: + """ + weights = RegNet_X_32GF_Weights.verify(weights) + + params = BlockParams.from_init_params(depth=23, w_0=320, w_a=69.86, w_m=2.0, group_width=168, **kwargs) + return _regnet(params, weights, progress, **kwargs) diff --git a/lib/python3.12/site-packages/torchvision/models/resnet.py b/lib/python3.12/site-packages/torchvision/models/resnet.py new file mode 100644 index 0000000000000000000000000000000000000000..83c0340cef74d9cb4c1dc92c38f4b3024be1f731 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/resnet.py @@ -0,0 +1,985 @@ +from functools import partial +from typing import Any, Callable, List, Optional, Type, Union + +import torch +import torch.nn as nn +from torch import Tensor + +from ..transforms._presets import ImageClassification +from ..utils import _log_api_usage_once +from ._api import register_model, Weights, WeightsEnum +from ._meta import _IMAGENET_CATEGORIES +from ._utils import _ovewrite_named_param, handle_legacy_interface + + +__all__ = [ + "ResNet", + "ResNet18_Weights", + "ResNet34_Weights", + "ResNet50_Weights", + "ResNet101_Weights", + "ResNet152_Weights", + "ResNeXt50_32X4D_Weights", + "ResNeXt101_32X8D_Weights", + "ResNeXt101_64X4D_Weights", + "Wide_ResNet50_2_Weights", + "Wide_ResNet101_2_Weights", + "resnet18", + "resnet34", + "resnet50", + "resnet101", + "resnet152", + "resnext50_32x4d", + "resnext101_32x8d", + "resnext101_64x4d", + "wide_resnet50_2", + "wide_resnet101_2", +] + + +def conv3x3(in_planes: int, out_planes: int, stride: int = 1, groups: int = 1, dilation: int = 1) -> nn.Conv2d: + """3x3 convolution with padding""" + return nn.Conv2d( + in_planes, + out_planes, + kernel_size=3, + stride=stride, + padding=dilation, + groups=groups, + bias=False, + dilation=dilation, + ) + + +def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d: + """1x1 convolution""" + return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) + + +class BasicBlock(nn.Module): + expansion: int = 1 + + def __init__( + self, + inplanes: int, + planes: int, + stride: int = 1, + downsample: Optional[nn.Module] = None, + groups: int = 1, + base_width: int = 64, + dilation: int = 1, + norm_layer: Optional[Callable[..., nn.Module]] = None, + ) -> None: + super().__init__() + if norm_layer is None: + norm_layer = nn.BatchNorm2d + if groups != 1 or base_width != 64: + raise ValueError("BasicBlock only supports groups=1 and base_width=64") + if dilation > 1: + raise NotImplementedError("Dilation > 1 not supported in BasicBlock") + # Both self.conv1 and self.downsample layers downsample the input when stride != 1 + self.conv1 = conv3x3(inplanes, planes, stride) + self.bn1 = norm_layer(planes) + self.relu = nn.ReLU(inplace=True) + self.conv2 = conv3x3(planes, planes) + self.bn2 = norm_layer(planes) + self.downsample = downsample + self.stride = stride + + def forward(self, x: Tensor) -> Tensor: + identity = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + + if self.downsample is not None: + identity = self.downsample(x) + + out += identity + out = self.relu(out) + + return out + + +class Bottleneck(nn.Module): + # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2) + # while original implementation places the stride at the first 1x1 convolution(self.conv1) + # according to "Deep residual learning for image recognition" https://arxiv.org/abs/1512.03385. + # This variant is also known as ResNet V1.5 and improves accuracy according to + # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch. + + expansion: int = 4 + + def __init__( + self, + inplanes: int, + planes: int, + stride: int = 1, + downsample: Optional[nn.Module] = None, + groups: int = 1, + base_width: int = 64, + dilation: int = 1, + norm_layer: Optional[Callable[..., nn.Module]] = None, + ) -> None: + super().__init__() + if norm_layer is None: + norm_layer = nn.BatchNorm2d + width = int(planes * (base_width / 64.0)) * groups + # Both self.conv2 and self.downsample layers downsample the input when stride != 1 + self.conv1 = conv1x1(inplanes, width) + self.bn1 = norm_layer(width) + self.conv2 = conv3x3(width, width, stride, groups, dilation) + self.bn2 = norm_layer(width) + self.conv3 = conv1x1(width, planes * self.expansion) + self.bn3 = norm_layer(planes * self.expansion) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x: Tensor) -> Tensor: + identity = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + if self.downsample is not None: + identity = self.downsample(x) + + out += identity + out = self.relu(out) + + return out + + +class ResNet(nn.Module): + def __init__( + self, + block: Type[Union[BasicBlock, Bottleneck]], + layers: List[int], + num_classes: int = 1000, + zero_init_residual: bool = False, + groups: int = 1, + width_per_group: int = 64, + replace_stride_with_dilation: Optional[List[bool]] = None, + norm_layer: Optional[Callable[..., nn.Module]] = None, + ) -> None: + super().__init__() + _log_api_usage_once(self) + if norm_layer is None: + norm_layer = nn.BatchNorm2d + self._norm_layer = norm_layer + + self.inplanes = 64 + self.dilation = 1 + if replace_stride_with_dilation is None: + # each element in the tuple indicates if we should replace + # the 2x2 stride with a dilated convolution instead + replace_stride_with_dilation = [False, False, False] + if len(replace_stride_with_dilation) != 3: + raise ValueError( + "replace_stride_with_dilation should be None " + f"or a 3-element tuple, got {replace_stride_with_dilation}" + ) + self.groups = groups + self.base_width = width_per_group + self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False) + self.bn1 = norm_layer(self.inplanes) + self.relu = nn.ReLU(inplace=True) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.layer1 = self._make_layer(block, 64, layers[0]) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2, dilate=replace_stride_with_dilation[0]) + self.layer3 = self._make_layer(block, 256, layers[2], stride=2, dilate=replace_stride_with_dilation[1]) + self.layer4 = self._make_layer(block, 512, layers[3], stride=2, dilate=replace_stride_with_dilation[2]) + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.fc = nn.Linear(512 * block.expansion, num_classes) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") + elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + # Zero-initialize the last BN in each residual branch, + # so that the residual branch starts with zeros, and each residual block behaves like an identity. + # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677 + if zero_init_residual: + for m in self.modules(): + if isinstance(m, Bottleneck) and m.bn3.weight is not None: + nn.init.constant_(m.bn3.weight, 0) # type: ignore[arg-type] + elif isinstance(m, BasicBlock) and m.bn2.weight is not None: + nn.init.constant_(m.bn2.weight, 0) # type: ignore[arg-type] + + def _make_layer( + self, + block: Type[Union[BasicBlock, Bottleneck]], + planes: int, + blocks: int, + stride: int = 1, + dilate: bool = False, + ) -> nn.Sequential: + norm_layer = self._norm_layer + downsample = None + previous_dilation = self.dilation + if dilate: + self.dilation *= stride + stride = 1 + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + conv1x1(self.inplanes, planes * block.expansion, stride), + norm_layer(planes * block.expansion), + ) + + layers = [] + layers.append( + block( + self.inplanes, planes, stride, downsample, self.groups, self.base_width, previous_dilation, norm_layer + ) + ) + self.inplanes = planes * block.expansion + for _ in range(1, blocks): + layers.append( + block( + self.inplanes, + planes, + groups=self.groups, + base_width=self.base_width, + dilation=self.dilation, + norm_layer=norm_layer, + ) + ) + + return nn.Sequential(*layers) + + def _forward_impl(self, x: Tensor) -> Tensor: + # See note [TorchScript super()] + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.maxpool(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + + x = self.avgpool(x) + x = torch.flatten(x, 1) + x = self.fc(x) + + return x + + def forward(self, x: Tensor) -> Tensor: + return self._forward_impl(x) + + +def _resnet( + block: Type[Union[BasicBlock, Bottleneck]], + layers: List[int], + weights: Optional[WeightsEnum], + progress: bool, + **kwargs: Any, +) -> ResNet: + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + + model = ResNet(block, layers, **kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +_COMMON_META = { + "min_size": (1, 1), + "categories": _IMAGENET_CATEGORIES, +} + + +class ResNet18_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/resnet18-f37072fd.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 11689512, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnet", + "_metrics": { + "ImageNet-1K": { + "acc@1": 69.758, + "acc@5": 89.078, + } + }, + "_ops": 1.814, + "_file_size": 44.661, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class ResNet34_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/resnet34-b627a593.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 21797672, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnet", + "_metrics": { + "ImageNet-1K": { + "acc@1": 73.314, + "acc@5": 91.420, + } + }, + "_ops": 3.664, + "_file_size": 83.275, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class ResNet50_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/resnet50-0676ba61.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 25557032, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnet", + "_metrics": { + "ImageNet-1K": { + "acc@1": 76.130, + "acc@5": 92.862, + } + }, + "_ops": 4.089, + "_file_size": 97.781, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/resnet50-11ad3fa6.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 25557032, + "recipe": "https://github.com/pytorch/vision/issues/3995#issuecomment-1013906621", + "_metrics": { + "ImageNet-1K": { + "acc@1": 80.858, + "acc@5": 95.434, + } + }, + "_ops": 4.089, + "_file_size": 97.79, + "_docs": """ + These weights improve upon the results of the original paper by using TorchVision's `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class ResNet101_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/resnet101-63fe2227.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 44549160, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnet", + "_metrics": { + "ImageNet-1K": { + "acc@1": 77.374, + "acc@5": 93.546, + } + }, + "_ops": 7.801, + "_file_size": 170.511, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/resnet101-cd907fc2.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 44549160, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe", + "_metrics": { + "ImageNet-1K": { + "acc@1": 81.886, + "acc@5": 95.780, + } + }, + "_ops": 7.801, + "_file_size": 170.53, + "_docs": """ + These weights improve upon the results of the original paper by using TorchVision's `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class ResNet152_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/resnet152-394f9c45.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 60192808, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnet", + "_metrics": { + "ImageNet-1K": { + "acc@1": 78.312, + "acc@5": 94.046, + } + }, + "_ops": 11.514, + "_file_size": 230.434, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/resnet152-f82ba261.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 60192808, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe", + "_metrics": { + "ImageNet-1K": { + "acc@1": 82.284, + "acc@5": 96.002, + } + }, + "_ops": 11.514, + "_file_size": 230.474, + "_docs": """ + These weights improve upon the results of the original paper by using TorchVision's `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class ResNeXt50_32X4D_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 25028904, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnext", + "_metrics": { + "ImageNet-1K": { + "acc@1": 77.618, + "acc@5": 93.698, + } + }, + "_ops": 4.23, + "_file_size": 95.789, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/resnext50_32x4d-1a0047aa.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 25028904, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe", + "_metrics": { + "ImageNet-1K": { + "acc@1": 81.198, + "acc@5": 95.340, + } + }, + "_ops": 4.23, + "_file_size": 95.833, + "_docs": """ + These weights improve upon the results of the original paper by using TorchVision's `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class ResNeXt101_32X8D_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 88791336, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnext", + "_metrics": { + "ImageNet-1K": { + "acc@1": 79.312, + "acc@5": 94.526, + } + }, + "_ops": 16.414, + "_file_size": 339.586, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/resnext101_32x8d-110c445d.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 88791336, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe-with-fixres", + "_metrics": { + "ImageNet-1K": { + "acc@1": 82.834, + "acc@5": 96.228, + } + }, + "_ops": 16.414, + "_file_size": 339.673, + "_docs": """ + These weights improve upon the results of the original paper by using TorchVision's `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class ResNeXt101_64X4D_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/resnext101_64x4d-173b62eb.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 83455272, + "recipe": "https://github.com/pytorch/vision/pull/5935", + "_metrics": { + "ImageNet-1K": { + "acc@1": 83.246, + "acc@5": 96.454, + } + }, + "_ops": 15.46, + "_file_size": 319.318, + "_docs": """ + These weights were trained from scratch by using TorchVision's `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class Wide_ResNet50_2_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 68883240, + "recipe": "https://github.com/pytorch/vision/pull/912#issue-445437439", + "_metrics": { + "ImageNet-1K": { + "acc@1": 78.468, + "acc@5": 94.086, + } + }, + "_ops": 11.398, + "_file_size": 131.82, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/wide_resnet50_2-9ba9bcbe.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 68883240, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe-with-fixres", + "_metrics": { + "ImageNet-1K": { + "acc@1": 81.602, + "acc@5": 95.758, + } + }, + "_ops": 11.398, + "_file_size": 263.124, + "_docs": """ + These weights improve upon the results of the original paper by using TorchVision's `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +class Wide_ResNet101_2_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 126886696, + "recipe": "https://github.com/pytorch/vision/pull/912#issue-445437439", + "_metrics": { + "ImageNet-1K": { + "acc@1": 78.848, + "acc@5": 94.284, + } + }, + "_ops": 22.753, + "_file_size": 242.896, + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", + }, + ) + IMAGENET1K_V2 = Weights( + url="https://download.pytorch.org/models/wide_resnet101_2-d733dc28.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "num_params": 126886696, + "recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe", + "_metrics": { + "ImageNet-1K": { + "acc@1": 82.510, + "acc@5": 96.020, + } + }, + "_ops": 22.753, + "_file_size": 484.747, + "_docs": """ + These weights improve upon the results of the original paper by using TorchVision's `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V2 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ResNet18_Weights.IMAGENET1K_V1)) +def resnet18(*, weights: Optional[ResNet18_Weights] = None, progress: bool = True, **kwargs: Any) -> ResNet: + """ResNet-18 from `Deep Residual Learning for Image Recognition `__. + + Args: + weights (:class:`~torchvision.models.ResNet18_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.ResNet18_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.resnet.ResNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.ResNet18_Weights + :members: + """ + weights = ResNet18_Weights.verify(weights) + + return _resnet(BasicBlock, [2, 2, 2, 2], weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ResNet34_Weights.IMAGENET1K_V1)) +def resnet34(*, weights: Optional[ResNet34_Weights] = None, progress: bool = True, **kwargs: Any) -> ResNet: + """ResNet-34 from `Deep Residual Learning for Image Recognition `__. + + Args: + weights (:class:`~torchvision.models.ResNet34_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.ResNet34_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.resnet.ResNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.ResNet34_Weights + :members: + """ + weights = ResNet34_Weights.verify(weights) + + return _resnet(BasicBlock, [3, 4, 6, 3], weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ResNet50_Weights.IMAGENET1K_V1)) +def resnet50(*, weights: Optional[ResNet50_Weights] = None, progress: bool = True, **kwargs: Any) -> ResNet: + """ResNet-50 from `Deep Residual Learning for Image Recognition `__. + + .. note:: + The bottleneck of TorchVision places the stride for downsampling to the second 3x3 + convolution while the original paper places it to the first 1x1 convolution. + This variant improves the accuracy and is known as `ResNet V1.5 + `_. + + Args: + weights (:class:`~torchvision.models.ResNet50_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.ResNet50_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.resnet.ResNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.ResNet50_Weights + :members: + """ + weights = ResNet50_Weights.verify(weights) + + return _resnet(Bottleneck, [3, 4, 6, 3], weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ResNet101_Weights.IMAGENET1K_V1)) +def resnet101(*, weights: Optional[ResNet101_Weights] = None, progress: bool = True, **kwargs: Any) -> ResNet: + """ResNet-101 from `Deep Residual Learning for Image Recognition `__. + + .. note:: + The bottleneck of TorchVision places the stride for downsampling to the second 3x3 + convolution while the original paper places it to the first 1x1 convolution. + This variant improves the accuracy and is known as `ResNet V1.5 + `_. + + Args: + weights (:class:`~torchvision.models.ResNet101_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.ResNet101_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.resnet.ResNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.ResNet101_Weights + :members: + """ + weights = ResNet101_Weights.verify(weights) + + return _resnet(Bottleneck, [3, 4, 23, 3], weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ResNet152_Weights.IMAGENET1K_V1)) +def resnet152(*, weights: Optional[ResNet152_Weights] = None, progress: bool = True, **kwargs: Any) -> ResNet: + """ResNet-152 from `Deep Residual Learning for Image Recognition `__. + + .. note:: + The bottleneck of TorchVision places the stride for downsampling to the second 3x3 + convolution while the original paper places it to the first 1x1 convolution. + This variant improves the accuracy and is known as `ResNet V1.5 + `_. + + Args: + weights (:class:`~torchvision.models.ResNet152_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.ResNet152_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.resnet.ResNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.ResNet152_Weights + :members: + """ + weights = ResNet152_Weights.verify(weights) + + return _resnet(Bottleneck, [3, 8, 36, 3], weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ResNeXt50_32X4D_Weights.IMAGENET1K_V1)) +def resnext50_32x4d( + *, weights: Optional[ResNeXt50_32X4D_Weights] = None, progress: bool = True, **kwargs: Any +) -> ResNet: + """ResNeXt-50 32x4d model from + `Aggregated Residual Transformation for Deep Neural Networks `_. + + Args: + weights (:class:`~torchvision.models.ResNeXt50_32X4D_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.ResNext50_32X4D_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.resnet.ResNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + .. autoclass:: torchvision.models.ResNeXt50_32X4D_Weights + :members: + """ + weights = ResNeXt50_32X4D_Weights.verify(weights) + + _ovewrite_named_param(kwargs, "groups", 32) + _ovewrite_named_param(kwargs, "width_per_group", 4) + return _resnet(Bottleneck, [3, 4, 6, 3], weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ResNeXt101_32X8D_Weights.IMAGENET1K_V1)) +def resnext101_32x8d( + *, weights: Optional[ResNeXt101_32X8D_Weights] = None, progress: bool = True, **kwargs: Any +) -> ResNet: + """ResNeXt-101 32x8d model from + `Aggregated Residual Transformation for Deep Neural Networks `_. + + Args: + weights (:class:`~torchvision.models.ResNeXt101_32X8D_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.ResNeXt101_32X8D_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.resnet.ResNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + .. autoclass:: torchvision.models.ResNeXt101_32X8D_Weights + :members: + """ + weights = ResNeXt101_32X8D_Weights.verify(weights) + + _ovewrite_named_param(kwargs, "groups", 32) + _ovewrite_named_param(kwargs, "width_per_group", 8) + return _resnet(Bottleneck, [3, 4, 23, 3], weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ResNeXt101_64X4D_Weights.IMAGENET1K_V1)) +def resnext101_64x4d( + *, weights: Optional[ResNeXt101_64X4D_Weights] = None, progress: bool = True, **kwargs: Any +) -> ResNet: + """ResNeXt-101 64x4d model from + `Aggregated Residual Transformation for Deep Neural Networks `_. + + Args: + weights (:class:`~torchvision.models.ResNeXt101_64X4D_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.ResNeXt101_64X4D_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.resnet.ResNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + .. autoclass:: torchvision.models.ResNeXt101_64X4D_Weights + :members: + """ + weights = ResNeXt101_64X4D_Weights.verify(weights) + + _ovewrite_named_param(kwargs, "groups", 64) + _ovewrite_named_param(kwargs, "width_per_group", 4) + return _resnet(Bottleneck, [3, 4, 23, 3], weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", Wide_ResNet50_2_Weights.IMAGENET1K_V1)) +def wide_resnet50_2( + *, weights: Optional[Wide_ResNet50_2_Weights] = None, progress: bool = True, **kwargs: Any +) -> ResNet: + """Wide ResNet-50-2 model from + `Wide Residual Networks `_. + + The model is the same as ResNet except for the bottleneck number of channels + which is twice larger in every block. The number of channels in outer 1x1 + convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048 + channels, and in Wide ResNet-50-2 has 2048-1024-2048. + + Args: + weights (:class:`~torchvision.models.Wide_ResNet50_2_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.Wide_ResNet50_2_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.resnet.ResNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + .. autoclass:: torchvision.models.Wide_ResNet50_2_Weights + :members: + """ + weights = Wide_ResNet50_2_Weights.verify(weights) + + _ovewrite_named_param(kwargs, "width_per_group", 64 * 2) + return _resnet(Bottleneck, [3, 4, 6, 3], weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", Wide_ResNet101_2_Weights.IMAGENET1K_V1)) +def wide_resnet101_2( + *, weights: Optional[Wide_ResNet101_2_Weights] = None, progress: bool = True, **kwargs: Any +) -> ResNet: + """Wide ResNet-101-2 model from + `Wide Residual Networks `_. + + The model is the same as ResNet except for the bottleneck number of channels + which is twice larger in every block. The number of channels in outer 1x1 + convolutions is the same, e.g. last block in ResNet-101 has 2048-512-2048 + channels, and in Wide ResNet-101-2 has 2048-1024-2048. + + Args: + weights (:class:`~torchvision.models.Wide_ResNet101_2_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.Wide_ResNet101_2_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.resnet.ResNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + .. autoclass:: torchvision.models.Wide_ResNet101_2_Weights + :members: + """ + weights = Wide_ResNet101_2_Weights.verify(weights) + + _ovewrite_named_param(kwargs, "width_per_group", 64 * 2) + return _resnet(Bottleneck, [3, 4, 23, 3], weights, progress, **kwargs) diff --git a/lib/python3.12/site-packages/torchvision/models/segmentation/__init__.py b/lib/python3.12/site-packages/torchvision/models/segmentation/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d6f37f958a131b76ce80306718b77d78bc3f045 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/segmentation/__init__.py @@ -0,0 +1,3 @@ +from .deeplabv3 import * +from .fcn import * +from .lraspp import * diff --git a/lib/python3.12/site-packages/torchvision/models/segmentation/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/segmentation/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4b635a017e8c1aad096a1a3136ec10eacd675e48 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/segmentation/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/segmentation/__pycache__/_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/segmentation/__pycache__/_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8f3b095c185813f1a94d7a1a567962abefec6e59 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/segmentation/__pycache__/_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/segmentation/__pycache__/deeplabv3.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/segmentation/__pycache__/deeplabv3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..343448185ddb9f2c05b95fc6f30f50d3ce1200fc Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/segmentation/__pycache__/deeplabv3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/segmentation/__pycache__/fcn.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/segmentation/__pycache__/fcn.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a64d9992e6965ca57554e97072b8ba879da3edf5 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/segmentation/__pycache__/fcn.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/segmentation/__pycache__/lraspp.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/segmentation/__pycache__/lraspp.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dd970b926ea4df0998031798452cdbc32cbcf9cf Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/segmentation/__pycache__/lraspp.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/segmentation/_utils.py b/lib/python3.12/site-packages/torchvision/models/segmentation/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..56560e9dab5c143699c918fa28236a902e530daf --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/segmentation/_utils.py @@ -0,0 +1,37 @@ +from collections import OrderedDict +from typing import Dict, Optional + +from torch import nn, Tensor +from torch.nn import functional as F + +from ...utils import _log_api_usage_once + + +class _SimpleSegmentationModel(nn.Module): + __constants__ = ["aux_classifier"] + + def __init__(self, backbone: nn.Module, classifier: nn.Module, aux_classifier: Optional[nn.Module] = None) -> None: + super().__init__() + _log_api_usage_once(self) + self.backbone = backbone + self.classifier = classifier + self.aux_classifier = aux_classifier + + def forward(self, x: Tensor) -> Dict[str, Tensor]: + input_shape = x.shape[-2:] + # contract: features is a dict of tensors + features = self.backbone(x) + + result = OrderedDict() + x = features["out"] + x = self.classifier(x) + x = F.interpolate(x, size=input_shape, mode="bilinear", align_corners=False) + result["out"] = x + + if self.aux_classifier is not None: + x = features["aux"] + x = self.aux_classifier(x) + x = F.interpolate(x, size=input_shape, mode="bilinear", align_corners=False) + result["aux"] = x + + return result diff --git a/lib/python3.12/site-packages/torchvision/models/segmentation/deeplabv3.py b/lib/python3.12/site-packages/torchvision/models/segmentation/deeplabv3.py new file mode 100644 index 0000000000000000000000000000000000000000..a92ddfe3b7af41a9ffd371d34cd459ba57965c53 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/segmentation/deeplabv3.py @@ -0,0 +1,390 @@ +from functools import partial +from typing import Any, Optional, Sequence + +import torch +from torch import nn +from torch.nn import functional as F + +from ...transforms._presets import SemanticSegmentation +from .._api import register_model, Weights, WeightsEnum +from .._meta import _VOC_CATEGORIES +from .._utils import _ovewrite_value_param, handle_legacy_interface, IntermediateLayerGetter +from ..mobilenetv3 import mobilenet_v3_large, MobileNet_V3_Large_Weights, MobileNetV3 +from ..resnet import ResNet, resnet101, ResNet101_Weights, resnet50, ResNet50_Weights +from ._utils import _SimpleSegmentationModel +from .fcn import FCNHead + + +__all__ = [ + "DeepLabV3", + "DeepLabV3_ResNet50_Weights", + "DeepLabV3_ResNet101_Weights", + "DeepLabV3_MobileNet_V3_Large_Weights", + "deeplabv3_mobilenet_v3_large", + "deeplabv3_resnet50", + "deeplabv3_resnet101", +] + + +class DeepLabV3(_SimpleSegmentationModel): + """ + Implements DeepLabV3 model from + `"Rethinking Atrous Convolution for Semantic Image Segmentation" + `_. + + Args: + backbone (nn.Module): the network used to compute the features for the model. + The backbone should return an OrderedDict[Tensor], with the key being + "out" for the last feature map used, and "aux" if an auxiliary classifier + is used. + classifier (nn.Module): module that takes the "out" element returned from + the backbone and returns a dense prediction. + aux_classifier (nn.Module, optional): auxiliary classifier used during training + """ + + pass + + +class DeepLabHead(nn.Sequential): + def __init__(self, in_channels: int, num_classes: int, atrous_rates: Sequence[int] = (12, 24, 36)) -> None: + super().__init__( + ASPP(in_channels, atrous_rates), + nn.Conv2d(256, 256, 3, padding=1, bias=False), + nn.BatchNorm2d(256), + nn.ReLU(), + nn.Conv2d(256, num_classes, 1), + ) + + +class ASPPConv(nn.Sequential): + def __init__(self, in_channels: int, out_channels: int, dilation: int) -> None: + modules = [ + nn.Conv2d(in_channels, out_channels, 3, padding=dilation, dilation=dilation, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(), + ] + super().__init__(*modules) + + +class ASPPPooling(nn.Sequential): + def __init__(self, in_channels: int, out_channels: int) -> None: + super().__init__( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(in_channels, out_channels, 1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + size = x.shape[-2:] + for mod in self: + x = mod(x) + return F.interpolate(x, size=size, mode="bilinear", align_corners=False) + + +class ASPP(nn.Module): + def __init__(self, in_channels: int, atrous_rates: Sequence[int], out_channels: int = 256) -> None: + super().__init__() + modules = [] + modules.append( + nn.Sequential(nn.Conv2d(in_channels, out_channels, 1, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU()) + ) + + rates = tuple(atrous_rates) + for rate in rates: + modules.append(ASPPConv(in_channels, out_channels, rate)) + + modules.append(ASPPPooling(in_channels, out_channels)) + + self.convs = nn.ModuleList(modules) + + self.project = nn.Sequential( + nn.Conv2d(len(self.convs) * out_channels, out_channels, 1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(), + nn.Dropout(0.5), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + _res = [] + for conv in self.convs: + _res.append(conv(x)) + res = torch.cat(_res, dim=1) + return self.project(res) + + +def _deeplabv3_resnet( + backbone: ResNet, + num_classes: int, + aux: Optional[bool], +) -> DeepLabV3: + return_layers = {"layer4": "out"} + if aux: + return_layers["layer3"] = "aux" + backbone = IntermediateLayerGetter(backbone, return_layers=return_layers) + + aux_classifier = FCNHead(1024, num_classes) if aux else None + classifier = DeepLabHead(2048, num_classes) + return DeepLabV3(backbone, classifier, aux_classifier) + + +_COMMON_META = { + "categories": _VOC_CATEGORIES, + "min_size": (1, 1), + "_docs": """ + These weights were trained on a subset of COCO, using only the 20 categories that are present in the Pascal VOC + dataset. + """, +} + + +class DeepLabV3_ResNet50_Weights(WeightsEnum): + COCO_WITH_VOC_LABELS_V1 = Weights( + url="https://download.pytorch.org/models/deeplabv3_resnet50_coco-cd0a2569.pth", + transforms=partial(SemanticSegmentation, resize_size=520), + meta={ + **_COMMON_META, + "num_params": 42004074, + "recipe": "https://github.com/pytorch/vision/tree/main/references/segmentation#deeplabv3_resnet50", + "_metrics": { + "COCO-val2017-VOC-labels": { + "miou": 66.4, + "pixel_acc": 92.4, + } + }, + "_ops": 178.722, + "_file_size": 160.515, + }, + ) + DEFAULT = COCO_WITH_VOC_LABELS_V1 + + +class DeepLabV3_ResNet101_Weights(WeightsEnum): + COCO_WITH_VOC_LABELS_V1 = Weights( + url="https://download.pytorch.org/models/deeplabv3_resnet101_coco-586e9e4e.pth", + transforms=partial(SemanticSegmentation, resize_size=520), + meta={ + **_COMMON_META, + "num_params": 60996202, + "recipe": "https://github.com/pytorch/vision/tree/main/references/segmentation#fcn_resnet101", + "_metrics": { + "COCO-val2017-VOC-labels": { + "miou": 67.4, + "pixel_acc": 92.4, + } + }, + "_ops": 258.743, + "_file_size": 233.217, + }, + ) + DEFAULT = COCO_WITH_VOC_LABELS_V1 + + +class DeepLabV3_MobileNet_V3_Large_Weights(WeightsEnum): + COCO_WITH_VOC_LABELS_V1 = Weights( + url="https://download.pytorch.org/models/deeplabv3_mobilenet_v3_large-fc3c493d.pth", + transforms=partial(SemanticSegmentation, resize_size=520), + meta={ + **_COMMON_META, + "num_params": 11029328, + "recipe": "https://github.com/pytorch/vision/tree/main/references/segmentation#deeplabv3_mobilenet_v3_large", + "_metrics": { + "COCO-val2017-VOC-labels": { + "miou": 60.3, + "pixel_acc": 91.2, + } + }, + "_ops": 10.452, + "_file_size": 42.301, + }, + ) + DEFAULT = COCO_WITH_VOC_LABELS_V1 + + +def _deeplabv3_mobilenetv3( + backbone: MobileNetV3, + num_classes: int, + aux: Optional[bool], +) -> DeepLabV3: + backbone = backbone.features + # Gather the indices of blocks which are strided. These are the locations of C1, ..., Cn-1 blocks. + # The first and last blocks are always included because they are the C0 (conv1) and Cn. + stage_indices = [0] + [i for i, b in enumerate(backbone) if getattr(b, "_is_cn", False)] + [len(backbone) - 1] + out_pos = stage_indices[-1] # use C5 which has output_stride = 16 + out_inplanes = backbone[out_pos].out_channels + aux_pos = stage_indices[-4] # use C2 here which has output_stride = 8 + aux_inplanes = backbone[aux_pos].out_channels + return_layers = {str(out_pos): "out"} + if aux: + return_layers[str(aux_pos)] = "aux" + backbone = IntermediateLayerGetter(backbone, return_layers=return_layers) + + aux_classifier = FCNHead(aux_inplanes, num_classes) if aux else None + classifier = DeepLabHead(out_inplanes, num_classes) + return DeepLabV3(backbone, classifier, aux_classifier) + + +@register_model() +@handle_legacy_interface( + weights=("pretrained", DeepLabV3_ResNet50_Weights.COCO_WITH_VOC_LABELS_V1), + weights_backbone=("pretrained_backbone", ResNet50_Weights.IMAGENET1K_V1), +) +def deeplabv3_resnet50( + *, + weights: Optional[DeepLabV3_ResNet50_Weights] = None, + progress: bool = True, + num_classes: Optional[int] = None, + aux_loss: Optional[bool] = None, + weights_backbone: Optional[ResNet50_Weights] = ResNet50_Weights.IMAGENET1K_V1, + **kwargs: Any, +) -> DeepLabV3: + """Constructs a DeepLabV3 model with a ResNet-50 backbone. + + .. betastatus:: segmentation module + + Reference: `Rethinking Atrous Convolution for Semantic Image Segmentation `__. + + Args: + weights (:class:`~torchvision.models.segmentation.DeepLabV3_ResNet50_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.segmentation.DeepLabV3_ResNet50_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + num_classes (int, optional): number of output classes of the model (including the background) + aux_loss (bool, optional): If True, it uses an auxiliary loss + weights_backbone (:class:`~torchvision.models.ResNet50_Weights`, optional): The pretrained weights for the + backbone + **kwargs: unused + + .. autoclass:: torchvision.models.segmentation.DeepLabV3_ResNet50_Weights + :members: + """ + weights = DeepLabV3_ResNet50_Weights.verify(weights) + weights_backbone = ResNet50_Weights.verify(weights_backbone) + + if weights is not None: + weights_backbone = None + num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"])) + aux_loss = _ovewrite_value_param("aux_loss", aux_loss, True) + elif num_classes is None: + num_classes = 21 + + backbone = resnet50(weights=weights_backbone, replace_stride_with_dilation=[False, True, True]) + model = _deeplabv3_resnet(backbone, num_classes, aux_loss) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +@register_model() +@handle_legacy_interface( + weights=("pretrained", DeepLabV3_ResNet101_Weights.COCO_WITH_VOC_LABELS_V1), + weights_backbone=("pretrained_backbone", ResNet101_Weights.IMAGENET1K_V1), +) +def deeplabv3_resnet101( + *, + weights: Optional[DeepLabV3_ResNet101_Weights] = None, + progress: bool = True, + num_classes: Optional[int] = None, + aux_loss: Optional[bool] = None, + weights_backbone: Optional[ResNet101_Weights] = ResNet101_Weights.IMAGENET1K_V1, + **kwargs: Any, +) -> DeepLabV3: + """Constructs a DeepLabV3 model with a ResNet-101 backbone. + + .. betastatus:: segmentation module + + Reference: `Rethinking Atrous Convolution for Semantic Image Segmentation `__. + + Args: + weights (:class:`~torchvision.models.segmentation.DeepLabV3_ResNet101_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.segmentation.DeepLabV3_ResNet101_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + num_classes (int, optional): number of output classes of the model (including the background) + aux_loss (bool, optional): If True, it uses an auxiliary loss + weights_backbone (:class:`~torchvision.models.ResNet101_Weights`, optional): The pretrained weights for the + backbone + **kwargs: unused + + .. autoclass:: torchvision.models.segmentation.DeepLabV3_ResNet101_Weights + :members: + """ + weights = DeepLabV3_ResNet101_Weights.verify(weights) + weights_backbone = ResNet101_Weights.verify(weights_backbone) + + if weights is not None: + weights_backbone = None + num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"])) + aux_loss = _ovewrite_value_param("aux_loss", aux_loss, True) + elif num_classes is None: + num_classes = 21 + + backbone = resnet101(weights=weights_backbone, replace_stride_with_dilation=[False, True, True]) + model = _deeplabv3_resnet(backbone, num_classes, aux_loss) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +@register_model() +@handle_legacy_interface( + weights=("pretrained", DeepLabV3_MobileNet_V3_Large_Weights.COCO_WITH_VOC_LABELS_V1), + weights_backbone=("pretrained_backbone", MobileNet_V3_Large_Weights.IMAGENET1K_V1), +) +def deeplabv3_mobilenet_v3_large( + *, + weights: Optional[DeepLabV3_MobileNet_V3_Large_Weights] = None, + progress: bool = True, + num_classes: Optional[int] = None, + aux_loss: Optional[bool] = None, + weights_backbone: Optional[MobileNet_V3_Large_Weights] = MobileNet_V3_Large_Weights.IMAGENET1K_V1, + **kwargs: Any, +) -> DeepLabV3: + """Constructs a DeepLabV3 model with a MobileNetV3-Large backbone. + + Reference: `Rethinking Atrous Convolution for Semantic Image Segmentation `__. + + Args: + weights (:class:`~torchvision.models.segmentation.DeepLabV3_MobileNet_V3_Large_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.segmentation.DeepLabV3_MobileNet_V3_Large_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + num_classes (int, optional): number of output classes of the model (including the background) + aux_loss (bool, optional): If True, it uses an auxiliary loss + weights_backbone (:class:`~torchvision.models.MobileNet_V3_Large_Weights`, optional): The pretrained weights + for the backbone + **kwargs: unused + + .. autoclass:: torchvision.models.segmentation.DeepLabV3_MobileNet_V3_Large_Weights + :members: + """ + weights = DeepLabV3_MobileNet_V3_Large_Weights.verify(weights) + weights_backbone = MobileNet_V3_Large_Weights.verify(weights_backbone) + + if weights is not None: + weights_backbone = None + num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"])) + aux_loss = _ovewrite_value_param("aux_loss", aux_loss, True) + elif num_classes is None: + num_classes = 21 + + backbone = mobilenet_v3_large(weights=weights_backbone, dilated=True) + model = _deeplabv3_mobilenetv3(backbone, num_classes, aux_loss) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model diff --git a/lib/python3.12/site-packages/torchvision/models/segmentation/fcn.py b/lib/python3.12/site-packages/torchvision/models/segmentation/fcn.py new file mode 100644 index 0000000000000000000000000000000000000000..fb2e242adac0e7430bab6155ae0347770e29fee9 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/segmentation/fcn.py @@ -0,0 +1,232 @@ +from functools import partial +from typing import Any, Optional + +from torch import nn + +from ...transforms._presets import SemanticSegmentation +from .._api import register_model, Weights, WeightsEnum +from .._meta import _VOC_CATEGORIES +from .._utils import _ovewrite_value_param, handle_legacy_interface, IntermediateLayerGetter +from ..resnet import ResNet, resnet101, ResNet101_Weights, resnet50, ResNet50_Weights +from ._utils import _SimpleSegmentationModel + + +__all__ = ["FCN", "FCN_ResNet50_Weights", "FCN_ResNet101_Weights", "fcn_resnet50", "fcn_resnet101"] + + +class FCN(_SimpleSegmentationModel): + """ + Implements FCN model from + `"Fully Convolutional Networks for Semantic Segmentation" + `_. + + Args: + backbone (nn.Module): the network used to compute the features for the model. + The backbone should return an OrderedDict[Tensor], with the key being + "out" for the last feature map used, and "aux" if an auxiliary classifier + is used. + classifier (nn.Module): module that takes the "out" element returned from + the backbone and returns a dense prediction. + aux_classifier (nn.Module, optional): auxiliary classifier used during training + """ + + pass + + +class FCNHead(nn.Sequential): + def __init__(self, in_channels: int, channels: int) -> None: + inter_channels = in_channels // 4 + layers = [ + nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False), + nn.BatchNorm2d(inter_channels), + nn.ReLU(), + nn.Dropout(0.1), + nn.Conv2d(inter_channels, channels, 1), + ] + + super().__init__(*layers) + + +_COMMON_META = { + "categories": _VOC_CATEGORIES, + "min_size": (1, 1), + "_docs": """ + These weights were trained on a subset of COCO, using only the 20 categories that are present in the Pascal VOC + dataset. + """, +} + + +class FCN_ResNet50_Weights(WeightsEnum): + COCO_WITH_VOC_LABELS_V1 = Weights( + url="https://download.pytorch.org/models/fcn_resnet50_coco-1167a1af.pth", + transforms=partial(SemanticSegmentation, resize_size=520), + meta={ + **_COMMON_META, + "num_params": 35322218, + "recipe": "https://github.com/pytorch/vision/tree/main/references/segmentation#fcn_resnet50", + "_metrics": { + "COCO-val2017-VOC-labels": { + "miou": 60.5, + "pixel_acc": 91.4, + } + }, + "_ops": 152.717, + "_file_size": 135.009, + }, + ) + DEFAULT = COCO_WITH_VOC_LABELS_V1 + + +class FCN_ResNet101_Weights(WeightsEnum): + COCO_WITH_VOC_LABELS_V1 = Weights( + url="https://download.pytorch.org/models/fcn_resnet101_coco-7ecb50ca.pth", + transforms=partial(SemanticSegmentation, resize_size=520), + meta={ + **_COMMON_META, + "num_params": 54314346, + "recipe": "https://github.com/pytorch/vision/tree/main/references/segmentation#deeplabv3_resnet101", + "_metrics": { + "COCO-val2017-VOC-labels": { + "miou": 63.7, + "pixel_acc": 91.9, + } + }, + "_ops": 232.738, + "_file_size": 207.711, + }, + ) + DEFAULT = COCO_WITH_VOC_LABELS_V1 + + +def _fcn_resnet( + backbone: ResNet, + num_classes: int, + aux: Optional[bool], +) -> FCN: + return_layers = {"layer4": "out"} + if aux: + return_layers["layer3"] = "aux" + backbone = IntermediateLayerGetter(backbone, return_layers=return_layers) + + aux_classifier = FCNHead(1024, num_classes) if aux else None + classifier = FCNHead(2048, num_classes) + return FCN(backbone, classifier, aux_classifier) + + +@register_model() +@handle_legacy_interface( + weights=("pretrained", FCN_ResNet50_Weights.COCO_WITH_VOC_LABELS_V1), + weights_backbone=("pretrained_backbone", ResNet50_Weights.IMAGENET1K_V1), +) +def fcn_resnet50( + *, + weights: Optional[FCN_ResNet50_Weights] = None, + progress: bool = True, + num_classes: Optional[int] = None, + aux_loss: Optional[bool] = None, + weights_backbone: Optional[ResNet50_Weights] = ResNet50_Weights.IMAGENET1K_V1, + **kwargs: Any, +) -> FCN: + """Fully-Convolutional Network model with a ResNet-50 backbone from the `Fully Convolutional + Networks for Semantic Segmentation `_ paper. + + .. betastatus:: segmentation module + + Args: + weights (:class:`~torchvision.models.segmentation.FCN_ResNet50_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.segmentation.FCN_ResNet50_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + num_classes (int, optional): number of output classes of the model (including the background). + aux_loss (bool, optional): If True, it uses an auxiliary loss. + weights_backbone (:class:`~torchvision.models.ResNet50_Weights`, optional): The pretrained + weights for the backbone. + **kwargs: parameters passed to the ``torchvision.models.segmentation.fcn.FCN`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.segmentation.FCN_ResNet50_Weights + :members: + """ + + weights = FCN_ResNet50_Weights.verify(weights) + weights_backbone = ResNet50_Weights.verify(weights_backbone) + + if weights is not None: + weights_backbone = None + num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"])) + aux_loss = _ovewrite_value_param("aux_loss", aux_loss, True) + elif num_classes is None: + num_classes = 21 + + backbone = resnet50(weights=weights_backbone, replace_stride_with_dilation=[False, True, True]) + model = _fcn_resnet(backbone, num_classes, aux_loss) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +@register_model() +@handle_legacy_interface( + weights=("pretrained", FCN_ResNet101_Weights.COCO_WITH_VOC_LABELS_V1), + weights_backbone=("pretrained_backbone", ResNet101_Weights.IMAGENET1K_V1), +) +def fcn_resnet101( + *, + weights: Optional[FCN_ResNet101_Weights] = None, + progress: bool = True, + num_classes: Optional[int] = None, + aux_loss: Optional[bool] = None, + weights_backbone: Optional[ResNet101_Weights] = ResNet101_Weights.IMAGENET1K_V1, + **kwargs: Any, +) -> FCN: + """Fully-Convolutional Network model with a ResNet-101 backbone from the `Fully Convolutional + Networks for Semantic Segmentation `_ paper. + + .. betastatus:: segmentation module + + Args: + weights (:class:`~torchvision.models.segmentation.FCN_ResNet101_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.segmentation.FCN_ResNet101_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + num_classes (int, optional): number of output classes of the model (including the background). + aux_loss (bool, optional): If True, it uses an auxiliary loss. + weights_backbone (:class:`~torchvision.models.ResNet101_Weights`, optional): The pretrained + weights for the backbone. + **kwargs: parameters passed to the ``torchvision.models.segmentation.fcn.FCN`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.segmentation.FCN_ResNet101_Weights + :members: + """ + + weights = FCN_ResNet101_Weights.verify(weights) + weights_backbone = ResNet101_Weights.verify(weights_backbone) + + if weights is not None: + weights_backbone = None + num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"])) + aux_loss = _ovewrite_value_param("aux_loss", aux_loss, True) + elif num_classes is None: + num_classes = 21 + + backbone = resnet101(weights=weights_backbone, replace_stride_with_dilation=[False, True, True]) + model = _fcn_resnet(backbone, num_classes, aux_loss) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model diff --git a/lib/python3.12/site-packages/torchvision/models/segmentation/lraspp.py b/lib/python3.12/site-packages/torchvision/models/segmentation/lraspp.py new file mode 100644 index 0000000000000000000000000000000000000000..70bced70fd37c3c681915492cea0c68c87cf0a7e --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/segmentation/lraspp.py @@ -0,0 +1,178 @@ +from collections import OrderedDict +from functools import partial +from typing import Any, Dict, Optional + +from torch import nn, Tensor +from torch.nn import functional as F + +from ...transforms._presets import SemanticSegmentation +from ...utils import _log_api_usage_once +from .._api import register_model, Weights, WeightsEnum +from .._meta import _VOC_CATEGORIES +from .._utils import _ovewrite_value_param, handle_legacy_interface, IntermediateLayerGetter +from ..mobilenetv3 import mobilenet_v3_large, MobileNet_V3_Large_Weights, MobileNetV3 + + +__all__ = ["LRASPP", "LRASPP_MobileNet_V3_Large_Weights", "lraspp_mobilenet_v3_large"] + + +class LRASPP(nn.Module): + """ + Implements a Lite R-ASPP Network for semantic segmentation from + `"Searching for MobileNetV3" + `_. + + Args: + backbone (nn.Module): the network used to compute the features for the model. + The backbone should return an OrderedDict[Tensor], with the key being + "high" for the high level feature map and "low" for the low level feature map. + low_channels (int): the number of channels of the low level features. + high_channels (int): the number of channels of the high level features. + num_classes (int, optional): number of output classes of the model (including the background). + inter_channels (int, optional): the number of channels for intermediate computations. + """ + + def __init__( + self, backbone: nn.Module, low_channels: int, high_channels: int, num_classes: int, inter_channels: int = 128 + ) -> None: + super().__init__() + _log_api_usage_once(self) + self.backbone = backbone + self.classifier = LRASPPHead(low_channels, high_channels, num_classes, inter_channels) + + def forward(self, input: Tensor) -> Dict[str, Tensor]: + features = self.backbone(input) + out = self.classifier(features) + out = F.interpolate(out, size=input.shape[-2:], mode="bilinear", align_corners=False) + + result = OrderedDict() + result["out"] = out + + return result + + +class LRASPPHead(nn.Module): + def __init__(self, low_channels: int, high_channels: int, num_classes: int, inter_channels: int) -> None: + super().__init__() + self.cbr = nn.Sequential( + nn.Conv2d(high_channels, inter_channels, 1, bias=False), + nn.BatchNorm2d(inter_channels), + nn.ReLU(inplace=True), + ) + self.scale = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(high_channels, inter_channels, 1, bias=False), + nn.Sigmoid(), + ) + self.low_classifier = nn.Conv2d(low_channels, num_classes, 1) + self.high_classifier = nn.Conv2d(inter_channels, num_classes, 1) + + def forward(self, input: Dict[str, Tensor]) -> Tensor: + low = input["low"] + high = input["high"] + + x = self.cbr(high) + s = self.scale(high) + x = x * s + x = F.interpolate(x, size=low.shape[-2:], mode="bilinear", align_corners=False) + + return self.low_classifier(low) + self.high_classifier(x) + + +def _lraspp_mobilenetv3(backbone: MobileNetV3, num_classes: int) -> LRASPP: + backbone = backbone.features + # Gather the indices of blocks which are strided. These are the locations of C1, ..., Cn-1 blocks. + # The first and last blocks are always included because they are the C0 (conv1) and Cn. + stage_indices = [0] + [i for i, b in enumerate(backbone) if getattr(b, "_is_cn", False)] + [len(backbone) - 1] + low_pos = stage_indices[-4] # use C2 here which has output_stride = 8 + high_pos = stage_indices[-1] # use C5 which has output_stride = 16 + low_channels = backbone[low_pos].out_channels + high_channels = backbone[high_pos].out_channels + backbone = IntermediateLayerGetter(backbone, return_layers={str(low_pos): "low", str(high_pos): "high"}) + + return LRASPP(backbone, low_channels, high_channels, num_classes) + + +class LRASPP_MobileNet_V3_Large_Weights(WeightsEnum): + COCO_WITH_VOC_LABELS_V1 = Weights( + url="https://download.pytorch.org/models/lraspp_mobilenet_v3_large-d234d4ea.pth", + transforms=partial(SemanticSegmentation, resize_size=520), + meta={ + "num_params": 3221538, + "categories": _VOC_CATEGORIES, + "min_size": (1, 1), + "recipe": "https://github.com/pytorch/vision/tree/main/references/segmentation#lraspp_mobilenet_v3_large", + "_metrics": { + "COCO-val2017-VOC-labels": { + "miou": 57.9, + "pixel_acc": 91.2, + } + }, + "_ops": 2.086, + "_file_size": 12.49, + "_docs": """ + These weights were trained on a subset of COCO, using only the 20 categories that are present in the + Pascal VOC dataset. + """, + }, + ) + DEFAULT = COCO_WITH_VOC_LABELS_V1 + + +@register_model() +@handle_legacy_interface( + weights=("pretrained", LRASPP_MobileNet_V3_Large_Weights.COCO_WITH_VOC_LABELS_V1), + weights_backbone=("pretrained_backbone", MobileNet_V3_Large_Weights.IMAGENET1K_V1), +) +def lraspp_mobilenet_v3_large( + *, + weights: Optional[LRASPP_MobileNet_V3_Large_Weights] = None, + progress: bool = True, + num_classes: Optional[int] = None, + weights_backbone: Optional[MobileNet_V3_Large_Weights] = MobileNet_V3_Large_Weights.IMAGENET1K_V1, + **kwargs: Any, +) -> LRASPP: + """Constructs a Lite R-ASPP Network model with a MobileNetV3-Large backbone from + `Searching for MobileNetV3 `_ paper. + + .. betastatus:: segmentation module + + Args: + weights (:class:`~torchvision.models.segmentation.LRASPP_MobileNet_V3_Large_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.segmentation.LRASPP_MobileNet_V3_Large_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + num_classes (int, optional): number of output classes of the model (including the background). + aux_loss (bool, optional): If True, it uses an auxiliary loss. + weights_backbone (:class:`~torchvision.models.MobileNet_V3_Large_Weights`, optional): The pretrained + weights for the backbone. + **kwargs: parameters passed to the ``torchvision.models.segmentation.LRASPP`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.segmentation.LRASPP_MobileNet_V3_Large_Weights + :members: + """ + if kwargs.pop("aux_loss", False): + raise NotImplementedError("This model does not use auxiliary loss") + + weights = LRASPP_MobileNet_V3_Large_Weights.verify(weights) + weights_backbone = MobileNet_V3_Large_Weights.verify(weights_backbone) + + if weights is not None: + weights_backbone = None + num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"])) + elif num_classes is None: + num_classes = 21 + + backbone = mobilenet_v3_large(weights=weights_backbone, dilated=True) + model = _lraspp_mobilenetv3(backbone, num_classes) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model diff --git a/lib/python3.12/site-packages/torchvision/models/shufflenetv2.py b/lib/python3.12/site-packages/torchvision/models/shufflenetv2.py new file mode 100644 index 0000000000000000000000000000000000000000..3f3322b7a88f183c15838308f39c38d36dea13c0 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/shufflenetv2.py @@ -0,0 +1,408 @@ +from functools import partial +from typing import Any, Callable, List, Optional + +import torch +import torch.nn as nn +from torch import Tensor + +from ..transforms._presets import ImageClassification +from ..utils import _log_api_usage_once +from ._api import register_model, Weights, WeightsEnum +from ._meta import _IMAGENET_CATEGORIES +from ._utils import _ovewrite_named_param, handle_legacy_interface + + +__all__ = [ + "ShuffleNetV2", + "ShuffleNet_V2_X0_5_Weights", + "ShuffleNet_V2_X1_0_Weights", + "ShuffleNet_V2_X1_5_Weights", + "ShuffleNet_V2_X2_0_Weights", + "shufflenet_v2_x0_5", + "shufflenet_v2_x1_0", + "shufflenet_v2_x1_5", + "shufflenet_v2_x2_0", +] + + +def channel_shuffle(x: Tensor, groups: int) -> Tensor: + batchsize, num_channels, height, width = x.size() + channels_per_group = num_channels // groups + + # reshape + x = x.view(batchsize, groups, channels_per_group, height, width) + + x = torch.transpose(x, 1, 2).contiguous() + + # flatten + x = x.view(batchsize, num_channels, height, width) + + return x + + +class InvertedResidual(nn.Module): + def __init__(self, inp: int, oup: int, stride: int) -> None: + super().__init__() + + if not (1 <= stride <= 3): + raise ValueError("illegal stride value") + self.stride = stride + + branch_features = oup // 2 + if (self.stride == 1) and (inp != branch_features << 1): + raise ValueError( + f"Invalid combination of stride {stride}, inp {inp} and oup {oup} values. If stride == 1 then inp should be equal to oup // 2 << 1." + ) + + if self.stride > 1: + self.branch1 = nn.Sequential( + self.depthwise_conv(inp, inp, kernel_size=3, stride=self.stride, padding=1), + nn.BatchNorm2d(inp), + nn.Conv2d(inp, branch_features, kernel_size=1, stride=1, padding=0, bias=False), + nn.BatchNorm2d(branch_features), + nn.ReLU(inplace=True), + ) + else: + self.branch1 = nn.Sequential() + + self.branch2 = nn.Sequential( + nn.Conv2d( + inp if (self.stride > 1) else branch_features, + branch_features, + kernel_size=1, + stride=1, + padding=0, + bias=False, + ), + nn.BatchNorm2d(branch_features), + nn.ReLU(inplace=True), + self.depthwise_conv(branch_features, branch_features, kernel_size=3, stride=self.stride, padding=1), + nn.BatchNorm2d(branch_features), + nn.Conv2d(branch_features, branch_features, kernel_size=1, stride=1, padding=0, bias=False), + nn.BatchNorm2d(branch_features), + nn.ReLU(inplace=True), + ) + + @staticmethod + def depthwise_conv( + i: int, o: int, kernel_size: int, stride: int = 1, padding: int = 0, bias: bool = False + ) -> nn.Conv2d: + return nn.Conv2d(i, o, kernel_size, stride, padding, bias=bias, groups=i) + + def forward(self, x: Tensor) -> Tensor: + if self.stride == 1: + x1, x2 = x.chunk(2, dim=1) + out = torch.cat((x1, self.branch2(x2)), dim=1) + else: + out = torch.cat((self.branch1(x), self.branch2(x)), dim=1) + + out = channel_shuffle(out, 2) + + return out + + +class ShuffleNetV2(nn.Module): + def __init__( + self, + stages_repeats: List[int], + stages_out_channels: List[int], + num_classes: int = 1000, + inverted_residual: Callable[..., nn.Module] = InvertedResidual, + ) -> None: + super().__init__() + _log_api_usage_once(self) + + if len(stages_repeats) != 3: + raise ValueError("expected stages_repeats as list of 3 positive ints") + if len(stages_out_channels) != 5: + raise ValueError("expected stages_out_channels as list of 5 positive ints") + self._stage_out_channels = stages_out_channels + + input_channels = 3 + output_channels = self._stage_out_channels[0] + self.conv1 = nn.Sequential( + nn.Conv2d(input_channels, output_channels, 3, 2, 1, bias=False), + nn.BatchNorm2d(output_channels), + nn.ReLU(inplace=True), + ) + input_channels = output_channels + + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + + # Static annotations for mypy + self.stage2: nn.Sequential + self.stage3: nn.Sequential + self.stage4: nn.Sequential + stage_names = [f"stage{i}" for i in [2, 3, 4]] + for name, repeats, output_channels in zip(stage_names, stages_repeats, self._stage_out_channels[1:]): + seq = [inverted_residual(input_channels, output_channels, 2)] + for i in range(repeats - 1): + seq.append(inverted_residual(output_channels, output_channels, 1)) + setattr(self, name, nn.Sequential(*seq)) + input_channels = output_channels + + output_channels = self._stage_out_channels[-1] + self.conv5 = nn.Sequential( + nn.Conv2d(input_channels, output_channels, 1, 1, 0, bias=False), + nn.BatchNorm2d(output_channels), + nn.ReLU(inplace=True), + ) + + self.fc = nn.Linear(output_channels, num_classes) + + def _forward_impl(self, x: Tensor) -> Tensor: + # See note [TorchScript super()] + x = self.conv1(x) + x = self.maxpool(x) + x = self.stage2(x) + x = self.stage3(x) + x = self.stage4(x) + x = self.conv5(x) + x = x.mean([2, 3]) # globalpool + x = self.fc(x) + return x + + def forward(self, x: Tensor) -> Tensor: + return self._forward_impl(x) + + +def _shufflenetv2( + weights: Optional[WeightsEnum], + progress: bool, + *args: Any, + **kwargs: Any, +) -> ShuffleNetV2: + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + + model = ShuffleNetV2(*args, **kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +_COMMON_META = { + "min_size": (1, 1), + "categories": _IMAGENET_CATEGORIES, + "recipe": "https://github.com/ericsun99/Shufflenet-v2-Pytorch", +} + + +class ShuffleNet_V2_X0_5_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + # Weights ported from https://github.com/ericsun99/Shufflenet-v2-Pytorch + url="https://download.pytorch.org/models/shufflenetv2_x0.5-f707e7126e.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 1366792, + "_metrics": { + "ImageNet-1K": { + "acc@1": 60.552, + "acc@5": 81.746, + } + }, + "_ops": 0.04, + "_file_size": 5.282, + "_docs": """These weights were trained from scratch to reproduce closely the results of the paper.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class ShuffleNet_V2_X1_0_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + # Weights ported from https://github.com/ericsun99/Shufflenet-v2-Pytorch + url="https://download.pytorch.org/models/shufflenetv2_x1-5666bf0f80.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 2278604, + "_metrics": { + "ImageNet-1K": { + "acc@1": 69.362, + "acc@5": 88.316, + } + }, + "_ops": 0.145, + "_file_size": 8.791, + "_docs": """These weights were trained from scratch to reproduce closely the results of the paper.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class ShuffleNet_V2_X1_5_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/shufflenetv2_x1_5-3c479a10.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "recipe": "https://github.com/pytorch/vision/pull/5906", + "num_params": 3503624, + "_metrics": { + "ImageNet-1K": { + "acc@1": 72.996, + "acc@5": 91.086, + } + }, + "_ops": 0.296, + "_file_size": 13.557, + "_docs": """ + These weights were trained from scratch by using TorchVision's `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class ShuffleNet_V2_X2_0_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/shufflenetv2_x2_0-8be3c8ee.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=232), + meta={ + **_COMMON_META, + "recipe": "https://github.com/pytorch/vision/pull/5906", + "num_params": 7393996, + "_metrics": { + "ImageNet-1K": { + "acc@1": 76.230, + "acc@5": 93.006, + } + }, + "_ops": 0.583, + "_file_size": 28.433, + "_docs": """ + These weights were trained from scratch by using TorchVision's `new training recipe + `_. + """, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ShuffleNet_V2_X0_5_Weights.IMAGENET1K_V1)) +def shufflenet_v2_x0_5( + *, weights: Optional[ShuffleNet_V2_X0_5_Weights] = None, progress: bool = True, **kwargs: Any +) -> ShuffleNetV2: + """ + Constructs a ShuffleNetV2 architecture with 0.5x output channels, as described in + `ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design + `__. + + Args: + weights (:class:`~torchvision.models.ShuffleNet_V2_X0_5_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.ShuffleNet_V2_X0_5_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.shufflenetv2.ShuffleNetV2`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.ShuffleNet_V2_X0_5_Weights + :members: + """ + weights = ShuffleNet_V2_X0_5_Weights.verify(weights) + + return _shufflenetv2(weights, progress, [4, 8, 4], [24, 48, 96, 192, 1024], **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ShuffleNet_V2_X1_0_Weights.IMAGENET1K_V1)) +def shufflenet_v2_x1_0( + *, weights: Optional[ShuffleNet_V2_X1_0_Weights] = None, progress: bool = True, **kwargs: Any +) -> ShuffleNetV2: + """ + Constructs a ShuffleNetV2 architecture with 1.0x output channels, as described in + `ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design + `__. + + Args: + weights (:class:`~torchvision.models.ShuffleNet_V2_X1_0_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.ShuffleNet_V2_X1_0_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.shufflenetv2.ShuffleNetV2`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.ShuffleNet_V2_X1_0_Weights + :members: + """ + weights = ShuffleNet_V2_X1_0_Weights.verify(weights) + + return _shufflenetv2(weights, progress, [4, 8, 4], [24, 116, 232, 464, 1024], **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ShuffleNet_V2_X1_5_Weights.IMAGENET1K_V1)) +def shufflenet_v2_x1_5( + *, weights: Optional[ShuffleNet_V2_X1_5_Weights] = None, progress: bool = True, **kwargs: Any +) -> ShuffleNetV2: + """ + Constructs a ShuffleNetV2 architecture with 1.5x output channels, as described in + `ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design + `__. + + Args: + weights (:class:`~torchvision.models.ShuffleNet_V2_X1_5_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.ShuffleNet_V2_X1_5_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.shufflenetv2.ShuffleNetV2`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.ShuffleNet_V2_X1_5_Weights + :members: + """ + weights = ShuffleNet_V2_X1_5_Weights.verify(weights) + + return _shufflenetv2(weights, progress, [4, 8, 4], [24, 176, 352, 704, 1024], **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ShuffleNet_V2_X2_0_Weights.IMAGENET1K_V1)) +def shufflenet_v2_x2_0( + *, weights: Optional[ShuffleNet_V2_X2_0_Weights] = None, progress: bool = True, **kwargs: Any +) -> ShuffleNetV2: + """ + Constructs a ShuffleNetV2 architecture with 2.0x output channels, as described in + `ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design + `__. + + Args: + weights (:class:`~torchvision.models.ShuffleNet_V2_X2_0_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.ShuffleNet_V2_X2_0_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.shufflenetv2.ShuffleNetV2`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.ShuffleNet_V2_X2_0_Weights + :members: + """ + weights = ShuffleNet_V2_X2_0_Weights.verify(weights) + + return _shufflenetv2(weights, progress, [4, 8, 4], [24, 244, 488, 976, 2048], **kwargs) diff --git a/lib/python3.12/site-packages/torchvision/models/squeezenet.py b/lib/python3.12/site-packages/torchvision/models/squeezenet.py new file mode 100644 index 0000000000000000000000000000000000000000..982b32107b09c280b4c7caa61e6b80be0cbf041e --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/squeezenet.py @@ -0,0 +1,223 @@ +from functools import partial +from typing import Any, Optional + +import torch +import torch.nn as nn +import torch.nn.init as init + +from ..transforms._presets import ImageClassification +from ..utils import _log_api_usage_once +from ._api import register_model, Weights, WeightsEnum +from ._meta import _IMAGENET_CATEGORIES +from ._utils import _ovewrite_named_param, handle_legacy_interface + + +__all__ = ["SqueezeNet", "SqueezeNet1_0_Weights", "SqueezeNet1_1_Weights", "squeezenet1_0", "squeezenet1_1"] + + +class Fire(nn.Module): + def __init__(self, inplanes: int, squeeze_planes: int, expand1x1_planes: int, expand3x3_planes: int) -> None: + super().__init__() + self.inplanes = inplanes + self.squeeze = nn.Conv2d(inplanes, squeeze_planes, kernel_size=1) + self.squeeze_activation = nn.ReLU(inplace=True) + self.expand1x1 = nn.Conv2d(squeeze_planes, expand1x1_planes, kernel_size=1) + self.expand1x1_activation = nn.ReLU(inplace=True) + self.expand3x3 = nn.Conv2d(squeeze_planes, expand3x3_planes, kernel_size=3, padding=1) + self.expand3x3_activation = nn.ReLU(inplace=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.squeeze_activation(self.squeeze(x)) + return torch.cat( + [self.expand1x1_activation(self.expand1x1(x)), self.expand3x3_activation(self.expand3x3(x))], 1 + ) + + +class SqueezeNet(nn.Module): + def __init__(self, version: str = "1_0", num_classes: int = 1000, dropout: float = 0.5) -> None: + super().__init__() + _log_api_usage_once(self) + self.num_classes = num_classes + if version == "1_0": + self.features = nn.Sequential( + nn.Conv2d(3, 96, kernel_size=7, stride=2), + nn.ReLU(inplace=True), + nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), + Fire(96, 16, 64, 64), + Fire(128, 16, 64, 64), + Fire(128, 32, 128, 128), + nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), + Fire(256, 32, 128, 128), + Fire(256, 48, 192, 192), + Fire(384, 48, 192, 192), + Fire(384, 64, 256, 256), + nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), + Fire(512, 64, 256, 256), + ) + elif version == "1_1": + self.features = nn.Sequential( + nn.Conv2d(3, 64, kernel_size=3, stride=2), + nn.ReLU(inplace=True), + nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), + Fire(64, 16, 64, 64), + Fire(128, 16, 64, 64), + nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), + Fire(128, 32, 128, 128), + Fire(256, 32, 128, 128), + nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), + Fire(256, 48, 192, 192), + Fire(384, 48, 192, 192), + Fire(384, 64, 256, 256), + Fire(512, 64, 256, 256), + ) + else: + # FIXME: Is this needed? SqueezeNet should only be called from the + # FIXME: squeezenet1_x() functions + # FIXME: This checking is not done for the other models + raise ValueError(f"Unsupported SqueezeNet version {version}: 1_0 or 1_1 expected") + + # Final convolution is initialized differently from the rest + final_conv = nn.Conv2d(512, self.num_classes, kernel_size=1) + self.classifier = nn.Sequential( + nn.Dropout(p=dropout), final_conv, nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d((1, 1)) + ) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + if m is final_conv: + init.normal_(m.weight, mean=0.0, std=0.01) + else: + init.kaiming_uniform_(m.weight) + if m.bias is not None: + init.constant_(m.bias, 0) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.features(x) + x = self.classifier(x) + return torch.flatten(x, 1) + + +def _squeezenet( + version: str, + weights: Optional[WeightsEnum], + progress: bool, + **kwargs: Any, +) -> SqueezeNet: + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + + model = SqueezeNet(version, **kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +_COMMON_META = { + "categories": _IMAGENET_CATEGORIES, + "recipe": "https://github.com/pytorch/vision/pull/49#issuecomment-277560717", + "_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""", +} + + +class SqueezeNet1_0_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/squeezenet1_0-b66bff10.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "min_size": (21, 21), + "num_params": 1248424, + "_metrics": { + "ImageNet-1K": { + "acc@1": 58.092, + "acc@5": 80.420, + } + }, + "_ops": 0.819, + "_file_size": 4.778, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class SqueezeNet1_1_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/squeezenet1_1-b8a52dc0.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "min_size": (17, 17), + "num_params": 1235496, + "_metrics": { + "ImageNet-1K": { + "acc@1": 58.178, + "acc@5": 80.624, + } + }, + "_ops": 0.349, + "_file_size": 4.729, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", SqueezeNet1_0_Weights.IMAGENET1K_V1)) +def squeezenet1_0( + *, weights: Optional[SqueezeNet1_0_Weights] = None, progress: bool = True, **kwargs: Any +) -> SqueezeNet: + """SqueezeNet model architecture from the `SqueezeNet: AlexNet-level + accuracy with 50x fewer parameters and <0.5MB model size + `_ paper. + + Args: + weights (:class:`~torchvision.models.SqueezeNet1_0_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.SqueezeNet1_0_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.squeezenet.SqueezeNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.SqueezeNet1_0_Weights + :members: + """ + weights = SqueezeNet1_0_Weights.verify(weights) + return _squeezenet("1_0", weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", SqueezeNet1_1_Weights.IMAGENET1K_V1)) +def squeezenet1_1( + *, weights: Optional[SqueezeNet1_1_Weights] = None, progress: bool = True, **kwargs: Any +) -> SqueezeNet: + """SqueezeNet 1.1 model from the `official SqueezeNet repo + `_. + + SqueezeNet 1.1 has 2.4x less computation and slightly fewer parameters + than SqueezeNet 1.0, without sacrificing accuracy. + + Args: + weights (:class:`~torchvision.models.SqueezeNet1_1_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.SqueezeNet1_1_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.squeezenet.SqueezeNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.SqueezeNet1_1_Weights + :members: + """ + weights = SqueezeNet1_1_Weights.verify(weights) + return _squeezenet("1_1", weights, progress, **kwargs) diff --git a/lib/python3.12/site-packages/torchvision/models/swin_transformer.py b/lib/python3.12/site-packages/torchvision/models/swin_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..2035f659bfc7f4e4f98d36cabff9be3802d49e59 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/swin_transformer.py @@ -0,0 +1,1033 @@ +import math +from functools import partial +from typing import Any, Callable, List, Optional + +import torch +import torch.nn.functional as F +from torch import nn, Tensor + +from ..ops.misc import MLP, Permute +from ..ops.stochastic_depth import StochasticDepth +from ..transforms._presets import ImageClassification, InterpolationMode +from ..utils import _log_api_usage_once +from ._api import register_model, Weights, WeightsEnum +from ._meta import _IMAGENET_CATEGORIES +from ._utils import _ovewrite_named_param, handle_legacy_interface + + +__all__ = [ + "SwinTransformer", + "Swin_T_Weights", + "Swin_S_Weights", + "Swin_B_Weights", + "Swin_V2_T_Weights", + "Swin_V2_S_Weights", + "Swin_V2_B_Weights", + "swin_t", + "swin_s", + "swin_b", + "swin_v2_t", + "swin_v2_s", + "swin_v2_b", +] + + +def _patch_merging_pad(x: torch.Tensor) -> torch.Tensor: + H, W, _ = x.shape[-3:] + x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2)) + x0 = x[..., 0::2, 0::2, :] # ... H/2 W/2 C + x1 = x[..., 1::2, 0::2, :] # ... H/2 W/2 C + x2 = x[..., 0::2, 1::2, :] # ... H/2 W/2 C + x3 = x[..., 1::2, 1::2, :] # ... H/2 W/2 C + x = torch.cat([x0, x1, x2, x3], -1) # ... H/2 W/2 4*C + return x + + +torch.fx.wrap("_patch_merging_pad") + + +def _get_relative_position_bias( + relative_position_bias_table: torch.Tensor, relative_position_index: torch.Tensor, window_size: List[int] +) -> torch.Tensor: + N = window_size[0] * window_size[1] + relative_position_bias = relative_position_bias_table[relative_position_index] # type: ignore[index] + relative_position_bias = relative_position_bias.view(N, N, -1) + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous().unsqueeze(0) + return relative_position_bias + + +torch.fx.wrap("_get_relative_position_bias") + + +class PatchMerging(nn.Module): + """Patch Merging Layer. + Args: + dim (int): Number of input channels. + norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. + """ + + def __init__(self, dim: int, norm_layer: Callable[..., nn.Module] = nn.LayerNorm): + super().__init__() + _log_api_usage_once(self) + self.dim = dim + self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) + self.norm = norm_layer(4 * dim) + + def forward(self, x: Tensor): + """ + Args: + x (Tensor): input tensor with expected layout of [..., H, W, C] + Returns: + Tensor with layout of [..., H/2, W/2, 2*C] + """ + x = _patch_merging_pad(x) + x = self.norm(x) + x = self.reduction(x) # ... H/2 W/2 2*C + return x + + +class PatchMergingV2(nn.Module): + """Patch Merging Layer for Swin Transformer V2. + Args: + dim (int): Number of input channels. + norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. + """ + + def __init__(self, dim: int, norm_layer: Callable[..., nn.Module] = nn.LayerNorm): + super().__init__() + _log_api_usage_once(self) + self.dim = dim + self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) + self.norm = norm_layer(2 * dim) # difference + + def forward(self, x: Tensor): + """ + Args: + x (Tensor): input tensor with expected layout of [..., H, W, C] + Returns: + Tensor with layout of [..., H/2, W/2, 2*C] + """ + x = _patch_merging_pad(x) + x = self.reduction(x) # ... H/2 W/2 2*C + x = self.norm(x) + return x + + +def shifted_window_attention( + input: Tensor, + qkv_weight: Tensor, + proj_weight: Tensor, + relative_position_bias: Tensor, + window_size: List[int], + num_heads: int, + shift_size: List[int], + attention_dropout: float = 0.0, + dropout: float = 0.0, + qkv_bias: Optional[Tensor] = None, + proj_bias: Optional[Tensor] = None, + logit_scale: Optional[torch.Tensor] = None, + training: bool = True, +) -> Tensor: + """ + Window based multi-head self attention (W-MSA) module with relative position bias. + It supports both of shifted and non-shifted window. + Args: + input (Tensor[N, H, W, C]): The input tensor or 4-dimensions. + qkv_weight (Tensor[in_dim, out_dim]): The weight tensor of query, key, value. + proj_weight (Tensor[out_dim, out_dim]): The weight tensor of projection. + relative_position_bias (Tensor): The learned relative position bias added to attention. + window_size (List[int]): Window size. + num_heads (int): Number of attention heads. + shift_size (List[int]): Shift size for shifted window attention. + attention_dropout (float): Dropout ratio of attention weight. Default: 0.0. + dropout (float): Dropout ratio of output. Default: 0.0. + qkv_bias (Tensor[out_dim], optional): The bias tensor of query, key, value. Default: None. + proj_bias (Tensor[out_dim], optional): The bias tensor of projection. Default: None. + logit_scale (Tensor[out_dim], optional): Logit scale of cosine attention for Swin Transformer V2. Default: None. + training (bool, optional): Training flag used by the dropout parameters. Default: True. + Returns: + Tensor[N, H, W, C]: The output tensor after shifted window attention. + """ + B, H, W, C = input.shape + # pad feature maps to multiples of window size + pad_r = (window_size[1] - W % window_size[1]) % window_size[1] + pad_b = (window_size[0] - H % window_size[0]) % window_size[0] + x = F.pad(input, (0, 0, 0, pad_r, 0, pad_b)) + _, pad_H, pad_W, _ = x.shape + + shift_size = shift_size.copy() + # If window size is larger than feature size, there is no need to shift window + if window_size[0] >= pad_H: + shift_size[0] = 0 + if window_size[1] >= pad_W: + shift_size[1] = 0 + + # cyclic shift + if sum(shift_size) > 0: + x = torch.roll(x, shifts=(-shift_size[0], -shift_size[1]), dims=(1, 2)) + + # partition windows + num_windows = (pad_H // window_size[0]) * (pad_W // window_size[1]) + x = x.view(B, pad_H // window_size[0], window_size[0], pad_W // window_size[1], window_size[1], C) + x = x.permute(0, 1, 3, 2, 4, 5).reshape(B * num_windows, window_size[0] * window_size[1], C) # B*nW, Ws*Ws, C + + # multi-head attention + if logit_scale is not None and qkv_bias is not None: + qkv_bias = qkv_bias.clone() + length = qkv_bias.numel() // 3 + qkv_bias[length : 2 * length].zero_() + qkv = F.linear(x, qkv_weight, qkv_bias) + qkv = qkv.reshape(x.size(0), x.size(1), 3, num_heads, C // num_heads).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] + if logit_scale is not None: + # cosine attention + attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1) + logit_scale = torch.clamp(logit_scale, max=math.log(100.0)).exp() + attn = attn * logit_scale + else: + q = q * (C // num_heads) ** -0.5 + attn = q.matmul(k.transpose(-2, -1)) + # add relative position bias + attn = attn + relative_position_bias + + if sum(shift_size) > 0: + # generate attention mask + attn_mask = x.new_zeros((pad_H, pad_W)) + h_slices = ((0, -window_size[0]), (-window_size[0], -shift_size[0]), (-shift_size[0], None)) + w_slices = ((0, -window_size[1]), (-window_size[1], -shift_size[1]), (-shift_size[1], None)) + count = 0 + for h in h_slices: + for w in w_slices: + attn_mask[h[0] : h[1], w[0] : w[1]] = count + count += 1 + attn_mask = attn_mask.view(pad_H // window_size[0], window_size[0], pad_W // window_size[1], window_size[1]) + attn_mask = attn_mask.permute(0, 2, 1, 3).reshape(num_windows, window_size[0] * window_size[1]) + attn_mask = attn_mask.unsqueeze(1) - attn_mask.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)) + attn = attn.view(x.size(0) // num_windows, num_windows, num_heads, x.size(1), x.size(1)) + attn = attn + attn_mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, num_heads, x.size(1), x.size(1)) + + attn = F.softmax(attn, dim=-1) + attn = F.dropout(attn, p=attention_dropout, training=training) + + x = attn.matmul(v).transpose(1, 2).reshape(x.size(0), x.size(1), C) + x = F.linear(x, proj_weight, proj_bias) + x = F.dropout(x, p=dropout, training=training) + + # reverse windows + x = x.view(B, pad_H // window_size[0], pad_W // window_size[1], window_size[0], window_size[1], C) + x = x.permute(0, 1, 3, 2, 4, 5).reshape(B, pad_H, pad_W, C) + + # reverse cyclic shift + if sum(shift_size) > 0: + x = torch.roll(x, shifts=(shift_size[0], shift_size[1]), dims=(1, 2)) + + # unpad features + x = x[:, :H, :W, :].contiguous() + return x + + +torch.fx.wrap("shifted_window_attention") + + +class ShiftedWindowAttention(nn.Module): + """ + See :func:`shifted_window_attention`. + """ + + def __init__( + self, + dim: int, + window_size: List[int], + shift_size: List[int], + num_heads: int, + qkv_bias: bool = True, + proj_bias: bool = True, + attention_dropout: float = 0.0, + dropout: float = 0.0, + ): + super().__init__() + if len(window_size) != 2 or len(shift_size) != 2: + raise ValueError("window_size and shift_size must be of length 2") + self.window_size = window_size + self.shift_size = shift_size + self.num_heads = num_heads + self.attention_dropout = attention_dropout + self.dropout = dropout + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.proj = nn.Linear(dim, dim, bias=proj_bias) + + self.define_relative_position_bias_table() + self.define_relative_position_index() + + def define_relative_position_bias_table(self): + # define a parameter table of relative position bias + self.relative_position_bias_table = nn.Parameter( + torch.zeros((2 * self.window_size[0] - 1) * (2 * self.window_size[1] - 1), self.num_heads) + ) # 2*Wh-1 * 2*Ww-1, nH + nn.init.trunc_normal_(self.relative_position_bias_table, std=0.02) + + def define_relative_position_index(self): + # get pair-wise relative position index for each token inside the window + coords_h = torch.arange(self.window_size[0]) + coords_w = torch.arange(self.window_size[1]) + coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij")) # 2, Wh, Ww + coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww + relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 + relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 + relative_position_index = relative_coords.sum(-1).flatten() # Wh*Ww*Wh*Ww + self.register_buffer("relative_position_index", relative_position_index) + + def get_relative_position_bias(self) -> torch.Tensor: + return _get_relative_position_bias( + self.relative_position_bias_table, self.relative_position_index, self.window_size # type: ignore[arg-type] + ) + + def forward(self, x: Tensor) -> Tensor: + """ + Args: + x (Tensor): Tensor with layout of [B, H, W, C] + Returns: + Tensor with same layout as input, i.e. [B, H, W, C] + """ + relative_position_bias = self.get_relative_position_bias() + return shifted_window_attention( + x, + self.qkv.weight, + self.proj.weight, + relative_position_bias, + self.window_size, + self.num_heads, + shift_size=self.shift_size, + attention_dropout=self.attention_dropout, + dropout=self.dropout, + qkv_bias=self.qkv.bias, + proj_bias=self.proj.bias, + training=self.training, + ) + + +class ShiftedWindowAttentionV2(ShiftedWindowAttention): + """ + See :func:`shifted_window_attention_v2`. + """ + + def __init__( + self, + dim: int, + window_size: List[int], + shift_size: List[int], + num_heads: int, + qkv_bias: bool = True, + proj_bias: bool = True, + attention_dropout: float = 0.0, + dropout: float = 0.0, + ): + super().__init__( + dim, + window_size, + shift_size, + num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + attention_dropout=attention_dropout, + dropout=dropout, + ) + + self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1)))) + # mlp to generate continuous relative position bias + self.cpb_mlp = nn.Sequential( + nn.Linear(2, 512, bias=True), nn.ReLU(inplace=True), nn.Linear(512, num_heads, bias=False) + ) + if qkv_bias: + length = self.qkv.bias.numel() // 3 + self.qkv.bias[length : 2 * length].data.zero_() + + def define_relative_position_bias_table(self): + # get relative_coords_table + relative_coords_h = torch.arange(-(self.window_size[0] - 1), self.window_size[0], dtype=torch.float32) + relative_coords_w = torch.arange(-(self.window_size[1] - 1), self.window_size[1], dtype=torch.float32) + relative_coords_table = torch.stack(torch.meshgrid([relative_coords_h, relative_coords_w], indexing="ij")) + relative_coords_table = relative_coords_table.permute(1, 2, 0).contiguous().unsqueeze(0) # 1, 2*Wh-1, 2*Ww-1, 2 + + relative_coords_table[:, :, :, 0] /= self.window_size[0] - 1 + relative_coords_table[:, :, :, 1] /= self.window_size[1] - 1 + + relative_coords_table *= 8 # normalize to -8, 8 + relative_coords_table = ( + torch.sign(relative_coords_table) * torch.log2(torch.abs(relative_coords_table) + 1.0) / 3.0 + ) + self.register_buffer("relative_coords_table", relative_coords_table) + + def get_relative_position_bias(self) -> torch.Tensor: + relative_position_bias = _get_relative_position_bias( + self.cpb_mlp(self.relative_coords_table).view(-1, self.num_heads), + self.relative_position_index, # type: ignore[arg-type] + self.window_size, + ) + relative_position_bias = 16 * torch.sigmoid(relative_position_bias) + return relative_position_bias + + def forward(self, x: Tensor): + """ + Args: + x (Tensor): Tensor with layout of [B, H, W, C] + Returns: + Tensor with same layout as input, i.e. [B, H, W, C] + """ + relative_position_bias = self.get_relative_position_bias() + return shifted_window_attention( + x, + self.qkv.weight, + self.proj.weight, + relative_position_bias, + self.window_size, + self.num_heads, + shift_size=self.shift_size, + attention_dropout=self.attention_dropout, + dropout=self.dropout, + qkv_bias=self.qkv.bias, + proj_bias=self.proj.bias, + logit_scale=self.logit_scale, + training=self.training, + ) + + +class SwinTransformerBlock(nn.Module): + """ + Swin Transformer Block. + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (List[int]): Window size. + shift_size (List[int]): Shift size for shifted window attention. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.0. + dropout (float): Dropout rate. Default: 0.0. + attention_dropout (float): Attention dropout rate. Default: 0.0. + stochastic_depth_prob: (float): Stochastic depth rate. Default: 0.0. + norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. + attn_layer (nn.Module): Attention layer. Default: ShiftedWindowAttention + """ + + def __init__( + self, + dim: int, + num_heads: int, + window_size: List[int], + shift_size: List[int], + mlp_ratio: float = 4.0, + dropout: float = 0.0, + attention_dropout: float = 0.0, + stochastic_depth_prob: float = 0.0, + norm_layer: Callable[..., nn.Module] = nn.LayerNorm, + attn_layer: Callable[..., nn.Module] = ShiftedWindowAttention, + ): + super().__init__() + _log_api_usage_once(self) + + self.norm1 = norm_layer(dim) + self.attn = attn_layer( + dim, + window_size, + shift_size, + num_heads, + attention_dropout=attention_dropout, + dropout=dropout, + ) + self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row") + self.norm2 = norm_layer(dim) + self.mlp = MLP(dim, [int(dim * mlp_ratio), dim], activation_layer=nn.GELU, inplace=None, dropout=dropout) + + for m in self.mlp.modules(): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.normal_(m.bias, std=1e-6) + + def forward(self, x: Tensor): + x = x + self.stochastic_depth(self.attn(self.norm1(x))) + x = x + self.stochastic_depth(self.mlp(self.norm2(x))) + return x + + +class SwinTransformerBlockV2(SwinTransformerBlock): + """ + Swin Transformer V2 Block. + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (List[int]): Window size. + shift_size (List[int]): Shift size for shifted window attention. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.0. + dropout (float): Dropout rate. Default: 0.0. + attention_dropout (float): Attention dropout rate. Default: 0.0. + stochastic_depth_prob: (float): Stochastic depth rate. Default: 0.0. + norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. + attn_layer (nn.Module): Attention layer. Default: ShiftedWindowAttentionV2. + """ + + def __init__( + self, + dim: int, + num_heads: int, + window_size: List[int], + shift_size: List[int], + mlp_ratio: float = 4.0, + dropout: float = 0.0, + attention_dropout: float = 0.0, + stochastic_depth_prob: float = 0.0, + norm_layer: Callable[..., nn.Module] = nn.LayerNorm, + attn_layer: Callable[..., nn.Module] = ShiftedWindowAttentionV2, + ): + super().__init__( + dim, + num_heads, + window_size, + shift_size, + mlp_ratio=mlp_ratio, + dropout=dropout, + attention_dropout=attention_dropout, + stochastic_depth_prob=stochastic_depth_prob, + norm_layer=norm_layer, + attn_layer=attn_layer, + ) + + def forward(self, x: Tensor): + # Here is the difference, we apply norm after the attention in V2. + # In V1 we applied norm before the attention. + x = x + self.stochastic_depth(self.norm1(self.attn(x))) + x = x + self.stochastic_depth(self.norm2(self.mlp(x))) + return x + + +class SwinTransformer(nn.Module): + """ + Implements Swin Transformer from the `"Swin Transformer: Hierarchical Vision Transformer using + Shifted Windows" `_ paper. + Args: + patch_size (List[int]): Patch size. + embed_dim (int): Patch embedding dimension. + depths (List(int)): Depth of each Swin Transformer layer. + num_heads (List(int)): Number of attention heads in different layers. + window_size (List[int]): Window size. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.0. + dropout (float): Dropout rate. Default: 0.0. + attention_dropout (float): Attention dropout rate. Default: 0.0. + stochastic_depth_prob (float): Stochastic depth rate. Default: 0.1. + num_classes (int): Number of classes for classification head. Default: 1000. + block (nn.Module, optional): SwinTransformer Block. Default: None. + norm_layer (nn.Module, optional): Normalization layer. Default: None. + downsample_layer (nn.Module): Downsample layer (patch merging). Default: PatchMerging. + """ + + def __init__( + self, + patch_size: List[int], + embed_dim: int, + depths: List[int], + num_heads: List[int], + window_size: List[int], + mlp_ratio: float = 4.0, + dropout: float = 0.0, + attention_dropout: float = 0.0, + stochastic_depth_prob: float = 0.1, + num_classes: int = 1000, + norm_layer: Optional[Callable[..., nn.Module]] = None, + block: Optional[Callable[..., nn.Module]] = None, + downsample_layer: Callable[..., nn.Module] = PatchMerging, + ): + super().__init__() + _log_api_usage_once(self) + self.num_classes = num_classes + + if block is None: + block = SwinTransformerBlock + if norm_layer is None: + norm_layer = partial(nn.LayerNorm, eps=1e-5) + + layers: List[nn.Module] = [] + # split image into non-overlapping patches + layers.append( + nn.Sequential( + nn.Conv2d( + 3, embed_dim, kernel_size=(patch_size[0], patch_size[1]), stride=(patch_size[0], patch_size[1]) + ), + Permute([0, 2, 3, 1]), + norm_layer(embed_dim), + ) + ) + + total_stage_blocks = sum(depths) + stage_block_id = 0 + # build SwinTransformer blocks + for i_stage in range(len(depths)): + stage: List[nn.Module] = [] + dim = embed_dim * 2**i_stage + for i_layer in range(depths[i_stage]): + # adjust stochastic depth probability based on the depth of the stage block + sd_prob = stochastic_depth_prob * float(stage_block_id) / (total_stage_blocks - 1) + stage.append( + block( + dim, + num_heads[i_stage], + window_size=window_size, + shift_size=[0 if i_layer % 2 == 0 else w // 2 for w in window_size], + mlp_ratio=mlp_ratio, + dropout=dropout, + attention_dropout=attention_dropout, + stochastic_depth_prob=sd_prob, + norm_layer=norm_layer, + ) + ) + stage_block_id += 1 + layers.append(nn.Sequential(*stage)) + # add patch merging layer + if i_stage < (len(depths) - 1): + layers.append(downsample_layer(dim, norm_layer)) + self.features = nn.Sequential(*layers) + + num_features = embed_dim * 2 ** (len(depths) - 1) + self.norm = norm_layer(num_features) + self.permute = Permute([0, 3, 1, 2]) # B H W C -> B C H W + self.avgpool = nn.AdaptiveAvgPool2d(1) + self.flatten = nn.Flatten(1) + self.head = nn.Linear(num_features, num_classes) + + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.zeros_(m.bias) + + def forward(self, x): + x = self.features(x) + x = self.norm(x) + x = self.permute(x) + x = self.avgpool(x) + x = self.flatten(x) + x = self.head(x) + return x + + +def _swin_transformer( + patch_size: List[int], + embed_dim: int, + depths: List[int], + num_heads: List[int], + window_size: List[int], + stochastic_depth_prob: float, + weights: Optional[WeightsEnum], + progress: bool, + **kwargs: Any, +) -> SwinTransformer: + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + + model = SwinTransformer( + patch_size=patch_size, + embed_dim=embed_dim, + depths=depths, + num_heads=num_heads, + window_size=window_size, + stochastic_depth_prob=stochastic_depth_prob, + **kwargs, + ) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +_COMMON_META = { + "categories": _IMAGENET_CATEGORIES, +} + + +class Swin_T_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/swin_t-704ceda3.pth", + transforms=partial( + ImageClassification, crop_size=224, resize_size=232, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_META, + "num_params": 28288354, + "min_size": (224, 224), + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#swintransformer", + "_metrics": { + "ImageNet-1K": { + "acc@1": 81.474, + "acc@5": 95.776, + } + }, + "_ops": 4.491, + "_file_size": 108.19, + "_docs": """These weights reproduce closely the results of the paper using a similar training recipe.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class Swin_S_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/swin_s-5e29d889.pth", + transforms=partial( + ImageClassification, crop_size=224, resize_size=246, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_META, + "num_params": 49606258, + "min_size": (224, 224), + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#swintransformer", + "_metrics": { + "ImageNet-1K": { + "acc@1": 83.196, + "acc@5": 96.360, + } + }, + "_ops": 8.741, + "_file_size": 189.786, + "_docs": """These weights reproduce closely the results of the paper using a similar training recipe.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class Swin_B_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/swin_b-68c6b09e.pth", + transforms=partial( + ImageClassification, crop_size=224, resize_size=238, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_META, + "num_params": 87768224, + "min_size": (224, 224), + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#swintransformer", + "_metrics": { + "ImageNet-1K": { + "acc@1": 83.582, + "acc@5": 96.640, + } + }, + "_ops": 15.431, + "_file_size": 335.364, + "_docs": """These weights reproduce closely the results of the paper using a similar training recipe.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class Swin_V2_T_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/swin_v2_t-b137f0e2.pth", + transforms=partial( + ImageClassification, crop_size=256, resize_size=260, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_META, + "num_params": 28351570, + "min_size": (256, 256), + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#swintransformer-v2", + "_metrics": { + "ImageNet-1K": { + "acc@1": 82.072, + "acc@5": 96.132, + } + }, + "_ops": 5.94, + "_file_size": 108.626, + "_docs": """These weights reproduce closely the results of the paper using a similar training recipe.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class Swin_V2_S_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/swin_v2_s-637d8ceb.pth", + transforms=partial( + ImageClassification, crop_size=256, resize_size=260, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_META, + "num_params": 49737442, + "min_size": (256, 256), + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#swintransformer-v2", + "_metrics": { + "ImageNet-1K": { + "acc@1": 83.712, + "acc@5": 96.816, + } + }, + "_ops": 11.546, + "_file_size": 190.675, + "_docs": """These weights reproduce closely the results of the paper using a similar training recipe.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class Swin_V2_B_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/swin_v2_b-781e5279.pth", + transforms=partial( + ImageClassification, crop_size=256, resize_size=272, interpolation=InterpolationMode.BICUBIC + ), + meta={ + **_COMMON_META, + "num_params": 87930848, + "min_size": (256, 256), + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#swintransformer-v2", + "_metrics": { + "ImageNet-1K": { + "acc@1": 84.112, + "acc@5": 96.864, + } + }, + "_ops": 20.325, + "_file_size": 336.372, + "_docs": """These weights reproduce closely the results of the paper using a similar training recipe.""", + }, + ) + DEFAULT = IMAGENET1K_V1 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", Swin_T_Weights.IMAGENET1K_V1)) +def swin_t(*, weights: Optional[Swin_T_Weights] = None, progress: bool = True, **kwargs: Any) -> SwinTransformer: + """ + Constructs a swin_tiny architecture from + `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows `_. + + Args: + weights (:class:`~torchvision.models.Swin_T_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.Swin_T_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.swin_transformer.SwinTransformer`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.Swin_T_Weights + :members: + """ + weights = Swin_T_Weights.verify(weights) + + return _swin_transformer( + patch_size=[4, 4], + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=[7, 7], + stochastic_depth_prob=0.2, + weights=weights, + progress=progress, + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", Swin_S_Weights.IMAGENET1K_V1)) +def swin_s(*, weights: Optional[Swin_S_Weights] = None, progress: bool = True, **kwargs: Any) -> SwinTransformer: + """ + Constructs a swin_small architecture from + `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows `_. + + Args: + weights (:class:`~torchvision.models.Swin_S_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.Swin_S_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.swin_transformer.SwinTransformer`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.Swin_S_Weights + :members: + """ + weights = Swin_S_Weights.verify(weights) + + return _swin_transformer( + patch_size=[4, 4], + embed_dim=96, + depths=[2, 2, 18, 2], + num_heads=[3, 6, 12, 24], + window_size=[7, 7], + stochastic_depth_prob=0.3, + weights=weights, + progress=progress, + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", Swin_B_Weights.IMAGENET1K_V1)) +def swin_b(*, weights: Optional[Swin_B_Weights] = None, progress: bool = True, **kwargs: Any) -> SwinTransformer: + """ + Constructs a swin_base architecture from + `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows `_. + + Args: + weights (:class:`~torchvision.models.Swin_B_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.Swin_B_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.swin_transformer.SwinTransformer`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.Swin_B_Weights + :members: + """ + weights = Swin_B_Weights.verify(weights) + + return _swin_transformer( + patch_size=[4, 4], + embed_dim=128, + depths=[2, 2, 18, 2], + num_heads=[4, 8, 16, 32], + window_size=[7, 7], + stochastic_depth_prob=0.5, + weights=weights, + progress=progress, + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", Swin_V2_T_Weights.IMAGENET1K_V1)) +def swin_v2_t(*, weights: Optional[Swin_V2_T_Weights] = None, progress: bool = True, **kwargs: Any) -> SwinTransformer: + """ + Constructs a swin_v2_tiny architecture from + `Swin Transformer V2: Scaling Up Capacity and Resolution `_. + + Args: + weights (:class:`~torchvision.models.Swin_V2_T_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.Swin_V2_T_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.swin_transformer.SwinTransformer`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.Swin_V2_T_Weights + :members: + """ + weights = Swin_V2_T_Weights.verify(weights) + + return _swin_transformer( + patch_size=[4, 4], + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=[8, 8], + stochastic_depth_prob=0.2, + weights=weights, + progress=progress, + block=SwinTransformerBlockV2, + downsample_layer=PatchMergingV2, + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", Swin_V2_S_Weights.IMAGENET1K_V1)) +def swin_v2_s(*, weights: Optional[Swin_V2_S_Weights] = None, progress: bool = True, **kwargs: Any) -> SwinTransformer: + """ + Constructs a swin_v2_small architecture from + `Swin Transformer V2: Scaling Up Capacity and Resolution `_. + + Args: + weights (:class:`~torchvision.models.Swin_V2_S_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.Swin_V2_S_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.swin_transformer.SwinTransformer`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.Swin_V2_S_Weights + :members: + """ + weights = Swin_V2_S_Weights.verify(weights) + + return _swin_transformer( + patch_size=[4, 4], + embed_dim=96, + depths=[2, 2, 18, 2], + num_heads=[3, 6, 12, 24], + window_size=[8, 8], + stochastic_depth_prob=0.3, + weights=weights, + progress=progress, + block=SwinTransformerBlockV2, + downsample_layer=PatchMergingV2, + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", Swin_V2_B_Weights.IMAGENET1K_V1)) +def swin_v2_b(*, weights: Optional[Swin_V2_B_Weights] = None, progress: bool = True, **kwargs: Any) -> SwinTransformer: + """ + Constructs a swin_v2_base architecture from + `Swin Transformer V2: Scaling Up Capacity and Resolution `_. + + Args: + weights (:class:`~torchvision.models.Swin_V2_B_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.Swin_V2_B_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.swin_transformer.SwinTransformer`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.Swin_V2_B_Weights + :members: + """ + weights = Swin_V2_B_Weights.verify(weights) + + return _swin_transformer( + patch_size=[4, 4], + embed_dim=128, + depths=[2, 2, 18, 2], + num_heads=[4, 8, 16, 32], + window_size=[8, 8], + stochastic_depth_prob=0.5, + weights=weights, + progress=progress, + block=SwinTransformerBlockV2, + downsample_layer=PatchMergingV2, + **kwargs, + ) diff --git a/lib/python3.12/site-packages/torchvision/models/vgg.py b/lib/python3.12/site-packages/torchvision/models/vgg.py new file mode 100644 index 0000000000000000000000000000000000000000..0a548570dd41fa87309e8aa4bdb373517f265060 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/vgg.py @@ -0,0 +1,511 @@ +from functools import partial +from typing import Any, cast, Dict, List, Optional, Union + +import torch +import torch.nn as nn + +from ..transforms._presets import ImageClassification +from ..utils import _log_api_usage_once +from ._api import register_model, Weights, WeightsEnum +from ._meta import _IMAGENET_CATEGORIES +from ._utils import _ovewrite_named_param, handle_legacy_interface + + +__all__ = [ + "VGG", + "VGG11_Weights", + "VGG11_BN_Weights", + "VGG13_Weights", + "VGG13_BN_Weights", + "VGG16_Weights", + "VGG16_BN_Weights", + "VGG19_Weights", + "VGG19_BN_Weights", + "vgg11", + "vgg11_bn", + "vgg13", + "vgg13_bn", + "vgg16", + "vgg16_bn", + "vgg19", + "vgg19_bn", +] + + +class VGG(nn.Module): + def __init__( + self, features: nn.Module, num_classes: int = 1000, init_weights: bool = True, dropout: float = 0.5 + ) -> None: + super().__init__() + _log_api_usage_once(self) + self.features = features + self.avgpool = nn.AdaptiveAvgPool2d((7, 7)) + self.classifier = nn.Sequential( + nn.Linear(512 * 7 * 7, 4096), + nn.ReLU(True), + nn.Dropout(p=dropout), + nn.Linear(4096, 4096), + nn.ReLU(True), + nn.Dropout(p=dropout), + nn.Linear(4096, num_classes), + ) + if init_weights: + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.BatchNorm2d): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Linear): + nn.init.normal_(m.weight, 0, 0.01) + nn.init.constant_(m.bias, 0) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.features(x) + x = self.avgpool(x) + x = torch.flatten(x, 1) + x = self.classifier(x) + return x + + +def make_layers(cfg: List[Union[str, int]], batch_norm: bool = False) -> nn.Sequential: + layers: List[nn.Module] = [] + in_channels = 3 + for v in cfg: + if v == "M": + layers += [nn.MaxPool2d(kernel_size=2, stride=2)] + else: + v = cast(int, v) + conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1) + if batch_norm: + layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)] + else: + layers += [conv2d, nn.ReLU(inplace=True)] + in_channels = v + return nn.Sequential(*layers) + + +cfgs: Dict[str, List[Union[str, int]]] = { + "A": [64, "M", 128, "M", 256, 256, "M", 512, 512, "M", 512, 512, "M"], + "B": [64, 64, "M", 128, 128, "M", 256, 256, "M", 512, 512, "M", 512, 512, "M"], + "D": [64, 64, "M", 128, 128, "M", 256, 256, 256, "M", 512, 512, 512, "M", 512, 512, 512, "M"], + "E": [64, 64, "M", 128, 128, "M", 256, 256, 256, 256, "M", 512, 512, 512, 512, "M", 512, 512, 512, 512, "M"], +} + + +def _vgg(cfg: str, batch_norm: bool, weights: Optional[WeightsEnum], progress: bool, **kwargs: Any) -> VGG: + if weights is not None: + kwargs["init_weights"] = False + if weights.meta["categories"] is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + model = VGG(make_layers(cfgs[cfg], batch_norm=batch_norm), **kwargs) + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + return model + + +_COMMON_META = { + "min_size": (32, 32), + "categories": _IMAGENET_CATEGORIES, + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#alexnet-and-vgg", + "_docs": """These weights were trained from scratch by using a simplified training recipe.""", +} + + +class VGG11_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/vgg11-8a719046.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 132863336, + "_metrics": { + "ImageNet-1K": { + "acc@1": 69.020, + "acc@5": 88.628, + } + }, + "_ops": 7.609, + "_file_size": 506.84, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class VGG11_BN_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/vgg11_bn-6002323d.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 132868840, + "_metrics": { + "ImageNet-1K": { + "acc@1": 70.370, + "acc@5": 89.810, + } + }, + "_ops": 7.609, + "_file_size": 506.881, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class VGG13_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/vgg13-19584684.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 133047848, + "_metrics": { + "ImageNet-1K": { + "acc@1": 69.928, + "acc@5": 89.246, + } + }, + "_ops": 11.308, + "_file_size": 507.545, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class VGG13_BN_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/vgg13_bn-abd245e5.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 133053736, + "_metrics": { + "ImageNet-1K": { + "acc@1": 71.586, + "acc@5": 90.374, + } + }, + "_ops": 11.308, + "_file_size": 507.59, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class VGG16_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/vgg16-397923af.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 138357544, + "_metrics": { + "ImageNet-1K": { + "acc@1": 71.592, + "acc@5": 90.382, + } + }, + "_ops": 15.47, + "_file_size": 527.796, + }, + ) + IMAGENET1K_FEATURES = Weights( + # Weights ported from https://github.com/amdegroot/ssd.pytorch/ + url="https://download.pytorch.org/models/vgg16_features-amdegroot-88682ab5.pth", + transforms=partial( + ImageClassification, + crop_size=224, + mean=(0.48235, 0.45882, 0.40784), + std=(1.0 / 255.0, 1.0 / 255.0, 1.0 / 255.0), + ), + meta={ + **_COMMON_META, + "num_params": 138357544, + "categories": None, + "recipe": "https://github.com/amdegroot/ssd.pytorch#training-ssd", + "_metrics": { + "ImageNet-1K": { + "acc@1": float("nan"), + "acc@5": float("nan"), + } + }, + "_ops": 15.47, + "_file_size": 527.802, + "_docs": """ + These weights can't be used for classification because they are missing values in the `classifier` + module. Only the `features` module has valid values and can be used for feature extraction. The weights + were trained using the original input standardization method as described in the paper. + """, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class VGG16_BN_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/vgg16_bn-6c64b313.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 138365992, + "_metrics": { + "ImageNet-1K": { + "acc@1": 73.360, + "acc@5": 91.516, + } + }, + "_ops": 15.47, + "_file_size": 527.866, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class VGG19_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/vgg19-dcbb9e9d.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 143667240, + "_metrics": { + "ImageNet-1K": { + "acc@1": 72.376, + "acc@5": 90.876, + } + }, + "_ops": 19.632, + "_file_size": 548.051, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class VGG19_BN_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/vgg19_bn-c79401a0.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 143678248, + "_metrics": { + "ImageNet-1K": { + "acc@1": 74.218, + "acc@5": 91.842, + } + }, + "_ops": 19.632, + "_file_size": 548.143, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", VGG11_Weights.IMAGENET1K_V1)) +def vgg11(*, weights: Optional[VGG11_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG: + """VGG-11 from `Very Deep Convolutional Networks for Large-Scale Image Recognition `__. + + Args: + weights (:class:`~torchvision.models.VGG11_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.VGG11_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.vgg.VGG`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.VGG11_Weights + :members: + """ + weights = VGG11_Weights.verify(weights) + + return _vgg("A", False, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", VGG11_BN_Weights.IMAGENET1K_V1)) +def vgg11_bn(*, weights: Optional[VGG11_BN_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG: + """VGG-11-BN from `Very Deep Convolutional Networks for Large-Scale Image Recognition `__. + + Args: + weights (:class:`~torchvision.models.VGG11_BN_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.VGG11_BN_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.vgg.VGG`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.VGG11_BN_Weights + :members: + """ + weights = VGG11_BN_Weights.verify(weights) + + return _vgg("A", True, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", VGG13_Weights.IMAGENET1K_V1)) +def vgg13(*, weights: Optional[VGG13_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG: + """VGG-13 from `Very Deep Convolutional Networks for Large-Scale Image Recognition `__. + + Args: + weights (:class:`~torchvision.models.VGG13_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.VGG13_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.vgg.VGG`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.VGG13_Weights + :members: + """ + weights = VGG13_Weights.verify(weights) + + return _vgg("B", False, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", VGG13_BN_Weights.IMAGENET1K_V1)) +def vgg13_bn(*, weights: Optional[VGG13_BN_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG: + """VGG-13-BN from `Very Deep Convolutional Networks for Large-Scale Image Recognition `__. + + Args: + weights (:class:`~torchvision.models.VGG13_BN_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.VGG13_BN_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.vgg.VGG`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.VGG13_BN_Weights + :members: + """ + weights = VGG13_BN_Weights.verify(weights) + + return _vgg("B", True, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", VGG16_Weights.IMAGENET1K_V1)) +def vgg16(*, weights: Optional[VGG16_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG: + """VGG-16 from `Very Deep Convolutional Networks for Large-Scale Image Recognition `__. + + Args: + weights (:class:`~torchvision.models.VGG16_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.VGG16_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.vgg.VGG`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.VGG16_Weights + :members: + """ + weights = VGG16_Weights.verify(weights) + + return _vgg("D", False, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", VGG16_BN_Weights.IMAGENET1K_V1)) +def vgg16_bn(*, weights: Optional[VGG16_BN_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG: + """VGG-16-BN from `Very Deep Convolutional Networks for Large-Scale Image Recognition `__. + + Args: + weights (:class:`~torchvision.models.VGG16_BN_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.VGG16_BN_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.vgg.VGG`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.VGG16_BN_Weights + :members: + """ + weights = VGG16_BN_Weights.verify(weights) + + return _vgg("D", True, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", VGG19_Weights.IMAGENET1K_V1)) +def vgg19(*, weights: Optional[VGG19_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG: + """VGG-19 from `Very Deep Convolutional Networks for Large-Scale Image Recognition `__. + + Args: + weights (:class:`~torchvision.models.VGG19_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.VGG19_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.vgg.VGG`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.VGG19_Weights + :members: + """ + weights = VGG19_Weights.verify(weights) + + return _vgg("E", False, weights, progress, **kwargs) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", VGG19_BN_Weights.IMAGENET1K_V1)) +def vgg19_bn(*, weights: Optional[VGG19_BN_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG: + """VGG-19_BN from `Very Deep Convolutional Networks for Large-Scale Image Recognition `__. + + Args: + weights (:class:`~torchvision.models.VGG19_BN_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.VGG19_BN_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.vgg.VGG`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.VGG19_BN_Weights + :members: + """ + weights = VGG19_BN_Weights.verify(weights) + + return _vgg("E", True, weights, progress, **kwargs) diff --git a/lib/python3.12/site-packages/torchvision/models/video/__init__.py b/lib/python3.12/site-packages/torchvision/models/video/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f1eedd3116001af22ec202d2ccec6eefad8090ae --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/video/__init__.py @@ -0,0 +1,4 @@ +from .mvit import * +from .resnet import * +from .s3d import * +from .swin_transformer import * diff --git a/lib/python3.12/site-packages/torchvision/models/video/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/video/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..080320be034fe8b8578dba24667a06026d8c0551 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/video/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/video/__pycache__/mvit.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/video/__pycache__/mvit.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..020e3eb66fd220b83c10c0fb8b82e2d53d20bba3 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/video/__pycache__/mvit.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/video/__pycache__/resnet.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/video/__pycache__/resnet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7ba7b627f6ef77b1d9945bf7edeefbc571f0bfc7 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/video/__pycache__/resnet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/video/__pycache__/s3d.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/video/__pycache__/s3d.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3bf14de889d7f4a9fabb4f0e9cf7f696375e973c Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/video/__pycache__/s3d.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/video/__pycache__/swin_transformer.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/models/video/__pycache__/swin_transformer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..83ec2a8ee747d429df81666f5f233db6e2aa6006 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/models/video/__pycache__/swin_transformer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/models/video/mvit.py b/lib/python3.12/site-packages/torchvision/models/video/mvit.py new file mode 100644 index 0000000000000000000000000000000000000000..159c12a4f3eac579f4e122741f57cc60f5cd0a23 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/video/mvit.py @@ -0,0 +1,897 @@ +import math +from dataclasses import dataclass +from functools import partial +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + +import torch +import torch.fx +import torch.nn as nn + +from ...ops import MLP, StochasticDepth +from ...transforms._presets import VideoClassification +from ...utils import _log_api_usage_once +from .._api import register_model, Weights, WeightsEnum +from .._meta import _KINETICS400_CATEGORIES +from .._utils import _ovewrite_named_param, handle_legacy_interface + + +__all__ = [ + "MViT", + "MViT_V1_B_Weights", + "mvit_v1_b", + "MViT_V2_S_Weights", + "mvit_v2_s", +] + + +@dataclass +class MSBlockConfig: + num_heads: int + input_channels: int + output_channels: int + kernel_q: List[int] + kernel_kv: List[int] + stride_q: List[int] + stride_kv: List[int] + + +def _prod(s: Sequence[int]) -> int: + product = 1 + for v in s: + product *= v + return product + + +def _unsqueeze(x: torch.Tensor, target_dim: int, expand_dim: int) -> Tuple[torch.Tensor, int]: + tensor_dim = x.dim() + if tensor_dim == target_dim - 1: + x = x.unsqueeze(expand_dim) + elif tensor_dim != target_dim: + raise ValueError(f"Unsupported input dimension {x.shape}") + return x, tensor_dim + + +def _squeeze(x: torch.Tensor, target_dim: int, expand_dim: int, tensor_dim: int) -> torch.Tensor: + if tensor_dim == target_dim - 1: + x = x.squeeze(expand_dim) + return x + + +torch.fx.wrap("_unsqueeze") +torch.fx.wrap("_squeeze") + + +class Pool(nn.Module): + def __init__( + self, + pool: nn.Module, + norm: Optional[nn.Module], + activation: Optional[nn.Module] = None, + norm_before_pool: bool = False, + ) -> None: + super().__init__() + self.pool = pool + layers = [] + if norm is not None: + layers.append(norm) + if activation is not None: + layers.append(activation) + self.norm_act = nn.Sequential(*layers) if layers else None + self.norm_before_pool = norm_before_pool + + def forward(self, x: torch.Tensor, thw: Tuple[int, int, int]) -> Tuple[torch.Tensor, Tuple[int, int, int]]: + x, tensor_dim = _unsqueeze(x, 4, 1) + + # Separate the class token and reshape the input + class_token, x = torch.tensor_split(x, indices=(1,), dim=2) + x = x.transpose(2, 3) + B, N, C = x.shape[:3] + x = x.reshape((B * N, C) + thw).contiguous() + + # normalizing prior pooling is useful when we use BN which can be absorbed to speed up inference + if self.norm_before_pool and self.norm_act is not None: + x = self.norm_act(x) + + # apply the pool on the input and add back the token + x = self.pool(x) + T, H, W = x.shape[2:] + x = x.reshape(B, N, C, -1).transpose(2, 3) + x = torch.cat((class_token, x), dim=2) + + if not self.norm_before_pool and self.norm_act is not None: + x = self.norm_act(x) + + x = _squeeze(x, 4, 1, tensor_dim) + return x, (T, H, W) + + +def _interpolate(embedding: torch.Tensor, d: int) -> torch.Tensor: + if embedding.shape[0] == d: + return embedding + + return ( + nn.functional.interpolate( + embedding.permute(1, 0).unsqueeze(0), + size=d, + mode="linear", + ) + .squeeze(0) + .permute(1, 0) + ) + + +def _add_rel_pos( + attn: torch.Tensor, + q: torch.Tensor, + q_thw: Tuple[int, int, int], + k_thw: Tuple[int, int, int], + rel_pos_h: torch.Tensor, + rel_pos_w: torch.Tensor, + rel_pos_t: torch.Tensor, +) -> torch.Tensor: + # Modified code from: https://github.com/facebookresearch/SlowFast/commit/1aebd71a2efad823d52b827a3deaf15a56cf4932 + q_t, q_h, q_w = q_thw + k_t, k_h, k_w = k_thw + dh = int(2 * max(q_h, k_h) - 1) + dw = int(2 * max(q_w, k_w) - 1) + dt = int(2 * max(q_t, k_t) - 1) + + # Scale up rel pos if shapes for q and k are different. + q_h_ratio = max(k_h / q_h, 1.0) + k_h_ratio = max(q_h / k_h, 1.0) + dist_h = torch.arange(q_h)[:, None] * q_h_ratio - (torch.arange(k_h)[None, :] + (1.0 - k_h)) * k_h_ratio + q_w_ratio = max(k_w / q_w, 1.0) + k_w_ratio = max(q_w / k_w, 1.0) + dist_w = torch.arange(q_w)[:, None] * q_w_ratio - (torch.arange(k_w)[None, :] + (1.0 - k_w)) * k_w_ratio + q_t_ratio = max(k_t / q_t, 1.0) + k_t_ratio = max(q_t / k_t, 1.0) + dist_t = torch.arange(q_t)[:, None] * q_t_ratio - (torch.arange(k_t)[None, :] + (1.0 - k_t)) * k_t_ratio + + # Interpolate rel pos if needed. + rel_pos_h = _interpolate(rel_pos_h, dh) + rel_pos_w = _interpolate(rel_pos_w, dw) + rel_pos_t = _interpolate(rel_pos_t, dt) + Rh = rel_pos_h[dist_h.long()] + Rw = rel_pos_w[dist_w.long()] + Rt = rel_pos_t[dist_t.long()] + + B, n_head, _, dim = q.shape + + r_q = q[:, :, 1:].reshape(B, n_head, q_t, q_h, q_w, dim) + rel_h_q = torch.einsum("bythwc,hkc->bythwk", r_q, Rh) # [B, H, q_t, qh, qw, k_h] + rel_w_q = torch.einsum("bythwc,wkc->bythwk", r_q, Rw) # [B, H, q_t, qh, qw, k_w] + # [B, H, q_t, q_h, q_w, dim] -> [q_t, B, H, q_h, q_w, dim] -> [q_t, B*H*q_h*q_w, dim] + r_q = r_q.permute(2, 0, 1, 3, 4, 5).reshape(q_t, B * n_head * q_h * q_w, dim) + # [q_t, B*H*q_h*q_w, dim] * [q_t, dim, k_t] = [q_t, B*H*q_h*q_w, k_t] -> [B*H*q_h*q_w, q_t, k_t] + rel_q_t = torch.matmul(r_q, Rt.transpose(1, 2)).transpose(0, 1) + # [B*H*q_h*q_w, q_t, k_t] -> [B, H, q_t, q_h, q_w, k_t] + rel_q_t = rel_q_t.view(B, n_head, q_h, q_w, q_t, k_t).permute(0, 1, 4, 2, 3, 5) + + # Combine rel pos. + rel_pos = ( + rel_h_q[:, :, :, :, :, None, :, None] + + rel_w_q[:, :, :, :, :, None, None, :] + + rel_q_t[:, :, :, :, :, :, None, None] + ).reshape(B, n_head, q_t * q_h * q_w, k_t * k_h * k_w) + + # Add it to attention + attn[:, :, 1:, 1:] += rel_pos + + return attn + + +def _add_shortcut(x: torch.Tensor, shortcut: torch.Tensor, residual_with_cls_embed: bool): + if residual_with_cls_embed: + x.add_(shortcut) + else: + x[:, :, 1:, :] += shortcut[:, :, 1:, :] + return x + + +torch.fx.wrap("_add_rel_pos") +torch.fx.wrap("_add_shortcut") + + +class MultiscaleAttention(nn.Module): + def __init__( + self, + input_size: List[int], + embed_dim: int, + output_dim: int, + num_heads: int, + kernel_q: List[int], + kernel_kv: List[int], + stride_q: List[int], + stride_kv: List[int], + residual_pool: bool, + residual_with_cls_embed: bool, + rel_pos_embed: bool, + dropout: float = 0.0, + norm_layer: Callable[..., nn.Module] = nn.LayerNorm, + ) -> None: + super().__init__() + self.embed_dim = embed_dim + self.output_dim = output_dim + self.num_heads = num_heads + self.head_dim = output_dim // num_heads + self.scaler = 1.0 / math.sqrt(self.head_dim) + self.residual_pool = residual_pool + self.residual_with_cls_embed = residual_with_cls_embed + + self.qkv = nn.Linear(embed_dim, 3 * output_dim) + layers: List[nn.Module] = [nn.Linear(output_dim, output_dim)] + if dropout > 0.0: + layers.append(nn.Dropout(dropout, inplace=True)) + self.project = nn.Sequential(*layers) + + self.pool_q: Optional[nn.Module] = None + if _prod(kernel_q) > 1 or _prod(stride_q) > 1: + padding_q = [int(q // 2) for q in kernel_q] + self.pool_q = Pool( + nn.Conv3d( + self.head_dim, + self.head_dim, + kernel_q, # type: ignore[arg-type] + stride=stride_q, # type: ignore[arg-type] + padding=padding_q, # type: ignore[arg-type] + groups=self.head_dim, + bias=False, + ), + norm_layer(self.head_dim), + ) + + self.pool_k: Optional[nn.Module] = None + self.pool_v: Optional[nn.Module] = None + if _prod(kernel_kv) > 1 or _prod(stride_kv) > 1: + padding_kv = [int(kv // 2) for kv in kernel_kv] + self.pool_k = Pool( + nn.Conv3d( + self.head_dim, + self.head_dim, + kernel_kv, # type: ignore[arg-type] + stride=stride_kv, # type: ignore[arg-type] + padding=padding_kv, # type: ignore[arg-type] + groups=self.head_dim, + bias=False, + ), + norm_layer(self.head_dim), + ) + self.pool_v = Pool( + nn.Conv3d( + self.head_dim, + self.head_dim, + kernel_kv, # type: ignore[arg-type] + stride=stride_kv, # type: ignore[arg-type] + padding=padding_kv, # type: ignore[arg-type] + groups=self.head_dim, + bias=False, + ), + norm_layer(self.head_dim), + ) + + self.rel_pos_h: Optional[nn.Parameter] = None + self.rel_pos_w: Optional[nn.Parameter] = None + self.rel_pos_t: Optional[nn.Parameter] = None + if rel_pos_embed: + size = max(input_size[1:]) + q_size = size // stride_q[1] if len(stride_q) > 0 else size + kv_size = size // stride_kv[1] if len(stride_kv) > 0 else size + spatial_dim = 2 * max(q_size, kv_size) - 1 + temporal_dim = 2 * input_size[0] - 1 + self.rel_pos_h = nn.Parameter(torch.zeros(spatial_dim, self.head_dim)) + self.rel_pos_w = nn.Parameter(torch.zeros(spatial_dim, self.head_dim)) + self.rel_pos_t = nn.Parameter(torch.zeros(temporal_dim, self.head_dim)) + nn.init.trunc_normal_(self.rel_pos_h, std=0.02) + nn.init.trunc_normal_(self.rel_pos_w, std=0.02) + nn.init.trunc_normal_(self.rel_pos_t, std=0.02) + + def forward(self, x: torch.Tensor, thw: Tuple[int, int, int]) -> Tuple[torch.Tensor, Tuple[int, int, int]]: + B, N, C = x.shape + q, k, v = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).transpose(1, 3).unbind(dim=2) + + if self.pool_k is not None: + k, k_thw = self.pool_k(k, thw) + else: + k_thw = thw + if self.pool_v is not None: + v = self.pool_v(v, thw)[0] + if self.pool_q is not None: + q, thw = self.pool_q(q, thw) + + attn = torch.matmul(self.scaler * q, k.transpose(2, 3)) + if self.rel_pos_h is not None and self.rel_pos_w is not None and self.rel_pos_t is not None: + attn = _add_rel_pos( + attn, + q, + thw, + k_thw, + self.rel_pos_h, + self.rel_pos_w, + self.rel_pos_t, + ) + attn = attn.softmax(dim=-1) + + x = torch.matmul(attn, v) + if self.residual_pool: + _add_shortcut(x, q, self.residual_with_cls_embed) + x = x.transpose(1, 2).reshape(B, -1, self.output_dim) + x = self.project(x) + + return x, thw + + +class MultiscaleBlock(nn.Module): + def __init__( + self, + input_size: List[int], + cnf: MSBlockConfig, + residual_pool: bool, + residual_with_cls_embed: bool, + rel_pos_embed: bool, + proj_after_attn: bool, + dropout: float = 0.0, + stochastic_depth_prob: float = 0.0, + norm_layer: Callable[..., nn.Module] = nn.LayerNorm, + ) -> None: + super().__init__() + self.proj_after_attn = proj_after_attn + + self.pool_skip: Optional[nn.Module] = None + if _prod(cnf.stride_q) > 1: + kernel_skip = [s + 1 if s > 1 else s for s in cnf.stride_q] + padding_skip = [int(k // 2) for k in kernel_skip] + self.pool_skip = Pool( + nn.MaxPool3d(kernel_skip, stride=cnf.stride_q, padding=padding_skip), None # type: ignore[arg-type] + ) + + attn_dim = cnf.output_channels if proj_after_attn else cnf.input_channels + + self.norm1 = norm_layer(cnf.input_channels) + self.norm2 = norm_layer(attn_dim) + self.needs_transposal = isinstance(self.norm1, nn.BatchNorm1d) + + self.attn = MultiscaleAttention( + input_size, + cnf.input_channels, + attn_dim, + cnf.num_heads, + kernel_q=cnf.kernel_q, + kernel_kv=cnf.kernel_kv, + stride_q=cnf.stride_q, + stride_kv=cnf.stride_kv, + rel_pos_embed=rel_pos_embed, + residual_pool=residual_pool, + residual_with_cls_embed=residual_with_cls_embed, + dropout=dropout, + norm_layer=norm_layer, + ) + self.mlp = MLP( + attn_dim, + [4 * attn_dim, cnf.output_channels], + activation_layer=nn.GELU, + dropout=dropout, + inplace=None, + ) + + self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row") + + self.project: Optional[nn.Module] = None + if cnf.input_channels != cnf.output_channels: + self.project = nn.Linear(cnf.input_channels, cnf.output_channels) + + def forward(self, x: torch.Tensor, thw: Tuple[int, int, int]) -> Tuple[torch.Tensor, Tuple[int, int, int]]: + x_norm1 = self.norm1(x.transpose(1, 2)).transpose(1, 2) if self.needs_transposal else self.norm1(x) + x_attn, thw_new = self.attn(x_norm1, thw) + x = x if self.project is None or not self.proj_after_attn else self.project(x_norm1) + x_skip = x if self.pool_skip is None else self.pool_skip(x, thw)[0] + x = x_skip + self.stochastic_depth(x_attn) + + x_norm2 = self.norm2(x.transpose(1, 2)).transpose(1, 2) if self.needs_transposal else self.norm2(x) + x_proj = x if self.project is None or self.proj_after_attn else self.project(x_norm2) + + return x_proj + self.stochastic_depth(self.mlp(x_norm2)), thw_new + + +class PositionalEncoding(nn.Module): + def __init__(self, embed_size: int, spatial_size: Tuple[int, int], temporal_size: int, rel_pos_embed: bool) -> None: + super().__init__() + self.spatial_size = spatial_size + self.temporal_size = temporal_size + + self.class_token = nn.Parameter(torch.zeros(embed_size)) + self.spatial_pos: Optional[nn.Parameter] = None + self.temporal_pos: Optional[nn.Parameter] = None + self.class_pos: Optional[nn.Parameter] = None + if not rel_pos_embed: + self.spatial_pos = nn.Parameter(torch.zeros(self.spatial_size[0] * self.spatial_size[1], embed_size)) + self.temporal_pos = nn.Parameter(torch.zeros(self.temporal_size, embed_size)) + self.class_pos = nn.Parameter(torch.zeros(embed_size)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + class_token = self.class_token.expand(x.size(0), -1).unsqueeze(1) + x = torch.cat((class_token, x), dim=1) + + if self.spatial_pos is not None and self.temporal_pos is not None and self.class_pos is not None: + hw_size, embed_size = self.spatial_pos.shape + pos_embedding = torch.repeat_interleave(self.temporal_pos, hw_size, dim=0) + pos_embedding.add_(self.spatial_pos.unsqueeze(0).expand(self.temporal_size, -1, -1).reshape(-1, embed_size)) + pos_embedding = torch.cat((self.class_pos.unsqueeze(0), pos_embedding), dim=0).unsqueeze(0) + x.add_(pos_embedding) + + return x + + +class MViT(nn.Module): + def __init__( + self, + spatial_size: Tuple[int, int], + temporal_size: int, + block_setting: Sequence[MSBlockConfig], + residual_pool: bool, + residual_with_cls_embed: bool, + rel_pos_embed: bool, + proj_after_attn: bool, + dropout: float = 0.5, + attention_dropout: float = 0.0, + stochastic_depth_prob: float = 0.0, + num_classes: int = 400, + block: Optional[Callable[..., nn.Module]] = None, + norm_layer: Optional[Callable[..., nn.Module]] = None, + patch_embed_kernel: Tuple[int, int, int] = (3, 7, 7), + patch_embed_stride: Tuple[int, int, int] = (2, 4, 4), + patch_embed_padding: Tuple[int, int, int] = (1, 3, 3), + ) -> None: + """ + MViT main class. + + Args: + spatial_size (tuple of ints): The spacial size of the input as ``(H, W)``. + temporal_size (int): The temporal size ``T`` of the input. + block_setting (sequence of MSBlockConfig): The Network structure. + residual_pool (bool): If True, use MViTv2 pooling residual connection. + residual_with_cls_embed (bool): If True, the addition on the residual connection will include + the class embedding. + rel_pos_embed (bool): If True, use MViTv2's relative positional embeddings. + proj_after_attn (bool): If True, apply the projection after the attention. + dropout (float): Dropout rate. Default: 0.0. + attention_dropout (float): Attention dropout rate. Default: 0.0. + stochastic_depth_prob: (float): Stochastic depth rate. Default: 0.0. + num_classes (int): The number of classes. + block (callable, optional): Module specifying the layer which consists of the attention and mlp. + norm_layer (callable, optional): Module specifying the normalization layer to use. + patch_embed_kernel (tuple of ints): The kernel of the convolution that patchifies the input. + patch_embed_stride (tuple of ints): The stride of the convolution that patchifies the input. + patch_embed_padding (tuple of ints): The padding of the convolution that patchifies the input. + """ + super().__init__() + # This implementation employs a different parameterization scheme than the one used at PyTorch Video: + # https://github.com/facebookresearch/pytorchvideo/blob/718d0a4/pytorchvideo/models/vision_transformers.py + # We remove any experimental configuration that didn't make it to the final variants of the models. To represent + # the configuration of the architecture we use the simplified form suggested at Table 1 of the paper. + _log_api_usage_once(self) + total_stage_blocks = len(block_setting) + if total_stage_blocks == 0: + raise ValueError("The configuration parameter can't be empty.") + + if block is None: + block = MultiscaleBlock + + if norm_layer is None: + norm_layer = partial(nn.LayerNorm, eps=1e-6) + + # Patch Embedding module + self.conv_proj = nn.Conv3d( + in_channels=3, + out_channels=block_setting[0].input_channels, + kernel_size=patch_embed_kernel, + stride=patch_embed_stride, + padding=patch_embed_padding, + ) + + input_size = [size // stride for size, stride in zip((temporal_size,) + spatial_size, self.conv_proj.stride)] + + # Spatio-Temporal Class Positional Encoding + self.pos_encoding = PositionalEncoding( + embed_size=block_setting[0].input_channels, + spatial_size=(input_size[1], input_size[2]), + temporal_size=input_size[0], + rel_pos_embed=rel_pos_embed, + ) + + # Encoder module + self.blocks = nn.ModuleList() + for stage_block_id, cnf in enumerate(block_setting): + # adjust stochastic depth probability based on the depth of the stage block + sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0) + + self.blocks.append( + block( + input_size=input_size, + cnf=cnf, + residual_pool=residual_pool, + residual_with_cls_embed=residual_with_cls_embed, + rel_pos_embed=rel_pos_embed, + proj_after_attn=proj_after_attn, + dropout=attention_dropout, + stochastic_depth_prob=sd_prob, + norm_layer=norm_layer, + ) + ) + + if len(cnf.stride_q) > 0: + input_size = [size // stride for size, stride in zip(input_size, cnf.stride_q)] + self.norm = norm_layer(block_setting[-1].output_channels) + + # Classifier module + self.head = nn.Sequential( + nn.Dropout(dropout, inplace=True), + nn.Linear(block_setting[-1].output_channels, num_classes), + ) + + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.trunc_normal_(m.weight, std=0.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0.0) + elif isinstance(m, nn.LayerNorm): + if m.weight is not None: + nn.init.constant_(m.weight, 1.0) + if m.bias is not None: + nn.init.constant_(m.bias, 0.0) + elif isinstance(m, PositionalEncoding): + for weights in m.parameters(): + nn.init.trunc_normal_(weights, std=0.02) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Convert if necessary (B, C, H, W) -> (B, C, 1, H, W) + x = _unsqueeze(x, 5, 2)[0] + # patchify and reshape: (B, C, T, H, W) -> (B, embed_channels[0], T', H', W') -> (B, THW', embed_channels[0]) + x = self.conv_proj(x) + x = x.flatten(2).transpose(1, 2) + + # add positional encoding + x = self.pos_encoding(x) + + # pass patches through the encoder + thw = (self.pos_encoding.temporal_size,) + self.pos_encoding.spatial_size + for block in self.blocks: + x, thw = block(x, thw) + x = self.norm(x) + + # classifier "token" as used by standard language architectures + x = x[:, 0] + x = self.head(x) + + return x + + +def _mvit( + block_setting: List[MSBlockConfig], + stochastic_depth_prob: float, + weights: Optional[WeightsEnum], + progress: bool, + **kwargs: Any, +) -> MViT: + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + assert weights.meta["min_size"][0] == weights.meta["min_size"][1] + _ovewrite_named_param(kwargs, "spatial_size", weights.meta["min_size"]) + _ovewrite_named_param(kwargs, "temporal_size", weights.meta["min_temporal_size"]) + spatial_size = kwargs.pop("spatial_size", (224, 224)) + temporal_size = kwargs.pop("temporal_size", 16) + + model = MViT( + spatial_size=spatial_size, + temporal_size=temporal_size, + block_setting=block_setting, + residual_pool=kwargs.pop("residual_pool", False), + residual_with_cls_embed=kwargs.pop("residual_with_cls_embed", True), + rel_pos_embed=kwargs.pop("rel_pos_embed", False), + proj_after_attn=kwargs.pop("proj_after_attn", False), + stochastic_depth_prob=stochastic_depth_prob, + **kwargs, + ) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +class MViT_V1_B_Weights(WeightsEnum): + KINETICS400_V1 = Weights( + url="https://download.pytorch.org/models/mvit_v1_b-dbeb1030.pth", + transforms=partial( + VideoClassification, + crop_size=(224, 224), + resize_size=(256,), + mean=(0.45, 0.45, 0.45), + std=(0.225, 0.225, 0.225), + ), + meta={ + "min_size": (224, 224), + "min_temporal_size": 16, + "categories": _KINETICS400_CATEGORIES, + "recipe": "https://github.com/facebookresearch/pytorchvideo/blob/main/docs/source/model_zoo.md", + "_docs": ( + "The weights were ported from the paper. The accuracies are estimated on video-level " + "with parameters `frame_rate=7.5`, `clips_per_video=5`, and `clip_len=16`" + ), + "num_params": 36610672, + "_metrics": { + "Kinetics-400": { + "acc@1": 78.477, + "acc@5": 93.582, + } + }, + "_ops": 70.599, + "_file_size": 139.764, + }, + ) + DEFAULT = KINETICS400_V1 + + +class MViT_V2_S_Weights(WeightsEnum): + KINETICS400_V1 = Weights( + url="https://download.pytorch.org/models/mvit_v2_s-ae3be167.pth", + transforms=partial( + VideoClassification, + crop_size=(224, 224), + resize_size=(256,), + mean=(0.45, 0.45, 0.45), + std=(0.225, 0.225, 0.225), + ), + meta={ + "min_size": (224, 224), + "min_temporal_size": 16, + "categories": _KINETICS400_CATEGORIES, + "recipe": "https://github.com/facebookresearch/SlowFast/blob/main/MODEL_ZOO.md", + "_docs": ( + "The weights were ported from the paper. The accuracies are estimated on video-level " + "with parameters `frame_rate=7.5`, `clips_per_video=5`, and `clip_len=16`" + ), + "num_params": 34537744, + "_metrics": { + "Kinetics-400": { + "acc@1": 80.757, + "acc@5": 94.665, + } + }, + "_ops": 64.224, + "_file_size": 131.884, + }, + ) + DEFAULT = KINETICS400_V1 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", MViT_V1_B_Weights.KINETICS400_V1)) +def mvit_v1_b(*, weights: Optional[MViT_V1_B_Weights] = None, progress: bool = True, **kwargs: Any) -> MViT: + """ + Constructs a base MViTV1 architecture from + `Multiscale Vision Transformers `__. + + .. betastatus:: video module + + Args: + weights (:class:`~torchvision.models.video.MViT_V1_B_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.video.MViT_V1_B_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.video.MViT`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.video.MViT_V1_B_Weights + :members: + """ + weights = MViT_V1_B_Weights.verify(weights) + + config: Dict[str, List] = { + "num_heads": [1, 2, 2, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 8, 8], + "input_channels": [96, 192, 192, 384, 384, 384, 384, 384, 384, 384, 384, 384, 384, 384, 768, 768], + "output_channels": [192, 192, 384, 384, 384, 384, 384, 384, 384, 384, 384, 384, 384, 768, 768, 768], + "kernel_q": [[], [3, 3, 3], [], [3, 3, 3], [], [], [], [], [], [], [], [], [], [], [3, 3, 3], []], + "kernel_kv": [ + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + ], + "stride_q": [[], [1, 2, 2], [], [1, 2, 2], [], [], [], [], [], [], [], [], [], [], [1, 2, 2], []], + "stride_kv": [ + [1, 8, 8], + [1, 4, 4], + [1, 4, 4], + [1, 2, 2], + [1, 2, 2], + [1, 2, 2], + [1, 2, 2], + [1, 2, 2], + [1, 2, 2], + [1, 2, 2], + [1, 2, 2], + [1, 2, 2], + [1, 2, 2], + [1, 2, 2], + [1, 1, 1], + [1, 1, 1], + ], + } + + block_setting = [] + for i in range(len(config["num_heads"])): + block_setting.append( + MSBlockConfig( + num_heads=config["num_heads"][i], + input_channels=config["input_channels"][i], + output_channels=config["output_channels"][i], + kernel_q=config["kernel_q"][i], + kernel_kv=config["kernel_kv"][i], + stride_q=config["stride_q"][i], + stride_kv=config["stride_kv"][i], + ) + ) + + return _mvit( + spatial_size=(224, 224), + temporal_size=16, + block_setting=block_setting, + residual_pool=False, + residual_with_cls_embed=False, + stochastic_depth_prob=kwargs.pop("stochastic_depth_prob", 0.2), + weights=weights, + progress=progress, + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", MViT_V2_S_Weights.KINETICS400_V1)) +def mvit_v2_s(*, weights: Optional[MViT_V2_S_Weights] = None, progress: bool = True, **kwargs: Any) -> MViT: + """Constructs a small MViTV2 architecture from + `Multiscale Vision Transformers `__ and + `MViTv2: Improved Multiscale Vision Transformers for Classification + and Detection `__. + + .. betastatus:: video module + + Args: + weights (:class:`~torchvision.models.video.MViT_V2_S_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.video.MViT_V2_S_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.video.MViT`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.video.MViT_V2_S_Weights + :members: + """ + weights = MViT_V2_S_Weights.verify(weights) + + config: Dict[str, List] = { + "num_heads": [1, 2, 2, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 8, 8], + "input_channels": [96, 96, 192, 192, 384, 384, 384, 384, 384, 384, 384, 384, 384, 384, 384, 768], + "output_channels": [96, 192, 192, 384, 384, 384, 384, 384, 384, 384, 384, 384, 384, 384, 768, 768], + "kernel_q": [ + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + ], + "kernel_kv": [ + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + ], + "stride_q": [ + [1, 1, 1], + [1, 2, 2], + [1, 1, 1], + [1, 2, 2], + [1, 1, 1], + [1, 1, 1], + [1, 1, 1], + [1, 1, 1], + [1, 1, 1], + [1, 1, 1], + [1, 1, 1], + [1, 1, 1], + [1, 1, 1], + [1, 1, 1], + [1, 2, 2], + [1, 1, 1], + ], + "stride_kv": [ + [1, 8, 8], + [1, 4, 4], + [1, 4, 4], + [1, 2, 2], + [1, 2, 2], + [1, 2, 2], + [1, 2, 2], + [1, 2, 2], + [1, 2, 2], + [1, 2, 2], + [1, 2, 2], + [1, 2, 2], + [1, 2, 2], + [1, 2, 2], + [1, 1, 1], + [1, 1, 1], + ], + } + + block_setting = [] + for i in range(len(config["num_heads"])): + block_setting.append( + MSBlockConfig( + num_heads=config["num_heads"][i], + input_channels=config["input_channels"][i], + output_channels=config["output_channels"][i], + kernel_q=config["kernel_q"][i], + kernel_kv=config["kernel_kv"][i], + stride_q=config["stride_q"][i], + stride_kv=config["stride_kv"][i], + ) + ) + + return _mvit( + spatial_size=(224, 224), + temporal_size=16, + block_setting=block_setting, + residual_pool=True, + residual_with_cls_embed=False, + rel_pos_embed=True, + proj_after_attn=True, + stochastic_depth_prob=kwargs.pop("stochastic_depth_prob", 0.2), + weights=weights, + progress=progress, + **kwargs, + ) diff --git a/lib/python3.12/site-packages/torchvision/models/video/resnet.py b/lib/python3.12/site-packages/torchvision/models/video/resnet.py new file mode 100644 index 0000000000000000000000000000000000000000..a1cb2884013c053118555344617e4b1efb8ddaab --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/video/resnet.py @@ -0,0 +1,503 @@ +from functools import partial +from typing import Any, Callable, List, Optional, Sequence, Tuple, Type, Union + +import torch.nn as nn +from torch import Tensor + +from ...transforms._presets import VideoClassification +from ...utils import _log_api_usage_once +from .._api import register_model, Weights, WeightsEnum +from .._meta import _KINETICS400_CATEGORIES +from .._utils import _ovewrite_named_param, handle_legacy_interface + + +__all__ = [ + "VideoResNet", + "R3D_18_Weights", + "MC3_18_Weights", + "R2Plus1D_18_Weights", + "r3d_18", + "mc3_18", + "r2plus1d_18", +] + + +class Conv3DSimple(nn.Conv3d): + def __init__( + self, in_planes: int, out_planes: int, midplanes: Optional[int] = None, stride: int = 1, padding: int = 1 + ) -> None: + + super().__init__( + in_channels=in_planes, + out_channels=out_planes, + kernel_size=(3, 3, 3), + stride=stride, + padding=padding, + bias=False, + ) + + @staticmethod + def get_downsample_stride(stride: int) -> Tuple[int, int, int]: + return stride, stride, stride + + +class Conv2Plus1D(nn.Sequential): + def __init__(self, in_planes: int, out_planes: int, midplanes: int, stride: int = 1, padding: int = 1) -> None: + super().__init__( + nn.Conv3d( + in_planes, + midplanes, + kernel_size=(1, 3, 3), + stride=(1, stride, stride), + padding=(0, padding, padding), + bias=False, + ), + nn.BatchNorm3d(midplanes), + nn.ReLU(inplace=True), + nn.Conv3d( + midplanes, out_planes, kernel_size=(3, 1, 1), stride=(stride, 1, 1), padding=(padding, 0, 0), bias=False + ), + ) + + @staticmethod + def get_downsample_stride(stride: int) -> Tuple[int, int, int]: + return stride, stride, stride + + +class Conv3DNoTemporal(nn.Conv3d): + def __init__( + self, in_planes: int, out_planes: int, midplanes: Optional[int] = None, stride: int = 1, padding: int = 1 + ) -> None: + + super().__init__( + in_channels=in_planes, + out_channels=out_planes, + kernel_size=(1, 3, 3), + stride=(1, stride, stride), + padding=(0, padding, padding), + bias=False, + ) + + @staticmethod + def get_downsample_stride(stride: int) -> Tuple[int, int, int]: + return 1, stride, stride + + +class BasicBlock(nn.Module): + + expansion = 1 + + def __init__( + self, + inplanes: int, + planes: int, + conv_builder: Callable[..., nn.Module], + stride: int = 1, + downsample: Optional[nn.Module] = None, + ) -> None: + midplanes = (inplanes * planes * 3 * 3 * 3) // (inplanes * 3 * 3 + 3 * planes) + + super().__init__() + self.conv1 = nn.Sequential( + conv_builder(inplanes, planes, midplanes, stride), nn.BatchNorm3d(planes), nn.ReLU(inplace=True) + ) + self.conv2 = nn.Sequential(conv_builder(planes, planes, midplanes), nn.BatchNorm3d(planes)) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x: Tensor) -> Tensor: + residual = x + + out = self.conv1(x) + out = self.conv2(out) + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class Bottleneck(nn.Module): + expansion = 4 + + def __init__( + self, + inplanes: int, + planes: int, + conv_builder: Callable[..., nn.Module], + stride: int = 1, + downsample: Optional[nn.Module] = None, + ) -> None: + + super().__init__() + midplanes = (inplanes * planes * 3 * 3 * 3) // (inplanes * 3 * 3 + 3 * planes) + + # 1x1x1 + self.conv1 = nn.Sequential( + nn.Conv3d(inplanes, planes, kernel_size=1, bias=False), nn.BatchNorm3d(planes), nn.ReLU(inplace=True) + ) + # Second kernel + self.conv2 = nn.Sequential( + conv_builder(planes, planes, midplanes, stride), nn.BatchNorm3d(planes), nn.ReLU(inplace=True) + ) + + # 1x1x1 + self.conv3 = nn.Sequential( + nn.Conv3d(planes, planes * self.expansion, kernel_size=1, bias=False), + nn.BatchNorm3d(planes * self.expansion), + ) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x: Tensor) -> Tensor: + residual = x + + out = self.conv1(x) + out = self.conv2(out) + out = self.conv3(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class BasicStem(nn.Sequential): + """The default conv-batchnorm-relu stem""" + + def __init__(self) -> None: + super().__init__( + nn.Conv3d(3, 64, kernel_size=(3, 7, 7), stride=(1, 2, 2), padding=(1, 3, 3), bias=False), + nn.BatchNorm3d(64), + nn.ReLU(inplace=True), + ) + + +class R2Plus1dStem(nn.Sequential): + """R(2+1)D stem is different than the default one as it uses separated 3D convolution""" + + def __init__(self) -> None: + super().__init__( + nn.Conv3d(3, 45, kernel_size=(1, 7, 7), stride=(1, 2, 2), padding=(0, 3, 3), bias=False), + nn.BatchNorm3d(45), + nn.ReLU(inplace=True), + nn.Conv3d(45, 64, kernel_size=(3, 1, 1), stride=(1, 1, 1), padding=(1, 0, 0), bias=False), + nn.BatchNorm3d(64), + nn.ReLU(inplace=True), + ) + + +class VideoResNet(nn.Module): + def __init__( + self, + block: Type[Union[BasicBlock, Bottleneck]], + conv_makers: Sequence[Type[Union[Conv3DSimple, Conv3DNoTemporal, Conv2Plus1D]]], + layers: List[int], + stem: Callable[..., nn.Module], + num_classes: int = 400, + zero_init_residual: bool = False, + ) -> None: + """Generic resnet video generator. + + Args: + block (Type[Union[BasicBlock, Bottleneck]]): resnet building block + conv_makers (List[Type[Union[Conv3DSimple, Conv3DNoTemporal, Conv2Plus1D]]]): generator + function for each layer + layers (List[int]): number of blocks per layer + stem (Callable[..., nn.Module]): module specifying the ResNet stem. + num_classes (int, optional): Dimension of the final FC layer. Defaults to 400. + zero_init_residual (bool, optional): Zero init bottleneck residual BN. Defaults to False. + """ + super().__init__() + _log_api_usage_once(self) + self.inplanes = 64 + + self.stem = stem() + + self.layer1 = self._make_layer(block, conv_makers[0], 64, layers[0], stride=1) + self.layer2 = self._make_layer(block, conv_makers[1], 128, layers[1], stride=2) + self.layer3 = self._make_layer(block, conv_makers[2], 256, layers[2], stride=2) + self.layer4 = self._make_layer(block, conv_makers[3], 512, layers[3], stride=2) + + self.avgpool = nn.AdaptiveAvgPool3d((1, 1, 1)) + self.fc = nn.Linear(512 * block.expansion, num_classes) + + # init weights + for m in self.modules(): + if isinstance(m, nn.Conv3d): + nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.BatchNorm3d): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Linear): + nn.init.normal_(m.weight, 0, 0.01) + nn.init.constant_(m.bias, 0) + + if zero_init_residual: + for m in self.modules(): + if isinstance(m, Bottleneck): + nn.init.constant_(m.bn3.weight, 0) # type: ignore[union-attr, arg-type] + + def forward(self, x: Tensor) -> Tensor: + x = self.stem(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + + x = self.avgpool(x) + # Flatten the layer to fc + x = x.flatten(1) + x = self.fc(x) + + return x + + def _make_layer( + self, + block: Type[Union[BasicBlock, Bottleneck]], + conv_builder: Type[Union[Conv3DSimple, Conv3DNoTemporal, Conv2Plus1D]], + planes: int, + blocks: int, + stride: int = 1, + ) -> nn.Sequential: + downsample = None + + if stride != 1 or self.inplanes != planes * block.expansion: + ds_stride = conv_builder.get_downsample_stride(stride) + downsample = nn.Sequential( + nn.Conv3d(self.inplanes, planes * block.expansion, kernel_size=1, stride=ds_stride, bias=False), + nn.BatchNorm3d(planes * block.expansion), + ) + layers = [] + layers.append(block(self.inplanes, planes, conv_builder, stride, downsample)) + + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes, conv_builder)) + + return nn.Sequential(*layers) + + +def _video_resnet( + block: Type[Union[BasicBlock, Bottleneck]], + conv_makers: Sequence[Type[Union[Conv3DSimple, Conv3DNoTemporal, Conv2Plus1D]]], + layers: List[int], + stem: Callable[..., nn.Module], + weights: Optional[WeightsEnum], + progress: bool, + **kwargs: Any, +) -> VideoResNet: + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + + model = VideoResNet(block, conv_makers, layers, stem, **kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +_COMMON_META = { + "min_size": (1, 1), + "categories": _KINETICS400_CATEGORIES, + "recipe": "https://github.com/pytorch/vision/tree/main/references/video_classification", + "_docs": ( + "The weights reproduce closely the accuracy of the paper. The accuracies are estimated on video-level " + "with parameters `frame_rate=15`, `clips_per_video=5`, and `clip_len=16`." + ), +} + + +class R3D_18_Weights(WeightsEnum): + KINETICS400_V1 = Weights( + url="https://download.pytorch.org/models/r3d_18-b3b3357e.pth", + transforms=partial(VideoClassification, crop_size=(112, 112), resize_size=(128, 171)), + meta={ + **_COMMON_META, + "num_params": 33371472, + "_metrics": { + "Kinetics-400": { + "acc@1": 63.200, + "acc@5": 83.479, + } + }, + "_ops": 40.697, + "_file_size": 127.359, + }, + ) + DEFAULT = KINETICS400_V1 + + +class MC3_18_Weights(WeightsEnum): + KINETICS400_V1 = Weights( + url="https://download.pytorch.org/models/mc3_18-a90a0ba3.pth", + transforms=partial(VideoClassification, crop_size=(112, 112), resize_size=(128, 171)), + meta={ + **_COMMON_META, + "num_params": 11695440, + "_metrics": { + "Kinetics-400": { + "acc@1": 63.960, + "acc@5": 84.130, + } + }, + "_ops": 43.343, + "_file_size": 44.672, + }, + ) + DEFAULT = KINETICS400_V1 + + +class R2Plus1D_18_Weights(WeightsEnum): + KINETICS400_V1 = Weights( + url="https://download.pytorch.org/models/r2plus1d_18-91a641e6.pth", + transforms=partial(VideoClassification, crop_size=(112, 112), resize_size=(128, 171)), + meta={ + **_COMMON_META, + "num_params": 31505325, + "_metrics": { + "Kinetics-400": { + "acc@1": 67.463, + "acc@5": 86.175, + } + }, + "_ops": 40.519, + "_file_size": 120.318, + }, + ) + DEFAULT = KINETICS400_V1 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", R3D_18_Weights.KINETICS400_V1)) +def r3d_18(*, weights: Optional[R3D_18_Weights] = None, progress: bool = True, **kwargs: Any) -> VideoResNet: + """Construct 18 layer Resnet3D model. + + .. betastatus:: video module + + Reference: `A Closer Look at Spatiotemporal Convolutions for Action Recognition `__. + + Args: + weights (:class:`~torchvision.models.video.R3D_18_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.video.R3D_18_Weights` + below for more details, and possible values. By default, no + pre-trained weights are used. + progress (bool): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.video.resnet.VideoResNet`` base class. + Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.video.R3D_18_Weights + :members: + """ + weights = R3D_18_Weights.verify(weights) + + return _video_resnet( + BasicBlock, + [Conv3DSimple] * 4, + [2, 2, 2, 2], + BasicStem, + weights, + progress, + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", MC3_18_Weights.KINETICS400_V1)) +def mc3_18(*, weights: Optional[MC3_18_Weights] = None, progress: bool = True, **kwargs: Any) -> VideoResNet: + """Construct 18 layer Mixed Convolution network as in + + .. betastatus:: video module + + Reference: `A Closer Look at Spatiotemporal Convolutions for Action Recognition `__. + + Args: + weights (:class:`~torchvision.models.video.MC3_18_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.video.MC3_18_Weights` + below for more details, and possible values. By default, no + pre-trained weights are used. + progress (bool): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.video.resnet.VideoResNet`` base class. + Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.video.MC3_18_Weights + :members: + """ + weights = MC3_18_Weights.verify(weights) + + return _video_resnet( + BasicBlock, + [Conv3DSimple] + [Conv3DNoTemporal] * 3, # type: ignore[list-item] + [2, 2, 2, 2], + BasicStem, + weights, + progress, + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", R2Plus1D_18_Weights.KINETICS400_V1)) +def r2plus1d_18(*, weights: Optional[R2Plus1D_18_Weights] = None, progress: bool = True, **kwargs: Any) -> VideoResNet: + """Construct 18 layer deep R(2+1)D network as in + + .. betastatus:: video module + + Reference: `A Closer Look at Spatiotemporal Convolutions for Action Recognition `__. + + Args: + weights (:class:`~torchvision.models.video.R2Plus1D_18_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.video.R2Plus1D_18_Weights` + below for more details, and possible values. By default, no + pre-trained weights are used. + progress (bool): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.video.resnet.VideoResNet`` base class. + Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.video.R2Plus1D_18_Weights + :members: + """ + weights = R2Plus1D_18_Weights.verify(weights) + + return _video_resnet( + BasicBlock, + [Conv2Plus1D] * 4, + [2, 2, 2, 2], + R2Plus1dStem, + weights, + progress, + **kwargs, + ) + + +# The dictionary below is internal implementation detail and will be removed in v0.15 +from .._utils import _ModelURLs + + +model_urls = _ModelURLs( + { + "r3d_18": R3D_18_Weights.KINETICS400_V1.url, + "mc3_18": MC3_18_Weights.KINETICS400_V1.url, + "r2plus1d_18": R2Plus1D_18_Weights.KINETICS400_V1.url, + } +) diff --git a/lib/python3.12/site-packages/torchvision/models/video/s3d.py b/lib/python3.12/site-packages/torchvision/models/video/s3d.py new file mode 100644 index 0000000000000000000000000000000000000000..4b202829b24fb1dc314452d38a521dfe6c8e446f --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/video/s3d.py @@ -0,0 +1,219 @@ +from functools import partial +from typing import Any, Callable, Optional + +import torch +from torch import nn +from torchvision.ops.misc import Conv3dNormActivation + +from ...transforms._presets import VideoClassification +from ...utils import _log_api_usage_once +from .._api import register_model, Weights, WeightsEnum +from .._meta import _KINETICS400_CATEGORIES +from .._utils import _ovewrite_named_param, handle_legacy_interface + + +__all__ = [ + "S3D", + "S3D_Weights", + "s3d", +] + + +class TemporalSeparableConv(nn.Sequential): + def __init__( + self, + in_planes: int, + out_planes: int, + kernel_size: int, + stride: int, + padding: int, + norm_layer: Callable[..., nn.Module], + ): + super().__init__( + Conv3dNormActivation( + in_planes, + out_planes, + kernel_size=(1, kernel_size, kernel_size), + stride=(1, stride, stride), + padding=(0, padding, padding), + bias=False, + norm_layer=norm_layer, + ), + Conv3dNormActivation( + out_planes, + out_planes, + kernel_size=(kernel_size, 1, 1), + stride=(stride, 1, 1), + padding=(padding, 0, 0), + bias=False, + norm_layer=norm_layer, + ), + ) + + +class SepInceptionBlock3D(nn.Module): + def __init__( + self, + in_planes: int, + b0_out: int, + b1_mid: int, + b1_out: int, + b2_mid: int, + b2_out: int, + b3_out: int, + norm_layer: Callable[..., nn.Module], + ): + super().__init__() + + self.branch0 = Conv3dNormActivation(in_planes, b0_out, kernel_size=1, stride=1, norm_layer=norm_layer) + self.branch1 = nn.Sequential( + Conv3dNormActivation(in_planes, b1_mid, kernel_size=1, stride=1, norm_layer=norm_layer), + TemporalSeparableConv(b1_mid, b1_out, kernel_size=3, stride=1, padding=1, norm_layer=norm_layer), + ) + self.branch2 = nn.Sequential( + Conv3dNormActivation(in_planes, b2_mid, kernel_size=1, stride=1, norm_layer=norm_layer), + TemporalSeparableConv(b2_mid, b2_out, kernel_size=3, stride=1, padding=1, norm_layer=norm_layer), + ) + self.branch3 = nn.Sequential( + nn.MaxPool3d(kernel_size=(3, 3, 3), stride=1, padding=1), + Conv3dNormActivation(in_planes, b3_out, kernel_size=1, stride=1, norm_layer=norm_layer), + ) + + def forward(self, x): + x0 = self.branch0(x) + x1 = self.branch1(x) + x2 = self.branch2(x) + x3 = self.branch3(x) + out = torch.cat((x0, x1, x2, x3), 1) + + return out + + +class S3D(nn.Module): + """S3D main class. + + Args: + num_class (int): number of classes for the classification task. + dropout (float): dropout probability. + norm_layer (Optional[Callable]): Module specifying the normalization layer to use. + + Inputs: + x (Tensor): batch of videos with dimensions (batch, channel, time, height, width) + """ + + def __init__( + self, + num_classes: int = 400, + dropout: float = 0.2, + norm_layer: Optional[Callable[..., torch.nn.Module]] = None, + ) -> None: + super().__init__() + _log_api_usage_once(self) + + if norm_layer is None: + norm_layer = partial(nn.BatchNorm3d, eps=0.001, momentum=0.001) + + self.features = nn.Sequential( + TemporalSeparableConv(3, 64, 7, 2, 3, norm_layer), + nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1)), + Conv3dNormActivation( + 64, + 64, + kernel_size=1, + stride=1, + norm_layer=norm_layer, + ), + TemporalSeparableConv(64, 192, 3, 1, 1, norm_layer), + nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1)), + SepInceptionBlock3D(192, 64, 96, 128, 16, 32, 32, norm_layer), + SepInceptionBlock3D(256, 128, 128, 192, 32, 96, 64, norm_layer), + nn.MaxPool3d(kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1)), + SepInceptionBlock3D(480, 192, 96, 208, 16, 48, 64, norm_layer), + SepInceptionBlock3D(512, 160, 112, 224, 24, 64, 64, norm_layer), + SepInceptionBlock3D(512, 128, 128, 256, 24, 64, 64, norm_layer), + SepInceptionBlock3D(512, 112, 144, 288, 32, 64, 64, norm_layer), + SepInceptionBlock3D(528, 256, 160, 320, 32, 128, 128, norm_layer), + nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2), padding=(0, 0, 0)), + SepInceptionBlock3D(832, 256, 160, 320, 32, 128, 128, norm_layer), + SepInceptionBlock3D(832, 384, 192, 384, 48, 128, 128, norm_layer), + ) + self.avgpool = nn.AvgPool3d(kernel_size=(2, 7, 7), stride=1) + self.classifier = nn.Sequential( + nn.Dropout(p=dropout), + nn.Conv3d(1024, num_classes, kernel_size=1, stride=1, bias=True), + ) + + def forward(self, x): + x = self.features(x) + x = self.avgpool(x) + x = self.classifier(x) + x = torch.mean(x, dim=(2, 3, 4)) + return x + + +class S3D_Weights(WeightsEnum): + KINETICS400_V1 = Weights( + url="https://download.pytorch.org/models/s3d-d76dad2f.pth", + transforms=partial( + VideoClassification, + crop_size=(224, 224), + resize_size=(256, 256), + ), + meta={ + "min_size": (224, 224), + "min_temporal_size": 14, + "categories": _KINETICS400_CATEGORIES, + "recipe": "https://github.com/pytorch/vision/tree/main/references/video_classification#s3d", + "_docs": ( + "The weights aim to approximate the accuracy of the paper. The accuracies are estimated on clip-level " + "with parameters `frame_rate=15`, `clips_per_video=1`, and `clip_len=128`." + ), + "num_params": 8320048, + "_metrics": { + "Kinetics-400": { + "acc@1": 68.368, + "acc@5": 88.050, + } + }, + "_ops": 17.979, + "_file_size": 31.972, + }, + ) + DEFAULT = KINETICS400_V1 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", S3D_Weights.KINETICS400_V1)) +def s3d(*, weights: Optional[S3D_Weights] = None, progress: bool = True, **kwargs: Any) -> S3D: + """Construct Separable 3D CNN model. + + Reference: `Rethinking Spatiotemporal Feature Learning `__. + + .. betastatus:: video module + + Args: + weights (:class:`~torchvision.models.video.S3D_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.video.S3D_Weights` + below for more details, and possible values. By default, no + pre-trained weights are used. + progress (bool): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.video.S3D`` base class. + Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.video.S3D_Weights + :members: + """ + weights = S3D_Weights.verify(weights) + + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + + model = S3D(**kwargs) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model diff --git a/lib/python3.12/site-packages/torchvision/models/video/swin_transformer.py b/lib/python3.12/site-packages/torchvision/models/video/swin_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..a8d87ffbe5af6caa1de0a3760fa5c506fdf8e231 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/video/swin_transformer.py @@ -0,0 +1,743 @@ +# Modified from 2d Swin Transformers in torchvision: +# https://github.com/pytorch/vision/blob/main/torchvision/models/swin_transformer.py + +from functools import partial +from typing import Any, Callable, List, Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn, Tensor + +from ...transforms._presets import VideoClassification + +from ...utils import _log_api_usage_once + +from .._api import register_model, Weights, WeightsEnum + +from .._meta import _KINETICS400_CATEGORIES +from .._utils import _ovewrite_named_param, handle_legacy_interface +from ..swin_transformer import PatchMerging, SwinTransformerBlock + +__all__ = [ + "SwinTransformer3d", + "Swin3D_T_Weights", + "Swin3D_S_Weights", + "Swin3D_B_Weights", + "swin3d_t", + "swin3d_s", + "swin3d_b", +] + + +def _get_window_and_shift_size( + shift_size: List[int], size_dhw: List[int], window_size: List[int] +) -> Tuple[List[int], List[int]]: + for i in range(3): + if size_dhw[i] <= window_size[i]: + # In this case, window_size will adapt to the input size, and no need to shift + window_size[i] = size_dhw[i] + shift_size[i] = 0 + + return window_size, shift_size + + +torch.fx.wrap("_get_window_and_shift_size") + + +def _get_relative_position_bias( + relative_position_bias_table: torch.Tensor, relative_position_index: torch.Tensor, window_size: List[int] +) -> Tensor: + window_vol = window_size[0] * window_size[1] * window_size[2] + # In 3d case we flatten the relative_position_bias + relative_position_bias = relative_position_bias_table[ + relative_position_index[:window_vol, :window_vol].flatten() # type: ignore[index] + ] + relative_position_bias = relative_position_bias.view(window_vol, window_vol, -1) + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous().unsqueeze(0) + return relative_position_bias + + +torch.fx.wrap("_get_relative_position_bias") + + +def _compute_pad_size_3d(size_dhw: Tuple[int, int, int], patch_size: Tuple[int, int, int]) -> Tuple[int, int, int]: + pad_size = [(patch_size[i] - size_dhw[i] % patch_size[i]) % patch_size[i] for i in range(3)] + return pad_size[0], pad_size[1], pad_size[2] + + +torch.fx.wrap("_compute_pad_size_3d") + + +def _compute_attention_mask_3d( + x: Tensor, + size_dhw: Tuple[int, int, int], + window_size: Tuple[int, int, int], + shift_size: Tuple[int, int, int], +) -> Tensor: + # generate attention mask + attn_mask = x.new_zeros(*size_dhw) + num_windows = (size_dhw[0] // window_size[0]) * (size_dhw[1] // window_size[1]) * (size_dhw[2] // window_size[2]) + slices = [ + ( + (0, -window_size[i]), + (-window_size[i], -shift_size[i]), + (-shift_size[i], None), + ) + for i in range(3) + ] + count = 0 + for d in slices[0]: + for h in slices[1]: + for w in slices[2]: + attn_mask[d[0] : d[1], h[0] : h[1], w[0] : w[1]] = count + count += 1 + + # Partition window on attn_mask + attn_mask = attn_mask.view( + size_dhw[0] // window_size[0], + window_size[0], + size_dhw[1] // window_size[1], + window_size[1], + size_dhw[2] // window_size[2], + window_size[2], + ) + attn_mask = attn_mask.permute(0, 2, 4, 1, 3, 5).reshape( + num_windows, window_size[0] * window_size[1] * window_size[2] + ) + attn_mask = attn_mask.unsqueeze(1) - attn_mask.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)) + return attn_mask + + +torch.fx.wrap("_compute_attention_mask_3d") + + +def shifted_window_attention_3d( + input: Tensor, + qkv_weight: Tensor, + proj_weight: Tensor, + relative_position_bias: Tensor, + window_size: List[int], + num_heads: int, + shift_size: List[int], + attention_dropout: float = 0.0, + dropout: float = 0.0, + qkv_bias: Optional[Tensor] = None, + proj_bias: Optional[Tensor] = None, + training: bool = True, +) -> Tensor: + """ + Window based multi-head self attention (W-MSA) module with relative position bias. + It supports both of shifted and non-shifted window. + Args: + input (Tensor[B, T, H, W, C]): The input tensor, 5-dimensions. + qkv_weight (Tensor[in_dim, out_dim]): The weight tensor of query, key, value. + proj_weight (Tensor[out_dim, out_dim]): The weight tensor of projection. + relative_position_bias (Tensor): The learned relative position bias added to attention. + window_size (List[int]): 3-dimensions window size, T, H, W . + num_heads (int): Number of attention heads. + shift_size (List[int]): Shift size for shifted window attention (T, H, W). + attention_dropout (float): Dropout ratio of attention weight. Default: 0.0. + dropout (float): Dropout ratio of output. Default: 0.0. + qkv_bias (Tensor[out_dim], optional): The bias tensor of query, key, value. Default: None. + proj_bias (Tensor[out_dim], optional): The bias tensor of projection. Default: None. + training (bool, optional): Training flag used by the dropout parameters. Default: True. + Returns: + Tensor[B, T, H, W, C]: The output tensor after shifted window attention. + """ + b, t, h, w, c = input.shape + # pad feature maps to multiples of window size + pad_size = _compute_pad_size_3d((t, h, w), (window_size[0], window_size[1], window_size[2])) + x = F.pad(input, (0, 0, 0, pad_size[2], 0, pad_size[1], 0, pad_size[0])) + _, tp, hp, wp, _ = x.shape + padded_size = (tp, hp, wp) + + # cyclic shift + if sum(shift_size) > 0: + x = torch.roll(x, shifts=(-shift_size[0], -shift_size[1], -shift_size[2]), dims=(1, 2, 3)) + + # partition windows + num_windows = ( + (padded_size[0] // window_size[0]) * (padded_size[1] // window_size[1]) * (padded_size[2] // window_size[2]) + ) + x = x.view( + b, + padded_size[0] // window_size[0], + window_size[0], + padded_size[1] // window_size[1], + window_size[1], + padded_size[2] // window_size[2], + window_size[2], + c, + ) + x = x.permute(0, 1, 3, 5, 2, 4, 6, 7).reshape( + b * num_windows, window_size[0] * window_size[1] * window_size[2], c + ) # B*nW, Wd*Wh*Ww, C + + # multi-head attention + qkv = F.linear(x, qkv_weight, qkv_bias) + qkv = qkv.reshape(x.size(0), x.size(1), 3, num_heads, c // num_heads).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] + q = q * (c // num_heads) ** -0.5 + attn = q.matmul(k.transpose(-2, -1)) + # add relative position bias + attn = attn + relative_position_bias + + if sum(shift_size) > 0: + # generate attention mask to handle shifted windows with varying size + attn_mask = _compute_attention_mask_3d( + x, + (padded_size[0], padded_size[1], padded_size[2]), + (window_size[0], window_size[1], window_size[2]), + (shift_size[0], shift_size[1], shift_size[2]), + ) + attn = attn.view(x.size(0) // num_windows, num_windows, num_heads, x.size(1), x.size(1)) + attn = attn + attn_mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, num_heads, x.size(1), x.size(1)) + + attn = F.softmax(attn, dim=-1) + attn = F.dropout(attn, p=attention_dropout, training=training) + + x = attn.matmul(v).transpose(1, 2).reshape(x.size(0), x.size(1), c) + x = F.linear(x, proj_weight, proj_bias) + x = F.dropout(x, p=dropout, training=training) + + # reverse windows + x = x.view( + b, + padded_size[0] // window_size[0], + padded_size[1] // window_size[1], + padded_size[2] // window_size[2], + window_size[0], + window_size[1], + window_size[2], + c, + ) + x = x.permute(0, 1, 4, 2, 5, 3, 6, 7).reshape(b, tp, hp, wp, c) + + # reverse cyclic shift + if sum(shift_size) > 0: + x = torch.roll(x, shifts=(shift_size[0], shift_size[1], shift_size[2]), dims=(1, 2, 3)) + + # unpad features + x = x[:, :t, :h, :w, :].contiguous() + return x + + +torch.fx.wrap("shifted_window_attention_3d") + + +class ShiftedWindowAttention3d(nn.Module): + """ + See :func:`shifted_window_attention_3d`. + """ + + def __init__( + self, + dim: int, + window_size: List[int], + shift_size: List[int], + num_heads: int, + qkv_bias: bool = True, + proj_bias: bool = True, + attention_dropout: float = 0.0, + dropout: float = 0.0, + ) -> None: + super().__init__() + if len(window_size) != 3 or len(shift_size) != 3: + raise ValueError("window_size and shift_size must be of length 2") + + self.window_size = window_size # Wd, Wh, Ww + self.shift_size = shift_size + self.num_heads = num_heads + self.attention_dropout = attention_dropout + self.dropout = dropout + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.proj = nn.Linear(dim, dim, bias=proj_bias) + + self.define_relative_position_bias_table() + self.define_relative_position_index() + + def define_relative_position_bias_table(self) -> None: + # define a parameter table of relative position bias + self.relative_position_bias_table = nn.Parameter( + torch.zeros( + (2 * self.window_size[0] - 1) * (2 * self.window_size[1] - 1) * (2 * self.window_size[2] - 1), + self.num_heads, + ) + ) # 2*Wd-1 * 2*Wh-1 * 2*Ww-1, nH + nn.init.trunc_normal_(self.relative_position_bias_table, std=0.02) + + def define_relative_position_index(self) -> None: + # get pair-wise relative position index for each token inside the window + coords_dhw = [torch.arange(self.window_size[i]) for i in range(3)] + coords = torch.stack( + torch.meshgrid(coords_dhw[0], coords_dhw[1], coords_dhw[2], indexing="ij") + ) # 3, Wd, Wh, Ww + coords_flatten = torch.flatten(coords, 1) # 3, Wd*Wh*Ww + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 3, Wd*Wh*Ww, Wd*Wh*Ww + relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wd*Wh*Ww, Wd*Wh*Ww, 3 + relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 2] += self.window_size[2] - 1 + + relative_coords[:, :, 0] *= (2 * self.window_size[1] - 1) * (2 * self.window_size[2] - 1) + relative_coords[:, :, 1] *= 2 * self.window_size[2] - 1 + # We don't flatten the relative_position_index here in 3d case. + relative_position_index = relative_coords.sum(-1) # Wd*Wh*Ww, Wd*Wh*Ww + self.register_buffer("relative_position_index", relative_position_index) + + def get_relative_position_bias(self, window_size: List[int]) -> torch.Tensor: + return _get_relative_position_bias(self.relative_position_bias_table, self.relative_position_index, window_size) # type: ignore + + def forward(self, x: Tensor) -> Tensor: + _, t, h, w, _ = x.shape + size_dhw = [t, h, w] + window_size, shift_size = self.window_size.copy(), self.shift_size.copy() + # Handle case where window_size is larger than the input tensor + window_size, shift_size = _get_window_and_shift_size(shift_size, size_dhw, window_size) + + relative_position_bias = self.get_relative_position_bias(window_size) + + return shifted_window_attention_3d( + x, + self.qkv.weight, + self.proj.weight, + relative_position_bias, + window_size, + self.num_heads, + shift_size=shift_size, + attention_dropout=self.attention_dropout, + dropout=self.dropout, + qkv_bias=self.qkv.bias, + proj_bias=self.proj.bias, + training=self.training, + ) + + +# Modified from: +# https://github.com/SwinTransformer/Video-Swin-Transformer/blob/master/mmaction/models/backbones/swin_transformer.py +class PatchEmbed3d(nn.Module): + """Video to Patch Embedding. + + Args: + patch_size (List[int]): Patch token size. + in_channels (int): Number of input channels. Default: 3 + embed_dim (int): Number of linear projection output channels. Default: 96. + norm_layer (nn.Module, optional): Normalization layer. Default: None + """ + + def __init__( + self, + patch_size: List[int], + in_channels: int = 3, + embed_dim: int = 96, + norm_layer: Optional[Callable[..., nn.Module]] = None, + ) -> None: + super().__init__() + _log_api_usage_once(self) + self.tuple_patch_size = (patch_size[0], patch_size[1], patch_size[2]) + + self.proj = nn.Conv3d( + in_channels, + embed_dim, + kernel_size=self.tuple_patch_size, + stride=self.tuple_patch_size, + ) + if norm_layer is not None: + self.norm = norm_layer(embed_dim) + else: + self.norm = nn.Identity() + + def forward(self, x: Tensor) -> Tensor: + """Forward function.""" + # padding + _, _, t, h, w = x.size() + pad_size = _compute_pad_size_3d((t, h, w), self.tuple_patch_size) + x = F.pad(x, (0, pad_size[2], 0, pad_size[1], 0, pad_size[0])) + x = self.proj(x) # B C T Wh Ww + x = x.permute(0, 2, 3, 4, 1) # B T Wh Ww C + if self.norm is not None: + x = self.norm(x) + return x + + +class SwinTransformer3d(nn.Module): + """ + Implements 3D Swin Transformer from the `"Video Swin Transformer" `_ paper. + Args: + patch_size (List[int]): Patch size. + embed_dim (int): Patch embedding dimension. + depths (List(int)): Depth of each Swin Transformer layer. + num_heads (List(int)): Number of attention heads in different layers. + window_size (List[int]): Window size. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.0. + dropout (float): Dropout rate. Default: 0.0. + attention_dropout (float): Attention dropout rate. Default: 0.0. + stochastic_depth_prob (float): Stochastic depth rate. Default: 0.1. + num_classes (int): Number of classes for classification head. Default: 400. + norm_layer (nn.Module, optional): Normalization layer. Default: None. + block (nn.Module, optional): SwinTransformer Block. Default: None. + downsample_layer (nn.Module): Downsample layer (patch merging). Default: PatchMerging. + patch_embed (nn.Module, optional): Patch Embedding layer. Default: None. + """ + + def __init__( + self, + patch_size: List[int], + embed_dim: int, + depths: List[int], + num_heads: List[int], + window_size: List[int], + mlp_ratio: float = 4.0, + dropout: float = 0.0, + attention_dropout: float = 0.0, + stochastic_depth_prob: float = 0.1, + num_classes: int = 400, + norm_layer: Optional[Callable[..., nn.Module]] = None, + block: Optional[Callable[..., nn.Module]] = None, + downsample_layer: Callable[..., nn.Module] = PatchMerging, + patch_embed: Optional[Callable[..., nn.Module]] = None, + ) -> None: + super().__init__() + _log_api_usage_once(self) + self.num_classes = num_classes + + if block is None: + block = partial(SwinTransformerBlock, attn_layer=ShiftedWindowAttention3d) + + if norm_layer is None: + norm_layer = partial(nn.LayerNorm, eps=1e-5) + + if patch_embed is None: + patch_embed = PatchEmbed3d + + # split image into non-overlapping patches + self.patch_embed = patch_embed(patch_size=patch_size, embed_dim=embed_dim, norm_layer=norm_layer) + self.pos_drop = nn.Dropout(p=dropout) + + layers: List[nn.Module] = [] + total_stage_blocks = sum(depths) + stage_block_id = 0 + # build SwinTransformer blocks + for i_stage in range(len(depths)): + stage: List[nn.Module] = [] + dim = embed_dim * 2**i_stage + for i_layer in range(depths[i_stage]): + # adjust stochastic depth probability based on the depth of the stage block + sd_prob = stochastic_depth_prob * float(stage_block_id) / (total_stage_blocks - 1) + stage.append( + block( + dim, + num_heads[i_stage], + window_size=window_size, + shift_size=[0 if i_layer % 2 == 0 else w // 2 for w in window_size], + mlp_ratio=mlp_ratio, + dropout=dropout, + attention_dropout=attention_dropout, + stochastic_depth_prob=sd_prob, + norm_layer=norm_layer, + attn_layer=ShiftedWindowAttention3d, + ) + ) + stage_block_id += 1 + layers.append(nn.Sequential(*stage)) + # add patch merging layer + if i_stage < (len(depths) - 1): + layers.append(downsample_layer(dim, norm_layer)) + self.features = nn.Sequential(*layers) + + self.num_features = embed_dim * 2 ** (len(depths) - 1) + self.norm = norm_layer(self.num_features) + self.avgpool = nn.AdaptiveAvgPool3d(1) + self.head = nn.Linear(self.num_features, num_classes) + + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.zeros_(m.bias) + + def forward(self, x: Tensor) -> Tensor: + # x: B C T H W + x = self.patch_embed(x) # B _T _H _W C + x = self.pos_drop(x) + x = self.features(x) # B _T _H _W C + x = self.norm(x) + x = x.permute(0, 4, 1, 2, 3) # B, C, _T, _H, _W + x = self.avgpool(x) + x = torch.flatten(x, 1) + x = self.head(x) + return x + + +def _swin_transformer3d( + patch_size: List[int], + embed_dim: int, + depths: List[int], + num_heads: List[int], + window_size: List[int], + stochastic_depth_prob: float, + weights: Optional[WeightsEnum], + progress: bool, + **kwargs: Any, +) -> SwinTransformer3d: + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + + model = SwinTransformer3d( + patch_size=patch_size, + embed_dim=embed_dim, + depths=depths, + num_heads=num_heads, + window_size=window_size, + stochastic_depth_prob=stochastic_depth_prob, + **kwargs, + ) + + if weights is not None: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +_COMMON_META = { + "categories": _KINETICS400_CATEGORIES, + "min_size": (1, 1), + "min_temporal_size": 1, +} + + +class Swin3D_T_Weights(WeightsEnum): + KINETICS400_V1 = Weights( + url="https://download.pytorch.org/models/swin3d_t-7615ae03.pth", + transforms=partial( + VideoClassification, + crop_size=(224, 224), + resize_size=(256,), + mean=(0.4850, 0.4560, 0.4060), + std=(0.2290, 0.2240, 0.2250), + ), + meta={ + **_COMMON_META, + "recipe": "https://github.com/SwinTransformer/Video-Swin-Transformer#kinetics-400", + "_docs": ( + "The weights were ported from the paper. The accuracies are estimated on video-level " + "with parameters `frame_rate=15`, `clips_per_video=12`, and `clip_len=32`" + ), + "num_params": 28158070, + "_metrics": { + "Kinetics-400": { + "acc@1": 77.715, + "acc@5": 93.519, + } + }, + "_ops": 43.882, + "_file_size": 121.543, + }, + ) + DEFAULT = KINETICS400_V1 + + +class Swin3D_S_Weights(WeightsEnum): + KINETICS400_V1 = Weights( + url="https://download.pytorch.org/models/swin3d_s-da41c237.pth", + transforms=partial( + VideoClassification, + crop_size=(224, 224), + resize_size=(256,), + mean=(0.4850, 0.4560, 0.4060), + std=(0.2290, 0.2240, 0.2250), + ), + meta={ + **_COMMON_META, + "recipe": "https://github.com/SwinTransformer/Video-Swin-Transformer#kinetics-400", + "_docs": ( + "The weights were ported from the paper. The accuracies are estimated on video-level " + "with parameters `frame_rate=15`, `clips_per_video=12`, and `clip_len=32`" + ), + "num_params": 49816678, + "_metrics": { + "Kinetics-400": { + "acc@1": 79.521, + "acc@5": 94.158, + } + }, + "_ops": 82.841, + "_file_size": 218.288, + }, + ) + DEFAULT = KINETICS400_V1 + + +class Swin3D_B_Weights(WeightsEnum): + KINETICS400_V1 = Weights( + url="https://download.pytorch.org/models/swin3d_b_1k-24f7c7c6.pth", + transforms=partial( + VideoClassification, + crop_size=(224, 224), + resize_size=(256,), + mean=(0.4850, 0.4560, 0.4060), + std=(0.2290, 0.2240, 0.2250), + ), + meta={ + **_COMMON_META, + "recipe": "https://github.com/SwinTransformer/Video-Swin-Transformer#kinetics-400", + "_docs": ( + "The weights were ported from the paper. The accuracies are estimated on video-level " + "with parameters `frame_rate=15`, `clips_per_video=12`, and `clip_len=32`" + ), + "num_params": 88048984, + "_metrics": { + "Kinetics-400": { + "acc@1": 79.427, + "acc@5": 94.386, + } + }, + "_ops": 140.667, + "_file_size": 364.134, + }, + ) + KINETICS400_IMAGENET22K_V1 = Weights( + url="https://download.pytorch.org/models/swin3d_b_22k-7c6ae6fa.pth", + transforms=partial( + VideoClassification, + crop_size=(224, 224), + resize_size=(256,), + mean=(0.4850, 0.4560, 0.4060), + std=(0.2290, 0.2240, 0.2250), + ), + meta={ + **_COMMON_META, + "recipe": "https://github.com/SwinTransformer/Video-Swin-Transformer#kinetics-400", + "_docs": ( + "The weights were ported from the paper. The accuracies are estimated on video-level " + "with parameters `frame_rate=15`, `clips_per_video=12`, and `clip_len=32`" + ), + "num_params": 88048984, + "_metrics": { + "Kinetics-400": { + "acc@1": 81.643, + "acc@5": 95.574, + } + }, + "_ops": 140.667, + "_file_size": 364.134, + }, + ) + DEFAULT = KINETICS400_V1 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", Swin3D_T_Weights.KINETICS400_V1)) +def swin3d_t(*, weights: Optional[Swin3D_T_Weights] = None, progress: bool = True, **kwargs: Any) -> SwinTransformer3d: + """ + Constructs a swin_tiny architecture from + `Video Swin Transformer `_. + + Args: + weights (:class:`~torchvision.models.video.Swin3D_T_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.video.Swin3D_T_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.video.swin_transformer.SwinTransformer`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.video.Swin3D_T_Weights + :members: + """ + weights = Swin3D_T_Weights.verify(weights) + + return _swin_transformer3d( + patch_size=[2, 4, 4], + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=[8, 7, 7], + stochastic_depth_prob=0.1, + weights=weights, + progress=progress, + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", Swin3D_S_Weights.KINETICS400_V1)) +def swin3d_s(*, weights: Optional[Swin3D_S_Weights] = None, progress: bool = True, **kwargs: Any) -> SwinTransformer3d: + """ + Constructs a swin_small architecture from + `Video Swin Transformer `_. + + Args: + weights (:class:`~torchvision.models.video.Swin3D_S_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.video.Swin3D_S_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.video.swin_transformer.SwinTransformer`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.video.Swin3D_S_Weights + :members: + """ + weights = Swin3D_S_Weights.verify(weights) + + return _swin_transformer3d( + patch_size=[2, 4, 4], + embed_dim=96, + depths=[2, 2, 18, 2], + num_heads=[3, 6, 12, 24], + window_size=[8, 7, 7], + stochastic_depth_prob=0.1, + weights=weights, + progress=progress, + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", Swin3D_B_Weights.KINETICS400_V1)) +def swin3d_b(*, weights: Optional[Swin3D_B_Weights] = None, progress: bool = True, **kwargs: Any) -> SwinTransformer3d: + """ + Constructs a swin_base architecture from + `Video Swin Transformer `_. + + Args: + weights (:class:`~torchvision.models.video.Swin3D_B_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.video.Swin3D_B_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.video.swin_transformer.SwinTransformer`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.video.Swin3D_B_Weights + :members: + """ + weights = Swin3D_B_Weights.verify(weights) + + return _swin_transformer3d( + patch_size=[2, 4, 4], + embed_dim=128, + depths=[2, 2, 18, 2], + num_heads=[4, 8, 16, 32], + window_size=[8, 7, 7], + stochastic_depth_prob=0.1, + weights=weights, + progress=progress, + **kwargs, + ) diff --git a/lib/python3.12/site-packages/torchvision/models/vision_transformer.py b/lib/python3.12/site-packages/torchvision/models/vision_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..f2983ef9db0a3adf232dd6fe1b90ce9417ce0853 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/models/vision_transformer.py @@ -0,0 +1,864 @@ +import math +from collections import OrderedDict +from functools import partial +from typing import Any, Callable, Dict, List, NamedTuple, Optional + +import torch +import torch.nn as nn + +from ..ops.misc import Conv2dNormActivation, MLP +from ..transforms._presets import ImageClassification, InterpolationMode +from ..utils import _log_api_usage_once +from ._api import register_model, Weights, WeightsEnum +from ._meta import _IMAGENET_CATEGORIES +from ._utils import _ovewrite_named_param, handle_legacy_interface + + +__all__ = [ + "VisionTransformer", + "ViT_B_16_Weights", + "ViT_B_32_Weights", + "ViT_L_16_Weights", + "ViT_L_32_Weights", + "ViT_H_14_Weights", + "vit_b_16", + "vit_b_32", + "vit_l_16", + "vit_l_32", + "vit_h_14", +] + + +class ConvStemConfig(NamedTuple): + out_channels: int + kernel_size: int + stride: int + norm_layer: Callable[..., nn.Module] = nn.BatchNorm2d + activation_layer: Callable[..., nn.Module] = nn.ReLU + + +class MLPBlock(MLP): + """Transformer MLP block.""" + + _version = 2 + + def __init__(self, in_dim: int, mlp_dim: int, dropout: float): + super().__init__(in_dim, [mlp_dim, in_dim], activation_layer=nn.GELU, inplace=None, dropout=dropout) + + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.normal_(m.bias, std=1e-6) + + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + version = local_metadata.get("version", None) + + if version is None or version < 2: + # Replacing legacy MLPBlock with MLP. See https://github.com/pytorch/vision/pull/6053 + for i in range(2): + for type in ["weight", "bias"]: + old_key = f"{prefix}linear_{i+1}.{type}" + new_key = f"{prefix}{3*i}.{type}" + if old_key in state_dict: + state_dict[new_key] = state_dict.pop(old_key) + + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + +class EncoderBlock(nn.Module): + """Transformer encoder block.""" + + def __init__( + self, + num_heads: int, + hidden_dim: int, + mlp_dim: int, + dropout: float, + attention_dropout: float, + norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6), + ): + super().__init__() + self.num_heads = num_heads + + # Attention block + self.ln_1 = norm_layer(hidden_dim) + self.self_attention = nn.MultiheadAttention(hidden_dim, num_heads, dropout=attention_dropout, batch_first=True) + self.dropout = nn.Dropout(dropout) + + # MLP block + self.ln_2 = norm_layer(hidden_dim) + self.mlp = MLPBlock(hidden_dim, mlp_dim, dropout) + + def forward(self, input: torch.Tensor): + torch._assert(input.dim() == 3, f"Expected (batch_size, seq_length, hidden_dim) got {input.shape}") + x = self.ln_1(input) + x, _ = self.self_attention(x, x, x, need_weights=False) + x = self.dropout(x) + x = x + input + + y = self.ln_2(x) + y = self.mlp(y) + return x + y + + +class Encoder(nn.Module): + """Transformer Model Encoder for sequence to sequence translation.""" + + def __init__( + self, + seq_length: int, + num_layers: int, + num_heads: int, + hidden_dim: int, + mlp_dim: int, + dropout: float, + attention_dropout: float, + norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6), + ): + super().__init__() + # Note that batch_size is on the first dim because + # we have batch_first=True in nn.MultiAttention() by default + self.pos_embedding = nn.Parameter(torch.empty(1, seq_length, hidden_dim).normal_(std=0.02)) # from BERT + self.dropout = nn.Dropout(dropout) + layers: OrderedDict[str, nn.Module] = OrderedDict() + for i in range(num_layers): + layers[f"encoder_layer_{i}"] = EncoderBlock( + num_heads, + hidden_dim, + mlp_dim, + dropout, + attention_dropout, + norm_layer, + ) + self.layers = nn.Sequential(layers) + self.ln = norm_layer(hidden_dim) + + def forward(self, input: torch.Tensor): + torch._assert(input.dim() == 3, f"Expected (batch_size, seq_length, hidden_dim) got {input.shape}") + input = input + self.pos_embedding + return self.ln(self.layers(self.dropout(input))) + + +class VisionTransformer(nn.Module): + """Vision Transformer as per https://arxiv.org/abs/2010.11929.""" + + def __init__( + self, + image_size: int, + patch_size: int, + num_layers: int, + num_heads: int, + hidden_dim: int, + mlp_dim: int, + dropout: float = 0.0, + attention_dropout: float = 0.0, + num_classes: int = 1000, + representation_size: Optional[int] = None, + norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6), + conv_stem_configs: Optional[List[ConvStemConfig]] = None, + ): + super().__init__() + _log_api_usage_once(self) + torch._assert(image_size % patch_size == 0, "Input shape indivisible by patch size!") + self.image_size = image_size + self.patch_size = patch_size + self.hidden_dim = hidden_dim + self.mlp_dim = mlp_dim + self.attention_dropout = attention_dropout + self.dropout = dropout + self.num_classes = num_classes + self.representation_size = representation_size + self.norm_layer = norm_layer + + if conv_stem_configs is not None: + # As per https://arxiv.org/abs/2106.14881 + seq_proj = nn.Sequential() + prev_channels = 3 + for i, conv_stem_layer_config in enumerate(conv_stem_configs): + seq_proj.add_module( + f"conv_bn_relu_{i}", + Conv2dNormActivation( + in_channels=prev_channels, + out_channels=conv_stem_layer_config.out_channels, + kernel_size=conv_stem_layer_config.kernel_size, + stride=conv_stem_layer_config.stride, + norm_layer=conv_stem_layer_config.norm_layer, + activation_layer=conv_stem_layer_config.activation_layer, + ), + ) + prev_channels = conv_stem_layer_config.out_channels + seq_proj.add_module( + "conv_last", nn.Conv2d(in_channels=prev_channels, out_channels=hidden_dim, kernel_size=1) + ) + self.conv_proj: nn.Module = seq_proj + else: + self.conv_proj = nn.Conv2d( + in_channels=3, out_channels=hidden_dim, kernel_size=patch_size, stride=patch_size + ) + + seq_length = (image_size // patch_size) ** 2 + + # Add a class token + self.class_token = nn.Parameter(torch.zeros(1, 1, hidden_dim)) + seq_length += 1 + + self.encoder = Encoder( + seq_length, + num_layers, + num_heads, + hidden_dim, + mlp_dim, + dropout, + attention_dropout, + norm_layer, + ) + self.seq_length = seq_length + + heads_layers: OrderedDict[str, nn.Module] = OrderedDict() + if representation_size is None: + heads_layers["head"] = nn.Linear(hidden_dim, num_classes) + else: + heads_layers["pre_logits"] = nn.Linear(hidden_dim, representation_size) + heads_layers["act"] = nn.Tanh() + heads_layers["head"] = nn.Linear(representation_size, num_classes) + + self.heads = nn.Sequential(heads_layers) + + if isinstance(self.conv_proj, nn.Conv2d): + # Init the patchify stem + fan_in = self.conv_proj.in_channels * self.conv_proj.kernel_size[0] * self.conv_proj.kernel_size[1] + nn.init.trunc_normal_(self.conv_proj.weight, std=math.sqrt(1 / fan_in)) + if self.conv_proj.bias is not None: + nn.init.zeros_(self.conv_proj.bias) + elif self.conv_proj.conv_last is not None and isinstance(self.conv_proj.conv_last, nn.Conv2d): + # Init the last 1x1 conv of the conv stem + nn.init.normal_( + self.conv_proj.conv_last.weight, mean=0.0, std=math.sqrt(2.0 / self.conv_proj.conv_last.out_channels) + ) + if self.conv_proj.conv_last.bias is not None: + nn.init.zeros_(self.conv_proj.conv_last.bias) + + if hasattr(self.heads, "pre_logits") and isinstance(self.heads.pre_logits, nn.Linear): + fan_in = self.heads.pre_logits.in_features + nn.init.trunc_normal_(self.heads.pre_logits.weight, std=math.sqrt(1 / fan_in)) + nn.init.zeros_(self.heads.pre_logits.bias) + + if isinstance(self.heads.head, nn.Linear): + nn.init.zeros_(self.heads.head.weight) + nn.init.zeros_(self.heads.head.bias) + + def _process_input(self, x: torch.Tensor) -> torch.Tensor: + n, c, h, w = x.shape + p = self.patch_size + torch._assert(h == self.image_size, f"Wrong image height! Expected {self.image_size} but got {h}!") + torch._assert(w == self.image_size, f"Wrong image width! Expected {self.image_size} but got {w}!") + n_h = h // p + n_w = w // p + + # (n, c, h, w) -> (n, hidden_dim, n_h, n_w) + x = self.conv_proj(x) + # (n, hidden_dim, n_h, n_w) -> (n, hidden_dim, (n_h * n_w)) + x = x.reshape(n, self.hidden_dim, n_h * n_w) + + # (n, hidden_dim, (n_h * n_w)) -> (n, (n_h * n_w), hidden_dim) + # The self attention layer expects inputs in the format (N, S, E) + # where S is the source sequence length, N is the batch size, E is the + # embedding dimension + x = x.permute(0, 2, 1) + + return x + + def forward(self, x: torch.Tensor): + # Reshape and permute the input tensor + x = self._process_input(x) + n = x.shape[0] + + # Expand the class token to the full batch + batch_class_token = self.class_token.expand(n, -1, -1) + x = torch.cat([batch_class_token, x], dim=1) + + x = self.encoder(x) + + # Classifier "token" as used by standard language architectures + x = x[:, 0] + + x = self.heads(x) + + return x + + +def _vision_transformer( + patch_size: int, + num_layers: int, + num_heads: int, + hidden_dim: int, + mlp_dim: int, + weights: Optional[WeightsEnum], + progress: bool, + **kwargs: Any, +) -> VisionTransformer: + if weights is not None: + _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) + assert weights.meta["min_size"][0] == weights.meta["min_size"][1] + _ovewrite_named_param(kwargs, "image_size", weights.meta["min_size"][0]) + image_size = kwargs.pop("image_size", 224) + + model = VisionTransformer( + image_size=image_size, + patch_size=patch_size, + num_layers=num_layers, + num_heads=num_heads, + hidden_dim=hidden_dim, + mlp_dim=mlp_dim, + **kwargs, + ) + + if weights: + model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) + + return model + + +_COMMON_META: Dict[str, Any] = { + "categories": _IMAGENET_CATEGORIES, +} + +_COMMON_SWAG_META = { + **_COMMON_META, + "recipe": "https://github.com/facebookresearch/SWAG", + "license": "https://github.com/facebookresearch/SWAG/blob/main/LICENSE", +} + + +class ViT_B_16_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/vit_b_16-c867db91.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 86567656, + "min_size": (224, 224), + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#vit_b_16", + "_metrics": { + "ImageNet-1K": { + "acc@1": 81.072, + "acc@5": 95.318, + } + }, + "_ops": 17.564, + "_file_size": 330.285, + "_docs": """ + These weights were trained from scratch by using a modified version of `DeIT + `_'s training recipe. + """, + }, + ) + IMAGENET1K_SWAG_E2E_V1 = Weights( + url="https://download.pytorch.org/models/vit_b_16_swag-9ac1b537.pth", + transforms=partial( + ImageClassification, + crop_size=384, + resize_size=384, + interpolation=InterpolationMode.BICUBIC, + ), + meta={ + **_COMMON_SWAG_META, + "num_params": 86859496, + "min_size": (384, 384), + "_metrics": { + "ImageNet-1K": { + "acc@1": 85.304, + "acc@5": 97.650, + } + }, + "_ops": 55.484, + "_file_size": 331.398, + "_docs": """ + These weights are learnt via transfer learning by end-to-end fine-tuning the original + `SWAG `_ weights on ImageNet-1K data. + """, + }, + ) + IMAGENET1K_SWAG_LINEAR_V1 = Weights( + url="https://download.pytorch.org/models/vit_b_16_lc_swag-4e70ced5.pth", + transforms=partial( + ImageClassification, + crop_size=224, + resize_size=224, + interpolation=InterpolationMode.BICUBIC, + ), + meta={ + **_COMMON_SWAG_META, + "recipe": "https://github.com/pytorch/vision/pull/5793", + "num_params": 86567656, + "min_size": (224, 224), + "_metrics": { + "ImageNet-1K": { + "acc@1": 81.886, + "acc@5": 96.180, + } + }, + "_ops": 17.564, + "_file_size": 330.285, + "_docs": """ + These weights are composed of the original frozen `SWAG `_ trunk + weights and a linear classifier learnt on top of them trained on ImageNet-1K data. + """, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class ViT_B_32_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/vit_b_32-d86f8d99.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 88224232, + "min_size": (224, 224), + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#vit_b_32", + "_metrics": { + "ImageNet-1K": { + "acc@1": 75.912, + "acc@5": 92.466, + } + }, + "_ops": 4.409, + "_file_size": 336.604, + "_docs": """ + These weights were trained from scratch by using a modified version of `DeIT + `_'s training recipe. + """, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class ViT_L_16_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/vit_l_16-852ce7e3.pth", + transforms=partial(ImageClassification, crop_size=224, resize_size=242), + meta={ + **_COMMON_META, + "num_params": 304326632, + "min_size": (224, 224), + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#vit_l_16", + "_metrics": { + "ImageNet-1K": { + "acc@1": 79.662, + "acc@5": 94.638, + } + }, + "_ops": 61.555, + "_file_size": 1161.023, + "_docs": """ + These weights were trained from scratch by using a modified version of TorchVision's + `new training recipe + `_. + """, + }, + ) + IMAGENET1K_SWAG_E2E_V1 = Weights( + url="https://download.pytorch.org/models/vit_l_16_swag-4f3808c9.pth", + transforms=partial( + ImageClassification, + crop_size=512, + resize_size=512, + interpolation=InterpolationMode.BICUBIC, + ), + meta={ + **_COMMON_SWAG_META, + "num_params": 305174504, + "min_size": (512, 512), + "_metrics": { + "ImageNet-1K": { + "acc@1": 88.064, + "acc@5": 98.512, + } + }, + "_ops": 361.986, + "_file_size": 1164.258, + "_docs": """ + These weights are learnt via transfer learning by end-to-end fine-tuning the original + `SWAG `_ weights on ImageNet-1K data. + """, + }, + ) + IMAGENET1K_SWAG_LINEAR_V1 = Weights( + url="https://download.pytorch.org/models/vit_l_16_lc_swag-4d563306.pth", + transforms=partial( + ImageClassification, + crop_size=224, + resize_size=224, + interpolation=InterpolationMode.BICUBIC, + ), + meta={ + **_COMMON_SWAG_META, + "recipe": "https://github.com/pytorch/vision/pull/5793", + "num_params": 304326632, + "min_size": (224, 224), + "_metrics": { + "ImageNet-1K": { + "acc@1": 85.146, + "acc@5": 97.422, + } + }, + "_ops": 61.555, + "_file_size": 1161.023, + "_docs": """ + These weights are composed of the original frozen `SWAG `_ trunk + weights and a linear classifier learnt on top of them trained on ImageNet-1K data. + """, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class ViT_L_32_Weights(WeightsEnum): + IMAGENET1K_V1 = Weights( + url="https://download.pytorch.org/models/vit_l_32-c7638314.pth", + transforms=partial(ImageClassification, crop_size=224), + meta={ + **_COMMON_META, + "num_params": 306535400, + "min_size": (224, 224), + "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#vit_l_32", + "_metrics": { + "ImageNet-1K": { + "acc@1": 76.972, + "acc@5": 93.07, + } + }, + "_ops": 15.378, + "_file_size": 1169.449, + "_docs": """ + These weights were trained from scratch by using a modified version of `DeIT + `_'s training recipe. + """, + }, + ) + DEFAULT = IMAGENET1K_V1 + + +class ViT_H_14_Weights(WeightsEnum): + IMAGENET1K_SWAG_E2E_V1 = Weights( + url="https://download.pytorch.org/models/vit_h_14_swag-80465313.pth", + transforms=partial( + ImageClassification, + crop_size=518, + resize_size=518, + interpolation=InterpolationMode.BICUBIC, + ), + meta={ + **_COMMON_SWAG_META, + "num_params": 633470440, + "min_size": (518, 518), + "_metrics": { + "ImageNet-1K": { + "acc@1": 88.552, + "acc@5": 98.694, + } + }, + "_ops": 1016.717, + "_file_size": 2416.643, + "_docs": """ + These weights are learnt via transfer learning by end-to-end fine-tuning the original + `SWAG `_ weights on ImageNet-1K data. + """, + }, + ) + IMAGENET1K_SWAG_LINEAR_V1 = Weights( + url="https://download.pytorch.org/models/vit_h_14_lc_swag-c1eb923e.pth", + transforms=partial( + ImageClassification, + crop_size=224, + resize_size=224, + interpolation=InterpolationMode.BICUBIC, + ), + meta={ + **_COMMON_SWAG_META, + "recipe": "https://github.com/pytorch/vision/pull/5793", + "num_params": 632045800, + "min_size": (224, 224), + "_metrics": { + "ImageNet-1K": { + "acc@1": 85.708, + "acc@5": 97.730, + } + }, + "_ops": 167.295, + "_file_size": 2411.209, + "_docs": """ + These weights are composed of the original frozen `SWAG `_ trunk + weights and a linear classifier learnt on top of them trained on ImageNet-1K data. + """, + }, + ) + DEFAULT = IMAGENET1K_SWAG_E2E_V1 + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ViT_B_16_Weights.IMAGENET1K_V1)) +def vit_b_16(*, weights: Optional[ViT_B_16_Weights] = None, progress: bool = True, **kwargs: Any) -> VisionTransformer: + """ + Constructs a vit_b_16 architecture from + `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale `_. + + Args: + weights (:class:`~torchvision.models.ViT_B_16_Weights`, optional): The pretrained + weights to use. See :class:`~torchvision.models.ViT_B_16_Weights` + below for more details and possible values. By default, no pre-trained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.vision_transformer.VisionTransformer`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.ViT_B_16_Weights + :members: + """ + weights = ViT_B_16_Weights.verify(weights) + + return _vision_transformer( + patch_size=16, + num_layers=12, + num_heads=12, + hidden_dim=768, + mlp_dim=3072, + weights=weights, + progress=progress, + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ViT_B_32_Weights.IMAGENET1K_V1)) +def vit_b_32(*, weights: Optional[ViT_B_32_Weights] = None, progress: bool = True, **kwargs: Any) -> VisionTransformer: + """ + Constructs a vit_b_32 architecture from + `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale `_. + + Args: + weights (:class:`~torchvision.models.ViT_B_32_Weights`, optional): The pretrained + weights to use. See :class:`~torchvision.models.ViT_B_32_Weights` + below for more details and possible values. By default, no pre-trained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.vision_transformer.VisionTransformer`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.ViT_B_32_Weights + :members: + """ + weights = ViT_B_32_Weights.verify(weights) + + return _vision_transformer( + patch_size=32, + num_layers=12, + num_heads=12, + hidden_dim=768, + mlp_dim=3072, + weights=weights, + progress=progress, + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ViT_L_16_Weights.IMAGENET1K_V1)) +def vit_l_16(*, weights: Optional[ViT_L_16_Weights] = None, progress: bool = True, **kwargs: Any) -> VisionTransformer: + """ + Constructs a vit_l_16 architecture from + `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale `_. + + Args: + weights (:class:`~torchvision.models.ViT_L_16_Weights`, optional): The pretrained + weights to use. See :class:`~torchvision.models.ViT_L_16_Weights` + below for more details and possible values. By default, no pre-trained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.vision_transformer.VisionTransformer`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.ViT_L_16_Weights + :members: + """ + weights = ViT_L_16_Weights.verify(weights) + + return _vision_transformer( + patch_size=16, + num_layers=24, + num_heads=16, + hidden_dim=1024, + mlp_dim=4096, + weights=weights, + progress=progress, + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", ViT_L_32_Weights.IMAGENET1K_V1)) +def vit_l_32(*, weights: Optional[ViT_L_32_Weights] = None, progress: bool = True, **kwargs: Any) -> VisionTransformer: + """ + Constructs a vit_l_32 architecture from + `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale `_. + + Args: + weights (:class:`~torchvision.models.ViT_L_32_Weights`, optional): The pretrained + weights to use. See :class:`~torchvision.models.ViT_L_32_Weights` + below for more details and possible values. By default, no pre-trained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.vision_transformer.VisionTransformer`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.ViT_L_32_Weights + :members: + """ + weights = ViT_L_32_Weights.verify(weights) + + return _vision_transformer( + patch_size=32, + num_layers=24, + num_heads=16, + hidden_dim=1024, + mlp_dim=4096, + weights=weights, + progress=progress, + **kwargs, + ) + + +@register_model() +@handle_legacy_interface(weights=("pretrained", None)) +def vit_h_14(*, weights: Optional[ViT_H_14_Weights] = None, progress: bool = True, **kwargs: Any) -> VisionTransformer: + """ + Constructs a vit_h_14 architecture from + `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale `_. + + Args: + weights (:class:`~torchvision.models.ViT_H_14_Weights`, optional): The pretrained + weights to use. See :class:`~torchvision.models.ViT_H_14_Weights` + below for more details and possible values. By default, no pre-trained weights are used. + progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.vision_transformer.VisionTransformer`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.ViT_H_14_Weights + :members: + """ + weights = ViT_H_14_Weights.verify(weights) + + return _vision_transformer( + patch_size=14, + num_layers=32, + num_heads=16, + hidden_dim=1280, + mlp_dim=5120, + weights=weights, + progress=progress, + **kwargs, + ) + + +def interpolate_embeddings( + image_size: int, + patch_size: int, + model_state: "OrderedDict[str, torch.Tensor]", + interpolation_mode: str = "bicubic", + reset_heads: bool = False, +) -> "OrderedDict[str, torch.Tensor]": + """This function helps interpolate positional embeddings during checkpoint loading, + especially when you want to apply a pre-trained model on images with different resolution. + + Args: + image_size (int): Image size of the new model. + patch_size (int): Patch size of the new model. + model_state (OrderedDict[str, torch.Tensor]): State dict of the pre-trained model. + interpolation_mode (str): The algorithm used for upsampling. Default: bicubic. + reset_heads (bool): If true, not copying the state of heads. Default: False. + + Returns: + OrderedDict[str, torch.Tensor]: A state dict which can be loaded into the new model. + """ + # Shape of pos_embedding is (1, seq_length, hidden_dim) + pos_embedding = model_state["encoder.pos_embedding"] + n, seq_length, hidden_dim = pos_embedding.shape + if n != 1: + raise ValueError(f"Unexpected position embedding shape: {pos_embedding.shape}") + + new_seq_length = (image_size // patch_size) ** 2 + 1 + + # Need to interpolate the weights for the position embedding. + # We do this by reshaping the positions embeddings to a 2d grid, performing + # an interpolation in the (h, w) space and then reshaping back to a 1d grid. + if new_seq_length != seq_length: + # The class token embedding shouldn't be interpolated, so we split it up. + seq_length -= 1 + new_seq_length -= 1 + pos_embedding_token = pos_embedding[:, :1, :] + pos_embedding_img = pos_embedding[:, 1:, :] + + # (1, seq_length, hidden_dim) -> (1, hidden_dim, seq_length) + pos_embedding_img = pos_embedding_img.permute(0, 2, 1) + seq_length_1d = int(math.sqrt(seq_length)) + if seq_length_1d * seq_length_1d != seq_length: + raise ValueError( + f"seq_length is not a perfect square! Instead got seq_length_1d * seq_length_1d = {seq_length_1d * seq_length_1d } and seq_length = {seq_length}" + ) + + # (1, hidden_dim, seq_length) -> (1, hidden_dim, seq_l_1d, seq_l_1d) + pos_embedding_img = pos_embedding_img.reshape(1, hidden_dim, seq_length_1d, seq_length_1d) + new_seq_length_1d = image_size // patch_size + + # Perform interpolation. + # (1, hidden_dim, seq_l_1d, seq_l_1d) -> (1, hidden_dim, new_seq_l_1d, new_seq_l_1d) + new_pos_embedding_img = nn.functional.interpolate( + pos_embedding_img, + size=new_seq_length_1d, + mode=interpolation_mode, + align_corners=True, + ) + + # (1, hidden_dim, new_seq_l_1d, new_seq_l_1d) -> (1, hidden_dim, new_seq_length) + new_pos_embedding_img = new_pos_embedding_img.reshape(1, hidden_dim, new_seq_length) + + # (1, hidden_dim, new_seq_length) -> (1, new_seq_length, hidden_dim) + new_pos_embedding_img = new_pos_embedding_img.permute(0, 2, 1) + new_pos_embedding = torch.cat([pos_embedding_token, new_pos_embedding_img], dim=1) + + model_state["encoder.pos_embedding"] = new_pos_embedding + + if reset_heads: + model_state_copy: "OrderedDict[str, torch.Tensor]" = OrderedDict() + for k, v in model_state.items(): + if not k.startswith("heads"): + model_state_copy[k] = v + model_state = model_state_copy + + return model_state diff --git a/lib/python3.12/site-packages/torchvision/ops/__init__.py b/lib/python3.12/site-packages/torchvision/ops/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..827505b842d4f1ad0e16dfe54ef28658364cc9ac --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/ops/__init__.py @@ -0,0 +1,73 @@ +from ._register_onnx_ops import _register_custom_op +from .boxes import ( + batched_nms, + box_area, + box_convert, + box_iou, + clip_boxes_to_image, + complete_box_iou, + distance_box_iou, + generalized_box_iou, + masks_to_boxes, + nms, + remove_small_boxes, +) +from .ciou_loss import complete_box_iou_loss +from .deform_conv import deform_conv2d, DeformConv2d +from .diou_loss import distance_box_iou_loss +from .drop_block import drop_block2d, drop_block3d, DropBlock2d, DropBlock3d +from .feature_pyramid_network import FeaturePyramidNetwork +from .focal_loss import sigmoid_focal_loss +from .giou_loss import generalized_box_iou_loss +from .misc import Conv2dNormActivation, Conv3dNormActivation, FrozenBatchNorm2d, MLP, Permute, SqueezeExcitation +from .poolers import MultiScaleRoIAlign +from .ps_roi_align import ps_roi_align, PSRoIAlign +from .ps_roi_pool import ps_roi_pool, PSRoIPool +from .roi_align import roi_align, RoIAlign +from .roi_pool import roi_pool, RoIPool +from .stochastic_depth import stochastic_depth, StochasticDepth + +_register_custom_op() + + +__all__ = [ + "masks_to_boxes", + "deform_conv2d", + "DeformConv2d", + "nms", + "batched_nms", + "remove_small_boxes", + "clip_boxes_to_image", + "box_convert", + "box_area", + "box_iou", + "generalized_box_iou", + "distance_box_iou", + "complete_box_iou", + "roi_align", + "RoIAlign", + "roi_pool", + "RoIPool", + "ps_roi_align", + "PSRoIAlign", + "ps_roi_pool", + "PSRoIPool", + "MultiScaleRoIAlign", + "FeaturePyramidNetwork", + "sigmoid_focal_loss", + "stochastic_depth", + "StochasticDepth", + "FrozenBatchNorm2d", + "Conv2dNormActivation", + "Conv3dNormActivation", + "SqueezeExcitation", + "MLP", + "Permute", + "generalized_box_iou_loss", + "distance_box_iou_loss", + "complete_box_iou_loss", + "drop_block2d", + "DropBlock2d", + "drop_block3d", + "DropBlock3d", +] diff --git a/lib/python3.12/site-packages/torchvision/ops/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/ops/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a6a79617c7621fb73152e97be1c8ad451b52294f Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/ops/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/ops/__pycache__/_box_convert.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/ops/__pycache__/_box_convert.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1b540cf679074208db35ad4ae53a63bd764d1072 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/ops/__pycache__/_box_convert.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/ops/__pycache__/_register_onnx_ops.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/ops/__pycache__/_register_onnx_ops.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8799a0e81034e9b6489047a2e2b60b8562bc3629 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/ops/__pycache__/_register_onnx_ops.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/ops/__pycache__/_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/ops/__pycache__/_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0a4f77049e3aa2905381a6567876bc7f8975df19 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/ops/__pycache__/_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/ops/__pycache__/boxes.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/ops/__pycache__/boxes.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..317b297f23fed60dd2b9afbdd2712ec713e3de9a Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/ops/__pycache__/boxes.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/ops/__pycache__/ciou_loss.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/ops/__pycache__/ciou_loss.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..68668aac32e9d0b1cbe2175b24f8b43821950175 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/ops/__pycache__/ciou_loss.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/ops/__pycache__/deform_conv.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/ops/__pycache__/deform_conv.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..62b83f40f1fb49fbd815e99e19ec2932e3b2cb6b Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/ops/__pycache__/deform_conv.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/ops/__pycache__/diou_loss.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/ops/__pycache__/diou_loss.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ce8259176351d3987078387fb8af5c96b2624d8b Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/ops/__pycache__/diou_loss.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/ops/__pycache__/drop_block.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/ops/__pycache__/drop_block.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4bf1f043f8d823ec02d86f83182f22f801b5fdac Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/ops/__pycache__/drop_block.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/ops/__pycache__/feature_pyramid_network.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/ops/__pycache__/feature_pyramid_network.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..687108258551e18b6b74b246834f76b1070f605b Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/ops/__pycache__/feature_pyramid_network.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/ops/__pycache__/focal_loss.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/ops/__pycache__/focal_loss.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8b7f2d577ace5ca905834619bdc1a055764aa552 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/ops/__pycache__/focal_loss.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/ops/__pycache__/giou_loss.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/ops/__pycache__/giou_loss.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..85a4f1768f7b2dafdf1b0702825b136da4102bfb Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/ops/__pycache__/giou_loss.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/ops/__pycache__/misc.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/ops/__pycache__/misc.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f3de8f84d5359fef16d4f051b94111e8dbb33d5f Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/ops/__pycache__/misc.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/ops/__pycache__/poolers.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/ops/__pycache__/poolers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eca2b29b2061eb6f33ddb5774ca7f833758c7679 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/ops/__pycache__/poolers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/ops/__pycache__/ps_roi_align.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/ops/__pycache__/ps_roi_align.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b40bb0a7021aec5b04e011cae6fa34d7141e7927 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/ops/__pycache__/ps_roi_align.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/ops/__pycache__/ps_roi_pool.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/ops/__pycache__/ps_roi_pool.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..422bab0524e96cbb4973bbe063ef0803670ef501 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/ops/__pycache__/ps_roi_pool.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/ops/__pycache__/roi_align.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/ops/__pycache__/roi_align.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b8813a6455ac294a8e24299c033c3b997b6232a8 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/ops/__pycache__/roi_align.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/ops/__pycache__/roi_pool.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/ops/__pycache__/roi_pool.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d84f98bf49a053886b2a11f5366dd807bfe73d14 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/ops/__pycache__/roi_pool.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/ops/__pycache__/stochastic_depth.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/ops/__pycache__/stochastic_depth.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..72b845e304e499b4eecf2fe33da606a630832ea5 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/ops/__pycache__/stochastic_depth.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/ops/_box_convert.py b/lib/python3.12/site-packages/torchvision/ops/_box_convert.py new file mode 100644 index 0000000000000000000000000000000000000000..124bdd0bcc6527d04e3c9c1f710b803f7dfbbb71 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/ops/_box_convert.py @@ -0,0 +1,81 @@ +import torch +from torch import Tensor + + +def _box_cxcywh_to_xyxy(boxes: Tensor) -> Tensor: + """ + Converts bounding boxes from (cx, cy, w, h) format to (x1, y1, x2, y2) format. + (cx, cy) refers to center of bounding box + (w, h) are width and height of bounding box + Args: + boxes (Tensor[N, 4]): boxes in (cx, cy, w, h) format which will be converted. + + Returns: + boxes (Tensor(N, 4)): boxes in (x1, y1, x2, y2) format. + """ + # We need to change all 4 of them so some temporary variable is needed. + cx, cy, w, h = boxes.unbind(-1) + x1 = cx - 0.5 * w + y1 = cy - 0.5 * h + x2 = cx + 0.5 * w + y2 = cy + 0.5 * h + + boxes = torch.stack((x1, y1, x2, y2), dim=-1) + + return boxes + + +def _box_xyxy_to_cxcywh(boxes: Tensor) -> Tensor: + """ + Converts bounding boxes from (x1, y1, x2, y2) format to (cx, cy, w, h) format. + (x1, y1) refer to top left of bounding box + (x2, y2) refer to bottom right of bounding box + Args: + boxes (Tensor[N, 4]): boxes in (x1, y1, x2, y2) format which will be converted. + + Returns: + boxes (Tensor(N, 4)): boxes in (cx, cy, w, h) format. + """ + x1, y1, x2, y2 = boxes.unbind(-1) + cx = (x1 + x2) / 2 + cy = (y1 + y2) / 2 + w = x2 - x1 + h = y2 - y1 + + boxes = torch.stack((cx, cy, w, h), dim=-1) + + return boxes + + +def _box_xywh_to_xyxy(boxes: Tensor) -> Tensor: + """ + Converts bounding boxes from (x, y, w, h) format to (x1, y1, x2, y2) format. + (x, y) refers to top left of bounding box. + (w, h) refers to width and height of box. + Args: + boxes (Tensor[N, 4]): boxes in (x, y, w, h) which will be converted. + + Returns: + boxes (Tensor[N, 4]): boxes in (x1, y1, x2, y2) format. + """ + x, y, w, h = boxes.unbind(-1) + boxes = torch.stack([x, y, x + w, y + h], dim=-1) + return boxes + + +def _box_xyxy_to_xywh(boxes: Tensor) -> Tensor: + """ + Converts bounding boxes from (x1, y1, x2, y2) format to (x, y, w, h) format. + (x1, y1) refer to top left of bounding box + (x2, y2) refer to bottom right of bounding box + Args: + boxes (Tensor[N, 4]): boxes in (x1, y1, x2, y2) which will be converted. + + Returns: + boxes (Tensor[N, 4]): boxes in (x, y, w, h) format. + """ + x1, y1, x2, y2 = boxes.unbind(-1) + w = x2 - x1 # x2 - x1 + h = y2 - y1 # y2 - y1 + boxes = torch.stack((x1, y1, w, h), dim=-1) + return boxes diff --git a/lib/python3.12/site-packages/torchvision/ops/_register_onnx_ops.py b/lib/python3.12/site-packages/torchvision/ops/_register_onnx_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..5dd263a5d8ef497becc4aa39252a93c913b84880 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/ops/_register_onnx_ops.py @@ -0,0 +1,107 @@ +import sys +import warnings + +import torch +from torch.onnx import symbolic_opset11 as opset11 +from torch.onnx.symbolic_helper import parse_args + +_ONNX_OPSET_VERSION_11 = 11 +_ONNX_OPSET_VERSION_16 = 16 +BASE_ONNX_OPSET_VERSION = _ONNX_OPSET_VERSION_11 + + +@parse_args("v", "v", "f") +def symbolic_multi_label_nms(g, boxes, scores, iou_threshold): + boxes = opset11.unsqueeze(g, boxes, 0) + scores = opset11.unsqueeze(g, opset11.unsqueeze(g, scores, 0), 0) + max_output_per_class = g.op("Constant", value_t=torch.tensor([sys.maxsize], dtype=torch.long)) + iou_threshold = g.op("Constant", value_t=torch.tensor([iou_threshold], dtype=torch.float)) + + # Cast boxes and scores to float32 in case they are float64 inputs + nms_out = g.op( + "NonMaxSuppression", + g.op("Cast", boxes, to_i=torch.onnx.TensorProtoDataType.FLOAT), + g.op("Cast", scores, to_i=torch.onnx.TensorProtoDataType.FLOAT), + max_output_per_class, + iou_threshold, + ) + return opset11.squeeze( + g, opset11.select(g, nms_out, 1, g.op("Constant", value_t=torch.tensor([2], dtype=torch.long))), 1 + ) + + +def _process_batch_indices_for_roi_align(g, rois): + indices = opset11.squeeze( + g, opset11.select(g, rois, 1, g.op("Constant", value_t=torch.tensor([0], dtype=torch.long))), 1 + ) + return g.op("Cast", indices, to_i=torch.onnx.TensorProtoDataType.INT64) + + +def _process_rois_for_roi_align(g, rois): + return opset11.select(g, rois, 1, g.op("Constant", value_t=torch.tensor([1, 2, 3, 4], dtype=torch.long))) + + +def _process_sampling_ratio_for_roi_align(g, sampling_ratio: int): + if sampling_ratio < 0: + warnings.warn( + "ONNX export for RoIAlign with a non-zero sampling_ratio is not supported. " + "The model will be exported with a sampling_ratio of 0." + ) + sampling_ratio = 0 + return sampling_ratio + + +@parse_args("v", "v", "f", "i", "i", "i", "i") +def roi_align_opset11(g, input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio, aligned): + batch_indices = _process_batch_indices_for_roi_align(g, rois) + rois = _process_rois_for_roi_align(g, rois) + if aligned: + warnings.warn( + "ROIAlign with aligned=True is only supported in opset >= 16. " + "Please export with opset 16 or higher, or use aligned=False." + ) + sampling_ratio = _process_sampling_ratio_for_roi_align(g, sampling_ratio) + return g.op( + "RoiAlign", + input, + rois, + batch_indices, + spatial_scale_f=spatial_scale, + output_height_i=pooled_height, + output_width_i=pooled_width, + sampling_ratio_i=sampling_ratio, + ) + + +@parse_args("v", "v", "f", "i", "i", "i", "i") +def roi_align_opset16(g, input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio, aligned): + batch_indices = _process_batch_indices_for_roi_align(g, rois) + rois = _process_rois_for_roi_align(g, rois) + coordinate_transformation_mode = "half_pixel" if aligned else "output_half_pixel" + sampling_ratio = _process_sampling_ratio_for_roi_align(g, sampling_ratio) + return g.op( + "RoiAlign", + input, + rois, + batch_indices, + coordinate_transformation_mode_s=coordinate_transformation_mode, + spatial_scale_f=spatial_scale, + output_height_i=pooled_height, + output_width_i=pooled_width, + sampling_ratio_i=sampling_ratio, + ) + + +@parse_args("v", "v", "f", "i", "i") +def roi_pool(g, input, rois, spatial_scale, pooled_height, pooled_width): + roi_pool = g.op( + "MaxRoiPool", input, rois, pooled_shape_i=(pooled_height, pooled_width), spatial_scale_f=spatial_scale + ) + return roi_pool, None + + +def _register_custom_op(): + torch.onnx.register_custom_op_symbolic("torchvision::nms", symbolic_multi_label_nms, _ONNX_OPSET_VERSION_11) + torch.onnx.register_custom_op_symbolic("torchvision::roi_align", roi_align_opset11, _ONNX_OPSET_VERSION_11) + torch.onnx.register_custom_op_symbolic("torchvision::roi_align", roi_align_opset16, _ONNX_OPSET_VERSION_16) + torch.onnx.register_custom_op_symbolic("torchvision::roi_pool", roi_pool, _ONNX_OPSET_VERSION_11) diff --git a/lib/python3.12/site-packages/torchvision/ops/_utils.py b/lib/python3.12/site-packages/torchvision/ops/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a6ca557a98b899b7c2a11ba0dca3d64730af4268 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/ops/_utils.py @@ -0,0 +1,106 @@ +from typing import List, Optional, Tuple, Union + +import torch +from torch import nn, Tensor + + +def _cat(tensors: List[Tensor], dim: int = 0) -> Tensor: + """ + Efficient version of torch.cat that avoids a copy if there is only a single element in a list + """ + # TODO add back the assert + # assert isinstance(tensors, (list, tuple)) + if len(tensors) == 1: + return tensors[0] + return torch.cat(tensors, dim) + + +def convert_boxes_to_roi_format(boxes: List[Tensor]) -> Tensor: + concat_boxes = _cat([b for b in boxes], dim=0) + temp = [] + for i, b in enumerate(boxes): + temp.append(torch.full_like(b[:, :1], i)) + ids = _cat(temp, dim=0) + rois = torch.cat([ids, concat_boxes], dim=1) + return rois + + +def check_roi_boxes_shape(boxes: Union[Tensor, List[Tensor]]): + if isinstance(boxes, (list, tuple)): + for _tensor in boxes: + torch._assert( + _tensor.size(1) == 4, "The shape of the tensor in the boxes list is not correct as List[Tensor[L, 4]]" + ) + elif isinstance(boxes, torch.Tensor): + torch._assert(boxes.size(1) == 5, "The boxes tensor shape is not correct as Tensor[K, 5]") + else: + torch._assert(False, "boxes is expected to be a Tensor[L, 5] or a List[Tensor[K, 4]]") + return + + +def split_normalization_params( + model: nn.Module, norm_classes: Optional[List[type]] = None +) -> Tuple[List[Tensor], List[Tensor]]: + # Adapted from https://github.com/facebookresearch/ClassyVision/blob/659d7f78/classy_vision/generic/util.py#L501 + if not norm_classes: + norm_classes = [ + nn.modules.batchnorm._BatchNorm, + nn.LayerNorm, + nn.GroupNorm, + nn.modules.instancenorm._InstanceNorm, + nn.LocalResponseNorm, + ] + + for t in norm_classes: + if not issubclass(t, nn.Module): + raise ValueError(f"Class {t} is not a subclass of nn.Module.") + + classes = tuple(norm_classes) + + norm_params = [] + other_params = [] + for module in model.modules(): + if next(module.children(), None): + other_params.extend(p for p in module.parameters(recurse=False) if p.requires_grad) + elif isinstance(module, classes): + norm_params.extend(p for p in module.parameters() if p.requires_grad) + else: + other_params.extend(p for p in module.parameters() if p.requires_grad) + return norm_params, other_params + + +def _upcast(t: Tensor) -> Tensor: + # Protects from numerical overflows in multiplications by upcasting to the equivalent higher type + if t.is_floating_point(): + return t if t.dtype in (torch.float32, torch.float64) else t.float() + else: + return t if t.dtype in (torch.int32, torch.int64) else t.int() + + +def _upcast_non_float(t: Tensor) -> Tensor: + # Protects from numerical overflows in multiplications by upcasting to the equivalent higher type + if t.dtype not in (torch.float32, torch.float64): + return t.float() + return t + + +def _loss_inter_union( + boxes1: torch.Tensor, + boxes2: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + + x1, y1, x2, y2 = boxes1.unbind(dim=-1) + x1g, y1g, x2g, y2g = boxes2.unbind(dim=-1) + + # Intersection keypoints + xkis1 = torch.max(x1, x1g) + ykis1 = torch.max(y1, y1g) + xkis2 = torch.min(x2, x2g) + ykis2 = torch.min(y2, y2g) + + intsctk = torch.zeros_like(x1) + mask = (ykis2 > ykis1) & (xkis2 > xkis1) + intsctk[mask] = (xkis2[mask] - xkis1[mask]) * (ykis2[mask] - ykis1[mask]) + unionk = (x2 - x1) * (y2 - y1) + (x2g - x1g) * (y2g - y1g) - intsctk + + return intsctk, unionk diff --git a/lib/python3.12/site-packages/torchvision/ops/boxes.py b/lib/python3.12/site-packages/torchvision/ops/boxes.py new file mode 100644 index 0000000000000000000000000000000000000000..96631278d48ea0d956d5d98a3e3fe1cbc4e51c19 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/ops/boxes.py @@ -0,0 +1,439 @@ +from typing import Tuple + +import torch +import torchvision +from torch import Tensor +from torchvision.extension import _assert_has_ops + +from ..utils import _log_api_usage_once +from ._box_convert import _box_cxcywh_to_xyxy, _box_xywh_to_xyxy, _box_xyxy_to_cxcywh, _box_xyxy_to_xywh +from ._utils import _upcast + + +def nms(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor: + """ + Performs non-maximum suppression (NMS) on the boxes according + to their intersection-over-union (IoU). + + NMS iteratively removes lower scoring boxes which have an + IoU greater than ``iou_threshold`` with another (higher scoring) + box. + + If multiple boxes have the exact same score and satisfy the IoU + criterion with respect to a reference box, the selected box is + not guaranteed to be the same between CPU and GPU. This is similar + to the behavior of argsort in PyTorch when repeated values are present. + + Args: + boxes (Tensor[N, 4])): boxes to perform NMS on. They + are expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and + ``0 <= y1 < y2``. + scores (Tensor[N]): scores for each one of the boxes + iou_threshold (float): discards all overlapping boxes with IoU > iou_threshold + + Returns: + Tensor: int64 tensor with the indices of the elements that have been kept + by NMS, sorted in decreasing order of scores + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(nms) + _assert_has_ops() + return torch.ops.torchvision.nms(boxes, scores, iou_threshold) + + +def batched_nms( + boxes: Tensor, + scores: Tensor, + idxs: Tensor, + iou_threshold: float, +) -> Tensor: + """ + Performs non-maximum suppression in a batched fashion. + + Each index value correspond to a category, and NMS + will not be applied between elements of different categories. + + Args: + boxes (Tensor[N, 4]): boxes where NMS will be performed. They + are expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and + ``0 <= y1 < y2``. + scores (Tensor[N]): scores for each one of the boxes + idxs (Tensor[N]): indices of the categories for each one of the boxes. + iou_threshold (float): discards all overlapping boxes with IoU > iou_threshold + + Returns: + Tensor: int64 tensor with the indices of the elements that have been kept by NMS, sorted + in decreasing order of scores + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(batched_nms) + # Benchmarks that drove the following thresholds are at + # https://github.com/pytorch/vision/issues/1311#issuecomment-781329339 + if boxes.numel() > (4000 if boxes.device.type == "cpu" else 20000) and not torchvision._is_tracing(): + return _batched_nms_vanilla(boxes, scores, idxs, iou_threshold) + else: + return _batched_nms_coordinate_trick(boxes, scores, idxs, iou_threshold) + + +@torch.jit._script_if_tracing +def _batched_nms_coordinate_trick( + boxes: Tensor, + scores: Tensor, + idxs: Tensor, + iou_threshold: float, +) -> Tensor: + # strategy: in order to perform NMS independently per class, + # we add an offset to all the boxes. The offset is dependent + # only on the class idx, and is large enough so that boxes + # from different classes do not overlap + if boxes.numel() == 0: + return torch.empty((0,), dtype=torch.int64, device=boxes.device) + max_coordinate = boxes.max() + offsets = idxs.to(boxes) * (max_coordinate + torch.tensor(1).to(boxes)) + boxes_for_nms = boxes + offsets[:, None] + keep = nms(boxes_for_nms, scores, iou_threshold) + return keep + + +@torch.jit._script_if_tracing +def _batched_nms_vanilla( + boxes: Tensor, + scores: Tensor, + idxs: Tensor, + iou_threshold: float, +) -> Tensor: + # Based on Detectron2 implementation, just manually call nms() on each class independently + keep_mask = torch.zeros_like(scores, dtype=torch.bool) + for class_id in torch.unique(idxs): + curr_indices = torch.where(idxs == class_id)[0] + curr_keep_indices = nms(boxes[curr_indices], scores[curr_indices], iou_threshold) + keep_mask[curr_indices[curr_keep_indices]] = True + keep_indices = torch.where(keep_mask)[0] + return keep_indices[scores[keep_indices].sort(descending=True)[1]] + + +def remove_small_boxes(boxes: Tensor, min_size: float) -> Tensor: + """ + Remove every box from ``boxes`` which contains at least one side length + that is smaller than ``min_size``. + + .. note:: + For sanitizing a :class:`~torchvision.tv_tensors.BoundingBoxes` object, consider using + the transform :func:`~torchvision.transforms.v2.SanitizeBoundingBoxes` instead. + + Args: + boxes (Tensor[N, 4]): boxes in ``(x1, y1, x2, y2)`` format + with ``0 <= x1 < x2`` and ``0 <= y1 < y2``. + min_size (float): minimum size + + Returns: + Tensor[K]: indices of the boxes that have both sides + larger than ``min_size`` + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(remove_small_boxes) + ws, hs = boxes[:, 2] - boxes[:, 0], boxes[:, 3] - boxes[:, 1] + keep = (ws >= min_size) & (hs >= min_size) + keep = torch.where(keep)[0] + return keep + + +def clip_boxes_to_image(boxes: Tensor, size: Tuple[int, int]) -> Tensor: + """ + Clip boxes so that they lie inside an image of size ``size``. + + .. note:: + For clipping a :class:`~torchvision.tv_tensors.BoundingBoxes` object, consider using + the transform :func:`~torchvision.transforms.v2.ClampBoundingBoxes` instead. + + Args: + boxes (Tensor[N, 4]): boxes in ``(x1, y1, x2, y2)`` format + with ``0 <= x1 < x2`` and ``0 <= y1 < y2``. + size (Tuple[height, width]): size of the image + + Returns: + Tensor[N, 4]: clipped boxes + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(clip_boxes_to_image) + dim = boxes.dim() + boxes_x = boxes[..., 0::2] + boxes_y = boxes[..., 1::2] + height, width = size + + if torchvision._is_tracing(): + boxes_x = torch.max(boxes_x, torch.tensor(0, dtype=boxes.dtype, device=boxes.device)) + boxes_x = torch.min(boxes_x, torch.tensor(width, dtype=boxes.dtype, device=boxes.device)) + boxes_y = torch.max(boxes_y, torch.tensor(0, dtype=boxes.dtype, device=boxes.device)) + boxes_y = torch.min(boxes_y, torch.tensor(height, dtype=boxes.dtype, device=boxes.device)) + else: + boxes_x = boxes_x.clamp(min=0, max=width) + boxes_y = boxes_y.clamp(min=0, max=height) + + clipped_boxes = torch.stack((boxes_x, boxes_y), dim=dim) + return clipped_boxes.reshape(boxes.shape) + + +def box_convert(boxes: Tensor, in_fmt: str, out_fmt: str) -> Tensor: + """ + Converts :class:`torch.Tensor` boxes from a given ``in_fmt`` to ``out_fmt``. + + .. note:: + For converting a :class:`torch.Tensor` or a :class:`~torchvision.tv_tensors.BoundingBoxes` object + between different formats, + consider using :func:`~torchvision.transforms.v2.functional.convert_bounding_box_format` instead. + Or see the corresponding transform :func:`~torchvision.transforms.v2.ConvertBoundingBoxFormat`. + + Supported ``in_fmt`` and ``out_fmt`` strings are: + + ``'xyxy'``: boxes are represented via corners, x1, y1 being top left and x2, y2 being bottom right. + This is the format that torchvision utilities expect. + + ``'xywh'``: boxes are represented via corner, width and height, x1, y2 being top left, w, h being width and height. + + ``'cxcywh'``: boxes are represented via centre, width and height, cx, cy being center of box, w, h + being width and height. + + Args: + boxes (Tensor[N, 4]): boxes which will be converted. + in_fmt (str): Input format of given boxes. Supported formats are ['xyxy', 'xywh', 'cxcywh']. + out_fmt (str): Output format of given boxes. Supported formats are ['xyxy', 'xywh', 'cxcywh'] + + Returns: + Tensor[N, 4]: Boxes into converted format. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(box_convert) + allowed_fmts = ("xyxy", "xywh", "cxcywh") + if in_fmt not in allowed_fmts or out_fmt not in allowed_fmts: + raise ValueError("Unsupported Bounding Box Conversions for given in_fmt and out_fmt") + + if in_fmt == out_fmt: + return boxes.clone() + + if in_fmt != "xyxy" and out_fmt != "xyxy": + # convert to xyxy and change in_fmt xyxy + if in_fmt == "xywh": + boxes = _box_xywh_to_xyxy(boxes) + elif in_fmt == "cxcywh": + boxes = _box_cxcywh_to_xyxy(boxes) + in_fmt = "xyxy" + + if in_fmt == "xyxy": + if out_fmt == "xywh": + boxes = _box_xyxy_to_xywh(boxes) + elif out_fmt == "cxcywh": + boxes = _box_xyxy_to_cxcywh(boxes) + elif out_fmt == "xyxy": + if in_fmt == "xywh": + boxes = _box_xywh_to_xyxy(boxes) + elif in_fmt == "cxcywh": + boxes = _box_cxcywh_to_xyxy(boxes) + return boxes + + +def box_area(boxes: Tensor) -> Tensor: + """ + Computes the area of a set of bounding boxes, which are specified by their + (x1, y1, x2, y2) coordinates. + + Args: + boxes (Tensor[N, 4]): boxes for which the area will be computed. They + are expected to be in (x1, y1, x2, y2) format with + ``0 <= x1 < x2`` and ``0 <= y1 < y2``. + + Returns: + Tensor[N]: the area for each box + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(box_area) + boxes = _upcast(boxes) + return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) + + +# implementation from https://github.com/kuangliu/torchcv/blob/master/torchcv/utils/box.py +# with slight modifications +def _box_inter_union(boxes1: Tensor, boxes2: Tensor) -> Tuple[Tensor, Tensor]: + area1 = box_area(boxes1) + area2 = box_area(boxes2) + + lt = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2] + rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2] + + wh = _upcast(rb - lt).clamp(min=0) # [N,M,2] + inter = wh[:, :, 0] * wh[:, :, 1] # [N,M] + + union = area1[:, None] + area2 - inter + + return inter, union + + +def box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor: + """ + Return intersection-over-union (Jaccard index) between two sets of boxes. + + Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with + ``0 <= x1 < x2`` and ``0 <= y1 < y2``. + + Args: + boxes1 (Tensor[N, 4]): first set of boxes + boxes2 (Tensor[M, 4]): second set of boxes + + Returns: + Tensor[N, M]: the NxM matrix containing the pairwise IoU values for every element in boxes1 and boxes2 + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(box_iou) + inter, union = _box_inter_union(boxes1, boxes2) + iou = inter / union + return iou + + +# Implementation adapted from https://github.com/facebookresearch/detr/blob/master/util/box_ops.py +def generalized_box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor: + """ + Return generalized intersection-over-union (Jaccard index) between two sets of boxes. + + Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with + ``0 <= x1 < x2`` and ``0 <= y1 < y2``. + + Args: + boxes1 (Tensor[N, 4]): first set of boxes + boxes2 (Tensor[M, 4]): second set of boxes + + Returns: + Tensor[N, M]: the NxM matrix containing the pairwise generalized IoU values + for every element in boxes1 and boxes2 + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(generalized_box_iou) + + inter, union = _box_inter_union(boxes1, boxes2) + iou = inter / union + + lti = torch.min(boxes1[:, None, :2], boxes2[:, :2]) + rbi = torch.max(boxes1[:, None, 2:], boxes2[:, 2:]) + + whi = _upcast(rbi - lti).clamp(min=0) # [N,M,2] + areai = whi[:, :, 0] * whi[:, :, 1] + + return iou - (areai - union) / areai + + +def complete_box_iou(boxes1: Tensor, boxes2: Tensor, eps: float = 1e-7) -> Tensor: + """ + Return complete intersection-over-union (Jaccard index) between two sets of boxes. + Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with + ``0 <= x1 < x2`` and ``0 <= y1 < y2``. + Args: + boxes1 (Tensor[N, 4]): first set of boxes + boxes2 (Tensor[M, 4]): second set of boxes + eps (float, optional): small number to prevent division by zero. Default: 1e-7 + Returns: + Tensor[N, M]: the NxM matrix containing the pairwise complete IoU values + for every element in boxes1 and boxes2 + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(complete_box_iou) + + boxes1 = _upcast(boxes1) + boxes2 = _upcast(boxes2) + + diou, iou = _box_diou_iou(boxes1, boxes2, eps) + + w_pred = boxes1[:, None, 2] - boxes1[:, None, 0] + h_pred = boxes1[:, None, 3] - boxes1[:, None, 1] + + w_gt = boxes2[:, 2] - boxes2[:, 0] + h_gt = boxes2[:, 3] - boxes2[:, 1] + + v = (4 / (torch.pi**2)) * torch.pow(torch.atan(w_pred / h_pred) - torch.atan(w_gt / h_gt), 2) + with torch.no_grad(): + alpha = v / (1 - iou + v + eps) + return diou - alpha * v + + +def distance_box_iou(boxes1: Tensor, boxes2: Tensor, eps: float = 1e-7) -> Tensor: + """ + Return distance intersection-over-union (Jaccard index) between two sets of boxes. + + Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with + ``0 <= x1 < x2`` and ``0 <= y1 < y2``. + + Args: + boxes1 (Tensor[N, 4]): first set of boxes + boxes2 (Tensor[M, 4]): second set of boxes + eps (float, optional): small number to prevent division by zero. Default: 1e-7 + + Returns: + Tensor[N, M]: the NxM matrix containing the pairwise distance IoU values + for every element in boxes1 and boxes2 + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(distance_box_iou) + + boxes1 = _upcast(boxes1) + boxes2 = _upcast(boxes2) + diou, _ = _box_diou_iou(boxes1, boxes2, eps=eps) + return diou + + +def _box_diou_iou(boxes1: Tensor, boxes2: Tensor, eps: float = 1e-7) -> Tuple[Tensor, Tensor]: + + iou = box_iou(boxes1, boxes2) + lti = torch.min(boxes1[:, None, :2], boxes2[:, :2]) + rbi = torch.max(boxes1[:, None, 2:], boxes2[:, 2:]) + whi = _upcast(rbi - lti).clamp(min=0) # [N,M,2] + diagonal_distance_squared = (whi[:, :, 0] ** 2) + (whi[:, :, 1] ** 2) + eps + # centers of boxes + x_p = (boxes1[:, 0] + boxes1[:, 2]) / 2 + y_p = (boxes1[:, 1] + boxes1[:, 3]) / 2 + x_g = (boxes2[:, 0] + boxes2[:, 2]) / 2 + y_g = (boxes2[:, 1] + boxes2[:, 3]) / 2 + # The distance between boxes' centers squared. + centers_distance_squared = (_upcast((x_p[:, None] - x_g[None, :])) ** 2) + ( + _upcast((y_p[:, None] - y_g[None, :])) ** 2 + ) + # The distance IoU is the IoU penalized by a normalized + # distance between boxes' centers squared. + return iou - (centers_distance_squared / diagonal_distance_squared), iou + + +def masks_to_boxes(masks: torch.Tensor) -> torch.Tensor: + """ + Compute the bounding boxes around the provided masks. + + Returns a [N, 4] tensor containing bounding boxes. The boxes are in ``(x1, y1, x2, y2)`` format with + ``0 <= x1 <= x2`` and ``0 <= y1 <= y2``. + + .. warning:: + + In most cases the output will guarantee ``x1 < x2`` and ``y1 < y2``. But + if the input is degenerate, e.g. if a mask is a single row or a single + column, then the output may have x1 = x2 or y1 = y2. + + Args: + masks (Tensor[N, H, W]): masks to transform where N is the number of masks + and (H, W) are the spatial dimensions. + + Returns: + Tensor[N, 4]: bounding boxes + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(masks_to_boxes) + if masks.numel() == 0: + return torch.zeros((0, 4), device=masks.device, dtype=torch.float) + + n = masks.shape[0] + + bounding_boxes = torch.zeros((n, 4), device=masks.device, dtype=torch.float) + + for index, mask in enumerate(masks): + y, x = torch.where(mask != 0) + + bounding_boxes[index, 0] = torch.min(x) + bounding_boxes[index, 1] = torch.min(y) + bounding_boxes[index, 2] = torch.max(x) + bounding_boxes[index, 3] = torch.max(y) + + return bounding_boxes diff --git a/lib/python3.12/site-packages/torchvision/ops/ciou_loss.py b/lib/python3.12/site-packages/torchvision/ops/ciou_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..75a1c4cb1f39f7b8ff33bdac3641ca71b6413d20 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/ops/ciou_loss.py @@ -0,0 +1,78 @@ +import torch + +from ..utils import _log_api_usage_once +from ._utils import _upcast_non_float +from .diou_loss import _diou_iou_loss + + +def complete_box_iou_loss( + boxes1: torch.Tensor, + boxes2: torch.Tensor, + reduction: str = "none", + eps: float = 1e-7, +) -> torch.Tensor: + + """ + Gradient-friendly IoU loss with an additional penalty that is non-zero when the + boxes do not overlap. This loss function considers important geometrical + factors such as overlap area, normalized central point distance and aspect ratio. + This loss is symmetric, so the boxes1 and boxes2 arguments are interchangeable. + + Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with + ``0 <= x1 < x2`` and ``0 <= y1 < y2``, and The two boxes should have the + same dimensions. + + Args: + boxes1 : (Tensor[N, 4] or Tensor[4]) first set of boxes + boxes2 : (Tensor[N, 4] or Tensor[4]) second set of boxes + reduction : (string, optional) Specifies the reduction to apply to the output: + ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: No reduction will be + applied to the output. ``'mean'``: The output will be averaged. + ``'sum'``: The output will be summed. Default: ``'none'`` + eps : (float): small number to prevent division by zero. Default: 1e-7 + + Returns: + Tensor: Loss tensor with the reduction option applied. + + Reference: + Zhaohui Zheng et al.: Complete Intersection over Union Loss: + https://arxiv.org/abs/1911.08287 + + """ + + # Original Implementation from https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/losses.py + + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(complete_box_iou_loss) + + boxes1 = _upcast_non_float(boxes1) + boxes2 = _upcast_non_float(boxes2) + + diou_loss, iou = _diou_iou_loss(boxes1, boxes2) + + x1, y1, x2, y2 = boxes1.unbind(dim=-1) + x1g, y1g, x2g, y2g = boxes2.unbind(dim=-1) + + # width and height of boxes + w_pred = x2 - x1 + h_pred = y2 - y1 + w_gt = x2g - x1g + h_gt = y2g - y1g + v = (4 / (torch.pi**2)) * torch.pow((torch.atan(w_gt / h_gt) - torch.atan(w_pred / h_pred)), 2) + with torch.no_grad(): + alpha = v / (1 - iou + v + eps) + + loss = diou_loss + alpha * v + + # Check reduction option and return loss accordingly + if reduction == "none": + pass + elif reduction == "mean": + loss = loss.mean() if loss.numel() > 0 else 0.0 * loss.sum() + elif reduction == "sum": + loss = loss.sum() + else: + raise ValueError( + f"Invalid Value for arg 'reduction': '{reduction} \n Supported reduction modes: 'none', 'mean', 'sum'" + ) + return loss diff --git a/lib/python3.12/site-packages/torchvision/ops/deform_conv.py b/lib/python3.12/site-packages/torchvision/ops/deform_conv.py new file mode 100644 index 0000000000000000000000000000000000000000..b3cc83332a0a780710cdd61f0930f81ca9ab18b9 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/ops/deform_conv.py @@ -0,0 +1,195 @@ +import math +from typing import Optional, Tuple + +import torch +from torch import nn, Tensor +from torch.nn import init +from torch.nn.modules.utils import _pair +from torch.nn.parameter import Parameter +from torchvision.extension import _assert_has_ops + +from ..utils import _log_api_usage_once + + +def deform_conv2d( + input: Tensor, + offset: Tensor, + weight: Tensor, + bias: Optional[Tensor] = None, + stride: Tuple[int, int] = (1, 1), + padding: Tuple[int, int] = (0, 0), + dilation: Tuple[int, int] = (1, 1), + mask: Optional[Tensor] = None, +) -> Tensor: + r""" + Performs Deformable Convolution v2, described in + `Deformable ConvNets v2: More Deformable, Better Results + `__ if :attr:`mask` is not ``None`` and + Performs Deformable Convolution, described in + `Deformable Convolutional Networks + `__ if :attr:`mask` is ``None``. + + Args: + input (Tensor[batch_size, in_channels, in_height, in_width]): input tensor + offset (Tensor[batch_size, 2 * offset_groups * kernel_height * kernel_width, out_height, out_width]): + offsets to be applied for each position in the convolution kernel. + weight (Tensor[out_channels, in_channels // groups, kernel_height, kernel_width]): convolution weights, + split into groups of size (in_channels // groups) + bias (Tensor[out_channels]): optional bias of shape (out_channels,). Default: None + stride (int or Tuple[int, int]): distance between convolution centers. Default: 1 + padding (int or Tuple[int, int]): height/width of padding of zeroes around + each image. Default: 0 + dilation (int or Tuple[int, int]): the spacing between kernel elements. Default: 1 + mask (Tensor[batch_size, offset_groups * kernel_height * kernel_width, out_height, out_width]): + masks to be applied for each position in the convolution kernel. Default: None + + Returns: + Tensor[batch_sz, out_channels, out_h, out_w]: result of convolution + + Examples:: + >>> input = torch.rand(4, 3, 10, 10) + >>> kh, kw = 3, 3 + >>> weight = torch.rand(5, 3, kh, kw) + >>> # offset and mask should have the same spatial size as the output + >>> # of the convolution. In this case, for an input of 10, stride of 1 + >>> # and kernel size of 3, without padding, the output size is 8 + >>> offset = torch.rand(4, 2 * kh * kw, 8, 8) + >>> mask = torch.rand(4, kh * kw, 8, 8) + >>> out = deform_conv2d(input, offset, weight, mask=mask) + >>> print(out.shape) + >>> # returns + >>> torch.Size([4, 5, 8, 8]) + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(deform_conv2d) + _assert_has_ops() + out_channels = weight.shape[0] + + use_mask = mask is not None + + if mask is None: + mask = torch.zeros((input.shape[0], 1), device=input.device, dtype=input.dtype) + + if bias is None: + bias = torch.zeros(out_channels, device=input.device, dtype=input.dtype) + + stride_h, stride_w = _pair(stride) + pad_h, pad_w = _pair(padding) + dil_h, dil_w = _pair(dilation) + weights_h, weights_w = weight.shape[-2:] + _, n_in_channels, _, _ = input.shape + + n_offset_grps = offset.shape[1] // (2 * weights_h * weights_w) + n_weight_grps = n_in_channels // weight.shape[1] + + if n_offset_grps == 0: + raise RuntimeError( + "the shape of the offset tensor at dimension 1 is not valid. It should " + "be a multiple of 2 * weight.size[2] * weight.size[3].\n" + f"Got offset.shape[1]={offset.shape[1]}, while 2 * weight.size[2] * weight.size[3]={2 * weights_h * weights_w}" + ) + + return torch.ops.torchvision.deform_conv2d( + input, + weight, + offset, + mask, + bias, + stride_h, + stride_w, + pad_h, + pad_w, + dil_h, + dil_w, + n_weight_grps, + n_offset_grps, + use_mask, + ) + + +class DeformConv2d(nn.Module): + """ + See :func:`deform_conv2d`. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + padding: int = 0, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + ): + super().__init__() + _log_api_usage_once(self) + + if in_channels % groups != 0: + raise ValueError("in_channels must be divisible by groups") + if out_channels % groups != 0: + raise ValueError("out_channels must be divisible by groups") + + self.in_channels = in_channels + self.out_channels = out_channels + self.kernel_size = _pair(kernel_size) + self.stride = _pair(stride) + self.padding = _pair(padding) + self.dilation = _pair(dilation) + self.groups = groups + + self.weight = Parameter( + torch.empty(out_channels, in_channels // groups, self.kernel_size[0], self.kernel_size[1]) + ) + + if bias: + self.bias = Parameter(torch.empty(out_channels)) + else: + self.register_parameter("bias", None) + + self.reset_parameters() + + def reset_parameters(self) -> None: + init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + + if self.bias is not None: + fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight) + bound = 1 / math.sqrt(fan_in) + init.uniform_(self.bias, -bound, bound) + + def forward(self, input: Tensor, offset: Tensor, mask: Optional[Tensor] = None) -> Tensor: + """ + Args: + input (Tensor[batch_size, in_channels, in_height, in_width]): input tensor + offset (Tensor[batch_size, 2 * offset_groups * kernel_height * kernel_width, out_height, out_width]): + offsets to be applied for each position in the convolution kernel. + mask (Tensor[batch_size, offset_groups * kernel_height * kernel_width, out_height, out_width]): + masks to be applied for each position in the convolution kernel. + """ + return deform_conv2d( + input, + offset, + self.weight, + self.bias, + stride=self.stride, + padding=self.padding, + dilation=self.dilation, + mask=mask, + ) + + def __repr__(self) -> str: + s = ( + f"{self.__class__.__name__}(" + f"{self.in_channels}" + f", {self.out_channels}" + f", kernel_size={self.kernel_size}" + f", stride={self.stride}" + ) + s += f", padding={self.padding}" if self.padding != (0, 0) else "" + s += f", dilation={self.dilation}" if self.dilation != (1, 1) else "" + s += f", groups={self.groups}" if self.groups != 1 else "" + s += ", bias=False" if self.bias is None else "" + s += ")" + + return s diff --git a/lib/python3.12/site-packages/torchvision/ops/diou_loss.py b/lib/python3.12/site-packages/torchvision/ops/diou_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..c64c6673a882f439f8ba9b0e25d1387d8d2f4284 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/ops/diou_loss.py @@ -0,0 +1,94 @@ +from typing import Tuple + +import torch + +from ..utils import _log_api_usage_once +from ._utils import _loss_inter_union, _upcast_non_float + + +def distance_box_iou_loss( + boxes1: torch.Tensor, + boxes2: torch.Tensor, + reduction: str = "none", + eps: float = 1e-7, +) -> torch.Tensor: + + """ + Gradient-friendly IoU loss with an additional penalty that is non-zero when the + distance between boxes' centers isn't zero. Indeed, for two exactly overlapping + boxes, the distance IoU is the same as the IoU loss. + This loss is symmetric, so the boxes1 and boxes2 arguments are interchangeable. + + Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with + ``0 <= x1 < x2`` and ``0 <= y1 < y2``, and The two boxes should have the + same dimensions. + + Args: + boxes1 (Tensor[N, 4]): first set of boxes + boxes2 (Tensor[N, 4]): second set of boxes + reduction (string, optional): Specifies the reduction to apply to the output: + ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: No reduction will be + applied to the output. ``'mean'``: The output will be averaged. + ``'sum'``: The output will be summed. Default: ``'none'`` + eps (float, optional): small number to prevent division by zero. Default: 1e-7 + + Returns: + Tensor: Loss tensor with the reduction option applied. + + Reference: + Zhaohui Zheng et al.: Distance Intersection over Union Loss: + https://arxiv.org/abs/1911.08287 + """ + + # Original Implementation from https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/losses.py + + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(distance_box_iou_loss) + + boxes1 = _upcast_non_float(boxes1) + boxes2 = _upcast_non_float(boxes2) + + loss, _ = _diou_iou_loss(boxes1, boxes2, eps) + + # Check reduction option and return loss accordingly + if reduction == "none": + pass + elif reduction == "mean": + loss = loss.mean() if loss.numel() > 0 else 0.0 * loss.sum() + elif reduction == "sum": + loss = loss.sum() + else: + raise ValueError( + f"Invalid Value for arg 'reduction': '{reduction} \n Supported reduction modes: 'none', 'mean', 'sum'" + ) + return loss + + +def _diou_iou_loss( + boxes1: torch.Tensor, + boxes2: torch.Tensor, + eps: float = 1e-7, +) -> Tuple[torch.Tensor, torch.Tensor]: + + intsct, union = _loss_inter_union(boxes1, boxes2) + iou = intsct / (union + eps) + # smallest enclosing box + x1, y1, x2, y2 = boxes1.unbind(dim=-1) + x1g, y1g, x2g, y2g = boxes2.unbind(dim=-1) + xc1 = torch.min(x1, x1g) + yc1 = torch.min(y1, y1g) + xc2 = torch.max(x2, x2g) + yc2 = torch.max(y2, y2g) + # The diagonal distance of the smallest enclosing box squared + diagonal_distance_squared = ((xc2 - xc1) ** 2) + ((yc2 - yc1) ** 2) + eps + # centers of boxes + x_p = (x2 + x1) / 2 + y_p = (y2 + y1) / 2 + x_g = (x1g + x2g) / 2 + y_g = (y1g + y2g) / 2 + # The distance between boxes' centers squared. + centers_distance_squared = ((x_p - x_g) ** 2) + ((y_p - y_g) ** 2) + # The distance IoU is the IoU penalized by a normalized + # distance between boxes' centers squared. + loss = 1 - iou + (centers_distance_squared / diagonal_distance_squared) + return loss, iou diff --git a/lib/python3.12/site-packages/torchvision/ops/drop_block.py b/lib/python3.12/site-packages/torchvision/ops/drop_block.py new file mode 100644 index 0000000000000000000000000000000000000000..e65496ea29ad1ee6e04bc19280a00a397a4f1040 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/ops/drop_block.py @@ -0,0 +1,155 @@ +import torch +import torch.fx +import torch.nn.functional as F +from torch import nn, Tensor + +from ..utils import _log_api_usage_once + + +def drop_block2d( + input: Tensor, p: float, block_size: int, inplace: bool = False, eps: float = 1e-06, training: bool = True +) -> Tensor: + """ + Implements DropBlock2d from `"DropBlock: A regularization method for convolutional networks" + `. + + Args: + input (Tensor[N, C, H, W]): The input tensor or 4-dimensions with the first one + being its batch i.e. a batch with ``N`` rows. + p (float): Probability of an element to be dropped. + block_size (int): Size of the block to drop. + inplace (bool): If set to ``True``, will do this operation in-place. Default: ``False``. + eps (float): A value added to the denominator for numerical stability. Default: 1e-6. + training (bool): apply dropblock if is ``True``. Default: ``True``. + + Returns: + Tensor[N, C, H, W]: The randomly zeroed tensor after dropblock. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(drop_block2d) + if p < 0.0 or p > 1.0: + raise ValueError(f"drop probability has to be between 0 and 1, but got {p}.") + if input.ndim != 4: + raise ValueError(f"input should be 4 dimensional. Got {input.ndim} dimensions.") + if not training or p == 0.0: + return input + + N, C, H, W = input.size() + block_size = min(block_size, W, H) + # compute the gamma of Bernoulli distribution + gamma = (p * H * W) / ((block_size**2) * ((H - block_size + 1) * (W - block_size + 1))) + noise = torch.empty((N, C, H - block_size + 1, W - block_size + 1), dtype=input.dtype, device=input.device) + noise.bernoulli_(gamma) + + noise = F.pad(noise, [block_size // 2] * 4, value=0) + noise = F.max_pool2d(noise, stride=(1, 1), kernel_size=(block_size, block_size), padding=block_size // 2) + noise = 1 - noise + normalize_scale = noise.numel() / (eps + noise.sum()) + if inplace: + input.mul_(noise).mul_(normalize_scale) + else: + input = input * noise * normalize_scale + return input + + +def drop_block3d( + input: Tensor, p: float, block_size: int, inplace: bool = False, eps: float = 1e-06, training: bool = True +) -> Tensor: + """ + Implements DropBlock3d from `"DropBlock: A regularization method for convolutional networks" + `. + + Args: + input (Tensor[N, C, D, H, W]): The input tensor or 5-dimensions with the first one + being its batch i.e. a batch with ``N`` rows. + p (float): Probability of an element to be dropped. + block_size (int): Size of the block to drop. + inplace (bool): If set to ``True``, will do this operation in-place. Default: ``False``. + eps (float): A value added to the denominator for numerical stability. Default: 1e-6. + training (bool): apply dropblock if is ``True``. Default: ``True``. + + Returns: + Tensor[N, C, D, H, W]: The randomly zeroed tensor after dropblock. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(drop_block3d) + if p < 0.0 or p > 1.0: + raise ValueError(f"drop probability has to be between 0 and 1, but got {p}.") + if input.ndim != 5: + raise ValueError(f"input should be 5 dimensional. Got {input.ndim} dimensions.") + if not training or p == 0.0: + return input + + N, C, D, H, W = input.size() + block_size = min(block_size, D, H, W) + # compute the gamma of Bernoulli distribution + gamma = (p * D * H * W) / ((block_size**3) * ((D - block_size + 1) * (H - block_size + 1) * (W - block_size + 1))) + noise = torch.empty( + (N, C, D - block_size + 1, H - block_size + 1, W - block_size + 1), dtype=input.dtype, device=input.device + ) + noise.bernoulli_(gamma) + + noise = F.pad(noise, [block_size // 2] * 6, value=0) + noise = F.max_pool3d( + noise, stride=(1, 1, 1), kernel_size=(block_size, block_size, block_size), padding=block_size // 2 + ) + noise = 1 - noise + normalize_scale = noise.numel() / (eps + noise.sum()) + if inplace: + input.mul_(noise).mul_(normalize_scale) + else: + input = input * noise * normalize_scale + return input + + +torch.fx.wrap("drop_block2d") + + +class DropBlock2d(nn.Module): + """ + See :func:`drop_block2d`. + """ + + def __init__(self, p: float, block_size: int, inplace: bool = False, eps: float = 1e-06) -> None: + super().__init__() + + self.p = p + self.block_size = block_size + self.inplace = inplace + self.eps = eps + + def forward(self, input: Tensor) -> Tensor: + """ + Args: + input (Tensor): Input feature map on which some areas will be randomly + dropped. + Returns: + Tensor: The tensor after DropBlock layer. + """ + return drop_block2d(input, self.p, self.block_size, self.inplace, self.eps, self.training) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}(p={self.p}, block_size={self.block_size}, inplace={self.inplace})" + return s + + +torch.fx.wrap("drop_block3d") + + +class DropBlock3d(DropBlock2d): + """ + See :func:`drop_block3d`. + """ + + def __init__(self, p: float, block_size: int, inplace: bool = False, eps: float = 1e-06) -> None: + super().__init__(p, block_size, inplace, eps) + + def forward(self, input: Tensor) -> Tensor: + """ + Args: + input (Tensor): Input feature map on which some areas will be randomly + dropped. + Returns: + Tensor: The tensor after DropBlock layer. + """ + return drop_block3d(input, self.p, self.block_size, self.inplace, self.eps, self.training) diff --git a/lib/python3.12/site-packages/torchvision/ops/feature_pyramid_network.py b/lib/python3.12/site-packages/torchvision/ops/feature_pyramid_network.py new file mode 100644 index 0000000000000000000000000000000000000000..2e7aef0e2fad773ac884d576f781a305e5e880d0 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/ops/feature_pyramid_network.py @@ -0,0 +1,250 @@ +from collections import OrderedDict +from typing import Callable, Dict, List, Optional, Tuple + +import torch.nn.functional as F +from torch import nn, Tensor + +from ..ops.misc import Conv2dNormActivation +from ..utils import _log_api_usage_once + + +class ExtraFPNBlock(nn.Module): + """ + Base class for the extra block in the FPN. + + Args: + results (List[Tensor]): the result of the FPN + x (List[Tensor]): the original feature maps + names (List[str]): the names for each one of the + original feature maps + + Returns: + results (List[Tensor]): the extended set of results + of the FPN + names (List[str]): the extended set of names for the results + """ + + def forward( + self, + results: List[Tensor], + x: List[Tensor], + names: List[str], + ) -> Tuple[List[Tensor], List[str]]: + pass + + +class FeaturePyramidNetwork(nn.Module): + """ + Module that adds a FPN from on top of a set of feature maps. This is based on + `"Feature Pyramid Network for Object Detection" `_. + + The feature maps are currently supposed to be in increasing depth + order. + + The input to the model is expected to be an OrderedDict[Tensor], containing + the feature maps on top of which the FPN will be added. + + Args: + in_channels_list (list[int]): number of channels for each feature map that + is passed to the module + out_channels (int): number of channels of the FPN representation + extra_blocks (ExtraFPNBlock or None): if provided, extra operations will + be performed. It is expected to take the fpn features, the original + features and the names of the original features as input, and returns + a new list of feature maps and their corresponding names + norm_layer (callable, optional): Module specifying the normalization layer to use. Default: None + + Examples:: + + >>> m = torchvision.ops.FeaturePyramidNetwork([10, 20, 30], 5) + >>> # get some dummy data + >>> x = OrderedDict() + >>> x['feat0'] = torch.rand(1, 10, 64, 64) + >>> x['feat2'] = torch.rand(1, 20, 16, 16) + >>> x['feat3'] = torch.rand(1, 30, 8, 8) + >>> # compute the FPN on top of x + >>> output = m(x) + >>> print([(k, v.shape) for k, v in output.items()]) + >>> # returns + >>> [('feat0', torch.Size([1, 5, 64, 64])), + >>> ('feat2', torch.Size([1, 5, 16, 16])), + >>> ('feat3', torch.Size([1, 5, 8, 8]))] + + """ + + _version = 2 + + def __init__( + self, + in_channels_list: List[int], + out_channels: int, + extra_blocks: Optional[ExtraFPNBlock] = None, + norm_layer: Optional[Callable[..., nn.Module]] = None, + ): + super().__init__() + _log_api_usage_once(self) + self.inner_blocks = nn.ModuleList() + self.layer_blocks = nn.ModuleList() + for in_channels in in_channels_list: + if in_channels == 0: + raise ValueError("in_channels=0 is currently not supported") + inner_block_module = Conv2dNormActivation( + in_channels, out_channels, kernel_size=1, padding=0, norm_layer=norm_layer, activation_layer=None + ) + layer_block_module = Conv2dNormActivation( + out_channels, out_channels, kernel_size=3, norm_layer=norm_layer, activation_layer=None + ) + self.inner_blocks.append(inner_block_module) + self.layer_blocks.append(layer_block_module) + + # initialize parameters now to avoid modifying the initialization of top_blocks + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_uniform_(m.weight, a=1) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + if extra_blocks is not None: + if not isinstance(extra_blocks, ExtraFPNBlock): + raise TypeError(f"extra_blocks should be of type ExtraFPNBlock not {type(extra_blocks)}") + self.extra_blocks = extra_blocks + + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + version = local_metadata.get("version", None) + + if version is None or version < 2: + num_blocks = len(self.inner_blocks) + for block in ["inner_blocks", "layer_blocks"]: + for i in range(num_blocks): + for type in ["weight", "bias"]: + old_key = f"{prefix}{block}.{i}.{type}" + new_key = f"{prefix}{block}.{i}.0.{type}" + if old_key in state_dict: + state_dict[new_key] = state_dict.pop(old_key) + + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + def get_result_from_inner_blocks(self, x: Tensor, idx: int) -> Tensor: + """ + This is equivalent to self.inner_blocks[idx](x), + but torchscript doesn't support this yet + """ + num_blocks = len(self.inner_blocks) + if idx < 0: + idx += num_blocks + out = x + for i, module in enumerate(self.inner_blocks): + if i == idx: + out = module(x) + return out + + def get_result_from_layer_blocks(self, x: Tensor, idx: int) -> Tensor: + """ + This is equivalent to self.layer_blocks[idx](x), + but torchscript doesn't support this yet + """ + num_blocks = len(self.layer_blocks) + if idx < 0: + idx += num_blocks + out = x + for i, module in enumerate(self.layer_blocks): + if i == idx: + out = module(x) + return out + + def forward(self, x: Dict[str, Tensor]) -> Dict[str, Tensor]: + """ + Computes the FPN for a set of feature maps. + + Args: + x (OrderedDict[Tensor]): feature maps for each feature level. + + Returns: + results (OrderedDict[Tensor]): feature maps after FPN layers. + They are ordered from the highest resolution first. + """ + # unpack OrderedDict into two lists for easier handling + names = list(x.keys()) + x = list(x.values()) + + last_inner = self.get_result_from_inner_blocks(x[-1], -1) + results = [] + results.append(self.get_result_from_layer_blocks(last_inner, -1)) + + for idx in range(len(x) - 2, -1, -1): + inner_lateral = self.get_result_from_inner_blocks(x[idx], idx) + feat_shape = inner_lateral.shape[-2:] + inner_top_down = F.interpolate(last_inner, size=feat_shape, mode="nearest") + last_inner = inner_lateral + inner_top_down + results.insert(0, self.get_result_from_layer_blocks(last_inner, idx)) + + if self.extra_blocks is not None: + results, names = self.extra_blocks(results, x, names) + + # make it back an OrderedDict + out = OrderedDict([(k, v) for k, v in zip(names, results)]) + + return out + + +class LastLevelMaxPool(ExtraFPNBlock): + """ + Applies a max_pool2d (not actual max_pool2d, we just subsample) on top of the last feature map + """ + + def forward( + self, + x: List[Tensor], + y: List[Tensor], + names: List[str], + ) -> Tuple[List[Tensor], List[str]]: + names.append("pool") + # Use max pooling to simulate stride 2 subsampling + x.append(F.max_pool2d(x[-1], kernel_size=1, stride=2, padding=0)) + return x, names + + +class LastLevelP6P7(ExtraFPNBlock): + """ + This module is used in RetinaNet to generate extra layers, P6 and P7. + """ + + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.p6 = nn.Conv2d(in_channels, out_channels, 3, 2, 1) + self.p7 = nn.Conv2d(out_channels, out_channels, 3, 2, 1) + for module in [self.p6, self.p7]: + nn.init.kaiming_uniform_(module.weight, a=1) + nn.init.constant_(module.bias, 0) + self.use_P5 = in_channels == out_channels + + def forward( + self, + p: List[Tensor], + c: List[Tensor], + names: List[str], + ) -> Tuple[List[Tensor], List[str]]: + p5, c5 = p[-1], c[-1] + x = p5 if self.use_P5 else c5 + p6 = self.p6(x) + p7 = self.p7(F.relu(p6)) + p.extend([p6, p7]) + names.extend(["p6", "p7"]) + return p, names diff --git a/lib/python3.12/site-packages/torchvision/ops/focal_loss.py b/lib/python3.12/site-packages/torchvision/ops/focal_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..08c282555fca01cbee78d79c2a0fa7bb8a9df570 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/ops/focal_loss.py @@ -0,0 +1,58 @@ +import torch +import torch.nn.functional as F + +from ..utils import _log_api_usage_once + + +def sigmoid_focal_loss( + inputs: torch.Tensor, + targets: torch.Tensor, + alpha: float = 0.25, + gamma: float = 2, + reduction: str = "none", +) -> torch.Tensor: + """ + Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002. + + Args: + inputs (Tensor): A float tensor of arbitrary shape. + The predictions for each example. + targets (Tensor): A float tensor with the same shape as inputs. Stores the binary + classification label for each element in inputs + (0 for the negative class and 1 for the positive class). + alpha (float): Weighting factor in range (0,1) to balance + positive vs negative examples or -1 for ignore. Default: ``0.25``. + gamma (float): Exponent of the modulating factor (1 - p_t) to + balance easy vs hard examples. Default: ``2``. + reduction (string): ``'none'`` | ``'mean'`` | ``'sum'`` + ``'none'``: No reduction will be applied to the output. + ``'mean'``: The output will be averaged. + ``'sum'``: The output will be summed. Default: ``'none'``. + Returns: + Loss tensor with the reduction option applied. + """ + # Original implementation from https://github.com/facebookresearch/fvcore/blob/master/fvcore/nn/focal_loss.py + + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(sigmoid_focal_loss) + p = torch.sigmoid(inputs) + ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none") + p_t = p * targets + (1 - p) * (1 - targets) + loss = ce_loss * ((1 - p_t) ** gamma) + + if alpha >= 0: + alpha_t = alpha * targets + (1 - alpha) * (1 - targets) + loss = alpha_t * loss + + # Check reduction option and return loss accordingly + if reduction == "none": + pass + elif reduction == "mean": + loss = loss.mean() + elif reduction == "sum": + loss = loss.sum() + else: + raise ValueError( + f"Invalid Value for arg 'reduction': '{reduction} \n Supported reduction modes: 'none', 'mean', 'sum'" + ) + return loss diff --git a/lib/python3.12/site-packages/torchvision/ops/giou_loss.py b/lib/python3.12/site-packages/torchvision/ops/giou_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..ec8bc8852fe71258cb5a46fdf5428581a23b4c3e --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/ops/giou_loss.py @@ -0,0 +1,76 @@ +import torch + +from ..utils import _log_api_usage_once +from ._utils import _loss_inter_union, _upcast_non_float + + +def generalized_box_iou_loss( + boxes1: torch.Tensor, + boxes2: torch.Tensor, + reduction: str = "none", + eps: float = 1e-7, +) -> torch.Tensor: + + """ + Gradient-friendly IoU loss with an additional penalty that is non-zero when the + boxes do not overlap and scales with the size of their smallest enclosing box. + This loss is symmetric, so the boxes1 and boxes2 arguments are interchangeable. + + Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with + ``0 <= x1 < x2`` and ``0 <= y1 < y2``, and The two boxes should have the + same dimensions. + + Args: + boxes1 (Tensor[N, 4] or Tensor[4]): first set of boxes + boxes2 (Tensor[N, 4] or Tensor[4]): second set of boxes + reduction (string, optional): Specifies the reduction to apply to the output: + ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: No reduction will be + applied to the output. ``'mean'``: The output will be averaged. + ``'sum'``: The output will be summed. Default: ``'none'`` + eps (float): small number to prevent division by zero. Default: 1e-7 + + Returns: + Tensor: Loss tensor with the reduction option applied. + + Reference: + Hamid Rezatofighi et al.: Generalized Intersection over Union: + A Metric and A Loss for Bounding Box Regression: + https://arxiv.org/abs/1902.09630 + """ + + # Original implementation from https://github.com/facebookresearch/fvcore/blob/bfff2ef/fvcore/nn/giou_loss.py + + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(generalized_box_iou_loss) + + boxes1 = _upcast_non_float(boxes1) + boxes2 = _upcast_non_float(boxes2) + intsctk, unionk = _loss_inter_union(boxes1, boxes2) + iouk = intsctk / (unionk + eps) + + x1, y1, x2, y2 = boxes1.unbind(dim=-1) + x1g, y1g, x2g, y2g = boxes2.unbind(dim=-1) + + # smallest enclosing box + xc1 = torch.min(x1, x1g) + yc1 = torch.min(y1, y1g) + xc2 = torch.max(x2, x2g) + yc2 = torch.max(y2, y2g) + + area_c = (xc2 - xc1) * (yc2 - yc1) + miouk = iouk - ((area_c - unionk) / (area_c + eps)) + + loss = 1 - miouk + + # Check reduction option and return loss accordingly + if reduction == "none": + pass + elif reduction == "mean": + loss = loss.mean() if loss.numel() > 0 else 0.0 * loss.sum() + elif reduction == "sum": + loss = loss.sum() + else: + raise ValueError( + f"Invalid Value for arg 'reduction': '{reduction} \n Supported reduction modes: 'none', 'mean', 'sum'" + ) + return loss diff --git a/lib/python3.12/site-packages/torchvision/ops/misc.py b/lib/python3.12/site-packages/torchvision/ops/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..0bbea6bce43425dec79a655857a92b1f36883098 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/ops/misc.py @@ -0,0 +1,320 @@ +import warnings +from typing import Callable, List, Optional, Sequence, Tuple, Union + +import torch +from torch import Tensor + +from ..utils import _log_api_usage_once, _make_ntuple + + +interpolate = torch.nn.functional.interpolate + + +class FrozenBatchNorm2d(torch.nn.Module): + """ + BatchNorm2d where the batch statistics and the affine parameters are fixed + + Args: + num_features (int): Number of features ``C`` from an expected input of size ``(N, C, H, W)`` + eps (float): a value added to the denominator for numerical stability. Default: 1e-5 + """ + + def __init__( + self, + num_features: int, + eps: float = 1e-5, + ): + super().__init__() + _log_api_usage_once(self) + self.eps = eps + self.register_buffer("weight", torch.ones(num_features)) + self.register_buffer("bias", torch.zeros(num_features)) + self.register_buffer("running_mean", torch.zeros(num_features)) + self.register_buffer("running_var", torch.ones(num_features)) + + def _load_from_state_dict( + self, + state_dict: dict, + prefix: str, + local_metadata: dict, + strict: bool, + missing_keys: List[str], + unexpected_keys: List[str], + error_msgs: List[str], + ): + num_batches_tracked_key = prefix + "num_batches_tracked" + if num_batches_tracked_key in state_dict: + del state_dict[num_batches_tracked_key] + + super()._load_from_state_dict( + state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ) + + def forward(self, x: Tensor) -> Tensor: + # move reshapes to the beginning + # to make it fuser-friendly + w = self.weight.reshape(1, -1, 1, 1) + b = self.bias.reshape(1, -1, 1, 1) + rv = self.running_var.reshape(1, -1, 1, 1) + rm = self.running_mean.reshape(1, -1, 1, 1) + scale = w * (rv + self.eps).rsqrt() + bias = b - rm * scale + return x * scale + bias + + def __repr__(self) -> str: + return f"{self.__class__.__name__}({self.weight.shape[0]}, eps={self.eps})" + + +class ConvNormActivation(torch.nn.Sequential): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: Union[int, Tuple[int, ...]] = 3, + stride: Union[int, Tuple[int, ...]] = 1, + padding: Optional[Union[int, Tuple[int, ...], str]] = None, + groups: int = 1, + norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm2d, + activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU, + dilation: Union[int, Tuple[int, ...]] = 1, + inplace: Optional[bool] = True, + bias: Optional[bool] = None, + conv_layer: Callable[..., torch.nn.Module] = torch.nn.Conv2d, + ) -> None: + + if padding is None: + if isinstance(kernel_size, int) and isinstance(dilation, int): + padding = (kernel_size - 1) // 2 * dilation + else: + _conv_dim = len(kernel_size) if isinstance(kernel_size, Sequence) else len(dilation) + kernel_size = _make_ntuple(kernel_size, _conv_dim) + dilation = _make_ntuple(dilation, _conv_dim) + padding = tuple((kernel_size[i] - 1) // 2 * dilation[i] for i in range(_conv_dim)) + if bias is None: + bias = norm_layer is None + + layers = [ + conv_layer( + in_channels, + out_channels, + kernel_size, + stride, + padding, + dilation=dilation, + groups=groups, + bias=bias, + ) + ] + + if norm_layer is not None: + layers.append(norm_layer(out_channels)) + + if activation_layer is not None: + params = {} if inplace is None else {"inplace": inplace} + layers.append(activation_layer(**params)) + super().__init__(*layers) + _log_api_usage_once(self) + self.out_channels = out_channels + + if self.__class__ == ConvNormActivation: + warnings.warn( + "Don't use ConvNormActivation directly, please use Conv2dNormActivation and Conv3dNormActivation instead." + ) + + +class Conv2dNormActivation(ConvNormActivation): + """ + Configurable block used for Convolution2d-Normalization-Activation blocks. + + Args: + in_channels (int): Number of channels in the input image + out_channels (int): Number of channels produced by the Convolution-Normalization-Activation block + kernel_size: (int, optional): Size of the convolving kernel. Default: 3 + stride (int, optional): Stride of the convolution. Default: 1 + padding (int, tuple or str, optional): Padding added to all four sides of the input. Default: None, in which case it will be calculated as ``padding = (kernel_size - 1) // 2 * dilation`` + groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1 + norm_layer (Callable[..., torch.nn.Module], optional): Norm layer that will be stacked on top of the convolution layer. If ``None`` this layer won't be used. Default: ``torch.nn.BatchNorm2d`` + activation_layer (Callable[..., torch.nn.Module], optional): Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the conv layer. If ``None`` this layer won't be used. Default: ``torch.nn.ReLU`` + dilation (int): Spacing between kernel elements. Default: 1 + inplace (bool): Parameter for the activation layer, which can optionally do the operation in-place. Default ``True`` + bias (bool, optional): Whether to use bias in the convolution layer. By default, biases are included if ``norm_layer is None``. + + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: Union[int, Tuple[int, int]] = 3, + stride: Union[int, Tuple[int, int]] = 1, + padding: Optional[Union[int, Tuple[int, int], str]] = None, + groups: int = 1, + norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm2d, + activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU, + dilation: Union[int, Tuple[int, int]] = 1, + inplace: Optional[bool] = True, + bias: Optional[bool] = None, + ) -> None: + + super().__init__( + in_channels, + out_channels, + kernel_size, + stride, + padding, + groups, + norm_layer, + activation_layer, + dilation, + inplace, + bias, + torch.nn.Conv2d, + ) + + +class Conv3dNormActivation(ConvNormActivation): + """ + Configurable block used for Convolution3d-Normalization-Activation blocks. + + Args: + in_channels (int): Number of channels in the input video. + out_channels (int): Number of channels produced by the Convolution-Normalization-Activation block + kernel_size: (int, optional): Size of the convolving kernel. Default: 3 + stride (int, optional): Stride of the convolution. Default: 1 + padding (int, tuple or str, optional): Padding added to all four sides of the input. Default: None, in which case it will be calculated as ``padding = (kernel_size - 1) // 2 * dilation`` + groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1 + norm_layer (Callable[..., torch.nn.Module], optional): Norm layer that will be stacked on top of the convolution layer. If ``None`` this layer won't be used. Default: ``torch.nn.BatchNorm3d`` + activation_layer (Callable[..., torch.nn.Module], optional): Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the conv layer. If ``None`` this layer won't be used. Default: ``torch.nn.ReLU`` + dilation (int): Spacing between kernel elements. Default: 1 + inplace (bool): Parameter for the activation layer, which can optionally do the operation in-place. Default ``True`` + bias (bool, optional): Whether to use bias in the convolution layer. By default, biases are included if ``norm_layer is None``. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: Union[int, Tuple[int, int, int]] = 3, + stride: Union[int, Tuple[int, int, int]] = 1, + padding: Optional[Union[int, Tuple[int, int, int], str]] = None, + groups: int = 1, + norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm3d, + activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU, + dilation: Union[int, Tuple[int, int, int]] = 1, + inplace: Optional[bool] = True, + bias: Optional[bool] = None, + ) -> None: + + super().__init__( + in_channels, + out_channels, + kernel_size, + stride, + padding, + groups, + norm_layer, + activation_layer, + dilation, + inplace, + bias, + torch.nn.Conv3d, + ) + + +class SqueezeExcitation(torch.nn.Module): + """ + This block implements the Squeeze-and-Excitation block from https://arxiv.org/abs/1709.01507 (see Fig. 1). + Parameters ``activation``, and ``scale_activation`` correspond to ``delta`` and ``sigma`` in eq. 3. + + Args: + input_channels (int): Number of channels in the input image + squeeze_channels (int): Number of squeeze channels + activation (Callable[..., torch.nn.Module], optional): ``delta`` activation. Default: ``torch.nn.ReLU`` + scale_activation (Callable[..., torch.nn.Module]): ``sigma`` activation. Default: ``torch.nn.Sigmoid`` + """ + + def __init__( + self, + input_channels: int, + squeeze_channels: int, + activation: Callable[..., torch.nn.Module] = torch.nn.ReLU, + scale_activation: Callable[..., torch.nn.Module] = torch.nn.Sigmoid, + ) -> None: + super().__init__() + _log_api_usage_once(self) + self.avgpool = torch.nn.AdaptiveAvgPool2d(1) + self.fc1 = torch.nn.Conv2d(input_channels, squeeze_channels, 1) + self.fc2 = torch.nn.Conv2d(squeeze_channels, input_channels, 1) + self.activation = activation() + self.scale_activation = scale_activation() + + def _scale(self, input: Tensor) -> Tensor: + scale = self.avgpool(input) + scale = self.fc1(scale) + scale = self.activation(scale) + scale = self.fc2(scale) + return self.scale_activation(scale) + + def forward(self, input: Tensor) -> Tensor: + scale = self._scale(input) + return scale * input + + +class MLP(torch.nn.Sequential): + """This block implements the multi-layer perceptron (MLP) module. + + Args: + in_channels (int): Number of channels of the input + hidden_channels (List[int]): List of the hidden channel dimensions + norm_layer (Callable[..., torch.nn.Module], optional): Norm layer that will be stacked on top of the linear layer. If ``None`` this layer won't be used. Default: ``None`` + activation_layer (Callable[..., torch.nn.Module], optional): Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the linear layer. If ``None`` this layer won't be used. Default: ``torch.nn.ReLU`` + inplace (bool, optional): Parameter for the activation layer, which can optionally do the operation in-place. + Default is ``None``, which uses the respective default values of the ``activation_layer`` and Dropout layer. + bias (bool): Whether to use bias in the linear layer. Default ``True`` + dropout (float): The probability for the dropout layer. Default: 0.0 + """ + + def __init__( + self, + in_channels: int, + hidden_channels: List[int], + norm_layer: Optional[Callable[..., torch.nn.Module]] = None, + activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU, + inplace: Optional[bool] = None, + bias: bool = True, + dropout: float = 0.0, + ): + # The addition of `norm_layer` is inspired from the implementation of TorchMultimodal: + # https://github.com/facebookresearch/multimodal/blob/5dec8a/torchmultimodal/modules/layers/mlp.py + params = {} if inplace is None else {"inplace": inplace} + + layers = [] + in_dim = in_channels + for hidden_dim in hidden_channels[:-1]: + layers.append(torch.nn.Linear(in_dim, hidden_dim, bias=bias)) + if norm_layer is not None: + layers.append(norm_layer(hidden_dim)) + layers.append(activation_layer(**params)) + layers.append(torch.nn.Dropout(dropout, **params)) + in_dim = hidden_dim + + layers.append(torch.nn.Linear(in_dim, hidden_channels[-1], bias=bias)) + layers.append(torch.nn.Dropout(dropout, **params)) + + super().__init__(*layers) + _log_api_usage_once(self) + + +class Permute(torch.nn.Module): + """This module returns a view of the tensor input with its dimensions permuted. + + Args: + dims (List[int]): The desired ordering of dimensions + """ + + def __init__(self, dims: List[int]): + super().__init__() + self.dims = dims + + def forward(self, x: Tensor) -> Tensor: + return torch.permute(x, self.dims) diff --git a/lib/python3.12/site-packages/torchvision/ops/poolers.py b/lib/python3.12/site-packages/torchvision/ops/poolers.py new file mode 100644 index 0000000000000000000000000000000000000000..9cdd83a598bea0a044179269f8369b58401801d1 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/ops/poolers.py @@ -0,0 +1,327 @@ +from typing import Dict, List, Optional, Tuple, Union + +import torch +import torch.fx +import torchvision +from torch import nn, Tensor +from torchvision.ops.boxes import box_area + +from ..utils import _log_api_usage_once +from .roi_align import roi_align + + +# copying result_idx_in_level to a specific index in result[] +# is not supported by ONNX tracing yet. +# _onnx_merge_levels() is an implementation supported by ONNX +# that merges the levels to the right indices +@torch.jit.unused +def _onnx_merge_levels(levels: Tensor, unmerged_results: List[Tensor]) -> Tensor: + first_result = unmerged_results[0] + dtype, device = first_result.dtype, first_result.device + res = torch.zeros( + (levels.size(0), first_result.size(1), first_result.size(2), first_result.size(3)), dtype=dtype, device=device + ) + for level in range(len(unmerged_results)): + index = torch.where(levels == level)[0].view(-1, 1, 1, 1) + index = index.expand( + index.size(0), + unmerged_results[level].size(1), + unmerged_results[level].size(2), + unmerged_results[level].size(3), + ) + res = res.scatter(0, index, unmerged_results[level]) + return res + + +# TODO: (eellison) T54974082 https://github.com/pytorch/pytorch/issues/26744/pytorch/issues/26744 +def initLevelMapper( + k_min: int, + k_max: int, + canonical_scale: int = 224, + canonical_level: int = 4, + eps: float = 1e-6, +): + return LevelMapper(k_min, k_max, canonical_scale, canonical_level, eps) + + +class LevelMapper: + """Determine which FPN level each RoI in a set of RoIs should map to based + on the heuristic in the FPN paper. + + Args: + k_min (int) + k_max (int) + canonical_scale (int) + canonical_level (int) + eps (float) + """ + + def __init__( + self, + k_min: int, + k_max: int, + canonical_scale: int = 224, + canonical_level: int = 4, + eps: float = 1e-6, + ): + self.k_min = k_min + self.k_max = k_max + self.s0 = canonical_scale + self.lvl0 = canonical_level + self.eps = eps + + def __call__(self, boxlists: List[Tensor]) -> Tensor: + """ + Args: + boxlists (list[BoxList]) + """ + # Compute level ids + s = torch.sqrt(torch.cat([box_area(boxlist) for boxlist in boxlists])) + + # Eqn.(1) in FPN paper + target_lvls = torch.floor(self.lvl0 + torch.log2(s / self.s0) + torch.tensor(self.eps, dtype=s.dtype)) + target_lvls = torch.clamp(target_lvls, min=self.k_min, max=self.k_max) + return (target_lvls.to(torch.int64) - self.k_min).to(torch.int64) + + +def _convert_to_roi_format(boxes: List[Tensor]) -> Tensor: + concat_boxes = torch.cat(boxes, dim=0) + device, dtype = concat_boxes.device, concat_boxes.dtype + ids = torch.cat( + [torch.full_like(b[:, :1], i, dtype=dtype, layout=torch.strided, device=device) for i, b in enumerate(boxes)], + dim=0, + ) + rois = torch.cat([ids, concat_boxes], dim=1) + return rois + + +def _infer_scale(feature: Tensor, original_size: List[int]) -> float: + # assumption: the scale is of the form 2 ** (-k), with k integer + size = feature.shape[-2:] + possible_scales: List[float] = [] + for s1, s2 in zip(size, original_size): + approx_scale = float(s1) / float(s2) + scale = 2 ** float(torch.tensor(approx_scale).log2().round()) + possible_scales.append(scale) + return possible_scales[0] + + +@torch.fx.wrap +def _setup_scales( + features: List[Tensor], image_shapes: List[Tuple[int, int]], canonical_scale: int, canonical_level: int +) -> Tuple[List[float], LevelMapper]: + if not image_shapes: + raise ValueError("images list should not be empty") + max_x = 0 + max_y = 0 + for shape in image_shapes: + max_x = max(shape[0], max_x) + max_y = max(shape[1], max_y) + original_input_shape = (max_x, max_y) + + scales = [_infer_scale(feat, original_input_shape) for feat in features] + # get the levels in the feature map by leveraging the fact that the network always + # downsamples by a factor of 2 at each level. + lvl_min = -torch.log2(torch.tensor(scales[0], dtype=torch.float32)).item() + lvl_max = -torch.log2(torch.tensor(scales[-1], dtype=torch.float32)).item() + + map_levels = initLevelMapper( + int(lvl_min), + int(lvl_max), + canonical_scale=canonical_scale, + canonical_level=canonical_level, + ) + return scales, map_levels + + +@torch.fx.wrap +def _filter_input(x: Dict[str, Tensor], featmap_names: List[str]) -> List[Tensor]: + x_filtered = [] + for k, v in x.items(): + if k in featmap_names: + x_filtered.append(v) + return x_filtered + + +@torch.fx.wrap +def _multiscale_roi_align( + x_filtered: List[Tensor], + boxes: List[Tensor], + output_size: List[int], + sampling_ratio: int, + scales: Optional[List[float]], + mapper: Optional[LevelMapper], +) -> Tensor: + """ + Args: + x_filtered (List[Tensor]): List of input tensors. + boxes (List[Tensor[N, 4]]): boxes to be used to perform the pooling operation, in + (x1, y1, x2, y2) format and in the image reference size, not the feature map + reference. The coordinate must satisfy ``0 <= x1 < x2`` and ``0 <= y1 < y2``. + output_size (Union[List[Tuple[int, int]], List[int]]): size of the output + sampling_ratio (int): sampling ratio for ROIAlign + scales (Optional[List[float]]): If None, scales will be automatically inferred. Default value is None. + mapper (Optional[LevelMapper]): If none, mapper will be automatically inferred. Default value is None. + Returns: + result (Tensor) + """ + if scales is None or mapper is None: + raise ValueError("scales and mapper should not be None") + + num_levels = len(x_filtered) + rois = _convert_to_roi_format(boxes) + + if num_levels == 1: + return roi_align( + x_filtered[0], + rois, + output_size=output_size, + spatial_scale=scales[0], + sampling_ratio=sampling_ratio, + ) + + levels = mapper(boxes) + + num_rois = len(rois) + num_channels = x_filtered[0].shape[1] + + dtype, device = x_filtered[0].dtype, x_filtered[0].device + result = torch.zeros( + ( + num_rois, + num_channels, + ) + + output_size, + dtype=dtype, + device=device, + ) + + tracing_results = [] + for level, (per_level_feature, scale) in enumerate(zip(x_filtered, scales)): + idx_in_level = torch.where(levels == level)[0] + rois_per_level = rois[idx_in_level] + + result_idx_in_level = roi_align( + per_level_feature, + rois_per_level, + output_size=output_size, + spatial_scale=scale, + sampling_ratio=sampling_ratio, + ) + + if torchvision._is_tracing(): + tracing_results.append(result_idx_in_level.to(dtype)) + else: + # result and result_idx_in_level's dtypes are based on dtypes of different + # elements in x_filtered. x_filtered contains tensors output by different + # layers. When autocast is active, it may choose different dtypes for + # different layers' outputs. Therefore, we defensively match result's dtype + # before copying elements from result_idx_in_level in the following op. + # We need to cast manually (can't rely on autocast to cast for us) because + # the op acts on result in-place, and autocast only affects out-of-place ops. + result[idx_in_level] = result_idx_in_level.to(result.dtype) + + if torchvision._is_tracing(): + result = _onnx_merge_levels(levels, tracing_results) + + return result + + +class MultiScaleRoIAlign(nn.Module): + """ + Multi-scale RoIAlign pooling, which is useful for detection with or without FPN. + + It infers the scale of the pooling via the heuristics specified in eq. 1 + of the `Feature Pyramid Network paper `_. + They keyword-only parameters ``canonical_scale`` and ``canonical_level`` + correspond respectively to ``224`` and ``k0=4`` in eq. 1, and + have the following meaning: ``canonical_level`` is the target level of the pyramid from + which to pool a region of interest with ``w x h = canonical_scale x canonical_scale``. + + Args: + featmap_names (List[str]): the names of the feature maps that will be used + for the pooling. + output_size (List[Tuple[int, int]] or List[int]): output size for the pooled region + sampling_ratio (int): sampling ratio for ROIAlign + canonical_scale (int, optional): canonical_scale for LevelMapper + canonical_level (int, optional): canonical_level for LevelMapper + + Examples:: + + >>> m = torchvision.ops.MultiScaleRoIAlign(['feat1', 'feat3'], 3, 2) + >>> i = OrderedDict() + >>> i['feat1'] = torch.rand(1, 5, 64, 64) + >>> i['feat2'] = torch.rand(1, 5, 32, 32) # this feature won't be used in the pooling + >>> i['feat3'] = torch.rand(1, 5, 16, 16) + >>> # create some random bounding boxes + >>> boxes = torch.rand(6, 4) * 256; boxes[:, 2:] += boxes[:, :2] + >>> # original image size, before computing the feature maps + >>> image_sizes = [(512, 512)] + >>> output = m(i, [boxes], image_sizes) + >>> print(output.shape) + >>> torch.Size([6, 5, 3, 3]) + + """ + + __annotations__ = {"scales": Optional[List[float]], "map_levels": Optional[LevelMapper]} + + def __init__( + self, + featmap_names: List[str], + output_size: Union[int, Tuple[int], List[int]], + sampling_ratio: int, + *, + canonical_scale: int = 224, + canonical_level: int = 4, + ): + super().__init__() + _log_api_usage_once(self) + if isinstance(output_size, int): + output_size = (output_size, output_size) + self.featmap_names = featmap_names + self.sampling_ratio = sampling_ratio + self.output_size = tuple(output_size) + self.scales = None + self.map_levels = None + self.canonical_scale = canonical_scale + self.canonical_level = canonical_level + + def forward( + self, + x: Dict[str, Tensor], + boxes: List[Tensor], + image_shapes: List[Tuple[int, int]], + ) -> Tensor: + """ + Args: + x (OrderedDict[Tensor]): feature maps for each level. They are assumed to have + all the same number of channels, but they can have different sizes. + boxes (List[Tensor[N, 4]]): boxes to be used to perform the pooling operation, in + (x1, y1, x2, y2) format and in the image reference size, not the feature map + reference. The coordinate must satisfy ``0 <= x1 < x2`` and ``0 <= y1 < y2``. + image_shapes (List[Tuple[height, width]]): the sizes of each image before they + have been fed to a CNN to obtain feature maps. This allows us to infer the + scale factor for each one of the levels to be pooled. + Returns: + result (Tensor) + """ + x_filtered = _filter_input(x, self.featmap_names) + if self.scales is None or self.map_levels is None: + self.scales, self.map_levels = _setup_scales( + x_filtered, image_shapes, self.canonical_scale, self.canonical_level + ) + + return _multiscale_roi_align( + x_filtered, + boxes, + self.output_size, + self.sampling_ratio, + self.scales, + self.map_levels, + ) + + def __repr__(self) -> str: + return ( + f"{self.__class__.__name__}(featmap_names={self.featmap_names}, " + f"output_size={self.output_size}, sampling_ratio={self.sampling_ratio})" + ) diff --git a/lib/python3.12/site-packages/torchvision/ops/ps_roi_align.py b/lib/python3.12/site-packages/torchvision/ops/ps_roi_align.py new file mode 100644 index 0000000000000000000000000000000000000000..82809b8f8885667b28eccd22aca60d1dca02f3bf --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/ops/ps_roi_align.py @@ -0,0 +1,90 @@ +import torch +import torch.fx +from torch import nn, Tensor +from torch.nn.modules.utils import _pair +from torchvision.extension import _assert_has_ops + +from ..utils import _log_api_usage_once +from ._utils import check_roi_boxes_shape, convert_boxes_to_roi_format + + +@torch.fx.wrap +def ps_roi_align( + input: Tensor, + boxes: Tensor, + output_size: int, + spatial_scale: float = 1.0, + sampling_ratio: int = -1, +) -> Tensor: + """ + Performs Position-Sensitive Region of Interest (RoI) Align operator + mentioned in Light-Head R-CNN. + + Args: + input (Tensor[N, C, H, W]): The input tensor, i.e. a batch with ``N`` elements. Each element + contains ``C`` feature maps of dimensions ``H x W``. + boxes (Tensor[K, 5] or List[Tensor[L, 4]]): the box coordinates in (x1, y1, x2, y2) + format where the regions will be taken from. + The coordinate must satisfy ``0 <= x1 < x2`` and ``0 <= y1 < y2``. + If a single Tensor is passed, then the first column should + contain the index of the corresponding element in the batch, i.e. a number in ``[0, N - 1]``. + If a list of Tensors is passed, then each Tensor will correspond to the boxes for an element i + in the batch. + output_size (int or Tuple[int, int]): the size of the output (in bins or pixels) after the pooling + is performed, as (height, width). + spatial_scale (float): a scaling factor that maps the box coordinates to + the input coordinates. For example, if your boxes are defined on the scale + of a 224x224 image and your input is a 112x112 feature map (resulting from a 0.5x scaling of + the original image), you'll want to set this to 0.5. Default: 1.0 + sampling_ratio (int): number of sampling points in the interpolation grid + used to compute the output value of each pooled output bin. If > 0, + then exactly ``sampling_ratio x sampling_ratio`` sampling points per bin are used. If + <= 0, then an adaptive number of grid points are used (computed as + ``ceil(roi_width / output_width)``, and likewise for height). Default: -1 + + Returns: + Tensor[K, C / (output_size[0] * output_size[1]), output_size[0], output_size[1]]: The pooled RoIs + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(ps_roi_align) + _assert_has_ops() + check_roi_boxes_shape(boxes) + rois = boxes + output_size = _pair(output_size) + if not isinstance(rois, torch.Tensor): + rois = convert_boxes_to_roi_format(rois) + output, _ = torch.ops.torchvision.ps_roi_align( + input, rois, spatial_scale, output_size[0], output_size[1], sampling_ratio + ) + return output + + +class PSRoIAlign(nn.Module): + """ + See :func:`ps_roi_align`. + """ + + def __init__( + self, + output_size: int, + spatial_scale: float, + sampling_ratio: int, + ): + super().__init__() + _log_api_usage_once(self) + self.output_size = output_size + self.spatial_scale = spatial_scale + self.sampling_ratio = sampling_ratio + + def forward(self, input: Tensor, rois: Tensor) -> Tensor: + return ps_roi_align(input, rois, self.output_size, self.spatial_scale, self.sampling_ratio) + + def __repr__(self) -> str: + s = ( + f"{self.__class__.__name__}(" + f"output_size={self.output_size}" + f", spatial_scale={self.spatial_scale}" + f", sampling_ratio={self.sampling_ratio}" + f")" + ) + return s diff --git a/lib/python3.12/site-packages/torchvision/ops/ps_roi_pool.py b/lib/python3.12/site-packages/torchvision/ops/ps_roi_pool.py new file mode 100644 index 0000000000000000000000000000000000000000..15292dcad97490aaa740cdec2d0aedb31e5662eb --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/ops/ps_roi_pool.py @@ -0,0 +1,70 @@ +import torch +import torch.fx +from torch import nn, Tensor +from torch.nn.modules.utils import _pair +from torchvision.extension import _assert_has_ops + +from ..utils import _log_api_usage_once +from ._utils import check_roi_boxes_shape, convert_boxes_to_roi_format + + +@torch.fx.wrap +def ps_roi_pool( + input: Tensor, + boxes: Tensor, + output_size: int, + spatial_scale: float = 1.0, +) -> Tensor: + """ + Performs Position-Sensitive Region of Interest (RoI) Pool operator + described in R-FCN + + Args: + input (Tensor[N, C, H, W]): The input tensor, i.e. a batch with ``N`` elements. Each element + contains ``C`` feature maps of dimensions ``H x W``. + boxes (Tensor[K, 5] or List[Tensor[L, 4]]): the box coordinates in (x1, y1, x2, y2) + format where the regions will be taken from. + The coordinate must satisfy ``0 <= x1 < x2`` and ``0 <= y1 < y2``. + If a single Tensor is passed, then the first column should + contain the index of the corresponding element in the batch, i.e. a number in ``[0, N - 1]``. + If a list of Tensors is passed, then each Tensor will correspond to the boxes for an element i + in the batch. + output_size (int or Tuple[int, int]): the size of the output (in bins or pixels) after the pooling + is performed, as (height, width). + spatial_scale (float): a scaling factor that maps the box coordinates to + the input coordinates. For example, if your boxes are defined on the scale + of a 224x224 image and your input is a 112x112 feature map (resulting from a 0.5x scaling of + the original image), you'll want to set this to 0.5. Default: 1.0 + + Returns: + Tensor[K, C / (output_size[0] * output_size[1]), output_size[0], output_size[1]]: The pooled RoIs. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(ps_roi_pool) + _assert_has_ops() + check_roi_boxes_shape(boxes) + rois = boxes + output_size = _pair(output_size) + if not isinstance(rois, torch.Tensor): + rois = convert_boxes_to_roi_format(rois) + output, _ = torch.ops.torchvision.ps_roi_pool(input, rois, spatial_scale, output_size[0], output_size[1]) + return output + + +class PSRoIPool(nn.Module): + """ + See :func:`ps_roi_pool`. + """ + + def __init__(self, output_size: int, spatial_scale: float): + super().__init__() + _log_api_usage_once(self) + self.output_size = output_size + self.spatial_scale = spatial_scale + + def forward(self, input: Tensor, rois: Tensor) -> Tensor: + return ps_roi_pool(input, rois, self.output_size, self.spatial_scale) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}(output_size={self.output_size}, spatial_scale={self.spatial_scale})" + return s diff --git a/lib/python3.12/site-packages/torchvision/ops/roi_align.py b/lib/python3.12/site-packages/torchvision/ops/roi_align.py new file mode 100644 index 0000000000000000000000000000000000000000..8b616ca9161c64e3a1fe1004776e466c5118d17d --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/ops/roi_align.py @@ -0,0 +1,293 @@ +import functools +from typing import List, Union + +import torch +import torch.fx +from torch import nn, Tensor +from torch._dynamo.utils import is_compile_supported +from torch.jit.annotations import BroadcastingList2 +from torch.nn.modules.utils import _pair +from torchvision.extension import _assert_has_ops, _has_ops + +from ..utils import _log_api_usage_once +from ._utils import check_roi_boxes_shape, convert_boxes_to_roi_format + + +def lazy_compile(**compile_kwargs): + """Lazily wrap a function with torch.compile on the first call + + This avoids eagerly importing dynamo. + """ + + def decorate_fn(fn): + @functools.wraps(fn) + def compile_hook(*args, **kwargs): + compiled_fn = torch.compile(fn, **compile_kwargs) + globals()[fn.__name__] = functools.wraps(fn)(compiled_fn) + return compiled_fn(*args, **kwargs) + + return compile_hook + + return decorate_fn + + +# NB: all inputs are tensors +def _bilinear_interpolate( + input, # [N, C, H, W] + roi_batch_ind, # [K] + y, # [K, PH, IY] + x, # [K, PW, IX] + ymask, # [K, IY] + xmask, # [K, IX] +): + _, channels, height, width = input.size() + + # deal with inverse element out of feature map boundary + y = y.clamp(min=0) + x = x.clamp(min=0) + y_low = y.int() + x_low = x.int() + y_high = torch.where(y_low >= height - 1, height - 1, y_low + 1) + y_low = torch.where(y_low >= height - 1, height - 1, y_low) + y = torch.where(y_low >= height - 1, y.to(input.dtype), y) + + x_high = torch.where(x_low >= width - 1, width - 1, x_low + 1) + x_low = torch.where(x_low >= width - 1, width - 1, x_low) + x = torch.where(x_low >= width - 1, x.to(input.dtype), x) + + ly = y - y_low + lx = x - x_low + hy = 1.0 - ly + hx = 1.0 - lx + + # do bilinear interpolation, but respect the masking! + # TODO: It's possible the masking here is unnecessary if y and + # x were clamped appropriately; hard to tell + def masked_index( + y, # [K, PH, IY] + x, # [K, PW, IX] + ): + if ymask is not None: + assert xmask is not None + y = torch.where(ymask[:, None, :], y, 0) + x = torch.where(xmask[:, None, :], x, 0) + return input[ + roi_batch_ind[:, None, None, None, None, None], + torch.arange(channels, device=input.device)[None, :, None, None, None, None], + y[:, None, :, None, :, None], # prev [K, PH, IY] + x[:, None, None, :, None, :], # prev [K, PW, IX] + ] # [K, C, PH, PW, IY, IX] + + v1 = masked_index(y_low, x_low) + v2 = masked_index(y_low, x_high) + v3 = masked_index(y_high, x_low) + v4 = masked_index(y_high, x_high) + + # all ws preemptively [K, C, PH, PW, IY, IX] + def outer_prod(y, x): + return y[:, None, :, None, :, None] * x[:, None, None, :, None, :] + + w1 = outer_prod(hy, hx) + w2 = outer_prod(hy, lx) + w3 = outer_prod(ly, hx) + w4 = outer_prod(ly, lx) + + val = w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4 + return val + + +# TODO: this doesn't actually cache +# TODO: main library should make this easier to do +def maybe_cast(tensor): + if torch.is_autocast_enabled() and tensor.is_cuda and tensor.dtype != torch.double: + return tensor.float() + else: + return tensor + + +# This is a pure Python and differentiable implementation of roi_align. When +# run in eager mode, it uses a lot of memory, but when compiled it has +# acceptable memory usage. The main point of this implementation is that +# its backwards is deterministic. +# It is transcribed directly off of the roi_align CUDA kernel, see +# https://dev-discuss.pytorch.org/t/a-pure-python-implementation-of-roi-align-that-looks-just-like-its-cuda-kernel/1266 +@lazy_compile(dynamic=True) +def _roi_align(input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio, aligned): + orig_dtype = input.dtype + + input = maybe_cast(input) + rois = maybe_cast(rois) + + _, _, height, width = input.size() + + ph = torch.arange(pooled_height, device=input.device) # [PH] + pw = torch.arange(pooled_width, device=input.device) # [PW] + + # input: [N, C, H, W] + # rois: [K, 5] + + roi_batch_ind = rois[:, 0].int() # [K] + offset = 0.5 if aligned else 0.0 + roi_start_w = rois[:, 1] * spatial_scale - offset # [K] + roi_start_h = rois[:, 2] * spatial_scale - offset # [K] + roi_end_w = rois[:, 3] * spatial_scale - offset # [K] + roi_end_h = rois[:, 4] * spatial_scale - offset # [K] + + roi_width = roi_end_w - roi_start_w # [K] + roi_height = roi_end_h - roi_start_h # [K] + if not aligned: + roi_width = torch.clamp(roi_width, min=1.0) # [K] + roi_height = torch.clamp(roi_height, min=1.0) # [K] + + bin_size_h = roi_height / pooled_height # [K] + bin_size_w = roi_width / pooled_width # [K] + + exact_sampling = sampling_ratio > 0 + + roi_bin_grid_h = sampling_ratio if exact_sampling else torch.ceil(roi_height / pooled_height) # scalar or [K] + roi_bin_grid_w = sampling_ratio if exact_sampling else torch.ceil(roi_width / pooled_width) # scalar or [K] + + """ + iy, ix = dims(2) + """ + + if exact_sampling: + count = max(roi_bin_grid_h * roi_bin_grid_w, 1) # scalar + iy = torch.arange(roi_bin_grid_h, device=input.device) # [IY] + ix = torch.arange(roi_bin_grid_w, device=input.device) # [IX] + ymask = None + xmask = None + else: + count = torch.clamp(roi_bin_grid_h * roi_bin_grid_w, min=1) # [K] + # When doing adaptive sampling, the number of samples we need to do + # is data-dependent based on how big the ROIs are. This is a bit + # awkward because first-class dims can't actually handle this. + # So instead, we inefficiently suppose that we needed to sample ALL + # the points and mask out things that turned out to be unnecessary + iy = torch.arange(height, device=input.device) # [IY] + ix = torch.arange(width, device=input.device) # [IX] + ymask = iy[None, :] < roi_bin_grid_h[:, None] # [K, IY] + xmask = ix[None, :] < roi_bin_grid_w[:, None] # [K, IX] + + def from_K(t): + return t[:, None, None] + + y = ( + from_K(roi_start_h) + + ph[None, :, None] * from_K(bin_size_h) + + (iy[None, None, :] + 0.5).to(input.dtype) * from_K(bin_size_h / roi_bin_grid_h) + ) # [K, PH, IY] + x = ( + from_K(roi_start_w) + + pw[None, :, None] * from_K(bin_size_w) + + (ix[None, None, :] + 0.5).to(input.dtype) * from_K(bin_size_w / roi_bin_grid_w) + ) # [K, PW, IX] + val = _bilinear_interpolate(input, roi_batch_ind, y, x, ymask, xmask) # [K, C, PH, PW, IY, IX] + + # Mask out samples that weren't actually adaptively needed + if not exact_sampling: + val = torch.where(ymask[:, None, None, None, :, None], val, 0) + val = torch.where(xmask[:, None, None, None, None, :], val, 0) + + output = val.sum((-1, -2)) # remove IY, IX ~> [K, C, PH, PW] + if isinstance(count, torch.Tensor): + output /= count[:, None, None, None] + else: + output /= count + + output = output.to(orig_dtype) + + return output + + +@torch.fx.wrap +def roi_align( + input: Tensor, + boxes: Union[Tensor, List[Tensor]], + output_size: BroadcastingList2[int], + spatial_scale: float = 1.0, + sampling_ratio: int = -1, + aligned: bool = False, +) -> Tensor: + """ + Performs Region of Interest (RoI) Align operator with average pooling, as described in Mask R-CNN. + + Args: + input (Tensor[N, C, H, W]): The input tensor, i.e. a batch with ``N`` elements. Each element + contains ``C`` feature maps of dimensions ``H x W``. + If the tensor is quantized, we expect a batch size of ``N == 1``. + boxes (Tensor[K, 5] or List[Tensor[L, 4]]): the box coordinates in (x1, y1, x2, y2) + format where the regions will be taken from. + The coordinate must satisfy ``0 <= x1 < x2`` and ``0 <= y1 < y2``. + If a single Tensor is passed, then the first column should + contain the index of the corresponding element in the batch, i.e. a number in ``[0, N - 1]``. + If a list of Tensors is passed, then each Tensor will correspond to the boxes for an element i + in the batch. + output_size (int or Tuple[int, int]): the size of the output (in bins or pixels) after the pooling + is performed, as (height, width). + spatial_scale (float): a scaling factor that maps the box coordinates to + the input coordinates. For example, if your boxes are defined on the scale + of a 224x224 image and your input is a 112x112 feature map (resulting from a 0.5x scaling of + the original image), you'll want to set this to 0.5. Default: 1.0 + sampling_ratio (int): number of sampling points in the interpolation grid + used to compute the output value of each pooled output bin. If > 0, + then exactly ``sampling_ratio x sampling_ratio`` sampling points per bin are used. If + <= 0, then an adaptive number of grid points are used (computed as + ``ceil(roi_width / output_width)``, and likewise for height). Default: -1 + aligned (bool): If False, use the legacy implementation. + If True, pixel shift the box coordinates it by -0.5 for a better alignment with the two + neighboring pixel indices. This version is used in Detectron2 + + Returns: + Tensor[K, C, output_size[0], output_size[1]]: The pooled RoIs. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(roi_align) + check_roi_boxes_shape(boxes) + rois = boxes + output_size = _pair(output_size) + if not isinstance(rois, torch.Tensor): + rois = convert_boxes_to_roi_format(rois) + if not torch.jit.is_scripting(): + if ( + not _has_ops() or (torch.are_deterministic_algorithms_enabled() and (input.is_cuda or input.is_mps)) + ) and is_compile_supported(input.device.type): + return _roi_align(input, rois, spatial_scale, output_size[0], output_size[1], sampling_ratio, aligned) + _assert_has_ops() + return torch.ops.torchvision.roi_align( + input, rois, spatial_scale, output_size[0], output_size[1], sampling_ratio, aligned + ) + + +class RoIAlign(nn.Module): + """ + See :func:`roi_align`. + """ + + def __init__( + self, + output_size: BroadcastingList2[int], + spatial_scale: float, + sampling_ratio: int, + aligned: bool = False, + ): + super().__init__() + _log_api_usage_once(self) + self.output_size = output_size + self.spatial_scale = spatial_scale + self.sampling_ratio = sampling_ratio + self.aligned = aligned + + def forward(self, input: Tensor, rois: Union[Tensor, List[Tensor]]) -> Tensor: + return roi_align(input, rois, self.output_size, self.spatial_scale, self.sampling_ratio, self.aligned) + + def __repr__(self) -> str: + s = ( + f"{self.__class__.__name__}(" + f"output_size={self.output_size}" + f", spatial_scale={self.spatial_scale}" + f", sampling_ratio={self.sampling_ratio}" + f", aligned={self.aligned}" + f")" + ) + return s diff --git a/lib/python3.12/site-packages/torchvision/ops/roi_pool.py b/lib/python3.12/site-packages/torchvision/ops/roi_pool.py new file mode 100644 index 0000000000000000000000000000000000000000..96282418f0769bc91018a872dbd0a106742bf884 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/ops/roi_pool.py @@ -0,0 +1,72 @@ +from typing import List, Union + +import torch +import torch.fx +from torch import nn, Tensor +from torch.jit.annotations import BroadcastingList2 +from torch.nn.modules.utils import _pair +from torchvision.extension import _assert_has_ops + +from ..utils import _log_api_usage_once +from ._utils import check_roi_boxes_shape, convert_boxes_to_roi_format + + +@torch.fx.wrap +def roi_pool( + input: Tensor, + boxes: Union[Tensor, List[Tensor]], + output_size: BroadcastingList2[int], + spatial_scale: float = 1.0, +) -> Tensor: + """ + Performs Region of Interest (RoI) Pool operator described in Fast R-CNN + + Args: + input (Tensor[N, C, H, W]): The input tensor, i.e. a batch with ``N`` elements. Each element + contains ``C`` feature maps of dimensions ``H x W``. + boxes (Tensor[K, 5] or List[Tensor[L, 4]]): the box coordinates in (x1, y1, x2, y2) + format where the regions will be taken from. + The coordinate must satisfy ``0 <= x1 < x2`` and ``0 <= y1 < y2``. + If a single Tensor is passed, then the first column should + contain the index of the corresponding element in the batch, i.e. a number in ``[0, N - 1]``. + If a list of Tensors is passed, then each Tensor will correspond to the boxes for an element i + in the batch. + output_size (int or Tuple[int, int]): the size of the output after the cropping + is performed, as (height, width) + spatial_scale (float): a scaling factor that maps the box coordinates to + the input coordinates. For example, if your boxes are defined on the scale + of a 224x224 image and your input is a 112x112 feature map (resulting from a 0.5x scaling of + the original image), you'll want to set this to 0.5. Default: 1.0 + + Returns: + Tensor[K, C, output_size[0], output_size[1]]: The pooled RoIs. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(roi_pool) + _assert_has_ops() + check_roi_boxes_shape(boxes) + rois = boxes + output_size = _pair(output_size) + if not isinstance(rois, torch.Tensor): + rois = convert_boxes_to_roi_format(rois) + output, _ = torch.ops.torchvision.roi_pool(input, rois, spatial_scale, output_size[0], output_size[1]) + return output + + +class RoIPool(nn.Module): + """ + See :func:`roi_pool`. + """ + + def __init__(self, output_size: BroadcastingList2[int], spatial_scale: float): + super().__init__() + _log_api_usage_once(self) + self.output_size = output_size + self.spatial_scale = spatial_scale + + def forward(self, input: Tensor, rois: Union[Tensor, List[Tensor]]) -> Tensor: + return roi_pool(input, rois, self.output_size, self.spatial_scale) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}(output_size={self.output_size}, spatial_scale={self.spatial_scale})" + return s diff --git a/lib/python3.12/site-packages/torchvision/ops/stochastic_depth.py b/lib/python3.12/site-packages/torchvision/ops/stochastic_depth.py new file mode 100644 index 0000000000000000000000000000000000000000..ff8167b2315e941f7e31a0626eeec270d350a710 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/ops/stochastic_depth.py @@ -0,0 +1,66 @@ +import torch +import torch.fx +from torch import nn, Tensor + +from ..utils import _log_api_usage_once + + +def stochastic_depth(input: Tensor, p: float, mode: str, training: bool = True) -> Tensor: + """ + Implements the Stochastic Depth from `"Deep Networks with Stochastic Depth" + `_ used for randomly dropping residual + branches of residual architectures. + + Args: + input (Tensor[N, ...]): The input tensor or arbitrary dimensions with the first one + being its batch i.e. a batch with ``N`` rows. + p (float): probability of the input to be zeroed. + mode (str): ``"batch"`` or ``"row"``. + ``"batch"`` randomly zeroes the entire input, ``"row"`` zeroes + randomly selected rows from the batch. + training: apply stochastic depth if is ``True``. Default: ``True`` + + Returns: + Tensor[N, ...]: The randomly zeroed tensor. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(stochastic_depth) + if p < 0.0 or p > 1.0: + raise ValueError(f"drop probability has to be between 0 and 1, but got {p}") + if mode not in ["batch", "row"]: + raise ValueError(f"mode has to be either 'batch' or 'row', but got {mode}") + if not training or p == 0.0: + return input + + survival_rate = 1.0 - p + if mode == "row": + size = [input.shape[0]] + [1] * (input.ndim - 1) + else: + size = [1] * input.ndim + noise = torch.empty(size, dtype=input.dtype, device=input.device) + noise = noise.bernoulli_(survival_rate) + if survival_rate > 0.0: + noise.div_(survival_rate) + return input * noise + + +torch.fx.wrap("stochastic_depth") + + +class StochasticDepth(nn.Module): + """ + See :func:`stochastic_depth`. + """ + + def __init__(self, p: float, mode: str) -> None: + super().__init__() + _log_api_usage_once(self) + self.p = p + self.mode = mode + + def forward(self, input: Tensor) -> Tensor: + return stochastic_depth(input, self.p, self.mode, self.training) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}(p={self.p}, mode={self.mode})" + return s diff --git a/lib/python3.12/site-packages/torchvision/transforms/__init__.py b/lib/python3.12/site-packages/torchvision/transforms/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..77680a14f0d0599f4004a2ce5c299c0f5e13a0d5 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/__init__.py @@ -0,0 +1,2 @@ +from .transforms import * +from .autoaugment import * diff --git a/lib/python3.12/site-packages/torchvision/transforms/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1c00299f534e4e179884ea330d421cd15474c6a6 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/__pycache__/_functional_pil.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/__pycache__/_functional_pil.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e8d9be5e5b3004f11e7225398d7c4fd58c746f65 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/__pycache__/_functional_pil.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/__pycache__/_functional_tensor.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/__pycache__/_functional_tensor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1a568ed843c5df3b01ee92ffee3dfdebeb581259 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/__pycache__/_functional_tensor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/__pycache__/_functional_video.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/__pycache__/_functional_video.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2e1b26aeb68c9e5156bfe60278fdc80399d5ca38 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/__pycache__/_functional_video.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/__pycache__/_presets.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/__pycache__/_presets.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9bed02193e05248c316034e4d7f0fe5930a82ead Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/__pycache__/_presets.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/__pycache__/_transforms_video.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/__pycache__/_transforms_video.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..22855a7e7450014bf9fa53974fc593a1a73efb94 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/__pycache__/_transforms_video.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/__pycache__/autoaugment.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/__pycache__/autoaugment.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c7d6237f63e2ede1761c050319e0d680d546f6eb Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/__pycache__/autoaugment.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/__pycache__/functional.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/__pycache__/functional.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..766e8fe21dfb824c1b7730eb1ffd16ad11807d65 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/__pycache__/functional.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/_functional_pil.py b/lib/python3.12/site-packages/torchvision/transforms/_functional_pil.py new file mode 100644 index 0000000000000000000000000000000000000000..bd943b6ee1016b45bf2e17ec169fc45b1cc311f4 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/_functional_pil.py @@ -0,0 +1,393 @@ +import numbers +from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple, Union + +import numpy as np +import torch +from PIL import Image, ImageEnhance, ImageOps + +try: + import accimage +except ImportError: + accimage = None + + +@torch.jit.unused +def _is_pil_image(img: Any) -> bool: + if accimage is not None: + return isinstance(img, (Image.Image, accimage.Image)) + else: + return isinstance(img, Image.Image) + + +@torch.jit.unused +def get_dimensions(img: Any) -> List[int]: + if _is_pil_image(img): + if hasattr(img, "getbands"): + channels = len(img.getbands()) + else: + channels = img.channels + width, height = img.size + return [channels, height, width] + raise TypeError(f"Unexpected type {type(img)}") + + +@torch.jit.unused +def get_image_size(img: Any) -> List[int]: + if _is_pil_image(img): + return list(img.size) + raise TypeError(f"Unexpected type {type(img)}") + + +@torch.jit.unused +def get_image_num_channels(img: Any) -> int: + if _is_pil_image(img): + if hasattr(img, "getbands"): + return len(img.getbands()) + else: + return img.channels + raise TypeError(f"Unexpected type {type(img)}") + + +@torch.jit.unused +def hflip(img: Image.Image) -> Image.Image: + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + + return img.transpose(Image.FLIP_LEFT_RIGHT) + + +@torch.jit.unused +def vflip(img: Image.Image) -> Image.Image: + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + + return img.transpose(Image.FLIP_TOP_BOTTOM) + + +@torch.jit.unused +def adjust_brightness(img: Image.Image, brightness_factor: float) -> Image.Image: + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + + enhancer = ImageEnhance.Brightness(img) + img = enhancer.enhance(brightness_factor) + return img + + +@torch.jit.unused +def adjust_contrast(img: Image.Image, contrast_factor: float) -> Image.Image: + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + + enhancer = ImageEnhance.Contrast(img) + img = enhancer.enhance(contrast_factor) + return img + + +@torch.jit.unused +def adjust_saturation(img: Image.Image, saturation_factor: float) -> Image.Image: + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + + enhancer = ImageEnhance.Color(img) + img = enhancer.enhance(saturation_factor) + return img + + +@torch.jit.unused +def adjust_hue(img: Image.Image, hue_factor: float) -> Image.Image: + if not (-0.5 <= hue_factor <= 0.5): + raise ValueError(f"hue_factor ({hue_factor}) is not in [-0.5, 0.5].") + + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + + input_mode = img.mode + if input_mode in {"L", "1", "I", "F"}: + return img + + h, s, v = img.convert("HSV").split() + + np_h = np.array(h, dtype=np.uint8) + # This will over/underflow, as desired + np_h += np.int32(hue_factor * 255).astype(np.uint8) + + h = Image.fromarray(np_h, "L") + + img = Image.merge("HSV", (h, s, v)).convert(input_mode) + return img + + +@torch.jit.unused +def adjust_gamma( + img: Image.Image, + gamma: float, + gain: float = 1.0, +) -> Image.Image: + + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + + if gamma < 0: + raise ValueError("Gamma should be a non-negative real number") + + input_mode = img.mode + img = img.convert("RGB") + gamma_map = [int((255 + 1 - 1e-3) * gain * pow(ele / 255.0, gamma)) for ele in range(256)] * 3 + img = img.point(gamma_map) # use PIL's point-function to accelerate this part + + img = img.convert(input_mode) + return img + + +@torch.jit.unused +def pad( + img: Image.Image, + padding: Union[int, List[int], Tuple[int, ...]], + fill: Optional[Union[float, List[float], Tuple[float, ...]]] = 0, + padding_mode: Literal["constant", "edge", "reflect", "symmetric"] = "constant", +) -> Image.Image: + + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + + if not isinstance(padding, (numbers.Number, tuple, list)): + raise TypeError("Got inappropriate padding arg") + if fill is not None and not isinstance(fill, (numbers.Number, tuple, list)): + raise TypeError("Got inappropriate fill arg") + if not isinstance(padding_mode, str): + raise TypeError("Got inappropriate padding_mode arg") + + if isinstance(padding, list): + padding = tuple(padding) + + if isinstance(padding, tuple) and len(padding) not in [1, 2, 4]: + raise ValueError(f"Padding must be an int or a 1, 2, or 4 element tuple, not a {len(padding)} element tuple") + + if isinstance(padding, tuple) and len(padding) == 1: + # Compatibility with `functional_tensor.pad` + padding = padding[0] + + if padding_mode not in ["constant", "edge", "reflect", "symmetric"]: + raise ValueError("Padding mode should be either constant, edge, reflect or symmetric") + + if padding_mode == "constant": + opts = _parse_fill(fill, img, name="fill") + if img.mode == "P": + palette = img.getpalette() + image = ImageOps.expand(img, border=padding, **opts) + image.putpalette(palette) + return image + + return ImageOps.expand(img, border=padding, **opts) + else: + if isinstance(padding, int): + pad_left = pad_right = pad_top = pad_bottom = padding + if isinstance(padding, tuple) and len(padding) == 2: + pad_left = pad_right = padding[0] + pad_top = pad_bottom = padding[1] + if isinstance(padding, tuple) and len(padding) == 4: + pad_left = padding[0] + pad_top = padding[1] + pad_right = padding[2] + pad_bottom = padding[3] + + p = [pad_left, pad_top, pad_right, pad_bottom] + cropping = -np.minimum(p, 0) + + if cropping.any(): + crop_left, crop_top, crop_right, crop_bottom = cropping + img = img.crop((crop_left, crop_top, img.width - crop_right, img.height - crop_bottom)) + + pad_left, pad_top, pad_right, pad_bottom = np.maximum(p, 0) + + if img.mode == "P": + palette = img.getpalette() + img = np.asarray(img) + img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right)), mode=padding_mode) + img = Image.fromarray(img) + img.putpalette(palette) + return img + + img = np.asarray(img) + # RGB image + if len(img.shape) == 3: + img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right), (0, 0)), padding_mode) + # Grayscale image + if len(img.shape) == 2: + img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right)), padding_mode) + + return Image.fromarray(img) + + +@torch.jit.unused +def crop( + img: Image.Image, + top: int, + left: int, + height: int, + width: int, +) -> Image.Image: + + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + + return img.crop((left, top, left + width, top + height)) + + +@torch.jit.unused +def resize( + img: Image.Image, + size: Union[List[int], int], + interpolation: int = Image.BILINEAR, +) -> Image.Image: + + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + if not (isinstance(size, list) and len(size) == 2): + raise TypeError(f"Got inappropriate size arg: {size}") + + return img.resize(tuple(size[::-1]), interpolation) + + +@torch.jit.unused +def _parse_fill( + fill: Optional[Union[float, List[float], Tuple[float, ...]]], + img: Image.Image, + name: str = "fillcolor", +) -> Dict[str, Optional[Union[float, List[float], Tuple[float, ...]]]]: + + # Process fill color for affine transforms + num_channels = get_image_num_channels(img) + if fill is None: + fill = 0 + if isinstance(fill, (int, float)) and num_channels > 1: + fill = tuple([fill] * num_channels) + if isinstance(fill, (list, tuple)): + if len(fill) == 1: + fill = fill * num_channels + elif len(fill) != num_channels: + msg = "The number of elements in 'fill' does not match the number of channels of the image ({} != {})" + raise ValueError(msg.format(len(fill), num_channels)) + + fill = tuple(fill) # type: ignore[arg-type] + + if img.mode != "F": + if isinstance(fill, (list, tuple)): + fill = tuple(int(x) for x in fill) + else: + fill = int(fill) + + return {name: fill} + + +@torch.jit.unused +def affine( + img: Image.Image, + matrix: List[float], + interpolation: int = Image.NEAREST, + fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None, +) -> Image.Image: + + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + + output_size = img.size + opts = _parse_fill(fill, img) + return img.transform(output_size, Image.AFFINE, matrix, interpolation, **opts) + + +@torch.jit.unused +def rotate( + img: Image.Image, + angle: float, + interpolation: int = Image.NEAREST, + expand: bool = False, + center: Optional[Tuple[int, int]] = None, + fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None, +) -> Image.Image: + + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + + opts = _parse_fill(fill, img) + return img.rotate(angle, interpolation, expand, center, **opts) + + +@torch.jit.unused +def perspective( + img: Image.Image, + perspective_coeffs: List[float], + interpolation: int = Image.BICUBIC, + fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None, +) -> Image.Image: + + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + + opts = _parse_fill(fill, img) + + return img.transform(img.size, Image.PERSPECTIVE, perspective_coeffs, interpolation, **opts) + + +@torch.jit.unused +def to_grayscale(img: Image.Image, num_output_channels: int) -> Image.Image: + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + + if num_output_channels == 1: + img = img.convert("L") + elif num_output_channels == 3: + img = img.convert("L") + np_img = np.array(img, dtype=np.uint8) + np_img = np.dstack([np_img, np_img, np_img]) + img = Image.fromarray(np_img, "RGB") + else: + raise ValueError("num_output_channels should be either 1 or 3") + + return img + + +@torch.jit.unused +def invert(img: Image.Image) -> Image.Image: + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + return ImageOps.invert(img) + + +@torch.jit.unused +def posterize(img: Image.Image, bits: int) -> Image.Image: + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + return ImageOps.posterize(img, bits) + + +@torch.jit.unused +def solarize(img: Image.Image, threshold: int) -> Image.Image: + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + return ImageOps.solarize(img, threshold) + + +@torch.jit.unused +def adjust_sharpness(img: Image.Image, sharpness_factor: float) -> Image.Image: + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + + enhancer = ImageEnhance.Sharpness(img) + img = enhancer.enhance(sharpness_factor) + return img + + +@torch.jit.unused +def autocontrast(img: Image.Image) -> Image.Image: + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + return ImageOps.autocontrast(img) + + +@torch.jit.unused +def equalize(img: Image.Image) -> Image.Image: + if not _is_pil_image(img): + raise TypeError(f"img should be PIL Image. Got {type(img)}") + return ImageOps.equalize(img) diff --git a/lib/python3.12/site-packages/torchvision/transforms/_functional_tensor.py b/lib/python3.12/site-packages/torchvision/transforms/_functional_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..618bbfbab7c8f40b34216c211a4a39f2a87ff72d --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/_functional_tensor.py @@ -0,0 +1,962 @@ +import warnings +from typing import List, Optional, Tuple, Union + +import torch +from torch import Tensor +from torch.nn.functional import conv2d, grid_sample, interpolate, pad as torch_pad + + +def _is_tensor_a_torch_image(x: Tensor) -> bool: + return x.ndim >= 2 + + +def _assert_image_tensor(img: Tensor) -> None: + if not _is_tensor_a_torch_image(img): + raise TypeError("Tensor is not a torch image.") + + +def get_dimensions(img: Tensor) -> List[int]: + _assert_image_tensor(img) + channels = 1 if img.ndim == 2 else img.shape[-3] + height, width = img.shape[-2:] + return [channels, height, width] + + +def get_image_size(img: Tensor) -> List[int]: + # Returns (w, h) of tensor image + _assert_image_tensor(img) + return [img.shape[-1], img.shape[-2]] + + +def get_image_num_channels(img: Tensor) -> int: + _assert_image_tensor(img) + if img.ndim == 2: + return 1 + elif img.ndim > 2: + return img.shape[-3] + + raise TypeError(f"Input ndim should be 2 or more. Got {img.ndim}") + + +def _max_value(dtype: torch.dtype) -> int: + if dtype == torch.uint8: + return 255 + elif dtype == torch.int8: + return 127 + elif dtype == torch.int16: + return 32767 + elif dtype == torch.uint16: + return 65535 + elif dtype == torch.int32: + return 2147483647 + elif dtype == torch.int64: + return 9223372036854775807 + else: + # This is only here for completeness. This value is implicitly assumed in a lot of places so changing it is not + # easy. + return 1 + + +def _assert_channels(img: Tensor, permitted: List[int]) -> None: + c = get_dimensions(img)[0] + if c not in permitted: + raise TypeError(f"Input image tensor permitted channel values are {permitted}, but found {c}") + + +def convert_image_dtype(image: torch.Tensor, dtype: torch.dtype = torch.float) -> torch.Tensor: + if image.dtype == dtype: + return image + + if image.is_floating_point(): + + # TODO: replace with dtype.is_floating_point when torchscript supports it + if torch.tensor(0, dtype=dtype).is_floating_point(): + return image.to(dtype) + + # float to int + if (image.dtype == torch.float32 and dtype in (torch.int32, torch.int64)) or ( + image.dtype == torch.float64 and dtype == torch.int64 + ): + msg = f"The cast from {image.dtype} to {dtype} cannot be performed safely." + raise RuntimeError(msg) + + # https://github.com/pytorch/vision/pull/2078#issuecomment-612045321 + # For data in the range 0-1, (float * 255).to(uint) is only 255 + # when float is exactly 1.0. + # `max + 1 - epsilon` provides more evenly distributed mapping of + # ranges of floats to ints. + eps = 1e-3 + max_val = float(_max_value(dtype)) + result = image.mul(max_val + 1.0 - eps) + return result.to(dtype) + else: + input_max = float(_max_value(image.dtype)) + + # int to float + # TODO: replace with dtype.is_floating_point when torchscript supports it + if torch.tensor(0, dtype=dtype).is_floating_point(): + image = image.to(dtype) + return image / input_max + + output_max = float(_max_value(dtype)) + + # int to int + if input_max > output_max: + # factor should be forced to int for torch jit script + # otherwise factor is a float and image // factor can produce different results + factor = int((input_max + 1) // (output_max + 1)) + image = torch.div(image, factor, rounding_mode="floor") + return image.to(dtype) + else: + # factor should be forced to int for torch jit script + # otherwise factor is a float and image * factor can produce different results + factor = int((output_max + 1) // (input_max + 1)) + image = image.to(dtype) + return image * factor + + +def vflip(img: Tensor) -> Tensor: + _assert_image_tensor(img) + + return img.flip(-2) + + +def hflip(img: Tensor) -> Tensor: + _assert_image_tensor(img) + + return img.flip(-1) + + +def crop(img: Tensor, top: int, left: int, height: int, width: int) -> Tensor: + _assert_image_tensor(img) + + _, h, w = get_dimensions(img) + right = left + width + bottom = top + height + + if left < 0 or top < 0 or right > w or bottom > h: + padding_ltrb = [ + max(-left + min(0, right), 0), + max(-top + min(0, bottom), 0), + max(right - max(w, left), 0), + max(bottom - max(h, top), 0), + ] + return pad(img[..., max(top, 0) : bottom, max(left, 0) : right], padding_ltrb, fill=0) + return img[..., top:bottom, left:right] + + +def rgb_to_grayscale(img: Tensor, num_output_channels: int = 1) -> Tensor: + if img.ndim < 3: + raise TypeError(f"Input image tensor should have at least 3 dimensions, but found {img.ndim}") + _assert_channels(img, [1, 3]) + + if num_output_channels not in (1, 3): + raise ValueError("num_output_channels should be either 1 or 3") + + if img.shape[-3] == 3: + r, g, b = img.unbind(dim=-3) + # This implementation closely follows the TF one: + # https://github.com/tensorflow/tensorflow/blob/v2.3.0/tensorflow/python/ops/image_ops_impl.py#L2105-L2138 + l_img = (0.2989 * r + 0.587 * g + 0.114 * b).to(img.dtype) + l_img = l_img.unsqueeze(dim=-3) + else: + l_img = img.clone() + + if num_output_channels == 3: + return l_img.expand(img.shape) + + return l_img + + +def adjust_brightness(img: Tensor, brightness_factor: float) -> Tensor: + if brightness_factor < 0: + raise ValueError(f"brightness_factor ({brightness_factor}) is not non-negative.") + + _assert_image_tensor(img) + + _assert_channels(img, [1, 3]) + + return _blend(img, torch.zeros_like(img), brightness_factor) + + +def adjust_contrast(img: Tensor, contrast_factor: float) -> Tensor: + if contrast_factor < 0: + raise ValueError(f"contrast_factor ({contrast_factor}) is not non-negative.") + + _assert_image_tensor(img) + + _assert_channels(img, [3, 1]) + c = get_dimensions(img)[0] + dtype = img.dtype if torch.is_floating_point(img) else torch.float32 + if c == 3: + mean = torch.mean(rgb_to_grayscale(img).to(dtype), dim=(-3, -2, -1), keepdim=True) + else: + mean = torch.mean(img.to(dtype), dim=(-3, -2, -1), keepdim=True) + + return _blend(img, mean, contrast_factor) + + +def adjust_hue(img: Tensor, hue_factor: float) -> Tensor: + if not (-0.5 <= hue_factor <= 0.5): + raise ValueError(f"hue_factor ({hue_factor}) is not in [-0.5, 0.5].") + + if not (isinstance(img, torch.Tensor)): + raise TypeError("Input img should be Tensor image") + + _assert_image_tensor(img) + + _assert_channels(img, [1, 3]) + if get_dimensions(img)[0] == 1: # Match PIL behaviour + return img + + orig_dtype = img.dtype + img = convert_image_dtype(img, torch.float32) + + img = _rgb2hsv(img) + h, s, v = img.unbind(dim=-3) + h = (h + hue_factor) % 1.0 + img = torch.stack((h, s, v), dim=-3) + img_hue_adj = _hsv2rgb(img) + + return convert_image_dtype(img_hue_adj, orig_dtype) + + +def adjust_saturation(img: Tensor, saturation_factor: float) -> Tensor: + if saturation_factor < 0: + raise ValueError(f"saturation_factor ({saturation_factor}) is not non-negative.") + + _assert_image_tensor(img) + + _assert_channels(img, [1, 3]) + + if get_dimensions(img)[0] == 1: # Match PIL behaviour + return img + + return _blend(img, rgb_to_grayscale(img), saturation_factor) + + +def adjust_gamma(img: Tensor, gamma: float, gain: float = 1) -> Tensor: + if not isinstance(img, torch.Tensor): + raise TypeError("Input img should be a Tensor.") + + _assert_channels(img, [1, 3]) + + if gamma < 0: + raise ValueError("Gamma should be a non-negative real number") + + result = img + dtype = img.dtype + if not torch.is_floating_point(img): + result = convert_image_dtype(result, torch.float32) + + result = (gain * result**gamma).clamp(0, 1) + + result = convert_image_dtype(result, dtype) + return result + + +def _blend(img1: Tensor, img2: Tensor, ratio: float) -> Tensor: + ratio = float(ratio) + bound = _max_value(img1.dtype) + return (ratio * img1 + (1.0 - ratio) * img2).clamp(0, bound).to(img1.dtype) + + +def _rgb2hsv(img: Tensor) -> Tensor: + r, g, b = img.unbind(dim=-3) + + # Implementation is based on https://github.com/python-pillow/Pillow/blob/4174d4267616897df3746d315d5a2d0f82c656ee/ + # src/libImaging/Convert.c#L330 + maxc = torch.max(img, dim=-3).values + minc = torch.min(img, dim=-3).values + + # The algorithm erases S and H channel where `maxc = minc`. This avoids NaN + # from happening in the results, because + # + S channel has division by `maxc`, which is zero only if `maxc = minc` + # + H channel has division by `(maxc - minc)`. + # + # Instead of overwriting NaN afterwards, we just prevent it from occurring, so + # we don't need to deal with it in case we save the NaN in a buffer in + # backprop, if it is ever supported, but it doesn't hurt to do so. + eqc = maxc == minc + + cr = maxc - minc + # Since `eqc => cr = 0`, replacing denominator with 1 when `eqc` is fine. + ones = torch.ones_like(maxc) + s = cr / torch.where(eqc, ones, maxc) + # Note that `eqc => maxc = minc = r = g = b`. So the following calculation + # of `h` would reduce to `bc - gc + 2 + rc - bc + 4 + rc - bc = 6` so it + # would not matter what values `rc`, `gc`, and `bc` have here, and thus + # replacing denominator with 1 when `eqc` is fine. + cr_divisor = torch.where(eqc, ones, cr) + rc = (maxc - r) / cr_divisor + gc = (maxc - g) / cr_divisor + bc = (maxc - b) / cr_divisor + + hr = (maxc == r) * (bc - gc) + hg = ((maxc == g) & (maxc != r)) * (2.0 + rc - bc) + hb = ((maxc != g) & (maxc != r)) * (4.0 + gc - rc) + h = hr + hg + hb + h = torch.fmod((h / 6.0 + 1.0), 1.0) + return torch.stack((h, s, maxc), dim=-3) + + +def _hsv2rgb(img: Tensor) -> Tensor: + h, s, v = img.unbind(dim=-3) + i = torch.floor(h * 6.0) + f = (h * 6.0) - i + i = i.to(dtype=torch.int32) + + p = torch.clamp((v * (1.0 - s)), 0.0, 1.0) + q = torch.clamp((v * (1.0 - s * f)), 0.0, 1.0) + t = torch.clamp((v * (1.0 - s * (1.0 - f))), 0.0, 1.0) + i = i % 6 + + mask = i.unsqueeze(dim=-3) == torch.arange(6, device=i.device).view(-1, 1, 1) + + a1 = torch.stack((v, q, p, p, t, v), dim=-3) + a2 = torch.stack((t, v, v, q, p, p), dim=-3) + a3 = torch.stack((p, p, t, v, v, q), dim=-3) + a4 = torch.stack((a1, a2, a3), dim=-4) + + return torch.einsum("...ijk, ...xijk -> ...xjk", mask.to(dtype=img.dtype), a4) + + +def _pad_symmetric(img: Tensor, padding: List[int]) -> Tensor: + # padding is left, right, top, bottom + + # crop if needed + if padding[0] < 0 or padding[1] < 0 or padding[2] < 0 or padding[3] < 0: + neg_min_padding = [-min(x, 0) for x in padding] + crop_left, crop_right, crop_top, crop_bottom = neg_min_padding + img = img[..., crop_top : img.shape[-2] - crop_bottom, crop_left : img.shape[-1] - crop_right] + padding = [max(x, 0) for x in padding] + + in_sizes = img.size() + + _x_indices = [i for i in range(in_sizes[-1])] # [0, 1, 2, 3, ...] + left_indices = [i for i in range(padding[0] - 1, -1, -1)] # e.g. [3, 2, 1, 0] + right_indices = [-(i + 1) for i in range(padding[1])] # e.g. [-1, -2, -3] + x_indices = torch.tensor(left_indices + _x_indices + right_indices, device=img.device) + + _y_indices = [i for i in range(in_sizes[-2])] + top_indices = [i for i in range(padding[2] - 1, -1, -1)] + bottom_indices = [-(i + 1) for i in range(padding[3])] + y_indices = torch.tensor(top_indices + _y_indices + bottom_indices, device=img.device) + + ndim = img.ndim + if ndim == 3: + return img[:, y_indices[:, None], x_indices[None, :]] + elif ndim == 4: + return img[:, :, y_indices[:, None], x_indices[None, :]] + else: + raise RuntimeError("Symmetric padding of N-D tensors are not supported yet") + + +def _parse_pad_padding(padding: Union[int, List[int]]) -> List[int]: + if isinstance(padding, int): + if torch.jit.is_scripting(): + # This maybe unreachable + raise ValueError("padding can't be an int while torchscripting, set it as a list [value, ]") + pad_left = pad_right = pad_top = pad_bottom = padding + elif len(padding) == 1: + pad_left = pad_right = pad_top = pad_bottom = padding[0] + elif len(padding) == 2: + pad_left = pad_right = padding[0] + pad_top = pad_bottom = padding[1] + else: + pad_left = padding[0] + pad_top = padding[1] + pad_right = padding[2] + pad_bottom = padding[3] + + return [pad_left, pad_right, pad_top, pad_bottom] + + +def pad( + img: Tensor, padding: Union[int, List[int]], fill: Optional[Union[int, float]] = 0, padding_mode: str = "constant" +) -> Tensor: + _assert_image_tensor(img) + + if fill is None: + fill = 0 + + if not isinstance(padding, (int, tuple, list)): + raise TypeError("Got inappropriate padding arg") + if not isinstance(fill, (int, float)): + raise TypeError("Got inappropriate fill arg") + if not isinstance(padding_mode, str): + raise TypeError("Got inappropriate padding_mode arg") + + if isinstance(padding, tuple): + padding = list(padding) + + if isinstance(padding, list): + # TODO: Jit is failing on loading this op when scripted and saved + # https://github.com/pytorch/pytorch/issues/81100 + if len(padding) not in [1, 2, 4]: + raise ValueError( + f"Padding must be an int or a 1, 2, or 4 element tuple, not a {len(padding)} element tuple" + ) + + if padding_mode not in ["constant", "edge", "reflect", "symmetric"]: + raise ValueError("Padding mode should be either constant, edge, reflect or symmetric") + + p = _parse_pad_padding(padding) + + if padding_mode == "edge": + # remap padding_mode str + padding_mode = "replicate" + elif padding_mode == "symmetric": + # route to another implementation + return _pad_symmetric(img, p) + + need_squeeze = False + if img.ndim < 4: + img = img.unsqueeze(dim=0) + need_squeeze = True + + out_dtype = img.dtype + need_cast = False + if (padding_mode != "constant") and img.dtype not in (torch.float32, torch.float64): + # Here we temporarily cast input tensor to float + # until pytorch issue is resolved : + # https://github.com/pytorch/pytorch/issues/40763 + need_cast = True + img = img.to(torch.float32) + + if padding_mode in ("reflect", "replicate"): + img = torch_pad(img, p, mode=padding_mode) + else: + img = torch_pad(img, p, mode=padding_mode, value=float(fill)) + + if need_squeeze: + img = img.squeeze(dim=0) + + if need_cast: + img = img.to(out_dtype) + + return img + + +def resize( + img: Tensor, + size: List[int], + interpolation: str = "bilinear", + antialias: Optional[bool] = True, +) -> Tensor: + _assert_image_tensor(img) + + if isinstance(size, tuple): + size = list(size) + + if antialias is None: + antialias = False + + if antialias and interpolation not in ["bilinear", "bicubic"]: + # We manually set it to False to avoid an error downstream in interpolate() + # This behaviour is documented: the parameter is irrelevant for modes + # that are not bilinear or bicubic. We used to raise an error here, but + # now we don't as True is the default. + antialias = False + + img, need_cast, need_squeeze, out_dtype = _cast_squeeze_in(img, [torch.float32, torch.float64]) + + # Define align_corners to avoid warnings + align_corners = False if interpolation in ["bilinear", "bicubic"] else None + + img = interpolate(img, size=size, mode=interpolation, align_corners=align_corners, antialias=antialias) + + if interpolation == "bicubic" and out_dtype == torch.uint8: + img = img.clamp(min=0, max=255) + + img = _cast_squeeze_out(img, need_cast=need_cast, need_squeeze=need_squeeze, out_dtype=out_dtype) + + return img + + +def _assert_grid_transform_inputs( + img: Tensor, + matrix: Optional[List[float]], + interpolation: str, + fill: Optional[Union[int, float, List[float]]], + supported_interpolation_modes: List[str], + coeffs: Optional[List[float]] = None, +) -> None: + + if not (isinstance(img, torch.Tensor)): + raise TypeError("Input img should be Tensor") + + _assert_image_tensor(img) + + if matrix is not None and not isinstance(matrix, list): + raise TypeError("Argument matrix should be a list") + + if matrix is not None and len(matrix) != 6: + raise ValueError("Argument matrix should have 6 float values") + + if coeffs is not None and len(coeffs) != 8: + raise ValueError("Argument coeffs should have 8 float values") + + if fill is not None and not isinstance(fill, (int, float, tuple, list)): + warnings.warn("Argument fill should be either int, float, tuple or list") + + # Check fill + num_channels = get_dimensions(img)[0] + if fill is not None and isinstance(fill, (tuple, list)) and len(fill) > 1 and len(fill) != num_channels: + msg = ( + "The number of elements in 'fill' cannot broadcast to match the number of " + "channels of the image ({} != {})" + ) + raise ValueError(msg.format(len(fill), num_channels)) + + if interpolation not in supported_interpolation_modes: + raise ValueError(f"Interpolation mode '{interpolation}' is unsupported with Tensor input") + + +def _cast_squeeze_in(img: Tensor, req_dtypes: List[torch.dtype]) -> Tuple[Tensor, bool, bool, torch.dtype]: + need_squeeze = False + # make image NCHW + if img.ndim < 4: + img = img.unsqueeze(dim=0) + need_squeeze = True + + out_dtype = img.dtype + need_cast = False + if out_dtype not in req_dtypes: + need_cast = True + req_dtype = req_dtypes[0] + img = img.to(req_dtype) + return img, need_cast, need_squeeze, out_dtype + + +def _cast_squeeze_out(img: Tensor, need_cast: bool, need_squeeze: bool, out_dtype: torch.dtype) -> Tensor: + if need_squeeze: + img = img.squeeze(dim=0) + + if need_cast: + if out_dtype in (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64): + # it is better to round before cast + img = torch.round(img) + img = img.to(out_dtype) + + return img + + +def _apply_grid_transform( + img: Tensor, grid: Tensor, mode: str, fill: Optional[Union[int, float, List[float]]] +) -> Tensor: + + img, need_cast, need_squeeze, out_dtype = _cast_squeeze_in(img, [grid.dtype]) + + if img.shape[0] > 1: + # Apply same grid to a batch of images + grid = grid.expand(img.shape[0], grid.shape[1], grid.shape[2], grid.shape[3]) + + # Append a dummy mask for customized fill colors, should be faster than grid_sample() twice + if fill is not None: + mask = torch.ones((img.shape[0], 1, img.shape[2], img.shape[3]), dtype=img.dtype, device=img.device) + img = torch.cat((img, mask), dim=1) + + img = grid_sample(img, grid, mode=mode, padding_mode="zeros", align_corners=False) + + # Fill with required color + if fill is not None: + mask = img[:, -1:, :, :] # N * 1 * H * W + img = img[:, :-1, :, :] # N * C * H * W + mask = mask.expand_as(img) + fill_list, len_fill = (fill, len(fill)) if isinstance(fill, (tuple, list)) else ([float(fill)], 1) + fill_img = torch.tensor(fill_list, dtype=img.dtype, device=img.device).view(1, len_fill, 1, 1).expand_as(img) + if mode == "nearest": + mask = mask < 0.5 + img[mask] = fill_img[mask] + else: # 'bilinear' + img = img * mask + (1.0 - mask) * fill_img + + img = _cast_squeeze_out(img, need_cast, need_squeeze, out_dtype) + return img + + +def _gen_affine_grid( + theta: Tensor, + w: int, + h: int, + ow: int, + oh: int, +) -> Tensor: + # https://github.com/pytorch/pytorch/blob/74b65c32be68b15dc7c9e8bb62459efbfbde33d8/aten/src/ATen/native/ + # AffineGridGenerator.cpp#L18 + # Difference with AffineGridGenerator is that: + # 1) we normalize grid values after applying theta + # 2) we can normalize by other image size, such that it covers "extend" option like in PIL.Image.rotate + + d = 0.5 + base_grid = torch.empty(1, oh, ow, 3, dtype=theta.dtype, device=theta.device) + x_grid = torch.linspace(-ow * 0.5 + d, ow * 0.5 + d - 1, steps=ow, device=theta.device) + base_grid[..., 0].copy_(x_grid) + y_grid = torch.linspace(-oh * 0.5 + d, oh * 0.5 + d - 1, steps=oh, device=theta.device).unsqueeze_(-1) + base_grid[..., 1].copy_(y_grid) + base_grid[..., 2].fill_(1) + + rescaled_theta = theta.transpose(1, 2) / torch.tensor([0.5 * w, 0.5 * h], dtype=theta.dtype, device=theta.device) + output_grid = base_grid.view(1, oh * ow, 3).bmm(rescaled_theta) + return output_grid.view(1, oh, ow, 2) + + +def affine( + img: Tensor, + matrix: List[float], + interpolation: str = "nearest", + fill: Optional[Union[int, float, List[float]]] = None, +) -> Tensor: + _assert_grid_transform_inputs(img, matrix, interpolation, fill, ["nearest", "bilinear"]) + + dtype = img.dtype if torch.is_floating_point(img) else torch.float32 + theta = torch.tensor(matrix, dtype=dtype, device=img.device).reshape(1, 2, 3) + shape = img.shape + # grid will be generated on the same device as theta and img + grid = _gen_affine_grid(theta, w=shape[-1], h=shape[-2], ow=shape[-1], oh=shape[-2]) + return _apply_grid_transform(img, grid, interpolation, fill=fill) + + +def _compute_affine_output_size(matrix: List[float], w: int, h: int) -> Tuple[int, int]: + + # Inspired of PIL implementation: + # https://github.com/python-pillow/Pillow/blob/11de3318867e4398057373ee9f12dcb33db7335c/src/PIL/Image.py#L2054 + + # pts are Top-Left, Top-Right, Bottom-Left, Bottom-Right points. + # Points are shifted due to affine matrix torch convention about + # the center point. Center is (0, 0) for image center pivot point (w * 0.5, h * 0.5) + pts = torch.tensor( + [ + [-0.5 * w, -0.5 * h, 1.0], + [-0.5 * w, 0.5 * h, 1.0], + [0.5 * w, 0.5 * h, 1.0], + [0.5 * w, -0.5 * h, 1.0], + ] + ) + theta = torch.tensor(matrix, dtype=torch.float).view(2, 3) + new_pts = torch.matmul(pts, theta.T) + min_vals, _ = new_pts.min(dim=0) + max_vals, _ = new_pts.max(dim=0) + + # shift points to [0, w] and [0, h] interval to match PIL results + min_vals += torch.tensor((w * 0.5, h * 0.5)) + max_vals += torch.tensor((w * 0.5, h * 0.5)) + + # Truncate precision to 1e-4 to avoid ceil of Xe-15 to 1.0 + tol = 1e-4 + cmax = torch.ceil((max_vals / tol).trunc_() * tol) + cmin = torch.floor((min_vals / tol).trunc_() * tol) + size = cmax - cmin + return int(size[0]), int(size[1]) # w, h + + +def rotate( + img: Tensor, + matrix: List[float], + interpolation: str = "nearest", + expand: bool = False, + fill: Optional[Union[int, float, List[float]]] = None, +) -> Tensor: + _assert_grid_transform_inputs(img, matrix, interpolation, fill, ["nearest", "bilinear"]) + w, h = img.shape[-1], img.shape[-2] + ow, oh = _compute_affine_output_size(matrix, w, h) if expand else (w, h) + dtype = img.dtype if torch.is_floating_point(img) else torch.float32 + theta = torch.tensor(matrix, dtype=dtype, device=img.device).reshape(1, 2, 3) + # grid will be generated on the same device as theta and img + grid = _gen_affine_grid(theta, w=w, h=h, ow=ow, oh=oh) + + return _apply_grid_transform(img, grid, interpolation, fill=fill) + + +def _perspective_grid(coeffs: List[float], ow: int, oh: int, dtype: torch.dtype, device: torch.device) -> Tensor: + # https://github.com/python-pillow/Pillow/blob/4634eafe3c695a014267eefdce830b4a825beed7/ + # src/libImaging/Geometry.c#L394 + + # + # x_out = (coeffs[0] * x + coeffs[1] * y + coeffs[2]) / (coeffs[6] * x + coeffs[7] * y + 1) + # y_out = (coeffs[3] * x + coeffs[4] * y + coeffs[5]) / (coeffs[6] * x + coeffs[7] * y + 1) + # + theta1 = torch.tensor( + [[[coeffs[0], coeffs[1], coeffs[2]], [coeffs[3], coeffs[4], coeffs[5]]]], dtype=dtype, device=device + ) + theta2 = torch.tensor([[[coeffs[6], coeffs[7], 1.0], [coeffs[6], coeffs[7], 1.0]]], dtype=dtype, device=device) + + d = 0.5 + base_grid = torch.empty(1, oh, ow, 3, dtype=dtype, device=device) + x_grid = torch.linspace(d, ow * 1.0 + d - 1.0, steps=ow, device=device) + base_grid[..., 0].copy_(x_grid) + y_grid = torch.linspace(d, oh * 1.0 + d - 1.0, steps=oh, device=device).unsqueeze_(-1) + base_grid[..., 1].copy_(y_grid) + base_grid[..., 2].fill_(1) + + rescaled_theta1 = theta1.transpose(1, 2) / torch.tensor([0.5 * ow, 0.5 * oh], dtype=dtype, device=device) + output_grid1 = base_grid.view(1, oh * ow, 3).bmm(rescaled_theta1) + output_grid2 = base_grid.view(1, oh * ow, 3).bmm(theta2.transpose(1, 2)) + + output_grid = output_grid1 / output_grid2 - 1.0 + return output_grid.view(1, oh, ow, 2) + + +def perspective( + img: Tensor, + perspective_coeffs: List[float], + interpolation: str = "bilinear", + fill: Optional[Union[int, float, List[float]]] = None, +) -> Tensor: + if not (isinstance(img, torch.Tensor)): + raise TypeError("Input img should be Tensor.") + + _assert_image_tensor(img) + + _assert_grid_transform_inputs( + img, + matrix=None, + interpolation=interpolation, + fill=fill, + supported_interpolation_modes=["nearest", "bilinear"], + coeffs=perspective_coeffs, + ) + + ow, oh = img.shape[-1], img.shape[-2] + dtype = img.dtype if torch.is_floating_point(img) else torch.float32 + grid = _perspective_grid(perspective_coeffs, ow=ow, oh=oh, dtype=dtype, device=img.device) + return _apply_grid_transform(img, grid, interpolation, fill=fill) + + +def _get_gaussian_kernel1d(kernel_size: int, sigma: float, dtype: torch.dtype, device: torch.device) -> Tensor: + ksize_half = (kernel_size - 1) * 0.5 + + x = torch.linspace(-ksize_half, ksize_half, steps=kernel_size, dtype=dtype, device=device) + pdf = torch.exp(-0.5 * (x / sigma).pow(2)) + kernel1d = pdf / pdf.sum() + + return kernel1d + + +def _get_gaussian_kernel2d( + kernel_size: List[int], sigma: List[float], dtype: torch.dtype, device: torch.device +) -> Tensor: + kernel1d_x = _get_gaussian_kernel1d(kernel_size[0], sigma[0], dtype, device) + kernel1d_y = _get_gaussian_kernel1d(kernel_size[1], sigma[1], dtype, device) + kernel2d = torch.mm(kernel1d_y[:, None], kernel1d_x[None, :]) + return kernel2d + + +def gaussian_blur(img: Tensor, kernel_size: List[int], sigma: List[float]) -> Tensor: + if not (isinstance(img, torch.Tensor)): + raise TypeError(f"img should be Tensor. Got {type(img)}") + + _assert_image_tensor(img) + + dtype = img.dtype if torch.is_floating_point(img) else torch.float32 + kernel = _get_gaussian_kernel2d(kernel_size, sigma, dtype=dtype, device=img.device) + kernel = kernel.expand(img.shape[-3], 1, kernel.shape[0], kernel.shape[1]) + + img, need_cast, need_squeeze, out_dtype = _cast_squeeze_in(img, [kernel.dtype]) + + # padding = (left, right, top, bottom) + padding = [kernel_size[0] // 2, kernel_size[0] // 2, kernel_size[1] // 2, kernel_size[1] // 2] + img = torch_pad(img, padding, mode="reflect") + img = conv2d(img, kernel, groups=img.shape[-3]) + + img = _cast_squeeze_out(img, need_cast, need_squeeze, out_dtype) + return img + + +def invert(img: Tensor) -> Tensor: + + _assert_image_tensor(img) + + if img.ndim < 3: + raise TypeError(f"Input image tensor should have at least 3 dimensions, but found {img.ndim}") + + _assert_channels(img, [1, 3]) + + return _max_value(img.dtype) - img + + +def posterize(img: Tensor, bits: int) -> Tensor: + + _assert_image_tensor(img) + + if img.ndim < 3: + raise TypeError(f"Input image tensor should have at least 3 dimensions, but found {img.ndim}") + if img.dtype != torch.uint8: + raise TypeError(f"Only torch.uint8 image tensors are supported, but found {img.dtype}") + + _assert_channels(img, [1, 3]) + mask = -int(2 ** (8 - bits)) # JIT-friendly for: ~(2 ** (8 - bits) - 1) + return img & mask + + +def solarize(img: Tensor, threshold: float) -> Tensor: + + _assert_image_tensor(img) + + if img.ndim < 3: + raise TypeError(f"Input image tensor should have at least 3 dimensions, but found {img.ndim}") + + _assert_channels(img, [1, 3]) + + if threshold > _max_value(img.dtype): + raise TypeError("Threshold should be less than bound of img.") + + inverted_img = invert(img) + return torch.where(img >= threshold, inverted_img, img) + + +def _blurred_degenerate_image(img: Tensor) -> Tensor: + dtype = img.dtype if torch.is_floating_point(img) else torch.float32 + + kernel = torch.ones((3, 3), dtype=dtype, device=img.device) + kernel[1, 1] = 5.0 + kernel /= kernel.sum() + kernel = kernel.expand(img.shape[-3], 1, kernel.shape[0], kernel.shape[1]) + + result_tmp, need_cast, need_squeeze, out_dtype = _cast_squeeze_in(img, [kernel.dtype]) + result_tmp = conv2d(result_tmp, kernel, groups=result_tmp.shape[-3]) + result_tmp = _cast_squeeze_out(result_tmp, need_cast, need_squeeze, out_dtype) + + result = img.clone() + result[..., 1:-1, 1:-1] = result_tmp + + return result + + +def adjust_sharpness(img: Tensor, sharpness_factor: float) -> Tensor: + if sharpness_factor < 0: + raise ValueError(f"sharpness_factor ({sharpness_factor}) is not non-negative.") + + _assert_image_tensor(img) + + _assert_channels(img, [1, 3]) + + if img.size(-1) <= 2 or img.size(-2) <= 2: + return img + + return _blend(img, _blurred_degenerate_image(img), sharpness_factor) + + +def autocontrast(img: Tensor) -> Tensor: + + _assert_image_tensor(img) + + if img.ndim < 3: + raise TypeError(f"Input image tensor should have at least 3 dimensions, but found {img.ndim}") + + _assert_channels(img, [1, 3]) + + bound = _max_value(img.dtype) + dtype = img.dtype if torch.is_floating_point(img) else torch.float32 + + minimum = img.amin(dim=(-2, -1), keepdim=True).to(dtype) + maximum = img.amax(dim=(-2, -1), keepdim=True).to(dtype) + scale = bound / (maximum - minimum) + eq_idxs = torch.isfinite(scale).logical_not() + minimum[eq_idxs] = 0 + scale[eq_idxs] = 1 + + return ((img - minimum) * scale).clamp(0, bound).to(img.dtype) + + +def _scale_channel(img_chan: Tensor) -> Tensor: + # TODO: we should expect bincount to always be faster than histc, but this + # isn't always the case. Once + # https://github.com/pytorch/pytorch/issues/53194 is fixed, remove the if + # block and only use bincount. + if img_chan.is_cuda: + hist = torch.histc(img_chan.to(torch.float32), bins=256, min=0, max=255) + else: + hist = torch.bincount(img_chan.reshape(-1), minlength=256) + + nonzero_hist = hist[hist != 0] + step = torch.div(nonzero_hist[:-1].sum(), 255, rounding_mode="floor") + if step == 0: + return img_chan + + lut = torch.div(torch.cumsum(hist, 0) + torch.div(step, 2, rounding_mode="floor"), step, rounding_mode="floor") + lut = torch.nn.functional.pad(lut, [1, 0])[:-1].clamp(0, 255) + + return lut[img_chan.to(torch.int64)].to(torch.uint8) + + +def _equalize_single_image(img: Tensor) -> Tensor: + return torch.stack([_scale_channel(img[c]) for c in range(img.size(0))]) + + +def equalize(img: Tensor) -> Tensor: + + _assert_image_tensor(img) + + if not (3 <= img.ndim <= 4): + raise TypeError(f"Input image tensor should have 3 or 4 dimensions, but found {img.ndim}") + if img.dtype != torch.uint8: + raise TypeError(f"Only torch.uint8 image tensors are supported, but found {img.dtype}") + + _assert_channels(img, [1, 3]) + + if img.ndim == 3: + return _equalize_single_image(img) + + return torch.stack([_equalize_single_image(x) for x in img]) + + +def normalize(tensor: Tensor, mean: List[float], std: List[float], inplace: bool = False) -> Tensor: + _assert_image_tensor(tensor) + + if not tensor.is_floating_point(): + raise TypeError(f"Input tensor should be a float tensor. Got {tensor.dtype}.") + + if tensor.ndim < 3: + raise ValueError( + f"Expected tensor to be a tensor image of size (..., C, H, W). Got tensor.size() = {tensor.size()}" + ) + + if not inplace: + tensor = tensor.clone() + + dtype = tensor.dtype + mean = torch.as_tensor(mean, dtype=dtype, device=tensor.device) + std = torch.as_tensor(std, dtype=dtype, device=tensor.device) + if (std == 0).any(): + raise ValueError(f"std evaluated to zero after conversion to {dtype}, leading to division by zero.") + if mean.ndim == 1: + mean = mean.view(-1, 1, 1) + if std.ndim == 1: + std = std.view(-1, 1, 1) + return tensor.sub_(mean).div_(std) + + +def erase(img: Tensor, i: int, j: int, h: int, w: int, v: Tensor, inplace: bool = False) -> Tensor: + _assert_image_tensor(img) + + if not inplace: + img = img.clone() + + img[..., i : i + h, j : j + w] = v + return img + + +def _create_identity_grid(size: List[int]) -> Tensor: + hw_space = [torch.linspace((-s + 1) / s, (s - 1) / s, s) for s in size] + grid_y, grid_x = torch.meshgrid(hw_space, indexing="ij") + return torch.stack([grid_x, grid_y], -1).unsqueeze(0) # 1 x H x W x 2 + + +def elastic_transform( + img: Tensor, + displacement: Tensor, + interpolation: str = "bilinear", + fill: Optional[Union[int, float, List[float]]] = None, +) -> Tensor: + + if not (isinstance(img, torch.Tensor)): + raise TypeError(f"img should be Tensor. Got {type(img)}") + + size = list(img.shape[-2:]) + displacement = displacement.to(img.device) + + identity_grid = _create_identity_grid(size) + grid = identity_grid.to(img.device) + displacement + return _apply_grid_transform(img, grid, interpolation, fill) diff --git a/lib/python3.12/site-packages/torchvision/transforms/_functional_video.py b/lib/python3.12/site-packages/torchvision/transforms/_functional_video.py new file mode 100644 index 0000000000000000000000000000000000000000..91df7d42cd71fc554aba51fcf5e90db30e3c3851 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/_functional_video.py @@ -0,0 +1,114 @@ +import warnings + +import torch + + +warnings.warn( + "The 'torchvision.transforms._functional_video' module is deprecated since 0.12 and will be removed in the future. " + "Please use the 'torchvision.transforms.functional' module instead." +) + + +def _is_tensor_video_clip(clip): + if not torch.is_tensor(clip): + raise TypeError("clip should be Tensor. Got %s" % type(clip)) + + if not clip.ndimension() == 4: + raise ValueError("clip should be 4D. Got %dD" % clip.dim()) + + return True + + +def crop(clip, i, j, h, w): + """ + Args: + clip (torch.tensor): Video clip to be cropped. Size is (C, T, H, W) + """ + if len(clip.size()) != 4: + raise ValueError("clip should be a 4D tensor") + return clip[..., i : i + h, j : j + w] + + +def resize(clip, target_size, interpolation_mode): + if len(target_size) != 2: + raise ValueError(f"target size should be tuple (height, width), instead got {target_size}") + return torch.nn.functional.interpolate(clip, size=target_size, mode=interpolation_mode, align_corners=False) + + +def resized_crop(clip, i, j, h, w, size, interpolation_mode="bilinear"): + """ + Do spatial cropping and resizing to the video clip + Args: + clip (torch.tensor): Video clip to be cropped. Size is (C, T, H, W) + i (int): i in (i,j) i.e coordinates of the upper left corner. + j (int): j in (i,j) i.e coordinates of the upper left corner. + h (int): Height of the cropped region. + w (int): Width of the cropped region. + size (tuple(int, int)): height and width of resized clip + Returns: + clip (torch.tensor): Resized and cropped clip. Size is (C, T, H, W) + """ + if not _is_tensor_video_clip(clip): + raise ValueError("clip should be a 4D torch.tensor") + clip = crop(clip, i, j, h, w) + clip = resize(clip, size, interpolation_mode) + return clip + + +def center_crop(clip, crop_size): + if not _is_tensor_video_clip(clip): + raise ValueError("clip should be a 4D torch.tensor") + h, w = clip.size(-2), clip.size(-1) + th, tw = crop_size + if h < th or w < tw: + raise ValueError("height and width must be no smaller than crop_size") + + i = int(round((h - th) / 2.0)) + j = int(round((w - tw) / 2.0)) + return crop(clip, i, j, th, tw) + + +def to_tensor(clip): + """ + Convert tensor data type from uint8 to float, divide value by 255.0 and + permute the dimensions of clip tensor + Args: + clip (torch.tensor, dtype=torch.uint8): Size is (T, H, W, C) + Return: + clip (torch.tensor, dtype=torch.float): Size is (C, T, H, W) + """ + _is_tensor_video_clip(clip) + if not clip.dtype == torch.uint8: + raise TypeError("clip tensor should have data type uint8. Got %s" % str(clip.dtype)) + return clip.float().permute(3, 0, 1, 2) / 255.0 + + +def normalize(clip, mean, std, inplace=False): + """ + Args: + clip (torch.tensor): Video clip to be normalized. Size is (C, T, H, W) + mean (tuple): pixel RGB mean. Size is (3) + std (tuple): pixel standard deviation. Size is (3) + Returns: + normalized clip (torch.tensor): Size is (C, T, H, W) + """ + if not _is_tensor_video_clip(clip): + raise ValueError("clip should be a 4D torch.tensor") + if not inplace: + clip = clip.clone() + mean = torch.as_tensor(mean, dtype=clip.dtype, device=clip.device) + std = torch.as_tensor(std, dtype=clip.dtype, device=clip.device) + clip.sub_(mean[:, None, None, None]).div_(std[:, None, None, None]) + return clip + + +def hflip(clip): + """ + Args: + clip (torch.tensor): Video clip to be normalized. Size is (C, T, H, W) + Returns: + flipped clip (torch.tensor): Size is (C, T, H, W) + """ + if not _is_tensor_video_clip(clip): + raise ValueError("clip should be a 4D torch.tensor") + return clip.flip(-1) diff --git a/lib/python3.12/site-packages/torchvision/transforms/_presets.py b/lib/python3.12/site-packages/torchvision/transforms/_presets.py new file mode 100644 index 0000000000000000000000000000000000000000..fb6f4ad5ca543ddaf41dbdaa4161cd89193f5506 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/_presets.py @@ -0,0 +1,216 @@ +""" +This file is part of the private API. Please do not use directly these classes as they will be modified on +future versions without warning. The classes should be accessed only via the transforms argument of Weights. +""" +from typing import Optional, Tuple, Union + +import torch +from torch import nn, Tensor + +from . import functional as F, InterpolationMode + + +__all__ = [ + "ObjectDetection", + "ImageClassification", + "VideoClassification", + "SemanticSegmentation", + "OpticalFlow", +] + + +class ObjectDetection(nn.Module): + def forward(self, img: Tensor) -> Tensor: + if not isinstance(img, Tensor): + img = F.pil_to_tensor(img) + return F.convert_image_dtype(img, torch.float) + + def __repr__(self) -> str: + return self.__class__.__name__ + "()" + + def describe(self) -> str: + return ( + "Accepts ``PIL.Image``, batched ``(B, C, H, W)`` and single ``(C, H, W)`` image ``torch.Tensor`` objects. " + "The images are rescaled to ``[0.0, 1.0]``." + ) + + +class ImageClassification(nn.Module): + def __init__( + self, + *, + crop_size: int, + resize_size: int = 256, + mean: Tuple[float, ...] = (0.485, 0.456, 0.406), + std: Tuple[float, ...] = (0.229, 0.224, 0.225), + interpolation: InterpolationMode = InterpolationMode.BILINEAR, + antialias: Optional[bool] = True, + ) -> None: + super().__init__() + self.crop_size = [crop_size] + self.resize_size = [resize_size] + self.mean = list(mean) + self.std = list(std) + self.interpolation = interpolation + self.antialias = antialias + + def forward(self, img: Tensor) -> Tensor: + img = F.resize(img, self.resize_size, interpolation=self.interpolation, antialias=self.antialias) + img = F.center_crop(img, self.crop_size) + if not isinstance(img, Tensor): + img = F.pil_to_tensor(img) + img = F.convert_image_dtype(img, torch.float) + img = F.normalize(img, mean=self.mean, std=self.std) + return img + + def __repr__(self) -> str: + format_string = self.__class__.__name__ + "(" + format_string += f"\n crop_size={self.crop_size}" + format_string += f"\n resize_size={self.resize_size}" + format_string += f"\n mean={self.mean}" + format_string += f"\n std={self.std}" + format_string += f"\n interpolation={self.interpolation}" + format_string += "\n)" + return format_string + + def describe(self) -> str: + return ( + "Accepts ``PIL.Image``, batched ``(B, C, H, W)`` and single ``(C, H, W)`` image ``torch.Tensor`` objects. " + f"The images are resized to ``resize_size={self.resize_size}`` using ``interpolation={self.interpolation}``, " + f"followed by a central crop of ``crop_size={self.crop_size}``. Finally the values are first rescaled to " + f"``[0.0, 1.0]`` and then normalized using ``mean={self.mean}`` and ``std={self.std}``." + ) + + +class VideoClassification(nn.Module): + def __init__( + self, + *, + crop_size: Tuple[int, int], + resize_size: Union[Tuple[int], Tuple[int, int]], + mean: Tuple[float, ...] = (0.43216, 0.394666, 0.37645), + std: Tuple[float, ...] = (0.22803, 0.22145, 0.216989), + interpolation: InterpolationMode = InterpolationMode.BILINEAR, + ) -> None: + super().__init__() + self.crop_size = list(crop_size) + self.resize_size = list(resize_size) + self.mean = list(mean) + self.std = list(std) + self.interpolation = interpolation + + def forward(self, vid: Tensor) -> Tensor: + need_squeeze = False + if vid.ndim < 5: + vid = vid.unsqueeze(dim=0) + need_squeeze = True + + N, T, C, H, W = vid.shape + vid = vid.view(-1, C, H, W) + # We hard-code antialias=False to preserve results after we changed + # its default from None to True (see + # https://github.com/pytorch/vision/pull/7160) + # TODO: we could re-train the video models with antialias=True? + vid = F.resize(vid, self.resize_size, interpolation=self.interpolation, antialias=False) + vid = F.center_crop(vid, self.crop_size) + vid = F.convert_image_dtype(vid, torch.float) + vid = F.normalize(vid, mean=self.mean, std=self.std) + H, W = self.crop_size + vid = vid.view(N, T, C, H, W) + vid = vid.permute(0, 2, 1, 3, 4) # (N, T, C, H, W) => (N, C, T, H, W) + + if need_squeeze: + vid = vid.squeeze(dim=0) + return vid + + def __repr__(self) -> str: + format_string = self.__class__.__name__ + "(" + format_string += f"\n crop_size={self.crop_size}" + format_string += f"\n resize_size={self.resize_size}" + format_string += f"\n mean={self.mean}" + format_string += f"\n std={self.std}" + format_string += f"\n interpolation={self.interpolation}" + format_string += "\n)" + return format_string + + def describe(self) -> str: + return ( + "Accepts batched ``(B, T, C, H, W)`` and single ``(T, C, H, W)`` video frame ``torch.Tensor`` objects. " + f"The frames are resized to ``resize_size={self.resize_size}`` using ``interpolation={self.interpolation}``, " + f"followed by a central crop of ``crop_size={self.crop_size}``. Finally the values are first rescaled to " + f"``[0.0, 1.0]`` and then normalized using ``mean={self.mean}`` and ``std={self.std}``. Finally the output " + "dimensions are permuted to ``(..., C, T, H, W)`` tensors." + ) + + +class SemanticSegmentation(nn.Module): + def __init__( + self, + *, + resize_size: Optional[int], + mean: Tuple[float, ...] = (0.485, 0.456, 0.406), + std: Tuple[float, ...] = (0.229, 0.224, 0.225), + interpolation: InterpolationMode = InterpolationMode.BILINEAR, + antialias: Optional[bool] = True, + ) -> None: + super().__init__() + self.resize_size = [resize_size] if resize_size is not None else None + self.mean = list(mean) + self.std = list(std) + self.interpolation = interpolation + self.antialias = antialias + + def forward(self, img: Tensor) -> Tensor: + if isinstance(self.resize_size, list): + img = F.resize(img, self.resize_size, interpolation=self.interpolation, antialias=self.antialias) + if not isinstance(img, Tensor): + img = F.pil_to_tensor(img) + img = F.convert_image_dtype(img, torch.float) + img = F.normalize(img, mean=self.mean, std=self.std) + return img + + def __repr__(self) -> str: + format_string = self.__class__.__name__ + "(" + format_string += f"\n resize_size={self.resize_size}" + format_string += f"\n mean={self.mean}" + format_string += f"\n std={self.std}" + format_string += f"\n interpolation={self.interpolation}" + format_string += "\n)" + return format_string + + def describe(self) -> str: + return ( + "Accepts ``PIL.Image``, batched ``(B, C, H, W)`` and single ``(C, H, W)`` image ``torch.Tensor`` objects. " + f"The images are resized to ``resize_size={self.resize_size}`` using ``interpolation={self.interpolation}``. " + f"Finally the values are first rescaled to ``[0.0, 1.0]`` and then normalized using ``mean={self.mean}`` and " + f"``std={self.std}``." + ) + + +class OpticalFlow(nn.Module): + def forward(self, img1: Tensor, img2: Tensor) -> Tuple[Tensor, Tensor]: + if not isinstance(img1, Tensor): + img1 = F.pil_to_tensor(img1) + if not isinstance(img2, Tensor): + img2 = F.pil_to_tensor(img2) + + img1 = F.convert_image_dtype(img1, torch.float) + img2 = F.convert_image_dtype(img2, torch.float) + + # map [0, 1] into [-1, 1] + img1 = F.normalize(img1, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) + img2 = F.normalize(img2, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) + + img1 = img1.contiguous() + img2 = img2.contiguous() + + return img1, img2 + + def __repr__(self) -> str: + return self.__class__.__name__ + "()" + + def describe(self) -> str: + return ( + "Accepts ``PIL.Image``, batched ``(B, C, H, W)`` and single ``(C, H, W)`` image ``torch.Tensor`` objects. " + "The images are rescaled to ``[-1.0, 1.0]``." + ) diff --git a/lib/python3.12/site-packages/torchvision/transforms/_transforms_video.py b/lib/python3.12/site-packages/torchvision/transforms/_transforms_video.py new file mode 100644 index 0000000000000000000000000000000000000000..a04da4f74849805641e4c470f6b6b8d5f7000e3a --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/_transforms_video.py @@ -0,0 +1,174 @@ +#!/usr/bin/env python3 + +import numbers +import random +import warnings + +from torchvision.transforms import RandomCrop, RandomResizedCrop + +from . import _functional_video as F + + +__all__ = [ + "RandomCropVideo", + "RandomResizedCropVideo", + "CenterCropVideo", + "NormalizeVideo", + "ToTensorVideo", + "RandomHorizontalFlipVideo", +] + + +warnings.warn( + "The 'torchvision.transforms._transforms_video' module is deprecated since 0.12 and will be removed in the future. " + "Please use the 'torchvision.transforms' module instead." +) + + +class RandomCropVideo(RandomCrop): + def __init__(self, size): + if isinstance(size, numbers.Number): + self.size = (int(size), int(size)) + else: + self.size = size + + def __call__(self, clip): + """ + Args: + clip (torch.tensor): Video clip to be cropped. Size is (C, T, H, W) + Returns: + torch.tensor: randomly cropped/resized video clip. + size is (C, T, OH, OW) + """ + i, j, h, w = self.get_params(clip, self.size) + return F.crop(clip, i, j, h, w) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(size={self.size})" + + +class RandomResizedCropVideo(RandomResizedCrop): + def __init__( + self, + size, + scale=(0.08, 1.0), + ratio=(3.0 / 4.0, 4.0 / 3.0), + interpolation_mode="bilinear", + ): + if isinstance(size, tuple): + if len(size) != 2: + raise ValueError(f"size should be tuple (height, width), instead got {size}") + self.size = size + else: + self.size = (size, size) + + self.interpolation_mode = interpolation_mode + self.scale = scale + self.ratio = ratio + + def __call__(self, clip): + """ + Args: + clip (torch.tensor): Video clip to be cropped. Size is (C, T, H, W) + Returns: + torch.tensor: randomly cropped/resized video clip. + size is (C, T, H, W) + """ + i, j, h, w = self.get_params(clip, self.scale, self.ratio) + return F.resized_crop(clip, i, j, h, w, self.size, self.interpolation_mode) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}, scale={self.scale}, ratio={self.ratio})" + + +class CenterCropVideo: + def __init__(self, crop_size): + if isinstance(crop_size, numbers.Number): + self.crop_size = (int(crop_size), int(crop_size)) + else: + self.crop_size = crop_size + + def __call__(self, clip): + """ + Args: + clip (torch.tensor): Video clip to be cropped. Size is (C, T, H, W) + Returns: + torch.tensor: central cropping of video clip. Size is + (C, T, crop_size, crop_size) + """ + return F.center_crop(clip, self.crop_size) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(crop_size={self.crop_size})" + + +class NormalizeVideo: + """ + Normalize the video clip by mean subtraction and division by standard deviation + Args: + mean (3-tuple): pixel RGB mean + std (3-tuple): pixel RGB standard deviation + inplace (boolean): whether do in-place normalization + """ + + def __init__(self, mean, std, inplace=False): + self.mean = mean + self.std = std + self.inplace = inplace + + def __call__(self, clip): + """ + Args: + clip (torch.tensor): video clip to be normalized. Size is (C, T, H, W) + """ + return F.normalize(clip, self.mean, self.std, self.inplace) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(mean={self.mean}, std={self.std}, inplace={self.inplace})" + + +class ToTensorVideo: + """ + Convert tensor data type from uint8 to float, divide value by 255.0 and + permute the dimensions of clip tensor + """ + + def __init__(self): + pass + + def __call__(self, clip): + """ + Args: + clip (torch.tensor, dtype=torch.uint8): Size is (T, H, W, C) + Return: + clip (torch.tensor, dtype=torch.float): Size is (C, T, H, W) + """ + return F.to_tensor(clip) + + def __repr__(self) -> str: + return self.__class__.__name__ + + +class RandomHorizontalFlipVideo: + """ + Flip the video clip along the horizontal direction with a given probability + Args: + p (float): probability of the clip being flipped. Default value is 0.5 + """ + + def __init__(self, p=0.5): + self.p = p + + def __call__(self, clip): + """ + Args: + clip (torch.tensor): Size is (C, T, H, W) + Return: + clip (torch.tensor): Size is (C, T, H, W) + """ + if random.random() < self.p: + clip = F.hflip(clip) + return clip + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(p={self.p})" diff --git a/lib/python3.12/site-packages/torchvision/transforms/autoaugment.py b/lib/python3.12/site-packages/torchvision/transforms/autoaugment.py new file mode 100644 index 0000000000000000000000000000000000000000..9dbbe91e741093b01ff8491ba8b39d9b6f578103 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/autoaugment.py @@ -0,0 +1,615 @@ +import math +from enum import Enum +from typing import Dict, List, Optional, Tuple + +import torch +from torch import Tensor + +from . import functional as F, InterpolationMode + +__all__ = ["AutoAugmentPolicy", "AutoAugment", "RandAugment", "TrivialAugmentWide", "AugMix"] + + +def _apply_op( + img: Tensor, op_name: str, magnitude: float, interpolation: InterpolationMode, fill: Optional[List[float]] +): + if op_name == "ShearX": + # magnitude should be arctan(magnitude) + # official autoaug: (1, level, 0, 0, 1, 0) + # https://github.com/tensorflow/models/blob/dd02069717128186b88afa8d857ce57d17957f03/research/autoaugment/augmentation_transforms.py#L290 + # compared to + # torchvision: (1, tan(level), 0, 0, 1, 0) + # https://github.com/pytorch/vision/blob/0c2373d0bba3499e95776e7936e207d8a1676e65/torchvision/transforms/functional.py#L976 + img = F.affine( + img, + angle=0.0, + translate=[0, 0], + scale=1.0, + shear=[math.degrees(math.atan(magnitude)), 0.0], + interpolation=interpolation, + fill=fill, + center=[0, 0], + ) + elif op_name == "ShearY": + # magnitude should be arctan(magnitude) + # See above + img = F.affine( + img, + angle=0.0, + translate=[0, 0], + scale=1.0, + shear=[0.0, math.degrees(math.atan(magnitude))], + interpolation=interpolation, + fill=fill, + center=[0, 0], + ) + elif op_name == "TranslateX": + img = F.affine( + img, + angle=0.0, + translate=[int(magnitude), 0], + scale=1.0, + interpolation=interpolation, + shear=[0.0, 0.0], + fill=fill, + ) + elif op_name == "TranslateY": + img = F.affine( + img, + angle=0.0, + translate=[0, int(magnitude)], + scale=1.0, + interpolation=interpolation, + shear=[0.0, 0.0], + fill=fill, + ) + elif op_name == "Rotate": + img = F.rotate(img, magnitude, interpolation=interpolation, fill=fill) + elif op_name == "Brightness": + img = F.adjust_brightness(img, 1.0 + magnitude) + elif op_name == "Color": + img = F.adjust_saturation(img, 1.0 + magnitude) + elif op_name == "Contrast": + img = F.adjust_contrast(img, 1.0 + magnitude) + elif op_name == "Sharpness": + img = F.adjust_sharpness(img, 1.0 + magnitude) + elif op_name == "Posterize": + img = F.posterize(img, int(magnitude)) + elif op_name == "Solarize": + img = F.solarize(img, magnitude) + elif op_name == "AutoContrast": + img = F.autocontrast(img) + elif op_name == "Equalize": + img = F.equalize(img) + elif op_name == "Invert": + img = F.invert(img) + elif op_name == "Identity": + pass + else: + raise ValueError(f"The provided operator {op_name} is not recognized.") + return img + + +class AutoAugmentPolicy(Enum): + """AutoAugment policies learned on different datasets. + Available policies are IMAGENET, CIFAR10 and SVHN. + """ + + IMAGENET = "imagenet" + CIFAR10 = "cifar10" + SVHN = "svhn" + + +# FIXME: Eliminate copy-pasted code for fill standardization and _augmentation_space() by moving stuff on a base class +class AutoAugment(torch.nn.Module): + r"""AutoAugment data augmentation method based on + `"AutoAugment: Learning Augmentation Strategies from Data" `_. + If the image is torch Tensor, it should be of type torch.uint8, and it is expected + to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + + Args: + policy (AutoAugmentPolicy): Desired policy enum defined by + :class:`torchvision.transforms.autoaugment.AutoAugmentPolicy`. Default is ``AutoAugmentPolicy.IMAGENET``. + interpolation (InterpolationMode): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. + fill (sequence or number, optional): Pixel fill value for the area outside the transformed + image. If given a number, the value is used for all bands respectively. + """ + + def __init__( + self, + policy: AutoAugmentPolicy = AutoAugmentPolicy.IMAGENET, + interpolation: InterpolationMode = InterpolationMode.NEAREST, + fill: Optional[List[float]] = None, + ) -> None: + super().__init__() + self.policy = policy + self.interpolation = interpolation + self.fill = fill + self.policies = self._get_policies(policy) + + def _get_policies( + self, policy: AutoAugmentPolicy + ) -> List[Tuple[Tuple[str, float, Optional[int]], Tuple[str, float, Optional[int]]]]: + if policy == AutoAugmentPolicy.IMAGENET: + return [ + (("Posterize", 0.4, 8), ("Rotate", 0.6, 9)), + (("Solarize", 0.6, 5), ("AutoContrast", 0.6, None)), + (("Equalize", 0.8, None), ("Equalize", 0.6, None)), + (("Posterize", 0.6, 7), ("Posterize", 0.6, 6)), + (("Equalize", 0.4, None), ("Solarize", 0.2, 4)), + (("Equalize", 0.4, None), ("Rotate", 0.8, 8)), + (("Solarize", 0.6, 3), ("Equalize", 0.6, None)), + (("Posterize", 0.8, 5), ("Equalize", 1.0, None)), + (("Rotate", 0.2, 3), ("Solarize", 0.6, 8)), + (("Equalize", 0.6, None), ("Posterize", 0.4, 6)), + (("Rotate", 0.8, 8), ("Color", 0.4, 0)), + (("Rotate", 0.4, 9), ("Equalize", 0.6, None)), + (("Equalize", 0.0, None), ("Equalize", 0.8, None)), + (("Invert", 0.6, None), ("Equalize", 1.0, None)), + (("Color", 0.6, 4), ("Contrast", 1.0, 8)), + (("Rotate", 0.8, 8), ("Color", 1.0, 2)), + (("Color", 0.8, 8), ("Solarize", 0.8, 7)), + (("Sharpness", 0.4, 7), ("Invert", 0.6, None)), + (("ShearX", 0.6, 5), ("Equalize", 1.0, None)), + (("Color", 0.4, 0), ("Equalize", 0.6, None)), + (("Equalize", 0.4, None), ("Solarize", 0.2, 4)), + (("Solarize", 0.6, 5), ("AutoContrast", 0.6, None)), + (("Invert", 0.6, None), ("Equalize", 1.0, None)), + (("Color", 0.6, 4), ("Contrast", 1.0, 8)), + (("Equalize", 0.8, None), ("Equalize", 0.6, None)), + ] + elif policy == AutoAugmentPolicy.CIFAR10: + return [ + (("Invert", 0.1, None), ("Contrast", 0.2, 6)), + (("Rotate", 0.7, 2), ("TranslateX", 0.3, 9)), + (("Sharpness", 0.8, 1), ("Sharpness", 0.9, 3)), + (("ShearY", 0.5, 8), ("TranslateY", 0.7, 9)), + (("AutoContrast", 0.5, None), ("Equalize", 0.9, None)), + (("ShearY", 0.2, 7), ("Posterize", 0.3, 7)), + (("Color", 0.4, 3), ("Brightness", 0.6, 7)), + (("Sharpness", 0.3, 9), ("Brightness", 0.7, 9)), + (("Equalize", 0.6, None), ("Equalize", 0.5, None)), + (("Contrast", 0.6, 7), ("Sharpness", 0.6, 5)), + (("Color", 0.7, 7), ("TranslateX", 0.5, 8)), + (("Equalize", 0.3, None), ("AutoContrast", 0.4, None)), + (("TranslateY", 0.4, 3), ("Sharpness", 0.2, 6)), + (("Brightness", 0.9, 6), ("Color", 0.2, 8)), + (("Solarize", 0.5, 2), ("Invert", 0.0, None)), + (("Equalize", 0.2, None), ("AutoContrast", 0.6, None)), + (("Equalize", 0.2, None), ("Equalize", 0.6, None)), + (("Color", 0.9, 9), ("Equalize", 0.6, None)), + (("AutoContrast", 0.8, None), ("Solarize", 0.2, 8)), + (("Brightness", 0.1, 3), ("Color", 0.7, 0)), + (("Solarize", 0.4, 5), ("AutoContrast", 0.9, None)), + (("TranslateY", 0.9, 9), ("TranslateY", 0.7, 9)), + (("AutoContrast", 0.9, None), ("Solarize", 0.8, 3)), + (("Equalize", 0.8, None), ("Invert", 0.1, None)), + (("TranslateY", 0.7, 9), ("AutoContrast", 0.9, None)), + ] + elif policy == AutoAugmentPolicy.SVHN: + return [ + (("ShearX", 0.9, 4), ("Invert", 0.2, None)), + (("ShearY", 0.9, 8), ("Invert", 0.7, None)), + (("Equalize", 0.6, None), ("Solarize", 0.6, 6)), + (("Invert", 0.9, None), ("Equalize", 0.6, None)), + (("Equalize", 0.6, None), ("Rotate", 0.9, 3)), + (("ShearX", 0.9, 4), ("AutoContrast", 0.8, None)), + (("ShearY", 0.9, 8), ("Invert", 0.4, None)), + (("ShearY", 0.9, 5), ("Solarize", 0.2, 6)), + (("Invert", 0.9, None), ("AutoContrast", 0.8, None)), + (("Equalize", 0.6, None), ("Rotate", 0.9, 3)), + (("ShearX", 0.9, 4), ("Solarize", 0.3, 3)), + (("ShearY", 0.8, 8), ("Invert", 0.7, None)), + (("Equalize", 0.9, None), ("TranslateY", 0.6, 6)), + (("Invert", 0.9, None), ("Equalize", 0.6, None)), + (("Contrast", 0.3, 3), ("Rotate", 0.8, 4)), + (("Invert", 0.8, None), ("TranslateY", 0.0, 2)), + (("ShearY", 0.7, 6), ("Solarize", 0.4, 8)), + (("Invert", 0.6, None), ("Rotate", 0.8, 4)), + (("ShearY", 0.3, 7), ("TranslateX", 0.9, 3)), + (("ShearX", 0.1, 6), ("Invert", 0.6, None)), + (("Solarize", 0.7, 2), ("TranslateY", 0.6, 7)), + (("ShearY", 0.8, 4), ("Invert", 0.8, None)), + (("ShearX", 0.7, 9), ("TranslateY", 0.8, 3)), + (("ShearY", 0.8, 5), ("AutoContrast", 0.7, None)), + (("ShearX", 0.7, 2), ("Invert", 0.1, None)), + ] + else: + raise ValueError(f"The provided policy {policy} is not recognized.") + + def _augmentation_space(self, num_bins: int, image_size: Tuple[int, int]) -> Dict[str, Tuple[Tensor, bool]]: + return { + # op_name: (magnitudes, signed) + "ShearX": (torch.linspace(0.0, 0.3, num_bins), True), + "ShearY": (torch.linspace(0.0, 0.3, num_bins), True), + "TranslateX": (torch.linspace(0.0, 150.0 / 331.0 * image_size[1], num_bins), True), + "TranslateY": (torch.linspace(0.0, 150.0 / 331.0 * image_size[0], num_bins), True), + "Rotate": (torch.linspace(0.0, 30.0, num_bins), True), + "Brightness": (torch.linspace(0.0, 0.9, num_bins), True), + "Color": (torch.linspace(0.0, 0.9, num_bins), True), + "Contrast": (torch.linspace(0.0, 0.9, num_bins), True), + "Sharpness": (torch.linspace(0.0, 0.9, num_bins), True), + "Posterize": (8 - (torch.arange(num_bins) / ((num_bins - 1) / 4)).round().int(), False), + "Solarize": (torch.linspace(255.0, 0.0, num_bins), False), + "AutoContrast": (torch.tensor(0.0), False), + "Equalize": (torch.tensor(0.0), False), + "Invert": (torch.tensor(0.0), False), + } + + @staticmethod + def get_params(transform_num: int) -> Tuple[int, Tensor, Tensor]: + """Get parameters for autoaugment transformation + + Returns: + params required by the autoaugment transformation + """ + policy_id = int(torch.randint(transform_num, (1,)).item()) + probs = torch.rand((2,)) + signs = torch.randint(2, (2,)) + + return policy_id, probs, signs + + def forward(self, img: Tensor) -> Tensor: + """ + img (PIL Image or Tensor): Image to be transformed. + + Returns: + PIL Image or Tensor: AutoAugmented image. + """ + fill = self.fill + channels, height, width = F.get_dimensions(img) + if isinstance(img, Tensor): + if isinstance(fill, (int, float)): + fill = [float(fill)] * channels + elif fill is not None: + fill = [float(f) for f in fill] + + transform_id, probs, signs = self.get_params(len(self.policies)) + + op_meta = self._augmentation_space(10, (height, width)) + for i, (op_name, p, magnitude_id) in enumerate(self.policies[transform_id]): + if probs[i] <= p: + magnitudes, signed = op_meta[op_name] + magnitude = float(magnitudes[magnitude_id].item()) if magnitude_id is not None else 0.0 + if signed and signs[i] == 0: + magnitude *= -1.0 + img = _apply_op(img, op_name, magnitude, interpolation=self.interpolation, fill=fill) + + return img + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(policy={self.policy}, fill={self.fill})" + + +class RandAugment(torch.nn.Module): + r"""RandAugment data augmentation method based on + `"RandAugment: Practical automated data augmentation with a reduced search space" + `_. + If the image is torch Tensor, it should be of type torch.uint8, and it is expected + to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + + Args: + num_ops (int): Number of augmentation transformations to apply sequentially. + magnitude (int): Magnitude for all the transformations. + num_magnitude_bins (int): The number of different magnitude values. + interpolation (InterpolationMode): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. + fill (sequence or number, optional): Pixel fill value for the area outside the transformed + image. If given a number, the value is used for all bands respectively. + """ + + def __init__( + self, + num_ops: int = 2, + magnitude: int = 9, + num_magnitude_bins: int = 31, + interpolation: InterpolationMode = InterpolationMode.NEAREST, + fill: Optional[List[float]] = None, + ) -> None: + super().__init__() + self.num_ops = num_ops + self.magnitude = magnitude + self.num_magnitude_bins = num_magnitude_bins + self.interpolation = interpolation + self.fill = fill + + def _augmentation_space(self, num_bins: int, image_size: Tuple[int, int]) -> Dict[str, Tuple[Tensor, bool]]: + return { + # op_name: (magnitudes, signed) + "Identity": (torch.tensor(0.0), False), + "ShearX": (torch.linspace(0.0, 0.3, num_bins), True), + "ShearY": (torch.linspace(0.0, 0.3, num_bins), True), + "TranslateX": (torch.linspace(0.0, 150.0 / 331.0 * image_size[1], num_bins), True), + "TranslateY": (torch.linspace(0.0, 150.0 / 331.0 * image_size[0], num_bins), True), + "Rotate": (torch.linspace(0.0, 30.0, num_bins), True), + "Brightness": (torch.linspace(0.0, 0.9, num_bins), True), + "Color": (torch.linspace(0.0, 0.9, num_bins), True), + "Contrast": (torch.linspace(0.0, 0.9, num_bins), True), + "Sharpness": (torch.linspace(0.0, 0.9, num_bins), True), + "Posterize": (8 - (torch.arange(num_bins) / ((num_bins - 1) / 4)).round().int(), False), + "Solarize": (torch.linspace(255.0, 0.0, num_bins), False), + "AutoContrast": (torch.tensor(0.0), False), + "Equalize": (torch.tensor(0.0), False), + } + + def forward(self, img: Tensor) -> Tensor: + """ + img (PIL Image or Tensor): Image to be transformed. + + Returns: + PIL Image or Tensor: Transformed image. + """ + fill = self.fill + channels, height, width = F.get_dimensions(img) + if isinstance(img, Tensor): + if isinstance(fill, (int, float)): + fill = [float(fill)] * channels + elif fill is not None: + fill = [float(f) for f in fill] + + op_meta = self._augmentation_space(self.num_magnitude_bins, (height, width)) + for _ in range(self.num_ops): + op_index = int(torch.randint(len(op_meta), (1,)).item()) + op_name = list(op_meta.keys())[op_index] + magnitudes, signed = op_meta[op_name] + magnitude = float(magnitudes[self.magnitude].item()) if magnitudes.ndim > 0 else 0.0 + if signed and torch.randint(2, (1,)): + magnitude *= -1.0 + img = _apply_op(img, op_name, magnitude, interpolation=self.interpolation, fill=fill) + + return img + + def __repr__(self) -> str: + s = ( + f"{self.__class__.__name__}(" + f"num_ops={self.num_ops}" + f", magnitude={self.magnitude}" + f", num_magnitude_bins={self.num_magnitude_bins}" + f", interpolation={self.interpolation}" + f", fill={self.fill}" + f")" + ) + return s + + +class TrivialAugmentWide(torch.nn.Module): + r"""Dataset-independent data-augmentation with TrivialAugment Wide, as described in + `"TrivialAugment: Tuning-free Yet State-of-the-Art Data Augmentation" `_. + If the image is torch Tensor, it should be of type torch.uint8, and it is expected + to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + + Args: + num_magnitude_bins (int): The number of different magnitude values. + interpolation (InterpolationMode): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. + fill (sequence or number, optional): Pixel fill value for the area outside the transformed + image. If given a number, the value is used for all bands respectively. + """ + + def __init__( + self, + num_magnitude_bins: int = 31, + interpolation: InterpolationMode = InterpolationMode.NEAREST, + fill: Optional[List[float]] = None, + ) -> None: + super().__init__() + self.num_magnitude_bins = num_magnitude_bins + self.interpolation = interpolation + self.fill = fill + + def _augmentation_space(self, num_bins: int) -> Dict[str, Tuple[Tensor, bool]]: + return { + # op_name: (magnitudes, signed) + "Identity": (torch.tensor(0.0), False), + "ShearX": (torch.linspace(0.0, 0.99, num_bins), True), + "ShearY": (torch.linspace(0.0, 0.99, num_bins), True), + "TranslateX": (torch.linspace(0.0, 32.0, num_bins), True), + "TranslateY": (torch.linspace(0.0, 32.0, num_bins), True), + "Rotate": (torch.linspace(0.0, 135.0, num_bins), True), + "Brightness": (torch.linspace(0.0, 0.99, num_bins), True), + "Color": (torch.linspace(0.0, 0.99, num_bins), True), + "Contrast": (torch.linspace(0.0, 0.99, num_bins), True), + "Sharpness": (torch.linspace(0.0, 0.99, num_bins), True), + "Posterize": (8 - (torch.arange(num_bins) / ((num_bins - 1) / 6)).round().int(), False), + "Solarize": (torch.linspace(255.0, 0.0, num_bins), False), + "AutoContrast": (torch.tensor(0.0), False), + "Equalize": (torch.tensor(0.0), False), + } + + def forward(self, img: Tensor) -> Tensor: + """ + img (PIL Image or Tensor): Image to be transformed. + + Returns: + PIL Image or Tensor: Transformed image. + """ + fill = self.fill + channels, height, width = F.get_dimensions(img) + if isinstance(img, Tensor): + if isinstance(fill, (int, float)): + fill = [float(fill)] * channels + elif fill is not None: + fill = [float(f) for f in fill] + + op_meta = self._augmentation_space(self.num_magnitude_bins) + op_index = int(torch.randint(len(op_meta), (1,)).item()) + op_name = list(op_meta.keys())[op_index] + magnitudes, signed = op_meta[op_name] + magnitude = ( + float(magnitudes[torch.randint(len(magnitudes), (1,), dtype=torch.long)].item()) + if magnitudes.ndim > 0 + else 0.0 + ) + if signed and torch.randint(2, (1,)): + magnitude *= -1.0 + + return _apply_op(img, op_name, magnitude, interpolation=self.interpolation, fill=fill) + + def __repr__(self) -> str: + s = ( + f"{self.__class__.__name__}(" + f"num_magnitude_bins={self.num_magnitude_bins}" + f", interpolation={self.interpolation}" + f", fill={self.fill}" + f")" + ) + return s + + +class AugMix(torch.nn.Module): + r"""AugMix data augmentation method based on + `"AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty" `_. + If the image is torch Tensor, it should be of type torch.uint8, and it is expected + to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + + Args: + severity (int): The severity of base augmentation operators. Default is ``3``. + mixture_width (int): The number of augmentation chains. Default is ``3``. + chain_depth (int): The depth of augmentation chains. A negative value denotes stochastic depth sampled from the interval [1, 3]. + Default is ``-1``. + alpha (float): The hyperparameter for the probability distributions. Default is ``1.0``. + all_ops (bool): Use all operations (including brightness, contrast, color and sharpness). Default is ``True``. + interpolation (InterpolationMode): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. + fill (sequence or number, optional): Pixel fill value for the area outside the transformed + image. If given a number, the value is used for all bands respectively. + """ + + def __init__( + self, + severity: int = 3, + mixture_width: int = 3, + chain_depth: int = -1, + alpha: float = 1.0, + all_ops: bool = True, + interpolation: InterpolationMode = InterpolationMode.BILINEAR, + fill: Optional[List[float]] = None, + ) -> None: + super().__init__() + self._PARAMETER_MAX = 10 + if not (1 <= severity <= self._PARAMETER_MAX): + raise ValueError(f"The severity must be between [1, {self._PARAMETER_MAX}]. Got {severity} instead.") + self.severity = severity + self.mixture_width = mixture_width + self.chain_depth = chain_depth + self.alpha = alpha + self.all_ops = all_ops + self.interpolation = interpolation + self.fill = fill + + def _augmentation_space(self, num_bins: int, image_size: Tuple[int, int]) -> Dict[str, Tuple[Tensor, bool]]: + s = { + # op_name: (magnitudes, signed) + "ShearX": (torch.linspace(0.0, 0.3, num_bins), True), + "ShearY": (torch.linspace(0.0, 0.3, num_bins), True), + "TranslateX": (torch.linspace(0.0, image_size[1] / 3.0, num_bins), True), + "TranslateY": (torch.linspace(0.0, image_size[0] / 3.0, num_bins), True), + "Rotate": (torch.linspace(0.0, 30.0, num_bins), True), + "Posterize": (4 - (torch.arange(num_bins) / ((num_bins - 1) / 4)).round().int(), False), + "Solarize": (torch.linspace(255.0, 0.0, num_bins), False), + "AutoContrast": (torch.tensor(0.0), False), + "Equalize": (torch.tensor(0.0), False), + } + if self.all_ops: + s.update( + { + "Brightness": (torch.linspace(0.0, 0.9, num_bins), True), + "Color": (torch.linspace(0.0, 0.9, num_bins), True), + "Contrast": (torch.linspace(0.0, 0.9, num_bins), True), + "Sharpness": (torch.linspace(0.0, 0.9, num_bins), True), + } + ) + return s + + @torch.jit.unused + def _pil_to_tensor(self, img) -> Tensor: + return F.pil_to_tensor(img) + + @torch.jit.unused + def _tensor_to_pil(self, img: Tensor): + return F.to_pil_image(img) + + def _sample_dirichlet(self, params: Tensor) -> Tensor: + # Must be on a separate method so that we can overwrite it in tests. + return torch._sample_dirichlet(params) + + def forward(self, orig_img: Tensor) -> Tensor: + """ + img (PIL Image or Tensor): Image to be transformed. + + Returns: + PIL Image or Tensor: Transformed image. + """ + fill = self.fill + channels, height, width = F.get_dimensions(orig_img) + if isinstance(orig_img, Tensor): + img = orig_img + if isinstance(fill, (int, float)): + fill = [float(fill)] * channels + elif fill is not None: + fill = [float(f) for f in fill] + else: + img = self._pil_to_tensor(orig_img) + + op_meta = self._augmentation_space(self._PARAMETER_MAX, (height, width)) + + orig_dims = list(img.shape) + batch = img.view([1] * max(4 - img.ndim, 0) + orig_dims) + batch_dims = [batch.size(0)] + [1] * (batch.ndim - 1) + + # Sample the beta weights for combining the original and augmented image. To get Beta, we use a Dirichlet + # with 2 parameters. The 1st column stores the weights of the original and the 2nd the ones of augmented image. + m = self._sample_dirichlet( + torch.tensor([self.alpha, self.alpha], device=batch.device).expand(batch_dims[0], -1) + ) + + # Sample the mixing weights and combine them with the ones sampled from Beta for the augmented images. + combined_weights = self._sample_dirichlet( + torch.tensor([self.alpha] * self.mixture_width, device=batch.device).expand(batch_dims[0], -1) + ) * m[:, 1].view([batch_dims[0], -1]) + + mix = m[:, 0].view(batch_dims) * batch + for i in range(self.mixture_width): + aug = batch + depth = self.chain_depth if self.chain_depth > 0 else int(torch.randint(low=1, high=4, size=(1,)).item()) + for _ in range(depth): + op_index = int(torch.randint(len(op_meta), (1,)).item()) + op_name = list(op_meta.keys())[op_index] + magnitudes, signed = op_meta[op_name] + magnitude = ( + float(magnitudes[torch.randint(self.severity, (1,), dtype=torch.long)].item()) + if magnitudes.ndim > 0 + else 0.0 + ) + if signed and torch.randint(2, (1,)): + magnitude *= -1.0 + aug = _apply_op(aug, op_name, magnitude, interpolation=self.interpolation, fill=fill) + mix.add_(combined_weights[:, i].view(batch_dims) * aug) + mix = mix.view(orig_dims).to(dtype=img.dtype) + + if not isinstance(orig_img, Tensor): + return self._tensor_to_pil(mix) + return mix + + def __repr__(self) -> str: + s = ( + f"{self.__class__.__name__}(" + f"severity={self.severity}" + f", mixture_width={self.mixture_width}" + f", chain_depth={self.chain_depth}" + f", alpha={self.alpha}" + f", all_ops={self.all_ops}" + f", interpolation={self.interpolation}" + f", fill={self.fill}" + f")" + ) + return s diff --git a/lib/python3.12/site-packages/torchvision/transforms/functional.py b/lib/python3.12/site-packages/torchvision/transforms/functional.py new file mode 100644 index 0000000000000000000000000000000000000000..8efe2a8878a06b37d7ede9496b213076c1f59c01 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/functional.py @@ -0,0 +1,1586 @@ +import math +import numbers +import sys +import warnings +from enum import Enum +from typing import Any, List, Optional, Tuple, Union + +import numpy as np +import torch +from PIL import Image +from PIL.Image import Image as PILImage +from torch import Tensor + +try: + import accimage +except ImportError: + accimage = None + +from ..utils import _log_api_usage_once +from . import _functional_pil as F_pil, _functional_tensor as F_t + + +class InterpolationMode(Enum): + """Interpolation modes + Available interpolation methods are ``nearest``, ``nearest-exact``, ``bilinear``, ``bicubic``, ``box``, ``hamming``, + and ``lanczos``. + """ + + NEAREST = "nearest" + NEAREST_EXACT = "nearest-exact" + BILINEAR = "bilinear" + BICUBIC = "bicubic" + # For PIL compatibility + BOX = "box" + HAMMING = "hamming" + LANCZOS = "lanczos" + + +# TODO: Once torchscript supports Enums with staticmethod +# this can be put into InterpolationMode as staticmethod +def _interpolation_modes_from_int(i: int) -> InterpolationMode: + inverse_modes_mapping = { + 0: InterpolationMode.NEAREST, + 2: InterpolationMode.BILINEAR, + 3: InterpolationMode.BICUBIC, + 4: InterpolationMode.BOX, + 5: InterpolationMode.HAMMING, + 1: InterpolationMode.LANCZOS, + } + return inverse_modes_mapping[i] + + +pil_modes_mapping = { + InterpolationMode.NEAREST: 0, + InterpolationMode.BILINEAR: 2, + InterpolationMode.BICUBIC: 3, + InterpolationMode.NEAREST_EXACT: 0, + InterpolationMode.BOX: 4, + InterpolationMode.HAMMING: 5, + InterpolationMode.LANCZOS: 1, +} + +_is_pil_image = F_pil._is_pil_image + + +def get_dimensions(img: Tensor) -> List[int]: + """Returns the dimensions of an image as [channels, height, width]. + + Args: + img (PIL Image or Tensor): The image to be checked. + + Returns: + List[int]: The image dimensions. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(get_dimensions) + if isinstance(img, torch.Tensor): + return F_t.get_dimensions(img) + + return F_pil.get_dimensions(img) + + +def get_image_size(img: Tensor) -> List[int]: + """Returns the size of an image as [width, height]. + + Args: + img (PIL Image or Tensor): The image to be checked. + + Returns: + List[int]: The image size. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(get_image_size) + if isinstance(img, torch.Tensor): + return F_t.get_image_size(img) + + return F_pil.get_image_size(img) + + +def get_image_num_channels(img: Tensor) -> int: + """Returns the number of channels of an image. + + Args: + img (PIL Image or Tensor): The image to be checked. + + Returns: + int: The number of channels. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(get_image_num_channels) + if isinstance(img, torch.Tensor): + return F_t.get_image_num_channels(img) + + return F_pil.get_image_num_channels(img) + + +@torch.jit.unused +def _is_numpy(img: Any) -> bool: + return isinstance(img, np.ndarray) + + +@torch.jit.unused +def _is_numpy_image(img: Any) -> bool: + return img.ndim in {2, 3} + + +def to_tensor(pic: Union[PILImage, np.ndarray]) -> Tensor: + """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor. + This function does not support torchscript. + + See :class:`~torchvision.transforms.ToTensor` for more details. + + Args: + pic (PIL Image or numpy.ndarray): Image to be converted to tensor. + + Returns: + Tensor: Converted image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(to_tensor) + if not (F_pil._is_pil_image(pic) or _is_numpy(pic)): + raise TypeError(f"pic should be PIL Image or ndarray. Got {type(pic)}") + + if _is_numpy(pic) and not _is_numpy_image(pic): + raise ValueError(f"pic should be 2/3 dimensional. Got {pic.ndim} dimensions.") + + default_float_dtype = torch.get_default_dtype() + + if isinstance(pic, np.ndarray): + # handle numpy array + if pic.ndim == 2: + pic = pic[:, :, None] + + img = torch.from_numpy(pic.transpose((2, 0, 1))).contiguous() + # backward compatibility + if isinstance(img, torch.ByteTensor): + return img.to(dtype=default_float_dtype).div(255) + else: + return img + + if accimage is not None and isinstance(pic, accimage.Image): + nppic = np.zeros([pic.channels, pic.height, pic.width], dtype=np.float32) + pic.copyto(nppic) + return torch.from_numpy(nppic).to(dtype=default_float_dtype) + + # handle PIL Image + mode_to_nptype = {"I": np.int32, "I;16" if sys.byteorder == "little" else "I;16B": np.int16, "F": np.float32} + img = torch.from_numpy(np.array(pic, mode_to_nptype.get(pic.mode, np.uint8), copy=True)) + + if pic.mode == "1": + img = 255 * img + img = img.view(pic.size[1], pic.size[0], F_pil.get_image_num_channels(pic)) + # put it from HWC to CHW format + img = img.permute((2, 0, 1)).contiguous() + if isinstance(img, torch.ByteTensor): + return img.to(dtype=default_float_dtype).div(255) + else: + return img + + +def pil_to_tensor(pic: Any) -> Tensor: + """Convert a ``PIL Image`` to a tensor of the same type. + This function does not support torchscript. + + See :class:`~torchvision.transforms.PILToTensor` for more details. + + .. note:: + + A deep copy of the underlying array is performed. + + Args: + pic (PIL Image): Image to be converted to tensor. + + Returns: + Tensor: Converted image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(pil_to_tensor) + if not F_pil._is_pil_image(pic): + raise TypeError(f"pic should be PIL Image. Got {type(pic)}") + + if accimage is not None and isinstance(pic, accimage.Image): + # accimage format is always uint8 internally, so always return uint8 here + nppic = np.zeros([pic.channels, pic.height, pic.width], dtype=np.uint8) + pic.copyto(nppic) + return torch.as_tensor(nppic) + + # handle PIL Image + img = torch.as_tensor(np.array(pic, copy=True)) + img = img.view(pic.size[1], pic.size[0], F_pil.get_image_num_channels(pic)) + # put it from HWC to CHW format + img = img.permute((2, 0, 1)) + return img + + +def convert_image_dtype(image: torch.Tensor, dtype: torch.dtype = torch.float) -> torch.Tensor: + """Convert a tensor image to the given ``dtype`` and scale the values accordingly + This function does not support PIL Image. + + Args: + image (torch.Tensor): Image to be converted + dtype (torch.dtype): Desired data type of the output + + Returns: + Tensor: Converted image + + .. note:: + + When converting from a smaller to a larger integer ``dtype`` the maximum values are **not** mapped exactly. + If converted back and forth, this mismatch has no effect. + + Raises: + RuntimeError: When trying to cast :class:`torch.float32` to :class:`torch.int32` or :class:`torch.int64` as + well as for trying to cast :class:`torch.float64` to :class:`torch.int64`. These conversions might lead to + overflow errors since the floating point ``dtype`` cannot store consecutive integers over the whole range + of the integer ``dtype``. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(convert_image_dtype) + if not isinstance(image, torch.Tensor): + raise TypeError("Input img should be Tensor Image") + + return F_t.convert_image_dtype(image, dtype) + + +def to_pil_image(pic, mode=None): + """Convert a tensor or an ndarray to PIL Image. This function does not support torchscript. + + See :class:`~torchvision.transforms.ToPILImage` for more details. + + Args: + pic (Tensor or numpy.ndarray): Image to be converted to PIL Image. + mode (`PIL.Image mode`_): color space and pixel depth of input data (optional). + + .. _PIL.Image mode: https://pillow.readthedocs.io/en/latest/handbook/concepts.html#concept-modes + + Returns: + PIL Image: Image converted to PIL Image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(to_pil_image) + + if isinstance(pic, torch.Tensor): + if pic.ndim == 3: + pic = pic.permute((1, 2, 0)) + pic = pic.numpy(force=True) + elif not isinstance(pic, np.ndarray): + raise TypeError(f"pic should be Tensor or ndarray. Got {type(pic)}.") + + if pic.ndim == 2: + # if 2D image, add channel dimension (HWC) + pic = np.expand_dims(pic, 2) + if pic.ndim != 3: + raise ValueError(f"pic should be 2/3 dimensional. Got {pic.ndim} dimensions.") + + if pic.shape[-1] > 4: + raise ValueError(f"pic should not have > 4 channels. Got {pic.shape[-1]} channels.") + + npimg = pic + + if np.issubdtype(npimg.dtype, np.floating) and mode != "F": + npimg = (npimg * 255).astype(np.uint8) + + if npimg.shape[2] == 1: + expected_mode = None + npimg = npimg[:, :, 0] + if npimg.dtype == np.uint8: + expected_mode = "L" + elif npimg.dtype == np.int16: + expected_mode = "I;16" if sys.byteorder == "little" else "I;16B" + elif npimg.dtype == np.int32: + expected_mode = "I" + elif npimg.dtype == np.float32: + expected_mode = "F" + if mode is not None and mode != expected_mode: + raise ValueError(f"Incorrect mode ({mode}) supplied for input type {np.dtype}. Should be {expected_mode}") + mode = expected_mode + + elif npimg.shape[2] == 2: + permitted_2_channel_modes = ["LA"] + if mode is not None and mode not in permitted_2_channel_modes: + raise ValueError(f"Only modes {permitted_2_channel_modes} are supported for 2D inputs") + + if mode is None and npimg.dtype == np.uint8: + mode = "LA" + + elif npimg.shape[2] == 4: + permitted_4_channel_modes = ["RGBA", "CMYK", "RGBX"] + if mode is not None and mode not in permitted_4_channel_modes: + raise ValueError(f"Only modes {permitted_4_channel_modes} are supported for 4D inputs") + + if mode is None and npimg.dtype == np.uint8: + mode = "RGBA" + else: + permitted_3_channel_modes = ["RGB", "YCbCr", "HSV"] + if mode is not None and mode not in permitted_3_channel_modes: + raise ValueError(f"Only modes {permitted_3_channel_modes} are supported for 3D inputs") + if mode is None and npimg.dtype == np.uint8: + mode = "RGB" + + if mode is None: + raise TypeError(f"Input type {npimg.dtype} is not supported") + + return Image.fromarray(npimg, mode=mode) + + +def normalize(tensor: Tensor, mean: List[float], std: List[float], inplace: bool = False) -> Tensor: + """Normalize a float tensor image with mean and standard deviation. + This transform does not support PIL Image. + + .. note:: + This transform acts out of place by default, i.e., it does not mutates the input tensor. + + See :class:`~torchvision.transforms.Normalize` for more details. + + Args: + tensor (Tensor): Float tensor image of size (C, H, W) or (B, C, H, W) to be normalized. + mean (sequence): Sequence of means for each channel. + std (sequence): Sequence of standard deviations for each channel. + inplace(bool,optional): Bool to make this operation inplace. + + Returns: + Tensor: Normalized Tensor image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(normalize) + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"img should be Tensor Image. Got {type(tensor)}") + + return F_t.normalize(tensor, mean=mean, std=std, inplace=inplace) + + +def _compute_resized_output_size( + image_size: Tuple[int, int], + size: Optional[List[int]], + max_size: Optional[int] = None, + allow_size_none: bool = False, # only True in v2 +) -> List[int]: + h, w = image_size + short, long = (w, h) if w <= h else (h, w) + if size is None: + if not allow_size_none: + raise ValueError("This should never happen!!") + if not isinstance(max_size, int): + raise ValueError(f"max_size must be an integer when size is None, but got {max_size} instead.") + new_short, new_long = int(max_size * short / long), max_size + new_w, new_h = (new_short, new_long) if w <= h else (new_long, new_short) + elif len(size) == 1: # specified size only for the smallest edge + requested_new_short = size if isinstance(size, int) else size[0] + new_short, new_long = requested_new_short, int(requested_new_short * long / short) + + if max_size is not None: + if max_size <= requested_new_short: + raise ValueError( + f"max_size = {max_size} must be strictly greater than the requested " + f"size for the smaller edge size = {size}" + ) + if new_long > max_size: + new_short, new_long = int(max_size * new_short / new_long), max_size + + new_w, new_h = (new_short, new_long) if w <= h else (new_long, new_short) + else: # specified both h and w + new_w, new_h = size[1], size[0] + return [new_h, new_w] + + +def resize( + img: Tensor, + size: List[int], + interpolation: InterpolationMode = InterpolationMode.BILINEAR, + max_size: Optional[int] = None, + antialias: Optional[bool] = True, +) -> Tensor: + r"""Resize the input image to the given size. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions + + Args: + img (PIL Image or Tensor): Image to be resized. + size (sequence or int): Desired output size. If size is a sequence like + (h, w), the output size will be matched to this. If size is an int, + the smaller edge of the image will be matched to this number maintaining + the aspect ratio. i.e, if height > width, then image will be rescaled to + :math:`\left(\text{size} \times \frac{\text{height}}{\text{width}}, \text{size}\right)`. + + .. note:: + In torchscript mode size as single int is not supported, use a sequence of length 1: ``[size, ]``. + interpolation (InterpolationMode): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. + Default is ``InterpolationMode.BILINEAR``. If input is Tensor, only ``InterpolationMode.NEAREST``, + ``InterpolationMode.NEAREST_EXACT``, ``InterpolationMode.BILINEAR`` and ``InterpolationMode.BICUBIC`` are + supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + max_size (int, optional): The maximum allowed for the longer edge of + the resized image. If the longer edge of the image is greater + than ``max_size`` after being resized according to ``size``, + ``size`` will be overruled so that the longer edge is equal to + ``max_size``. + As a result, the smaller edge may be shorter than ``size``. This + is only supported if ``size`` is an int (or a sequence of length + 1 in torchscript mode). + antialias (bool, optional): Whether to apply antialiasing. + It only affects **tensors** with bilinear or bicubic modes and it is + ignored otherwise: on PIL images, antialiasing is always applied on + bilinear or bicubic modes; on other modes (for PIL images and + tensors), antialiasing makes no sense and this parameter is ignored. + Possible values are: + + - ``True`` (default): will apply antialiasing for bilinear or bicubic modes. + Other mode aren't affected. This is probably what you want to use. + - ``False``: will not apply antialiasing for tensors on any mode. PIL + images are still antialiased on bilinear or bicubic modes, because + PIL doesn't support no antialias. + - ``None``: equivalent to ``False`` for tensors and ``True`` for + PIL images. This value exists for legacy reasons and you probably + don't want to use it unless you really know what you are doing. + + The default value changed from ``None`` to ``True`` in + v0.17, for the PIL and Tensor backends to be consistent. + + Returns: + PIL Image or Tensor: Resized image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(resize) + + if isinstance(interpolation, int): + interpolation = _interpolation_modes_from_int(interpolation) + elif not isinstance(interpolation, InterpolationMode): + raise TypeError( + "Argument interpolation should be a InterpolationMode or a corresponding Pillow integer constant" + ) + + if isinstance(size, (list, tuple)): + if len(size) not in [1, 2]: + raise ValueError( + f"Size must be an int or a 1 or 2 element tuple/list, not a {len(size)} element tuple/list" + ) + if max_size is not None and len(size) != 1: + raise ValueError( + "max_size should only be passed if size specifies the length of the smaller edge, " + "i.e. size should be an int or a sequence of length 1 in torchscript mode." + ) + + _, image_height, image_width = get_dimensions(img) + if isinstance(size, int): + size = [size] + output_size = _compute_resized_output_size((image_height, image_width), size, max_size) + + if [image_height, image_width] == output_size: + return img + + if not isinstance(img, torch.Tensor): + if antialias is False: + warnings.warn("Anti-alias option is always applied for PIL Image input. Argument antialias is ignored.") + pil_interpolation = pil_modes_mapping[interpolation] + return F_pil.resize(img, size=output_size, interpolation=pil_interpolation) + + return F_t.resize(img, size=output_size, interpolation=interpolation.value, antialias=antialias) + + +def pad(img: Tensor, padding: List[int], fill: Union[int, float] = 0, padding_mode: str = "constant") -> Tensor: + r"""Pad the given image on all sides with the given "pad" value. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means at most 2 leading dimensions for mode reflect and symmetric, + at most 3 leading dimensions for mode edge, + and an arbitrary number of leading dimensions for mode constant + + Args: + img (PIL Image or Tensor): Image to be padded. + padding (int or sequence): Padding on each border. If a single int is provided this + is used to pad all borders. If sequence of length 2 is provided this is the padding + on left/right and top/bottom respectively. If a sequence of length 4 is provided + this is the padding for the left, top, right and bottom borders respectively. + + .. note:: + In torchscript mode padding as single int is not supported, use a sequence of + length 1: ``[padding, ]``. + fill (number or tuple): Pixel fill value for constant fill. Default is 0. + If a tuple of length 3, it is used to fill R, G, B channels respectively. + This value is only used when the padding_mode is constant. + Only number is supported for torch Tensor. + Only int or tuple value is supported for PIL Image. + padding_mode (str): Type of padding. Should be: constant, edge, reflect or symmetric. + Default is constant. + + - constant: pads with a constant value, this value is specified with fill + + - edge: pads with the last value at the edge of the image. + If input a 5D torch Tensor, the last 3 dimensions will be padded instead of the last 2 + + - reflect: pads with reflection of image without repeating the last value on the edge. + For example, padding [1, 2, 3, 4] with 2 elements on both sides in reflect mode + will result in [3, 2, 1, 2, 3, 4, 3, 2] + + - symmetric: pads with reflection of image repeating the last value on the edge. + For example, padding [1, 2, 3, 4] with 2 elements on both sides in symmetric mode + will result in [2, 1, 1, 2, 3, 4, 4, 3] + + Returns: + PIL Image or Tensor: Padded image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(pad) + if not isinstance(img, torch.Tensor): + return F_pil.pad(img, padding=padding, fill=fill, padding_mode=padding_mode) + + return F_t.pad(img, padding=padding, fill=fill, padding_mode=padding_mode) + + +def crop(img: Tensor, top: int, left: int, height: int, width: int) -> Tensor: + """Crop the given image at specified location and output size. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions. + If image size is smaller than output size along any edge, image is padded with 0 and then cropped. + + Args: + img (PIL Image or Tensor): Image to be cropped. (0,0) denotes the top left corner of the image. + top (int): Vertical component of the top left corner of the crop box. + left (int): Horizontal component of the top left corner of the crop box. + height (int): Height of the crop box. + width (int): Width of the crop box. + + Returns: + PIL Image or Tensor: Cropped image. + """ + + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(crop) + if not isinstance(img, torch.Tensor): + return F_pil.crop(img, top, left, height, width) + + return F_t.crop(img, top, left, height, width) + + +def center_crop(img: Tensor, output_size: List[int]) -> Tensor: + """Crops the given image at the center. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions. + If image size is smaller than output size along any edge, image is padded with 0 and then center cropped. + + Args: + img (PIL Image or Tensor): Image to be cropped. + output_size (sequence or int): (height, width) of the crop box. If int or sequence with single int, + it is used for both directions. + + Returns: + PIL Image or Tensor: Cropped image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(center_crop) + if isinstance(output_size, numbers.Number): + output_size = (int(output_size), int(output_size)) + elif isinstance(output_size, (tuple, list)) and len(output_size) == 1: + output_size = (output_size[0], output_size[0]) + + _, image_height, image_width = get_dimensions(img) + crop_height, crop_width = output_size + + if crop_width > image_width or crop_height > image_height: + padding_ltrb = [ + (crop_width - image_width) // 2 if crop_width > image_width else 0, + (crop_height - image_height) // 2 if crop_height > image_height else 0, + (crop_width - image_width + 1) // 2 if crop_width > image_width else 0, + (crop_height - image_height + 1) // 2 if crop_height > image_height else 0, + ] + img = pad(img, padding_ltrb, fill=0) # PIL uses fill value 0 + _, image_height, image_width = get_dimensions(img) + if crop_width == image_width and crop_height == image_height: + return img + + crop_top = int(round((image_height - crop_height) / 2.0)) + crop_left = int(round((image_width - crop_width) / 2.0)) + return crop(img, crop_top, crop_left, crop_height, crop_width) + + +def resized_crop( + img: Tensor, + top: int, + left: int, + height: int, + width: int, + size: List[int], + interpolation: InterpolationMode = InterpolationMode.BILINEAR, + antialias: Optional[bool] = True, +) -> Tensor: + """Crop the given image and resize it to desired size. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions + + Notably used in :class:`~torchvision.transforms.RandomResizedCrop`. + + Args: + img (PIL Image or Tensor): Image to be cropped. (0,0) denotes the top left corner of the image. + top (int): Vertical component of the top left corner of the crop box. + left (int): Horizontal component of the top left corner of the crop box. + height (int): Height of the crop box. + width (int): Width of the crop box. + size (sequence or int): Desired output size. Same semantics as ``resize``. + interpolation (InterpolationMode): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. + Default is ``InterpolationMode.BILINEAR``. If input is Tensor, only ``InterpolationMode.NEAREST``, + ``InterpolationMode.NEAREST_EXACT``, ``InterpolationMode.BILINEAR`` and ``InterpolationMode.BICUBIC`` are + supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + antialias (bool, optional): Whether to apply antialiasing. + It only affects **tensors** with bilinear or bicubic modes and it is + ignored otherwise: on PIL images, antialiasing is always applied on + bilinear or bicubic modes; on other modes (for PIL images and + tensors), antialiasing makes no sense and this parameter is ignored. + Possible values are: + + - ``True`` (default): will apply antialiasing for bilinear or bicubic modes. + Other mode aren't affected. This is probably what you want to use. + - ``False``: will not apply antialiasing for tensors on any mode. PIL + images are still antialiased on bilinear or bicubic modes, because + PIL doesn't support no antialias. + - ``None``: equivalent to ``False`` for tensors and ``True`` for + PIL images. This value exists for legacy reasons and you probably + don't want to use it unless you really know what you are doing. + + The default value changed from ``None`` to ``True`` in + v0.17, for the PIL and Tensor backends to be consistent. + Returns: + PIL Image or Tensor: Cropped image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(resized_crop) + img = crop(img, top, left, height, width) + img = resize(img, size, interpolation, antialias=antialias) + return img + + +def hflip(img: Tensor) -> Tensor: + """Horizontally flip the given image. + + Args: + img (PIL Image or Tensor): Image to be flipped. If img + is a Tensor, it is expected to be in [..., H, W] format, + where ... means it can have an arbitrary number of leading + dimensions. + + Returns: + PIL Image or Tensor: Horizontally flipped image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(hflip) + if not isinstance(img, torch.Tensor): + return F_pil.hflip(img) + + return F_t.hflip(img) + + +def _get_perspective_coeffs(startpoints: List[List[int]], endpoints: List[List[int]]) -> List[float]: + """Helper function to get the coefficients (a, b, c, d, e, f, g, h) for the perspective transforms. + + In Perspective Transform each pixel (x, y) in the original image gets transformed as, + (x, y) -> ( (ax + by + c) / (gx + hy + 1), (dx + ey + f) / (gx + hy + 1) ) + + Args: + startpoints (list of list of ints): List containing four lists of two integers corresponding to four corners + ``[top-left, top-right, bottom-right, bottom-left]`` of the original image. + endpoints (list of list of ints): List containing four lists of two integers corresponding to four corners + ``[top-left, top-right, bottom-right, bottom-left]`` of the transformed image. + + Returns: + octuple (a, b, c, d, e, f, g, h) for transforming each pixel. + """ + if len(startpoints) != 4 or len(endpoints) != 4: + raise ValueError( + f"Please provide exactly four corners, got {len(startpoints)} startpoints and {len(endpoints)} endpoints." + ) + a_matrix = torch.zeros(2 * len(startpoints), 8, dtype=torch.float64) + + for i, (p1, p2) in enumerate(zip(endpoints, startpoints)): + a_matrix[2 * i, :] = torch.tensor([p1[0], p1[1], 1, 0, 0, 0, -p2[0] * p1[0], -p2[0] * p1[1]]) + a_matrix[2 * i + 1, :] = torch.tensor([0, 0, 0, p1[0], p1[1], 1, -p2[1] * p1[0], -p2[1] * p1[1]]) + + b_matrix = torch.tensor(startpoints, dtype=torch.float64).view(8) + # do least squares in double precision to prevent numerical issues + res = torch.linalg.lstsq(a_matrix, b_matrix, driver="gels").solution.to(torch.float32) + + output: List[float] = res.tolist() + return output + + +def perspective( + img: Tensor, + startpoints: List[List[int]], + endpoints: List[List[int]], + interpolation: InterpolationMode = InterpolationMode.BILINEAR, + fill: Optional[List[float]] = None, +) -> Tensor: + """Perform perspective transform of the given image. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions. + + Args: + img (PIL Image or Tensor): Image to be transformed. + startpoints (list of list of ints): List containing four lists of two integers corresponding to four corners + ``[top-left, top-right, bottom-right, bottom-left]`` of the original image. + endpoints (list of list of ints): List containing four lists of two integers corresponding to four corners + ``[top-left, top-right, bottom-right, bottom-left]`` of the transformed image. + interpolation (InterpolationMode): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + fill (sequence or number, optional): Pixel fill value for the area outside the transformed + image. If given a number, the value is used for all bands respectively. + + .. note:: + In torchscript mode single int/float value is not supported, please use a sequence + of length 1: ``[value, ]``. + + Returns: + PIL Image or Tensor: transformed Image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(perspective) + + coeffs = _get_perspective_coeffs(startpoints, endpoints) + + if isinstance(interpolation, int): + interpolation = _interpolation_modes_from_int(interpolation) + elif not isinstance(interpolation, InterpolationMode): + raise TypeError( + "Argument interpolation should be a InterpolationMode or a corresponding Pillow integer constant" + ) + + if not isinstance(img, torch.Tensor): + pil_interpolation = pil_modes_mapping[interpolation] + return F_pil.perspective(img, coeffs, interpolation=pil_interpolation, fill=fill) + + return F_t.perspective(img, coeffs, interpolation=interpolation.value, fill=fill) + + +def vflip(img: Tensor) -> Tensor: + """Vertically flip the given image. + + Args: + img (PIL Image or Tensor): Image to be flipped. If img + is a Tensor, it is expected to be in [..., H, W] format, + where ... means it can have an arbitrary number of leading + dimensions. + + Returns: + PIL Image or Tensor: Vertically flipped image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(vflip) + if not isinstance(img, torch.Tensor): + return F_pil.vflip(img) + + return F_t.vflip(img) + + +def five_crop(img: Tensor, size: List[int]) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + """Crop the given image into four corners and the central crop. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions + + .. Note:: + This transform returns a tuple of images and there may be a + mismatch in the number of inputs and targets your ``Dataset`` returns. + + Args: + img (PIL Image or Tensor): Image to be cropped. + size (sequence or int): Desired output size of the crop. If size is an + int instead of sequence like (h, w), a square crop (size, size) is + made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). + + Returns: + tuple: tuple (tl, tr, bl, br, center) + Corresponding top left, top right, bottom left, bottom right and center crop. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(five_crop) + if isinstance(size, numbers.Number): + size = (int(size), int(size)) + elif isinstance(size, (tuple, list)) and len(size) == 1: + size = (size[0], size[0]) + + if len(size) != 2: + raise ValueError("Please provide only two dimensions (h, w) for size.") + + _, image_height, image_width = get_dimensions(img) + crop_height, crop_width = size + if crop_width > image_width or crop_height > image_height: + msg = "Requested crop size {} is bigger than input size {}" + raise ValueError(msg.format(size, (image_height, image_width))) + + tl = crop(img, 0, 0, crop_height, crop_width) + tr = crop(img, 0, image_width - crop_width, crop_height, crop_width) + bl = crop(img, image_height - crop_height, 0, crop_height, crop_width) + br = crop(img, image_height - crop_height, image_width - crop_width, crop_height, crop_width) + + center = center_crop(img, [crop_height, crop_width]) + + return tl, tr, bl, br, center + + +def ten_crop( + img: Tensor, size: List[int], vertical_flip: bool = False +) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + """Generate ten cropped images from the given image. + Crop the given image into four corners and the central crop plus the + flipped version of these (horizontal flipping is used by default). + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions + + .. Note:: + This transform returns a tuple of images and there may be a + mismatch in the number of inputs and targets your ``Dataset`` returns. + + Args: + img (PIL Image or Tensor): Image to be cropped. + size (sequence or int): Desired output size of the crop. If size is an + int instead of sequence like (h, w), a square crop (size, size) is + made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). + vertical_flip (bool): Use vertical flipping instead of horizontal + + Returns: + tuple: tuple (tl, tr, bl, br, center, tl_flip, tr_flip, bl_flip, br_flip, center_flip) + Corresponding top left, top right, bottom left, bottom right and + center crop and same for the flipped image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(ten_crop) + if isinstance(size, numbers.Number): + size = (int(size), int(size)) + elif isinstance(size, (tuple, list)) and len(size) == 1: + size = (size[0], size[0]) + + if len(size) != 2: + raise ValueError("Please provide only two dimensions (h, w) for size.") + + first_five = five_crop(img, size) + + if vertical_flip: + img = vflip(img) + else: + img = hflip(img) + + second_five = five_crop(img, size) + return first_five + second_five + + +def adjust_brightness(img: Tensor, brightness_factor: float) -> Tensor: + """Adjust brightness of an image. + + Args: + img (PIL Image or Tensor): Image to be adjusted. + If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, + where ... means it can have an arbitrary number of leading dimensions. + brightness_factor (float): How much to adjust the brightness. Can be + any non-negative number. 0 gives a black image, 1 gives the + original image while 2 increases the brightness by a factor of 2. + + Returns: + PIL Image or Tensor: Brightness adjusted image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(adjust_brightness) + if not isinstance(img, torch.Tensor): + return F_pil.adjust_brightness(img, brightness_factor) + + return F_t.adjust_brightness(img, brightness_factor) + + +def adjust_contrast(img: Tensor, contrast_factor: float) -> Tensor: + """Adjust contrast of an image. + + Args: + img (PIL Image or Tensor): Image to be adjusted. + If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, + where ... means it can have an arbitrary number of leading dimensions. + contrast_factor (float): How much to adjust the contrast. Can be any + non-negative number. 0 gives a solid gray image, 1 gives the + original image while 2 increases the contrast by a factor of 2. + + Returns: + PIL Image or Tensor: Contrast adjusted image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(adjust_contrast) + if not isinstance(img, torch.Tensor): + return F_pil.adjust_contrast(img, contrast_factor) + + return F_t.adjust_contrast(img, contrast_factor) + + +def adjust_saturation(img: Tensor, saturation_factor: float) -> Tensor: + """Adjust color saturation of an image. + + Args: + img (PIL Image or Tensor): Image to be adjusted. + If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, + where ... means it can have an arbitrary number of leading dimensions. + saturation_factor (float): How much to adjust the saturation. 0 will + give a black and white image, 1 will give the original image while + 2 will enhance the saturation by a factor of 2. + + Returns: + PIL Image or Tensor: Saturation adjusted image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(adjust_saturation) + if not isinstance(img, torch.Tensor): + return F_pil.adjust_saturation(img, saturation_factor) + + return F_t.adjust_saturation(img, saturation_factor) + + +def adjust_hue(img: Tensor, hue_factor: float) -> Tensor: + """Adjust hue of an image. + + The image hue is adjusted by converting the image to HSV and + cyclically shifting the intensities in the hue channel (H). + The image is then converted back to original image mode. + + `hue_factor` is the amount of shift in H channel and must be in the + interval `[-0.5, 0.5]`. + + See `Hue`_ for more details. + + .. _Hue: https://en.wikipedia.org/wiki/Hue + + Args: + img (PIL Image or Tensor): Image to be adjusted. + If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, + where ... means it can have an arbitrary number of leading dimensions. + If img is PIL Image mode "1", "I", "F" and modes with transparency (alpha channel) are not supported. + Note: the pixel values of the input image has to be non-negative for conversion to HSV space; + thus it does not work if you normalize your image to an interval with negative values, + or use an interpolation that generates negative values before using this function. + hue_factor (float): How much to shift the hue channel. Should be in + [-0.5, 0.5]. 0.5 and -0.5 give complete reversal of hue channel in + HSV space in positive and negative direction respectively. + 0 means no shift. Therefore, both -0.5 and 0.5 will give an image + with complementary colors while 0 gives the original image. + + Returns: + PIL Image or Tensor: Hue adjusted image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(adjust_hue) + if not isinstance(img, torch.Tensor): + return F_pil.adjust_hue(img, hue_factor) + + return F_t.adjust_hue(img, hue_factor) + + +def adjust_gamma(img: Tensor, gamma: float, gain: float = 1) -> Tensor: + r"""Perform gamma correction on an image. + + Also known as Power Law Transform. Intensities in RGB mode are adjusted + based on the following equation: + + .. math:: + I_{\text{out}} = 255 \times \text{gain} \times \left(\frac{I_{\text{in}}}{255}\right)^{\gamma} + + See `Gamma Correction`_ for more details. + + .. _Gamma Correction: https://en.wikipedia.org/wiki/Gamma_correction + + Args: + img (PIL Image or Tensor): PIL Image to be adjusted. + If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, + where ... means it can have an arbitrary number of leading dimensions. + If img is PIL Image, modes with transparency (alpha channel) are not supported. + gamma (float): Non negative real number, same as :math:`\gamma` in the equation. + gamma larger than 1 make the shadows darker, + while gamma smaller than 1 make dark regions lighter. + gain (float): The constant multiplier. + Returns: + PIL Image or Tensor: Gamma correction adjusted image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(adjust_gamma) + if not isinstance(img, torch.Tensor): + return F_pil.adjust_gamma(img, gamma, gain) + + return F_t.adjust_gamma(img, gamma, gain) + + +def _get_inverse_affine_matrix( + center: List[float], angle: float, translate: List[float], scale: float, shear: List[float], inverted: bool = True +) -> List[float]: + # Helper method to compute inverse matrix for affine transformation + + # Pillow requires inverse affine transformation matrix: + # Affine matrix is : M = T * C * RotateScaleShear * C^-1 + # + # where T is translation matrix: [1, 0, tx | 0, 1, ty | 0, 0, 1] + # C is translation matrix to keep center: [1, 0, cx | 0, 1, cy | 0, 0, 1] + # RotateScaleShear is rotation with scale and shear matrix + # + # RotateScaleShear(a, s, (sx, sy)) = + # = R(a) * S(s) * SHy(sy) * SHx(sx) + # = [ s*cos(a - sy)/cos(sy), s*(-cos(a - sy)*tan(sx)/cos(sy) - sin(a)), 0 ] + # [ s*sin(a - sy)/cos(sy), s*(-sin(a - sy)*tan(sx)/cos(sy) + cos(a)), 0 ] + # [ 0 , 0 , 1 ] + # where R is a rotation matrix, S is a scaling matrix, and SHx and SHy are the shears: + # SHx(s) = [1, -tan(s)] and SHy(s) = [1 , 0] + # [0, 1 ] [-tan(s), 1] + # + # Thus, the inverse is M^-1 = C * RotateScaleShear^-1 * C^-1 * T^-1 + + rot = math.radians(angle) + sx = math.radians(shear[0]) + sy = math.radians(shear[1]) + + cx, cy = center + tx, ty = translate + + # RSS without scaling + a = math.cos(rot - sy) / math.cos(sy) + b = -math.cos(rot - sy) * math.tan(sx) / math.cos(sy) - math.sin(rot) + c = math.sin(rot - sy) / math.cos(sy) + d = -math.sin(rot - sy) * math.tan(sx) / math.cos(sy) + math.cos(rot) + + if inverted: + # Inverted rotation matrix with scale and shear + # det([[a, b], [c, d]]) == 1, since det(rotation) = 1 and det(shear) = 1 + matrix = [d, -b, 0.0, -c, a, 0.0] + matrix = [x / scale for x in matrix] + # Apply inverse of translation and of center translation: RSS^-1 * C^-1 * T^-1 + matrix[2] += matrix[0] * (-cx - tx) + matrix[1] * (-cy - ty) + matrix[5] += matrix[3] * (-cx - tx) + matrix[4] * (-cy - ty) + # Apply center translation: C * RSS^-1 * C^-1 * T^-1 + matrix[2] += cx + matrix[5] += cy + else: + matrix = [a, b, 0.0, c, d, 0.0] + matrix = [x * scale for x in matrix] + # Apply inverse of center translation: RSS * C^-1 + matrix[2] += matrix[0] * (-cx) + matrix[1] * (-cy) + matrix[5] += matrix[3] * (-cx) + matrix[4] * (-cy) + # Apply translation and center : T * C * RSS * C^-1 + matrix[2] += cx + tx + matrix[5] += cy + ty + + return matrix + + +def rotate( + img: Tensor, + angle: float, + interpolation: InterpolationMode = InterpolationMode.NEAREST, + expand: bool = False, + center: Optional[List[int]] = None, + fill: Optional[List[float]] = None, +) -> Tensor: + """Rotate the image by angle. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions. + + Args: + img (PIL Image or Tensor): image to be rotated. + angle (number): rotation angle value in degrees, counter-clockwise. + interpolation (InterpolationMode): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + expand (bool, optional): Optional expansion flag. + If true, expands the output image to make it large enough to hold the entire rotated image. + If false or omitted, make the output image the same size as the input image. + Note that the expand flag assumes rotation around the center and no translation. + center (sequence, optional): Optional center of rotation. Origin is the upper left corner. + Default is the center of the image. + fill (sequence or number, optional): Pixel fill value for the area outside the transformed + image. If given a number, the value is used for all bands respectively. + + .. note:: + In torchscript mode single int/float value is not supported, please use a sequence + of length 1: ``[value, ]``. + Returns: + PIL Image or Tensor: Rotated image. + + .. _filters: https://pillow.readthedocs.io/en/latest/handbook/concepts.html#filters + + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(rotate) + + if isinstance(interpolation, int): + interpolation = _interpolation_modes_from_int(interpolation) + elif not isinstance(interpolation, InterpolationMode): + raise TypeError( + "Argument interpolation should be a InterpolationMode or a corresponding Pillow integer constant" + ) + + if not isinstance(angle, (int, float)): + raise TypeError("Argument angle should be int or float") + + if center is not None and not isinstance(center, (list, tuple)): + raise TypeError("Argument center should be a sequence") + + if not isinstance(img, torch.Tensor): + pil_interpolation = pil_modes_mapping[interpolation] + return F_pil.rotate(img, angle=angle, interpolation=pil_interpolation, expand=expand, center=center, fill=fill) + + center_f = [0.0, 0.0] + if center is not None: + _, height, width = get_dimensions(img) + # Center values should be in pixel coordinates but translated such that (0, 0) corresponds to image center. + center_f = [1.0 * (c - s * 0.5) for c, s in zip(center, [width, height])] + + # due to current incoherence of rotation angle direction between affine and rotate implementations + # we need to set -angle. + matrix = _get_inverse_affine_matrix(center_f, -angle, [0.0, 0.0], 1.0, [0.0, 0.0]) + return F_t.rotate(img, matrix=matrix, interpolation=interpolation.value, expand=expand, fill=fill) + + +def affine( + img: Tensor, + angle: float, + translate: List[int], + scale: float, + shear: List[float], + interpolation: InterpolationMode = InterpolationMode.NEAREST, + fill: Optional[List[float]] = None, + center: Optional[List[int]] = None, +) -> Tensor: + """Apply affine transformation on the image keeping image center invariant. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions. + + Args: + img (PIL Image or Tensor): image to transform. + angle (number): rotation angle in degrees between -180 and 180, clockwise direction. + translate (sequence of integers): horizontal and vertical translations (post-rotation translation) + scale (float): overall scale + shear (float or sequence): shear angle value in degrees between -180 to 180, clockwise direction. + If a sequence is specified, the first value corresponds to a shear parallel to the x-axis, while + the second value corresponds to a shear parallel to the y-axis. + interpolation (InterpolationMode): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + fill (sequence or number, optional): Pixel fill value for the area outside the transformed + image. If given a number, the value is used for all bands respectively. + + .. note:: + In torchscript mode single int/float value is not supported, please use a sequence + of length 1: ``[value, ]``. + center (sequence, optional): Optional center of rotation. Origin is the upper left corner. + Default is the center of the image. + + Returns: + PIL Image or Tensor: Transformed image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(affine) + + if isinstance(interpolation, int): + interpolation = _interpolation_modes_from_int(interpolation) + elif not isinstance(interpolation, InterpolationMode): + raise TypeError( + "Argument interpolation should be a InterpolationMode or a corresponding Pillow integer constant" + ) + + if not isinstance(angle, (int, float)): + raise TypeError("Argument angle should be int or float") + + if not isinstance(translate, (list, tuple)): + raise TypeError("Argument translate should be a sequence") + + if len(translate) != 2: + raise ValueError("Argument translate should be a sequence of length 2") + + if scale <= 0.0: + raise ValueError("Argument scale should be positive") + + if not isinstance(shear, (numbers.Number, (list, tuple))): + raise TypeError("Shear should be either a single value or a sequence of two values") + + if isinstance(angle, int): + angle = float(angle) + + if isinstance(translate, tuple): + translate = list(translate) + + if isinstance(shear, numbers.Number): + shear = [shear, 0.0] + + if isinstance(shear, tuple): + shear = list(shear) + + if len(shear) == 1: + shear = [shear[0], shear[0]] + + if len(shear) != 2: + raise ValueError(f"Shear should be a sequence containing two values. Got {shear}") + + if center is not None and not isinstance(center, (list, tuple)): + raise TypeError("Argument center should be a sequence") + + _, height, width = get_dimensions(img) + if not isinstance(img, torch.Tensor): + # center = (width * 0.5 + 0.5, height * 0.5 + 0.5) + # it is visually better to estimate the center without 0.5 offset + # otherwise image rotated by 90 degrees is shifted vs output image of torch.rot90 or F_t.affine + if center is None: + center = [width * 0.5, height * 0.5] + matrix = _get_inverse_affine_matrix(center, angle, translate, scale, shear) + pil_interpolation = pil_modes_mapping[interpolation] + return F_pil.affine(img, matrix=matrix, interpolation=pil_interpolation, fill=fill) + + center_f = [0.0, 0.0] + if center is not None: + _, height, width = get_dimensions(img) + # Center values should be in pixel coordinates but translated such that (0, 0) corresponds to image center. + center_f = [1.0 * (c - s * 0.5) for c, s in zip(center, [width, height])] + + translate_f = [1.0 * t for t in translate] + matrix = _get_inverse_affine_matrix(center_f, angle, translate_f, scale, shear) + return F_t.affine(img, matrix=matrix, interpolation=interpolation.value, fill=fill) + + +# Looks like to_grayscale() is a stand-alone functional that is never called +# from the transform classes. Perhaps it's still here for BC? I can't be +# bothered to dig. +@torch.jit.unused +def to_grayscale(img, num_output_channels=1): + """Convert PIL image of any mode (RGB, HSV, LAB, etc) to grayscale version of image. + This transform does not support torch Tensor. + + Args: + img (PIL Image): PIL Image to be converted to grayscale. + num_output_channels (int): number of channels of the output image. Value can be 1 or 3. Default is 1. + + Returns: + PIL Image: Grayscale version of the image. + + - if num_output_channels = 1 : returned image is single channel + - if num_output_channels = 3 : returned image is 3 channel with r = g = b + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(to_grayscale) + if isinstance(img, Image.Image): + return F_pil.to_grayscale(img, num_output_channels) + + raise TypeError("Input should be PIL Image") + + +def rgb_to_grayscale(img: Tensor, num_output_channels: int = 1) -> Tensor: + """Convert RGB image to grayscale version of image. + If the image is torch Tensor, it is expected + to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions + + Note: + Please, note that this method supports only RGB images as input. For inputs in other color spaces, + please, consider using :meth:`~torchvision.transforms.functional.to_grayscale` with PIL Image. + + Args: + img (PIL Image or Tensor): RGB Image to be converted to grayscale. + num_output_channels (int): number of channels of the output image. Value can be 1 or 3. Default, 1. + + Returns: + PIL Image or Tensor: Grayscale version of the image. + + - if num_output_channels = 1 : returned image is single channel + - if num_output_channels = 3 : returned image is 3 channel with r = g = b + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(rgb_to_grayscale) + if not isinstance(img, torch.Tensor): + return F_pil.to_grayscale(img, num_output_channels) + + return F_t.rgb_to_grayscale(img, num_output_channels) + + +def erase(img: Tensor, i: int, j: int, h: int, w: int, v: Tensor, inplace: bool = False) -> Tensor: + """Erase the input Tensor Image with given value. + This transform does not support PIL Image. + + Args: + img (Tensor Image): Tensor image of size (C, H, W) to be erased + i (int): i in (i,j) i.e coordinates of the upper left corner. + j (int): j in (i,j) i.e coordinates of the upper left corner. + h (int): Height of the erased region. + w (int): Width of the erased region. + v: Erasing value. + inplace(bool, optional): For in-place operations. By default, is set False. + + Returns: + Tensor Image: Erased image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(erase) + if not isinstance(img, torch.Tensor): + raise TypeError(f"img should be Tensor Image. Got {type(img)}") + + return F_t.erase(img, i, j, h, w, v, inplace=inplace) + + +def gaussian_blur(img: Tensor, kernel_size: List[int], sigma: Optional[List[float]] = None) -> Tensor: + """Performs Gaussian blurring on the image by given kernel + + The convolution will be using reflection padding corresponding to the kernel size, to maintain the input shape. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means at most one leading dimension. + + Args: + img (PIL Image or Tensor): Image to be blurred + kernel_size (sequence of ints or int): Gaussian kernel size. Can be a sequence of integers + like ``(kx, ky)`` or a single integer for square kernels. + + .. note:: + In torchscript mode kernel_size as single int is not supported, use a sequence of + length 1: ``[ksize, ]``. + sigma (sequence of floats or float, optional): Gaussian kernel standard deviation. Can be a + sequence of floats like ``(sigma_x, sigma_y)`` or a single float to define the + same sigma in both X/Y directions. If None, then it is computed using + ``kernel_size`` as ``sigma = 0.3 * ((kernel_size - 1) * 0.5 - 1) + 0.8``. + Default, None. + + .. note:: + In torchscript mode sigma as single float is + not supported, use a sequence of length 1: ``[sigma, ]``. + + Returns: + PIL Image or Tensor: Gaussian Blurred version of the image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(gaussian_blur) + if not isinstance(kernel_size, (int, list, tuple)): + raise TypeError(f"kernel_size should be int or a sequence of integers. Got {type(kernel_size)}") + if isinstance(kernel_size, int): + kernel_size = [kernel_size, kernel_size] + if len(kernel_size) != 2: + raise ValueError(f"If kernel_size is a sequence its length should be 2. Got {len(kernel_size)}") + for ksize in kernel_size: + if ksize % 2 == 0 or ksize < 0: + raise ValueError(f"kernel_size should have odd and positive integers. Got {kernel_size}") + + if sigma is None: + sigma = [ksize * 0.15 + 0.35 for ksize in kernel_size] + + if sigma is not None and not isinstance(sigma, (int, float, list, tuple)): + raise TypeError(f"sigma should be either float or sequence of floats. Got {type(sigma)}") + if isinstance(sigma, (int, float)): + sigma = [float(sigma), float(sigma)] + if isinstance(sigma, (list, tuple)) and len(sigma) == 1: + sigma = [sigma[0], sigma[0]] + if len(sigma) != 2: + raise ValueError(f"If sigma is a sequence, its length should be 2. Got {len(sigma)}") + for s in sigma: + if s <= 0.0: + raise ValueError(f"sigma should have positive values. Got {sigma}") + + t_img = img + if not isinstance(img, torch.Tensor): + if not F_pil._is_pil_image(img): + raise TypeError(f"img should be PIL Image or Tensor. Got {type(img)}") + + t_img = pil_to_tensor(img) + + output = F_t.gaussian_blur(t_img, kernel_size, sigma) + + if not isinstance(img, torch.Tensor): + output = to_pil_image(output, mode=img.mode) + return output + + +def invert(img: Tensor) -> Tensor: + """Invert the colors of an RGB/grayscale image. + + Args: + img (PIL Image or Tensor): Image to have its colors inverted. + If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, + where ... means it can have an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + + Returns: + PIL Image or Tensor: Color inverted image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(invert) + if not isinstance(img, torch.Tensor): + return F_pil.invert(img) + + return F_t.invert(img) + + +def posterize(img: Tensor, bits: int) -> Tensor: + """Posterize an image by reducing the number of bits for each color channel. + + Args: + img (PIL Image or Tensor): Image to have its colors posterized. + If img is torch Tensor, it should be of type torch.uint8, and + it is expected to be in [..., 1 or 3, H, W] format, where ... means + it can have an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + bits (int): The number of bits to keep for each channel (0-8). + Returns: + PIL Image or Tensor: Posterized image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(posterize) + if not (0 <= bits <= 8): + raise ValueError(f"The number if bits should be between 0 and 8. Got {bits}") + + if not isinstance(img, torch.Tensor): + return F_pil.posterize(img, bits) + + return F_t.posterize(img, bits) + + +def solarize(img: Tensor, threshold: float) -> Tensor: + """Solarize an RGB/grayscale image by inverting all pixel values above a threshold. + + Args: + img (PIL Image or Tensor): Image to have its colors inverted. + If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, + where ... means it can have an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + threshold (float): All pixels equal or above this value are inverted. + Returns: + PIL Image or Tensor: Solarized image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(solarize) + if not isinstance(img, torch.Tensor): + return F_pil.solarize(img, threshold) + + return F_t.solarize(img, threshold) + + +def adjust_sharpness(img: Tensor, sharpness_factor: float) -> Tensor: + """Adjust the sharpness of an image. + + Args: + img (PIL Image or Tensor): Image to be adjusted. + If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, + where ... means it can have an arbitrary number of leading dimensions. + sharpness_factor (float): How much to adjust the sharpness. Can be + any non-negative number. 0 gives a blurred image, 1 gives the + original image while 2 increases the sharpness by a factor of 2. + + Returns: + PIL Image or Tensor: Sharpness adjusted image. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(adjust_sharpness) + if not isinstance(img, torch.Tensor): + return F_pil.adjust_sharpness(img, sharpness_factor) + + return F_t.adjust_sharpness(img, sharpness_factor) + + +def autocontrast(img: Tensor) -> Tensor: + """Maximize contrast of an image by remapping its + pixels per channel so that the lowest becomes black and the lightest + becomes white. + + Args: + img (PIL Image or Tensor): Image on which autocontrast is applied. + If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, + where ... means it can have an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + + Returns: + PIL Image or Tensor: An image that was autocontrasted. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(autocontrast) + if not isinstance(img, torch.Tensor): + return F_pil.autocontrast(img) + + return F_t.autocontrast(img) + + +def equalize(img: Tensor) -> Tensor: + """Equalize the histogram of an image by applying + a non-linear mapping to the input in order to create a uniform + distribution of grayscale values in the output. + + Args: + img (PIL Image or Tensor): Image on which equalize is applied. + If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, + where ... means it can have an arbitrary number of leading dimensions. + The tensor dtype must be ``torch.uint8`` and values are expected to be in ``[0, 255]``. + If img is PIL Image, it is expected to be in mode "P", "L" or "RGB". + + Returns: + PIL Image or Tensor: An image that was equalized. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(equalize) + if not isinstance(img, torch.Tensor): + return F_pil.equalize(img) + + return F_t.equalize(img) + + +def elastic_transform( + img: Tensor, + displacement: Tensor, + interpolation: InterpolationMode = InterpolationMode.BILINEAR, + fill: Optional[List[float]] = None, +) -> Tensor: + """Transform a tensor image with elastic transformations. + Given alpha and sigma, it will generate displacement + vectors for all pixels based on random offsets. Alpha controls the strength + and sigma controls the smoothness of the displacements. + The displacements are added to an identity grid and the resulting grid is + used to grid_sample from the image. + + Applications: + Randomly transforms the morphology of objects in images and produces a + see-through-water-like effect. + + Args: + img (PIL Image or Tensor): Image on which elastic_transform is applied. + If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, + where ... means it can have an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "P", "L" or "RGB". + displacement (Tensor): The displacement field. Expected shape is [1, H, W, 2]. + interpolation (InterpolationMode): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. + Default is ``InterpolationMode.BILINEAR``. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + fill (number or str or tuple): Pixel fill value for constant fill. Default is 0. + If a tuple of length 3, it is used to fill R, G, B channels respectively. + This value is only used when the padding_mode is constant. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(elastic_transform) + # Backward compatibility with integer value + if isinstance(interpolation, int): + warnings.warn( + "Argument interpolation should be of type InterpolationMode instead of int. " + "Please, use InterpolationMode enum." + ) + interpolation = _interpolation_modes_from_int(interpolation) + + if not isinstance(displacement, torch.Tensor): + raise TypeError("Argument displacement should be a Tensor") + + t_img = img + if not isinstance(img, torch.Tensor): + if not F_pil._is_pil_image(img): + raise TypeError(f"img should be PIL Image or Tensor. Got {type(img)}") + t_img = pil_to_tensor(img) + + shape = t_img.shape + shape = (1,) + shape[-2:] + (2,) + if shape != displacement.shape: + raise ValueError(f"Argument displacement shape should be {shape}, but given {displacement.shape}") + + # TODO: if image shape is [N1, N2, ..., C, H, W] and + # displacement is [1, H, W, 2] we need to reshape input image + # such grid_sampler takes internal code for 4D input + + output = F_t.elastic_transform( + t_img, + displacement, + interpolation=interpolation.value, + fill=fill, + ) + + if not isinstance(img, torch.Tensor): + output = to_pil_image(output, mode=img.mode) + return output diff --git a/lib/python3.12/site-packages/torchvision/transforms/transforms.py b/lib/python3.12/site-packages/torchvision/transforms/transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..07932390efeceb8e0ae1d1aada6faddd56b17461 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/transforms.py @@ -0,0 +1,2153 @@ +import math +import numbers +import random +import warnings +from collections.abc import Sequence +from typing import List, Optional, Tuple, Union + +import torch +from torch import Tensor + +try: + import accimage +except ImportError: + accimage = None + +from ..utils import _log_api_usage_once +from . import functional as F +from .functional import _interpolation_modes_from_int, InterpolationMode + +__all__ = [ + "Compose", + "ToTensor", + "PILToTensor", + "ConvertImageDtype", + "ToPILImage", + "Normalize", + "Resize", + "CenterCrop", + "Pad", + "Lambda", + "RandomApply", + "RandomChoice", + "RandomOrder", + "RandomCrop", + "RandomHorizontalFlip", + "RandomVerticalFlip", + "RandomResizedCrop", + "FiveCrop", + "TenCrop", + "LinearTransformation", + "ColorJitter", + "RandomRotation", + "RandomAffine", + "Grayscale", + "RandomGrayscale", + "RandomPerspective", + "RandomErasing", + "GaussianBlur", + "InterpolationMode", + "RandomInvert", + "RandomPosterize", + "RandomSolarize", + "RandomAdjustSharpness", + "RandomAutocontrast", + "RandomEqualize", + "ElasticTransform", +] + + +class Compose: + """Composes several transforms together. This transform does not support torchscript. + Please, see the note below. + + Args: + transforms (list of ``Transform`` objects): list of transforms to compose. + + Example: + >>> transforms.Compose([ + >>> transforms.CenterCrop(10), + >>> transforms.PILToTensor(), + >>> transforms.ConvertImageDtype(torch.float), + >>> ]) + + .. note:: + In order to script the transformations, please use ``torch.nn.Sequential`` as below. + + >>> transforms = torch.nn.Sequential( + >>> transforms.CenterCrop(10), + >>> transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), + >>> ) + >>> scripted_transforms = torch.jit.script(transforms) + + Make sure to use only scriptable transformations, i.e. that work with ``torch.Tensor``, does not require + `lambda` functions or ``PIL.Image``. + + """ + + def __init__(self, transforms): + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(self) + self.transforms = transforms + + def __call__(self, img): + for t in self.transforms: + img = t(img) + return img + + def __repr__(self) -> str: + format_string = self.__class__.__name__ + "(" + for t in self.transforms: + format_string += "\n" + format_string += f" {t}" + format_string += "\n)" + return format_string + + +class ToTensor: + """Convert a PIL Image or ndarray to tensor and scale the values accordingly. + + This transform does not support torchscript. + + Converts a PIL Image or numpy.ndarray (H x W x C) in the range + [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0] + if the PIL Image belongs to one of the modes (L, LA, P, I, F, RGB, YCbCr, RGBA, CMYK, 1) + or if the numpy.ndarray has dtype = np.uint8 + + In the other cases, tensors are returned without scaling. + + .. note:: + Because the input image is scaled to [0.0, 1.0], this transformation should not be used when + transforming target image masks. See the `references`_ for implementing the transforms for image masks. + + .. _references: https://github.com/pytorch/vision/tree/main/references/segmentation + """ + + def __init__(self) -> None: + _log_api_usage_once(self) + + def __call__(self, pic): + """ + Args: + pic (PIL Image or numpy.ndarray): Image to be converted to tensor. + + Returns: + Tensor: Converted image. + """ + return F.to_tensor(pic) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}()" + + +class PILToTensor: + """Convert a PIL Image to a tensor of the same type - this does not scale values. + + This transform does not support torchscript. + + Converts a PIL Image (H x W x C) to a Tensor of shape (C x H x W). + """ + + def __init__(self) -> None: + _log_api_usage_once(self) + + def __call__(self, pic): + """ + .. note:: + + A deep copy of the underlying array is performed. + + Args: + pic (PIL Image): Image to be converted to tensor. + + Returns: + Tensor: Converted image. + """ + return F.pil_to_tensor(pic) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}()" + + +class ConvertImageDtype(torch.nn.Module): + """Convert a tensor image to the given ``dtype`` and scale the values accordingly. + + This function does not support PIL Image. + + Args: + dtype (torch.dtype): Desired data type of the output + + .. note:: + + When converting from a smaller to a larger integer ``dtype`` the maximum values are **not** mapped exactly. + If converted back and forth, this mismatch has no effect. + + Raises: + RuntimeError: When trying to cast :class:`torch.float32` to :class:`torch.int32` or :class:`torch.int64` as + well as for trying to cast :class:`torch.float64` to :class:`torch.int64`. These conversions might lead to + overflow errors since the floating point ``dtype`` cannot store consecutive integers over the whole range + of the integer ``dtype``. + """ + + def __init__(self, dtype: torch.dtype) -> None: + super().__init__() + _log_api_usage_once(self) + self.dtype = dtype + + def forward(self, image): + return F.convert_image_dtype(image, self.dtype) + + +class ToPILImage: + """Convert a tensor or an ndarray to PIL Image + + This transform does not support torchscript. + + Converts a torch.*Tensor of shape C x H x W or a numpy ndarray of shape + H x W x C to a PIL Image while adjusting the value range depending on the ``mode``. + + Args: + mode (`PIL.Image mode`_): color space and pixel depth of input data (optional). + If ``mode`` is ``None`` (default) there are some assumptions made about the input data: + + - If the input has 4 channels, the ``mode`` is assumed to be ``RGBA``. + - If the input has 3 channels, the ``mode`` is assumed to be ``RGB``. + - If the input has 2 channels, the ``mode`` is assumed to be ``LA``. + - If the input has 1 channel, the ``mode`` is determined by the data type (i.e ``int``, ``float``, ``short``). + + .. _PIL.Image mode: https://pillow.readthedocs.io/en/latest/handbook/concepts.html#concept-modes + """ + + def __init__(self, mode=None): + _log_api_usage_once(self) + self.mode = mode + + def __call__(self, pic): + """ + Args: + pic (Tensor or numpy.ndarray): Image to be converted to PIL Image. + + Returns: + PIL Image: Image converted to PIL Image. + + """ + return F.to_pil_image(pic, self.mode) + + def __repr__(self) -> str: + format_string = self.__class__.__name__ + "(" + if self.mode is not None: + format_string += f"mode={self.mode}" + format_string += ")" + return format_string + + +class Normalize(torch.nn.Module): + """Normalize a tensor image with mean and standard deviation. + This transform does not support PIL Image. + Given mean: ``(mean[1],...,mean[n])`` and std: ``(std[1],..,std[n])`` for ``n`` + channels, this transform will normalize each channel of the input + ``torch.*Tensor`` i.e., + ``output[channel] = (input[channel] - mean[channel]) / std[channel]`` + + .. note:: + This transform acts out of place, i.e., it does not mutate the input tensor. + + Args: + mean (sequence): Sequence of means for each channel. + std (sequence): Sequence of standard deviations for each channel. + inplace(bool,optional): Bool to make this operation in-place. + + """ + + def __init__(self, mean, std, inplace=False): + super().__init__() + _log_api_usage_once(self) + self.mean = mean + self.std = std + self.inplace = inplace + + def forward(self, tensor: Tensor) -> Tensor: + """ + Args: + tensor (Tensor): Tensor image to be normalized. + + Returns: + Tensor: Normalized Tensor image. + """ + return F.normalize(tensor, self.mean, self.std, self.inplace) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(mean={self.mean}, std={self.std})" + + +class Resize(torch.nn.Module): + """Resize the input image to the given size. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means a maximum of two leading dimensions + + Args: + size (sequence or int): Desired output size. If size is a sequence like + (h, w), output size will be matched to this. If size is an int, + smaller edge of the image will be matched to this number. + i.e, if height > width, then image will be rescaled to + (size * height / width, size). + + .. note:: + In torchscript mode size as single int is not supported, use a sequence of length 1: ``[size, ]``. + interpolation (InterpolationMode): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.NEAREST_EXACT``, + ``InterpolationMode.BILINEAR`` and ``InterpolationMode.BICUBIC`` are supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + max_size (int, optional): The maximum allowed for the longer edge of + the resized image. If the longer edge of the image is greater + than ``max_size`` after being resized according to ``size``, + ``size`` will be overruled so that the longer edge is equal to + ``max_size``. + As a result, the smaller edge may be shorter than ``size``. This + is only supported if ``size`` is an int (or a sequence of length + 1 in torchscript mode). + antialias (bool, optional): Whether to apply antialiasing. + It only affects **tensors** with bilinear or bicubic modes and it is + ignored otherwise: on PIL images, antialiasing is always applied on + bilinear or bicubic modes; on other modes (for PIL images and + tensors), antialiasing makes no sense and this parameter is ignored. + Possible values are: + + - ``True`` (default): will apply antialiasing for bilinear or bicubic modes. + Other mode aren't affected. This is probably what you want to use. + - ``False``: will not apply antialiasing for tensors on any mode. PIL + images are still antialiased on bilinear or bicubic modes, because + PIL doesn't support no antialias. + - ``None``: equivalent to ``False`` for tensors and ``True`` for + PIL images. This value exists for legacy reasons and you probably + don't want to use it unless you really know what you are doing. + + The default value changed from ``None`` to ``True`` in + v0.17, for the PIL and Tensor backends to be consistent. + """ + + def __init__(self, size, interpolation=InterpolationMode.BILINEAR, max_size=None, antialias=True): + super().__init__() + _log_api_usage_once(self) + if not isinstance(size, (int, Sequence)): + raise TypeError(f"Size should be int or sequence. Got {type(size)}") + if isinstance(size, Sequence) and len(size) not in (1, 2): + raise ValueError("If size is a sequence, it should have 1 or 2 values") + self.size = size + self.max_size = max_size + + if isinstance(interpolation, int): + interpolation = _interpolation_modes_from_int(interpolation) + + self.interpolation = interpolation + self.antialias = antialias + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be scaled. + + Returns: + PIL Image or Tensor: Rescaled image. + """ + return F.resize(img, self.size, self.interpolation, self.max_size, self.antialias) + + def __repr__(self) -> str: + detail = f"(size={self.size}, interpolation={self.interpolation.value}, max_size={self.max_size}, antialias={self.antialias})" + return f"{self.__class__.__name__}{detail}" + + +class CenterCrop(torch.nn.Module): + """Crops the given image at the center. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions. + If image size is smaller than output size along any edge, image is padded with 0 and then center cropped. + + Args: + size (sequence or int): Desired output size of the crop. If size is an + int instead of sequence like (h, w), a square crop (size, size) is + made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). + """ + + def __init__(self, size): + super().__init__() + _log_api_usage_once(self) + self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.") + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be cropped. + + Returns: + PIL Image or Tensor: Cropped image. + """ + return F.center_crop(img, self.size) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(size={self.size})" + + +class Pad(torch.nn.Module): + """Pad the given image on all sides with the given "pad" value. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means at most 2 leading dimensions for mode reflect and symmetric, + at most 3 leading dimensions for mode edge, + and an arbitrary number of leading dimensions for mode constant + + Args: + padding (int or sequence): Padding on each border. If a single int is provided this + is used to pad all borders. If sequence of length 2 is provided this is the padding + on left/right and top/bottom respectively. If a sequence of length 4 is provided + this is the padding for the left, top, right and bottom borders respectively. + + .. note:: + In torchscript mode padding as single int is not supported, use a sequence of + length 1: ``[padding, ]``. + fill (number or tuple): Pixel fill value for constant fill. Default is 0. If a tuple of + length 3, it is used to fill R, G, B channels respectively. + This value is only used when the padding_mode is constant. + Only number is supported for torch Tensor. + Only int or tuple value is supported for PIL Image. + padding_mode (str): Type of padding. Should be: constant, edge, reflect or symmetric. + Default is constant. + + - constant: pads with a constant value, this value is specified with fill + + - edge: pads with the last value at the edge of the image. + If input a 5D torch Tensor, the last 3 dimensions will be padded instead of the last 2 + + - reflect: pads with reflection of image without repeating the last value on the edge. + For example, padding [1, 2, 3, 4] with 2 elements on both sides in reflect mode + will result in [3, 2, 1, 2, 3, 4, 3, 2] + + - symmetric: pads with reflection of image repeating the last value on the edge. + For example, padding [1, 2, 3, 4] with 2 elements on both sides in symmetric mode + will result in [2, 1, 1, 2, 3, 4, 4, 3] + """ + + def __init__(self, padding, fill=0, padding_mode="constant"): + super().__init__() + _log_api_usage_once(self) + if not isinstance(padding, (numbers.Number, tuple, list)): + raise TypeError("Got inappropriate padding arg") + + if not isinstance(fill, (numbers.Number, tuple, list)): + raise TypeError("Got inappropriate fill arg") + + if padding_mode not in ["constant", "edge", "reflect", "symmetric"]: + raise ValueError("Padding mode should be either constant, edge, reflect or symmetric") + + if isinstance(padding, Sequence) and len(padding) not in [1, 2, 4]: + raise ValueError( + f"Padding must be an int or a 1, 2, or 4 element tuple, not a {len(padding)} element tuple" + ) + + self.padding = padding + self.fill = fill + self.padding_mode = padding_mode + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be padded. + + Returns: + PIL Image or Tensor: Padded image. + """ + return F.pad(img, self.padding, self.fill, self.padding_mode) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(padding={self.padding}, fill={self.fill}, padding_mode={self.padding_mode})" + + +class Lambda: + """Apply a user-defined lambda as a transform. This transform does not support torchscript. + + Args: + lambd (function): Lambda/function to be used for transform. + """ + + def __init__(self, lambd): + _log_api_usage_once(self) + if not callable(lambd): + raise TypeError(f"Argument lambd should be callable, got {repr(type(lambd).__name__)}") + self.lambd = lambd + + def __call__(self, img): + return self.lambd(img) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}()" + + +class RandomTransforms: + """Base class for a list of transformations with randomness + + Args: + transforms (sequence): list of transformations + """ + + def __init__(self, transforms): + _log_api_usage_once(self) + if not isinstance(transforms, Sequence): + raise TypeError("Argument transforms should be a sequence") + self.transforms = transforms + + def __call__(self, *args, **kwargs): + raise NotImplementedError() + + def __repr__(self) -> str: + format_string = self.__class__.__name__ + "(" + for t in self.transforms: + format_string += "\n" + format_string += f" {t}" + format_string += "\n)" + return format_string + + +class RandomApply(torch.nn.Module): + """Apply randomly a list of transformations with a given probability. + + .. note:: + In order to script the transformation, please use ``torch.nn.ModuleList`` as input instead of list/tuple of + transforms as shown below: + + >>> transforms = transforms.RandomApply(torch.nn.ModuleList([ + >>> transforms.ColorJitter(), + >>> ]), p=0.3) + >>> scripted_transforms = torch.jit.script(transforms) + + Make sure to use only scriptable transformations, i.e. that work with ``torch.Tensor``, does not require + `lambda` functions or ``PIL.Image``. + + Args: + transforms (sequence or torch.nn.Module): list of transformations + p (float): probability + """ + + def __init__(self, transforms, p=0.5): + super().__init__() + _log_api_usage_once(self) + self.transforms = transforms + self.p = p + + def forward(self, img): + if self.p < torch.rand(1): + return img + for t in self.transforms: + img = t(img) + return img + + def __repr__(self) -> str: + format_string = self.__class__.__name__ + "(" + format_string += f"\n p={self.p}" + for t in self.transforms: + format_string += "\n" + format_string += f" {t}" + format_string += "\n)" + return format_string + + +class RandomOrder(RandomTransforms): + """Apply a list of transformations in a random order. This transform does not support torchscript.""" + + def __call__(self, img): + order = list(range(len(self.transforms))) + random.shuffle(order) + for i in order: + img = self.transforms[i](img) + return img + + +class RandomChoice(RandomTransforms): + """Apply single transformation randomly picked from a list. This transform does not support torchscript.""" + + def __init__(self, transforms, p=None): + super().__init__(transforms) + if p is not None and not isinstance(p, Sequence): + raise TypeError("Argument p should be a sequence") + self.p = p + + def __call__(self, *args): + t = random.choices(self.transforms, weights=self.p)[0] + return t(*args) + + def __repr__(self) -> str: + return f"{super().__repr__()}(p={self.p})" + + +class RandomCrop(torch.nn.Module): + """Crop the given image at a random location. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions, + but if non-constant padding is used, the input is expected to have at most 2 leading dimensions + + Args: + size (sequence or int): Desired output size of the crop. If size is an + int instead of sequence like (h, w), a square crop (size, size) is + made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). + padding (int or sequence, optional): Optional padding on each border + of the image. Default is None. If a single int is provided this + is used to pad all borders. If sequence of length 2 is provided this is the padding + on left/right and top/bottom respectively. If a sequence of length 4 is provided + this is the padding for the left, top, right and bottom borders respectively. + + .. note:: + In torchscript mode padding as single int is not supported, use a sequence of + length 1: ``[padding, ]``. + pad_if_needed (boolean): It will pad the image if smaller than the + desired size to avoid raising an exception. Since cropping is done + after padding, the padding seems to be done at a random offset. + fill (number or tuple): Pixel fill value for constant fill. Default is 0. If a tuple of + length 3, it is used to fill R, G, B channels respectively. + This value is only used when the padding_mode is constant. + Only number is supported for torch Tensor. + Only int or tuple value is supported for PIL Image. + padding_mode (str): Type of padding. Should be: constant, edge, reflect or symmetric. + Default is constant. + + - constant: pads with a constant value, this value is specified with fill + + - edge: pads with the last value at the edge of the image. + If input a 5D torch Tensor, the last 3 dimensions will be padded instead of the last 2 + + - reflect: pads with reflection of image without repeating the last value on the edge. + For example, padding [1, 2, 3, 4] with 2 elements on both sides in reflect mode + will result in [3, 2, 1, 2, 3, 4, 3, 2] + + - symmetric: pads with reflection of image repeating the last value on the edge. + For example, padding [1, 2, 3, 4] with 2 elements on both sides in symmetric mode + will result in [2, 1, 1, 2, 3, 4, 4, 3] + """ + + @staticmethod + def get_params(img: Tensor, output_size: Tuple[int, int]) -> Tuple[int, int, int, int]: + """Get parameters for ``crop`` for a random crop. + + Args: + img (PIL Image or Tensor): Image to be cropped. + output_size (tuple): Expected output size of the crop. + + Returns: + tuple: params (i, j, h, w) to be passed to ``crop`` for random crop. + """ + _, h, w = F.get_dimensions(img) + th, tw = output_size + + if h < th or w < tw: + raise ValueError(f"Required crop size {(th, tw)} is larger than input image size {(h, w)}") + + if w == tw and h == th: + return 0, 0, h, w + + i = torch.randint(0, h - th + 1, size=(1,)).item() + j = torch.randint(0, w - tw + 1, size=(1,)).item() + return i, j, th, tw + + def __init__(self, size, padding=None, pad_if_needed=False, fill=0, padding_mode="constant"): + super().__init__() + _log_api_usage_once(self) + + self.size = tuple(_setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")) + + self.padding = padding + self.pad_if_needed = pad_if_needed + self.fill = fill + self.padding_mode = padding_mode + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be cropped. + + Returns: + PIL Image or Tensor: Cropped image. + """ + if self.padding is not None: + img = F.pad(img, self.padding, self.fill, self.padding_mode) + + _, height, width = F.get_dimensions(img) + # pad the width if needed + if self.pad_if_needed and width < self.size[1]: + padding = [self.size[1] - width, 0] + img = F.pad(img, padding, self.fill, self.padding_mode) + # pad the height if needed + if self.pad_if_needed and height < self.size[0]: + padding = [0, self.size[0] - height] + img = F.pad(img, padding, self.fill, self.padding_mode) + + i, j, h, w = self.get_params(img, self.size) + + return F.crop(img, i, j, h, w) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(size={self.size}, padding={self.padding})" + + +class RandomHorizontalFlip(torch.nn.Module): + """Horizontally flip the given image randomly with a given probability. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means an arbitrary number of leading + dimensions + + Args: + p (float): probability of the image being flipped. Default value is 0.5 + """ + + def __init__(self, p=0.5): + super().__init__() + _log_api_usage_once(self) + self.p = p + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be flipped. + + Returns: + PIL Image or Tensor: Randomly flipped image. + """ + if torch.rand(1) < self.p: + return F.hflip(img) + return img + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(p={self.p})" + + +class RandomVerticalFlip(torch.nn.Module): + """Vertically flip the given image randomly with a given probability. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means an arbitrary number of leading + dimensions + + Args: + p (float): probability of the image being flipped. Default value is 0.5 + """ + + def __init__(self, p=0.5): + super().__init__() + _log_api_usage_once(self) + self.p = p + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be flipped. + + Returns: + PIL Image or Tensor: Randomly flipped image. + """ + if torch.rand(1) < self.p: + return F.vflip(img) + return img + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(p={self.p})" + + +class RandomPerspective(torch.nn.Module): + """Performs a random perspective transformation of the given image with a given probability. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions. + + Args: + distortion_scale (float): argument to control the degree of distortion and ranges from 0 to 1. + Default is 0.5. + p (float): probability of the image being transformed. Default is 0.5. + interpolation (InterpolationMode): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + fill (sequence or number): Pixel fill value for the area outside the transformed + image. Default is ``0``. If given a number, the value is used for all bands respectively. + """ + + def __init__(self, distortion_scale=0.5, p=0.5, interpolation=InterpolationMode.BILINEAR, fill=0): + super().__init__() + _log_api_usage_once(self) + self.p = p + + if isinstance(interpolation, int): + interpolation = _interpolation_modes_from_int(interpolation) + + self.interpolation = interpolation + self.distortion_scale = distortion_scale + + if fill is None: + fill = 0 + elif not isinstance(fill, (Sequence, numbers.Number)): + raise TypeError("Fill should be either a sequence or a number.") + + self.fill = fill + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be Perspectively transformed. + + Returns: + PIL Image or Tensor: Randomly transformed image. + """ + + fill = self.fill + channels, height, width = F.get_dimensions(img) + if isinstance(img, Tensor): + if isinstance(fill, (int, float)): + fill = [float(fill)] * channels + else: + fill = [float(f) for f in fill] + + if torch.rand(1) < self.p: + startpoints, endpoints = self.get_params(width, height, self.distortion_scale) + return F.perspective(img, startpoints, endpoints, self.interpolation, fill) + return img + + @staticmethod + def get_params(width: int, height: int, distortion_scale: float) -> Tuple[List[List[int]], List[List[int]]]: + """Get parameters for ``perspective`` for a random perspective transform. + + Args: + width (int): width of the image. + height (int): height of the image. + distortion_scale (float): argument to control the degree of distortion and ranges from 0 to 1. + + Returns: + List containing [top-left, top-right, bottom-right, bottom-left] of the original image, + List containing [top-left, top-right, bottom-right, bottom-left] of the transformed image. + """ + half_height = height // 2 + half_width = width // 2 + topleft = [ + int(torch.randint(0, int(distortion_scale * half_width) + 1, size=(1,)).item()), + int(torch.randint(0, int(distortion_scale * half_height) + 1, size=(1,)).item()), + ] + topright = [ + int(torch.randint(width - int(distortion_scale * half_width) - 1, width, size=(1,)).item()), + int(torch.randint(0, int(distortion_scale * half_height) + 1, size=(1,)).item()), + ] + botright = [ + int(torch.randint(width - int(distortion_scale * half_width) - 1, width, size=(1,)).item()), + int(torch.randint(height - int(distortion_scale * half_height) - 1, height, size=(1,)).item()), + ] + botleft = [ + int(torch.randint(0, int(distortion_scale * half_width) + 1, size=(1,)).item()), + int(torch.randint(height - int(distortion_scale * half_height) - 1, height, size=(1,)).item()), + ] + startpoints = [[0, 0], [width - 1, 0], [width - 1, height - 1], [0, height - 1]] + endpoints = [topleft, topright, botright, botleft] + return startpoints, endpoints + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(p={self.p})" + + +class RandomResizedCrop(torch.nn.Module): + """Crop a random portion of image and resize it to a given size. + + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions + + A crop of the original image is made: the crop has a random area (H * W) + and a random aspect ratio. This crop is finally resized to the given + size. This is popularly used to train the Inception networks. + + Args: + size (int or sequence): expected output size of the crop, for each edge. If size is an + int instead of sequence like (h, w), a square output size ``(size, size)`` is + made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). + + .. note:: + In torchscript mode size as single int is not supported, use a sequence of length 1: ``[size, ]``. + scale (tuple of float): Specifies the lower and upper bounds for the random area of the crop, + before resizing. The scale is defined with respect to the area of the original image. + ratio (tuple of float): lower and upper bounds for the random aspect ratio of the crop, before + resizing. + interpolation (InterpolationMode): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.NEAREST_EXACT``, + ``InterpolationMode.BILINEAR`` and ``InterpolationMode.BICUBIC`` are supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + antialias (bool, optional): Whether to apply antialiasing. + It only affects **tensors** with bilinear or bicubic modes and it is + ignored otherwise: on PIL images, antialiasing is always applied on + bilinear or bicubic modes; on other modes (for PIL images and + tensors), antialiasing makes no sense and this parameter is ignored. + Possible values are: + + - ``True`` (default): will apply antialiasing for bilinear or bicubic modes. + Other mode aren't affected. This is probably what you want to use. + - ``False``: will not apply antialiasing for tensors on any mode. PIL + images are still antialiased on bilinear or bicubic modes, because + PIL doesn't support no antialias. + - ``None``: equivalent to ``False`` for tensors and ``True`` for + PIL images. This value exists for legacy reasons and you probably + don't want to use it unless you really know what you are doing. + + The default value changed from ``None`` to ``True`` in + v0.17, for the PIL and Tensor backends to be consistent. + """ + + def __init__( + self, + size, + scale=(0.08, 1.0), + ratio=(3.0 / 4.0, 4.0 / 3.0), + interpolation=InterpolationMode.BILINEAR, + antialias: Optional[bool] = True, + ): + super().__init__() + _log_api_usage_once(self) + self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.") + + if not isinstance(scale, Sequence): + raise TypeError("Scale should be a sequence") + if not isinstance(ratio, Sequence): + raise TypeError("Ratio should be a sequence") + if (scale[0] > scale[1]) or (ratio[0] > ratio[1]): + warnings.warn("Scale and ratio should be of kind (min, max)") + + if isinstance(interpolation, int): + interpolation = _interpolation_modes_from_int(interpolation) + + self.interpolation = interpolation + self.antialias = antialias + self.scale = scale + self.ratio = ratio + + @staticmethod + def get_params(img: Tensor, scale: List[float], ratio: List[float]) -> Tuple[int, int, int, int]: + """Get parameters for ``crop`` for a random sized crop. + + Args: + img (PIL Image or Tensor): Input image. + scale (list): range of scale of the origin size cropped + ratio (list): range of aspect ratio of the origin aspect ratio cropped + + Returns: + tuple: params (i, j, h, w) to be passed to ``crop`` for a random + sized crop. + """ + _, height, width = F.get_dimensions(img) + area = height * width + + log_ratio = torch.log(torch.tensor(ratio)) + for _ in range(10): + target_area = area * torch.empty(1).uniform_(scale[0], scale[1]).item() + aspect_ratio = torch.exp(torch.empty(1).uniform_(log_ratio[0], log_ratio[1])).item() + + w = int(round(math.sqrt(target_area * aspect_ratio))) + h = int(round(math.sqrt(target_area / aspect_ratio))) + + if 0 < w <= width and 0 < h <= height: + i = torch.randint(0, height - h + 1, size=(1,)).item() + j = torch.randint(0, width - w + 1, size=(1,)).item() + return i, j, h, w + + # Fallback to central crop + in_ratio = float(width) / float(height) + if in_ratio < min(ratio): + w = width + h = int(round(w / min(ratio))) + elif in_ratio > max(ratio): + h = height + w = int(round(h * max(ratio))) + else: # whole image + w = width + h = height + i = (height - h) // 2 + j = (width - w) // 2 + return i, j, h, w + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be cropped and resized. + + Returns: + PIL Image or Tensor: Randomly cropped and resized image. + """ + i, j, h, w = self.get_params(img, self.scale, self.ratio) + return F.resized_crop(img, i, j, h, w, self.size, self.interpolation, antialias=self.antialias) + + def __repr__(self) -> str: + interpolate_str = self.interpolation.value + format_string = self.__class__.__name__ + f"(size={self.size}" + format_string += f", scale={tuple(round(s, 4) for s in self.scale)}" + format_string += f", ratio={tuple(round(r, 4) for r in self.ratio)}" + format_string += f", interpolation={interpolate_str}" + format_string += f", antialias={self.antialias})" + return format_string + + +class FiveCrop(torch.nn.Module): + """Crop the given image into four corners and the central crop. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means an arbitrary number of leading + dimensions + + .. Note:: + This transform returns a tuple of images and there may be a mismatch in the number of + inputs and targets your Dataset returns. See below for an example of how to deal with + this. + + Args: + size (sequence or int): Desired output size of the crop. If size is an ``int`` + instead of sequence like (h, w), a square crop of size (size, size) is made. + If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). + + Example: + >>> transform = Compose([ + >>> FiveCrop(size), # this is a list of PIL Images + >>> Lambda(lambda crops: torch.stack([PILToTensor()(crop) for crop in crops])) # returns a 4D tensor + >>> ]) + >>> #In your test loop you can do the following: + >>> input, target = batch # input is a 5d tensor, target is 2d + >>> bs, ncrops, c, h, w = input.size() + >>> result = model(input.view(-1, c, h, w)) # fuse batch size and ncrops + >>> result_avg = result.view(bs, ncrops, -1).mean(1) # avg over crops + """ + + def __init__(self, size): + super().__init__() + _log_api_usage_once(self) + self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.") + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be cropped. + + Returns: + tuple of 5 images. Image can be PIL Image or Tensor + """ + return F.five_crop(img, self.size) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(size={self.size})" + + +class TenCrop(torch.nn.Module): + """Crop the given image into four corners and the central crop plus the flipped version of + these (horizontal flipping is used by default). + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means an arbitrary number of leading + dimensions + + .. Note:: + This transform returns a tuple of images and there may be a mismatch in the number of + inputs and targets your Dataset returns. See below for an example of how to deal with + this. + + Args: + size (sequence or int): Desired output size of the crop. If size is an + int instead of sequence like (h, w), a square crop (size, size) is + made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). + vertical_flip (bool): Use vertical flipping instead of horizontal + + Example: + >>> transform = Compose([ + >>> TenCrop(size), # this is a tuple of PIL Images + >>> Lambda(lambda crops: torch.stack([PILToTensor()(crop) for crop in crops])) # returns a 4D tensor + >>> ]) + >>> #In your test loop you can do the following: + >>> input, target = batch # input is a 5d tensor, target is 2d + >>> bs, ncrops, c, h, w = input.size() + >>> result = model(input.view(-1, c, h, w)) # fuse batch size and ncrops + >>> result_avg = result.view(bs, ncrops, -1).mean(1) # avg over crops + """ + + def __init__(self, size, vertical_flip=False): + super().__init__() + _log_api_usage_once(self) + self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.") + self.vertical_flip = vertical_flip + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be cropped. + + Returns: + tuple of 10 images. Image can be PIL Image or Tensor + """ + return F.ten_crop(img, self.size, self.vertical_flip) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(size={self.size}, vertical_flip={self.vertical_flip})" + + +class LinearTransformation(torch.nn.Module): + """Transform a tensor image with a square transformation matrix and a mean_vector computed + offline. + This transform does not support PIL Image. + Given transformation_matrix and mean_vector, will flatten the torch.*Tensor and + subtract mean_vector from it which is then followed by computing the dot + product with the transformation matrix and then reshaping the tensor to its + original shape. + + Applications: + whitening transformation: Suppose X is a column vector zero-centered data. + Then compute the data covariance matrix [D x D] with torch.mm(X.t(), X), + perform SVD on this matrix and pass it as transformation_matrix. + + Args: + transformation_matrix (Tensor): tensor [D x D], D = C x H x W + mean_vector (Tensor): tensor [D], D = C x H x W + """ + + def __init__(self, transformation_matrix, mean_vector): + super().__init__() + _log_api_usage_once(self) + if transformation_matrix.size(0) != transformation_matrix.size(1): + raise ValueError( + "transformation_matrix should be square. Got " + f"{tuple(transformation_matrix.size())} rectangular matrix." + ) + + if mean_vector.size(0) != transformation_matrix.size(0): + raise ValueError( + f"mean_vector should have the same length {mean_vector.size(0)}" + f" as any one of the dimensions of the transformation_matrix [{tuple(transformation_matrix.size())}]" + ) + + if transformation_matrix.device != mean_vector.device: + raise ValueError( + f"Input tensors should be on the same device. Got {transformation_matrix.device} and {mean_vector.device}" + ) + + if transformation_matrix.dtype != mean_vector.dtype: + raise ValueError( + f"Input tensors should have the same dtype. Got {transformation_matrix.dtype} and {mean_vector.dtype}" + ) + + self.transformation_matrix = transformation_matrix + self.mean_vector = mean_vector + + def forward(self, tensor: Tensor) -> Tensor: + """ + Args: + tensor (Tensor): Tensor image to be whitened. + + Returns: + Tensor: Transformed image. + """ + shape = tensor.shape + n = shape[-3] * shape[-2] * shape[-1] + if n != self.transformation_matrix.shape[0]: + raise ValueError( + "Input tensor and transformation matrix have incompatible shape." + + f"[{shape[-3]} x {shape[-2]} x {shape[-1]}] != " + + f"{self.transformation_matrix.shape[0]}" + ) + + if tensor.device.type != self.mean_vector.device.type: + raise ValueError( + "Input tensor should be on the same device as transformation matrix and mean vector. " + f"Got {tensor.device} vs {self.mean_vector.device}" + ) + + flat_tensor = tensor.view(-1, n) - self.mean_vector + transformation_matrix = self.transformation_matrix.to(flat_tensor.dtype) + transformed_tensor = torch.mm(flat_tensor, transformation_matrix) + tensor = transformed_tensor.view(shape) + return tensor + + def __repr__(self) -> str: + s = ( + f"{self.__class__.__name__}(transformation_matrix=" + f"{self.transformation_matrix.tolist()}" + f", mean_vector={self.mean_vector.tolist()})" + ) + return s + + +class ColorJitter(torch.nn.Module): + """Randomly change the brightness, contrast, saturation and hue of an image. + If the image is torch Tensor, it is expected + to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. + If img is PIL Image, mode "1", "I", "F" and modes with transparency (alpha channel) are not supported. + + Args: + brightness (float or tuple of float (min, max)): How much to jitter brightness. + brightness_factor is chosen uniformly from [max(0, 1 - brightness), 1 + brightness] + or the given [min, max]. Should be non negative numbers. + contrast (float or tuple of float (min, max)): How much to jitter contrast. + contrast_factor is chosen uniformly from [max(0, 1 - contrast), 1 + contrast] + or the given [min, max]. Should be non-negative numbers. + saturation (float or tuple of float (min, max)): How much to jitter saturation. + saturation_factor is chosen uniformly from [max(0, 1 - saturation), 1 + saturation] + or the given [min, max]. Should be non negative numbers. + hue (float or tuple of float (min, max)): How much to jitter hue. + hue_factor is chosen uniformly from [-hue, hue] or the given [min, max]. + Should have 0<= hue <= 0.5 or -0.5 <= min <= max <= 0.5. + To jitter hue, the pixel values of the input image has to be non-negative for conversion to HSV space; + thus it does not work if you normalize your image to an interval with negative values, + or use an interpolation that generates negative values before using this function. + """ + + def __init__( + self, + brightness: Union[float, Tuple[float, float]] = 0, + contrast: Union[float, Tuple[float, float]] = 0, + saturation: Union[float, Tuple[float, float]] = 0, + hue: Union[float, Tuple[float, float]] = 0, + ) -> None: + super().__init__() + _log_api_usage_once(self) + self.brightness = self._check_input(brightness, "brightness") + self.contrast = self._check_input(contrast, "contrast") + self.saturation = self._check_input(saturation, "saturation") + self.hue = self._check_input(hue, "hue", center=0, bound=(-0.5, 0.5), clip_first_on_zero=False) + + @torch.jit.unused + def _check_input(self, value, name, center=1, bound=(0, float("inf")), clip_first_on_zero=True): + if isinstance(value, numbers.Number): + if value < 0: + raise ValueError(f"If {name} is a single number, it must be non negative.") + value = [center - float(value), center + float(value)] + if clip_first_on_zero: + value[0] = max(value[0], 0.0) + elif isinstance(value, (tuple, list)) and len(value) == 2: + value = [float(value[0]), float(value[1])] + else: + raise TypeError(f"{name} should be a single number or a list/tuple with length 2.") + + if not bound[0] <= value[0] <= value[1] <= bound[1]: + raise ValueError(f"{name} values should be between {bound}, but got {value}.") + + # if value is 0 or (1., 1.) for brightness/contrast/saturation + # or (0., 0.) for hue, do nothing + if value[0] == value[1] == center: + return None + else: + return tuple(value) + + @staticmethod + def get_params( + brightness: Optional[List[float]], + contrast: Optional[List[float]], + saturation: Optional[List[float]], + hue: Optional[List[float]], + ) -> Tuple[Tensor, Optional[float], Optional[float], Optional[float], Optional[float]]: + """Get the parameters for the randomized transform to be applied on image. + + Args: + brightness (tuple of float (min, max), optional): The range from which the brightness_factor is chosen + uniformly. Pass None to turn off the transformation. + contrast (tuple of float (min, max), optional): The range from which the contrast_factor is chosen + uniformly. Pass None to turn off the transformation. + saturation (tuple of float (min, max), optional): The range from which the saturation_factor is chosen + uniformly. Pass None to turn off the transformation. + hue (tuple of float (min, max), optional): The range from which the hue_factor is chosen uniformly. + Pass None to turn off the transformation. + + Returns: + tuple: The parameters used to apply the randomized transform + along with their random order. + """ + fn_idx = torch.randperm(4) + + b = None if brightness is None else float(torch.empty(1).uniform_(brightness[0], brightness[1])) + c = None if contrast is None else float(torch.empty(1).uniform_(contrast[0], contrast[1])) + s = None if saturation is None else float(torch.empty(1).uniform_(saturation[0], saturation[1])) + h = None if hue is None else float(torch.empty(1).uniform_(hue[0], hue[1])) + + return fn_idx, b, c, s, h + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Input image. + + Returns: + PIL Image or Tensor: Color jittered image. + """ + fn_idx, brightness_factor, contrast_factor, saturation_factor, hue_factor = self.get_params( + self.brightness, self.contrast, self.saturation, self.hue + ) + + for fn_id in fn_idx: + if fn_id == 0 and brightness_factor is not None: + img = F.adjust_brightness(img, brightness_factor) + elif fn_id == 1 and contrast_factor is not None: + img = F.adjust_contrast(img, contrast_factor) + elif fn_id == 2 and saturation_factor is not None: + img = F.adjust_saturation(img, saturation_factor) + elif fn_id == 3 and hue_factor is not None: + img = F.adjust_hue(img, hue_factor) + + return img + + def __repr__(self) -> str: + s = ( + f"{self.__class__.__name__}(" + f"brightness={self.brightness}" + f", contrast={self.contrast}" + f", saturation={self.saturation}" + f", hue={self.hue})" + ) + return s + + +class RandomRotation(torch.nn.Module): + """Rotate the image by angle. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions. + + Args: + degrees (sequence or number): Range of degrees to select from. + If degrees is a number instead of sequence like (min, max), the range of degrees + will be (-degrees, +degrees). + interpolation (InterpolationMode): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + expand (bool, optional): Optional expansion flag. + If true, expands the output to make it large enough to hold the entire rotated image. + If false or omitted, make the output image the same size as the input image. + Note that the expand flag assumes rotation around the center and no translation. + center (sequence, optional): Optional center of rotation, (x, y). Origin is the upper left corner. + Default is the center of the image. + fill (sequence or number): Pixel fill value for the area outside the rotated + image. Default is ``0``. If given a number, the value is used for all bands respectively. + + .. _filters: https://pillow.readthedocs.io/en/latest/handbook/concepts.html#filters + + """ + + def __init__(self, degrees, interpolation=InterpolationMode.NEAREST, expand=False, center=None, fill=0): + super().__init__() + _log_api_usage_once(self) + + if isinstance(interpolation, int): + interpolation = _interpolation_modes_from_int(interpolation) + + self.degrees = _setup_angle(degrees, name="degrees", req_sizes=(2,)) + + if center is not None: + _check_sequence_input(center, "center", req_sizes=(2,)) + + self.center = center + + self.interpolation = interpolation + self.expand = expand + + if fill is None: + fill = 0 + elif not isinstance(fill, (Sequence, numbers.Number)): + raise TypeError("Fill should be either a sequence or a number.") + + self.fill = fill + + @staticmethod + def get_params(degrees: List[float]) -> float: + """Get parameters for ``rotate`` for a random rotation. + + Returns: + float: angle parameter to be passed to ``rotate`` for random rotation. + """ + angle = float(torch.empty(1).uniform_(float(degrees[0]), float(degrees[1])).item()) + return angle + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be rotated. + + Returns: + PIL Image or Tensor: Rotated image. + """ + fill = self.fill + channels, _, _ = F.get_dimensions(img) + if isinstance(img, Tensor): + if isinstance(fill, (int, float)): + fill = [float(fill)] * channels + else: + fill = [float(f) for f in fill] + angle = self.get_params(self.degrees) + + return F.rotate(img, angle, self.interpolation, self.expand, self.center, fill) + + def __repr__(self) -> str: + interpolate_str = self.interpolation.value + format_string = self.__class__.__name__ + f"(degrees={self.degrees}" + format_string += f", interpolation={interpolate_str}" + format_string += f", expand={self.expand}" + if self.center is not None: + format_string += f", center={self.center}" + if self.fill is not None: + format_string += f", fill={self.fill}" + format_string += ")" + return format_string + + +class RandomAffine(torch.nn.Module): + """Random affine transformation of the image keeping center invariant. + If the image is torch Tensor, it is expected + to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions. + + Args: + degrees (sequence or number): Range of degrees to select from. + If degrees is a number instead of sequence like (min, max), the range of degrees + will be (-degrees, +degrees). Set to 0 to deactivate rotations. + translate (tuple, optional): tuple of maximum absolute fraction for horizontal + and vertical translations. For example translate=(a, b), then horizontal shift + is randomly sampled in the range -img_width * a < dx < img_width * a and vertical shift is + randomly sampled in the range -img_height * b < dy < img_height * b. Will not translate by default. + scale (tuple, optional): scaling factor interval, e.g (a, b), then scale is + randomly sampled from the range a <= scale <= b. Will keep original scale by default. + shear (sequence or number, optional): Range of degrees to select from. + If shear is a number, a shear parallel to the x-axis in the range (-shear, +shear) + will be applied. Else if shear is a sequence of 2 values a shear parallel to the x-axis in the + range (shear[0], shear[1]) will be applied. Else if shear is a sequence of 4 values, + an x-axis shear in (shear[0], shear[1]) and y-axis shear in (shear[2], shear[3]) will be applied. + Will not apply shear by default. + interpolation (InterpolationMode): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + fill (sequence or number): Pixel fill value for the area outside the transformed + image. Default is ``0``. If given a number, the value is used for all bands respectively. + center (sequence, optional): Optional center of rotation, (x, y). Origin is the upper left corner. + Default is the center of the image. + + .. _filters: https://pillow.readthedocs.io/en/latest/handbook/concepts.html#filters + + """ + + def __init__( + self, + degrees, + translate=None, + scale=None, + shear=None, + interpolation=InterpolationMode.NEAREST, + fill=0, + center=None, + ): + super().__init__() + _log_api_usage_once(self) + + if isinstance(interpolation, int): + interpolation = _interpolation_modes_from_int(interpolation) + + self.degrees = _setup_angle(degrees, name="degrees", req_sizes=(2,)) + + if translate is not None: + _check_sequence_input(translate, "translate", req_sizes=(2,)) + for t in translate: + if not (0.0 <= t <= 1.0): + raise ValueError("translation values should be between 0 and 1") + self.translate = translate + + if scale is not None: + _check_sequence_input(scale, "scale", req_sizes=(2,)) + for s in scale: + if s <= 0: + raise ValueError("scale values should be positive") + self.scale = scale + + if shear is not None: + self.shear = _setup_angle(shear, name="shear", req_sizes=(2, 4)) + else: + self.shear = shear + + self.interpolation = interpolation + + if fill is None: + fill = 0 + elif not isinstance(fill, (Sequence, numbers.Number)): + raise TypeError("Fill should be either a sequence or a number.") + + self.fill = fill + + if center is not None: + _check_sequence_input(center, "center", req_sizes=(2,)) + + self.center = center + + @staticmethod + def get_params( + degrees: List[float], + translate: Optional[List[float]], + scale_ranges: Optional[List[float]], + shears: Optional[List[float]], + img_size: List[int], + ) -> Tuple[float, Tuple[int, int], float, Tuple[float, float]]: + """Get parameters for affine transformation + + Returns: + params to be passed to the affine transformation + """ + angle = float(torch.empty(1).uniform_(float(degrees[0]), float(degrees[1])).item()) + if translate is not None: + max_dx = float(translate[0] * img_size[0]) + max_dy = float(translate[1] * img_size[1]) + tx = int(round(torch.empty(1).uniform_(-max_dx, max_dx).item())) + ty = int(round(torch.empty(1).uniform_(-max_dy, max_dy).item())) + translations = (tx, ty) + else: + translations = (0, 0) + + if scale_ranges is not None: + scale = float(torch.empty(1).uniform_(scale_ranges[0], scale_ranges[1]).item()) + else: + scale = 1.0 + + shear_x = shear_y = 0.0 + if shears is not None: + shear_x = float(torch.empty(1).uniform_(shears[0], shears[1]).item()) + if len(shears) == 4: + shear_y = float(torch.empty(1).uniform_(shears[2], shears[3]).item()) + + shear = (shear_x, shear_y) + + return angle, translations, scale, shear + + def forward(self, img): + """ + img (PIL Image or Tensor): Image to be transformed. + + Returns: + PIL Image or Tensor: Affine transformed image. + """ + fill = self.fill + channels, height, width = F.get_dimensions(img) + if isinstance(img, Tensor): + if isinstance(fill, (int, float)): + fill = [float(fill)] * channels + else: + fill = [float(f) for f in fill] + + img_size = [width, height] # flip for keeping BC on get_params call + + ret = self.get_params(self.degrees, self.translate, self.scale, self.shear, img_size) + + return F.affine(img, *ret, interpolation=self.interpolation, fill=fill, center=self.center) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}(degrees={self.degrees}" + s += f", translate={self.translate}" if self.translate is not None else "" + s += f", scale={self.scale}" if self.scale is not None else "" + s += f", shear={self.shear}" if self.shear is not None else "" + s += f", interpolation={self.interpolation.value}" if self.interpolation != InterpolationMode.NEAREST else "" + s += f", fill={self.fill}" if self.fill != 0 else "" + s += f", center={self.center}" if self.center is not None else "" + s += ")" + + return s + + +class Grayscale(torch.nn.Module): + """Convert image to grayscale. + If the image is torch Tensor, it is expected + to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions + + Args: + num_output_channels (int): (1 or 3) number of channels desired for output image + + Returns: + PIL Image: Grayscale version of the input. + + - If ``num_output_channels == 1`` : returned image is single channel + - If ``num_output_channels == 3`` : returned image is 3 channel with r == g == b + + """ + + def __init__(self, num_output_channels=1): + super().__init__() + _log_api_usage_once(self) + self.num_output_channels = num_output_channels + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be converted to grayscale. + + Returns: + PIL Image or Tensor: Grayscaled image. + """ + return F.rgb_to_grayscale(img, num_output_channels=self.num_output_channels) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(num_output_channels={self.num_output_channels})" + + +class RandomGrayscale(torch.nn.Module): + """Randomly convert image to grayscale with a probability of p (default 0.1). + If the image is torch Tensor, it is expected + to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions + + Args: + p (float): probability that image should be converted to grayscale. + + Returns: + PIL Image or Tensor: Grayscale version of the input image with probability p and unchanged + with probability (1-p). + - If input image is 1 channel: grayscale version is 1 channel + - If input image is 3 channel: grayscale version is 3 channel with r == g == b + + """ + + def __init__(self, p=0.1): + super().__init__() + _log_api_usage_once(self) + self.p = p + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be converted to grayscale. + + Returns: + PIL Image or Tensor: Randomly grayscaled image. + """ + num_output_channels, _, _ = F.get_dimensions(img) + if torch.rand(1) < self.p: + return F.rgb_to_grayscale(img, num_output_channels=num_output_channels) + return img + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(p={self.p})" + + +class RandomErasing(torch.nn.Module): + """Randomly selects a rectangle region in a torch.Tensor image and erases its pixels. + This transform does not support PIL Image. + 'Random Erasing Data Augmentation' by Zhong et al. See https://arxiv.org/abs/1708.04896 + + Args: + p: probability that the random erasing operation will be performed. + scale: range of proportion of erased area against input image. + ratio: range of aspect ratio of erased area. + value: erasing value. Default is 0. If a single int, it is used to + erase all pixels. If a tuple of length 3, it is used to erase + R, G, B channels respectively. + If a str of 'random', erasing each pixel with random values. + inplace: boolean to make this transform inplace. Default set to False. + + Returns: + Erased Image. + + Example: + >>> transform = transforms.Compose([ + >>> transforms.RandomHorizontalFlip(), + >>> transforms.PILToTensor(), + >>> transforms.ConvertImageDtype(torch.float), + >>> transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), + >>> transforms.RandomErasing(), + >>> ]) + """ + + def __init__(self, p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=0, inplace=False): + super().__init__() + _log_api_usage_once(self) + if not isinstance(value, (numbers.Number, str, tuple, list)): + raise TypeError("Argument value should be either a number or str or a sequence") + if isinstance(value, str) and value != "random": + raise ValueError("If value is str, it should be 'random'") + if not isinstance(scale, Sequence): + raise TypeError("Scale should be a sequence") + if not isinstance(ratio, Sequence): + raise TypeError("Ratio should be a sequence") + if (scale[0] > scale[1]) or (ratio[0] > ratio[1]): + warnings.warn("Scale and ratio should be of kind (min, max)") + if scale[0] < 0 or scale[1] > 1: + raise ValueError("Scale should be between 0 and 1") + if p < 0 or p > 1: + raise ValueError("Random erasing probability should be between 0 and 1") + + self.p = p + self.scale = scale + self.ratio = ratio + self.value = value + self.inplace = inplace + + @staticmethod + def get_params( + img: Tensor, scale: Tuple[float, float], ratio: Tuple[float, float], value: Optional[List[float]] = None + ) -> Tuple[int, int, int, int, Tensor]: + """Get parameters for ``erase`` for a random erasing. + + Args: + img (Tensor): Tensor image to be erased. + scale (sequence): range of proportion of erased area against input image. + ratio (sequence): range of aspect ratio of erased area. + value (list, optional): erasing value. If None, it is interpreted as "random" + (erasing each pixel with random values). If ``len(value)`` is 1, it is interpreted as a number, + i.e. ``value[0]``. + + Returns: + tuple: params (i, j, h, w, v) to be passed to ``erase`` for random erasing. + """ + img_c, img_h, img_w = img.shape[-3], img.shape[-2], img.shape[-1] + area = img_h * img_w + + log_ratio = torch.log(torch.tensor(ratio)) + for _ in range(10): + erase_area = area * torch.empty(1).uniform_(scale[0], scale[1]).item() + aspect_ratio = torch.exp(torch.empty(1).uniform_(log_ratio[0], log_ratio[1])).item() + + h = int(round(math.sqrt(erase_area * aspect_ratio))) + w = int(round(math.sqrt(erase_area / aspect_ratio))) + if not (h < img_h and w < img_w): + continue + + if value is None: + v = torch.empty([img_c, h, w], dtype=torch.float32).normal_() + else: + v = torch.tensor(value)[:, None, None] + + i = torch.randint(0, img_h - h + 1, size=(1,)).item() + j = torch.randint(0, img_w - w + 1, size=(1,)).item() + return i, j, h, w, v + + # Return original image + return 0, 0, img_h, img_w, img + + def forward(self, img): + """ + Args: + img (Tensor): Tensor image to be erased. + + Returns: + img (Tensor): Erased Tensor image. + """ + if torch.rand(1) < self.p: + + # cast self.value to script acceptable type + if isinstance(self.value, (int, float)): + value = [float(self.value)] + elif isinstance(self.value, str): + value = None + elif isinstance(self.value, (list, tuple)): + value = [float(v) for v in self.value] + else: + value = self.value + + if value is not None and not (len(value) in (1, img.shape[-3])): + raise ValueError( + "If value is a sequence, it should have either a single value or " + f"{img.shape[-3]} (number of input channels)" + ) + + x, y, h, w, v = self.get_params(img, scale=self.scale, ratio=self.ratio, value=value) + return F.erase(img, x, y, h, w, v, self.inplace) + return img + + def __repr__(self) -> str: + s = ( + f"{self.__class__.__name__}" + f"(p={self.p}, " + f"scale={self.scale}, " + f"ratio={self.ratio}, " + f"value={self.value}, " + f"inplace={self.inplace})" + ) + return s + + +class GaussianBlur(torch.nn.Module): + """Blurs image with randomly chosen Gaussian blur. + If the image is torch Tensor, it is expected + to have [..., C, H, W] shape, where ... means at most one leading dimension. + + Args: + kernel_size (int or sequence): Size of the Gaussian kernel. + sigma (float or tuple of float (min, max)): Standard deviation to be used for + creating kernel to perform blurring. If float, sigma is fixed. If it is tuple + of float (min, max), sigma is chosen uniformly at random to lie in the + given range. + + Returns: + PIL Image or Tensor: Gaussian blurred version of the input image. + + """ + + def __init__(self, kernel_size, sigma=(0.1, 2.0)): + super().__init__() + _log_api_usage_once(self) + self.kernel_size = _setup_size(kernel_size, "Kernel size should be a tuple/list of two integers") + for ks in self.kernel_size: + if ks <= 0 or ks % 2 == 0: + raise ValueError("Kernel size value should be an odd and positive number.") + + if isinstance(sigma, numbers.Number): + if sigma <= 0: + raise ValueError("If sigma is a single number, it must be positive.") + sigma = (sigma, sigma) + elif isinstance(sigma, Sequence) and len(sigma) == 2: + if not 0.0 < sigma[0] <= sigma[1]: + raise ValueError("sigma values should be positive and of the form (min, max).") + else: + raise ValueError("sigma should be a single number or a list/tuple with length 2.") + + self.sigma = sigma + + @staticmethod + def get_params(sigma_min: float, sigma_max: float) -> float: + """Choose sigma for random gaussian blurring. + + Args: + sigma_min (float): Minimum standard deviation that can be chosen for blurring kernel. + sigma_max (float): Maximum standard deviation that can be chosen for blurring kernel. + + Returns: + float: Standard deviation to be passed to calculate kernel for gaussian blurring. + """ + return torch.empty(1).uniform_(sigma_min, sigma_max).item() + + def forward(self, img: Tensor) -> Tensor: + """ + Args: + img (PIL Image or Tensor): image to be blurred. + + Returns: + PIL Image or Tensor: Gaussian blurred image + """ + sigma = self.get_params(self.sigma[0], self.sigma[1]) + return F.gaussian_blur(img, self.kernel_size, [sigma, sigma]) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}(kernel_size={self.kernel_size}, sigma={self.sigma})" + return s + + +def _setup_size(size, error_msg): + if isinstance(size, numbers.Number): + return int(size), int(size) + + if isinstance(size, Sequence) and len(size) == 1: + return size[0], size[0] + + if len(size) != 2: + raise ValueError(error_msg) + + return size + + +def _check_sequence_input(x, name, req_sizes): + msg = req_sizes[0] if len(req_sizes) < 2 else " or ".join([str(s) for s in req_sizes]) + if not isinstance(x, Sequence): + raise TypeError(f"{name} should be a sequence of length {msg}.") + if len(x) not in req_sizes: + raise ValueError(f"{name} should be a sequence of length {msg}.") + + +def _setup_angle(x, name, req_sizes=(2,)): + if isinstance(x, numbers.Number): + if x < 0: + raise ValueError(f"If {name} is a single number, it must be positive.") + x = [-x, x] + else: + _check_sequence_input(x, name, req_sizes) + + return [float(d) for d in x] + + +class RandomInvert(torch.nn.Module): + """Inverts the colors of the given image randomly with a given probability. + If img is a Tensor, it is expected to be in [..., 1 or 3, H, W] format, + where ... means it can have an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + + Args: + p (float): probability of the image being color inverted. Default value is 0.5 + """ + + def __init__(self, p=0.5): + super().__init__() + _log_api_usage_once(self) + self.p = p + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be inverted. + + Returns: + PIL Image or Tensor: Randomly color inverted image. + """ + if torch.rand(1).item() < self.p: + return F.invert(img) + return img + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(p={self.p})" + + +class RandomPosterize(torch.nn.Module): + """Posterize the image randomly with a given probability by reducing the + number of bits for each color channel. If the image is torch Tensor, it should be of type torch.uint8, + and it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + + Args: + bits (int): number of bits to keep for each channel (0-8) + p (float): probability of the image being posterized. Default value is 0.5 + """ + + def __init__(self, bits, p=0.5): + super().__init__() + _log_api_usage_once(self) + self.bits = bits + self.p = p + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be posterized. + + Returns: + PIL Image or Tensor: Randomly posterized image. + """ + if torch.rand(1).item() < self.p: + return F.posterize(img, self.bits) + return img + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(bits={self.bits},p={self.p})" + + +class RandomSolarize(torch.nn.Module): + """Solarize the image randomly with a given probability by inverting all pixel + values above a threshold. If img is a Tensor, it is expected to be in [..., 1 or 3, H, W] format, + where ... means it can have an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + + Args: + threshold (float): all pixels equal or above this value are inverted. + p (float): probability of the image being solarized. Default value is 0.5 + """ + + def __init__(self, threshold, p=0.5): + super().__init__() + _log_api_usage_once(self) + self.threshold = threshold + self.p = p + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be solarized. + + Returns: + PIL Image or Tensor: Randomly solarized image. + """ + if torch.rand(1).item() < self.p: + return F.solarize(img, self.threshold) + return img + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(threshold={self.threshold},p={self.p})" + + +class RandomAdjustSharpness(torch.nn.Module): + """Adjust the sharpness of the image randomly with a given probability. If the image is torch Tensor, + it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. + + Args: + sharpness_factor (float): How much to adjust the sharpness. Can be + any non-negative number. 0 gives a blurred image, 1 gives the + original image while 2 increases the sharpness by a factor of 2. + p (float): probability of the image being sharpened. Default value is 0.5 + """ + + def __init__(self, sharpness_factor, p=0.5): + super().__init__() + _log_api_usage_once(self) + self.sharpness_factor = sharpness_factor + self.p = p + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be sharpened. + + Returns: + PIL Image or Tensor: Randomly sharpened image. + """ + if torch.rand(1).item() < self.p: + return F.adjust_sharpness(img, self.sharpness_factor) + return img + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(sharpness_factor={self.sharpness_factor},p={self.p})" + + +class RandomAutocontrast(torch.nn.Module): + """Autocontrast the pixels of the given image randomly with a given probability. + If the image is torch Tensor, it is expected + to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + + Args: + p (float): probability of the image being autocontrasted. Default value is 0.5 + """ + + def __init__(self, p=0.5): + super().__init__() + _log_api_usage_once(self) + self.p = p + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be autocontrasted. + + Returns: + PIL Image or Tensor: Randomly autocontrasted image. + """ + if torch.rand(1).item() < self.p: + return F.autocontrast(img) + return img + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(p={self.p})" + + +class RandomEqualize(torch.nn.Module): + """Equalize the histogram of the given image randomly with a given probability. + If the image is torch Tensor, it is expected + to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "P", "L" or "RGB". + + Args: + p (float): probability of the image being equalized. Default value is 0.5 + """ + + def __init__(self, p=0.5): + super().__init__() + _log_api_usage_once(self) + self.p = p + + def forward(self, img): + """ + Args: + img (PIL Image or Tensor): Image to be equalized. + + Returns: + PIL Image or Tensor: Randomly equalized image. + """ + if torch.rand(1).item() < self.p: + return F.equalize(img) + return img + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(p={self.p})" + + +class ElasticTransform(torch.nn.Module): + """Transform a tensor image with elastic transformations. + Given alpha and sigma, it will generate displacement + vectors for all pixels based on random offsets. Alpha controls the strength + and sigma controls the smoothness of the displacements. + The displacements are added to an identity grid and the resulting grid is + used to grid_sample from the image. + + Applications: + Randomly transforms the morphology of objects in images and produces a + see-through-water-like effect. + + Args: + alpha (float or sequence of floats): Magnitude of displacements. Default is 50.0. + sigma (float or sequence of floats): Smoothness of displacements. Default is 5.0. + interpolation (InterpolationMode): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + fill (sequence or number): Pixel fill value for the area outside the transformed + image. Default is ``0``. If given a number, the value is used for all bands respectively. + + """ + + def __init__(self, alpha=50.0, sigma=5.0, interpolation=InterpolationMode.BILINEAR, fill=0): + super().__init__() + _log_api_usage_once(self) + if not isinstance(alpha, (float, Sequence)): + raise TypeError(f"alpha should be float or a sequence of floats. Got {type(alpha)}") + if isinstance(alpha, Sequence) and len(alpha) != 2: + raise ValueError(f"If alpha is a sequence its length should be 2. Got {len(alpha)}") + if isinstance(alpha, Sequence): + for element in alpha: + if not isinstance(element, float): + raise TypeError(f"alpha should be a sequence of floats. Got {type(element)}") + + if isinstance(alpha, float): + alpha = [float(alpha), float(alpha)] + if isinstance(alpha, (list, tuple)) and len(alpha) == 1: + alpha = [alpha[0], alpha[0]] + + self.alpha = alpha + + if not isinstance(sigma, (float, Sequence)): + raise TypeError(f"sigma should be float or a sequence of floats. Got {type(sigma)}") + if isinstance(sigma, Sequence) and len(sigma) != 2: + raise ValueError(f"If sigma is a sequence its length should be 2. Got {len(sigma)}") + if isinstance(sigma, Sequence): + for element in sigma: + if not isinstance(element, float): + raise TypeError(f"sigma should be a sequence of floats. Got {type(element)}") + + if isinstance(sigma, float): + sigma = [float(sigma), float(sigma)] + if isinstance(sigma, (list, tuple)) and len(sigma) == 1: + sigma = [sigma[0], sigma[0]] + + self.sigma = sigma + + if isinstance(interpolation, int): + interpolation = _interpolation_modes_from_int(interpolation) + self.interpolation = interpolation + + if isinstance(fill, (int, float)): + fill = [float(fill)] + elif isinstance(fill, (list, tuple)): + fill = [float(f) for f in fill] + else: + raise TypeError(f"fill should be int or float or a list or tuple of them. Got {type(fill)}") + self.fill = fill + + @staticmethod + def get_params(alpha: List[float], sigma: List[float], size: List[int]) -> Tensor: + dx = torch.rand([1, 1] + size) * 2 - 1 + if sigma[0] > 0.0: + kx = int(8 * sigma[0] + 1) + # if kernel size is even we have to make it odd + if kx % 2 == 0: + kx += 1 + dx = F.gaussian_blur(dx, [kx, kx], sigma) + dx = dx * alpha[0] / size[0] + + dy = torch.rand([1, 1] + size) * 2 - 1 + if sigma[1] > 0.0: + ky = int(8 * sigma[1] + 1) + # if kernel size is even we have to make it odd + if ky % 2 == 0: + ky += 1 + dy = F.gaussian_blur(dy, [ky, ky], sigma) + dy = dy * alpha[1] / size[1] + return torch.concat([dx, dy], 1).permute([0, 2, 3, 1]) # 1 x H x W x 2 + + def forward(self, tensor: Tensor) -> Tensor: + """ + Args: + tensor (PIL Image or Tensor): Image to be transformed. + + Returns: + PIL Image or Tensor: Transformed image. + """ + _, height, width = F.get_dimensions(tensor) + displacement = self.get_params(self.alpha, self.sigma, [height, width]) + return F.elastic_transform(tensor, displacement, self.interpolation, self.fill) + + def __repr__(self): + format_string = self.__class__.__name__ + format_string += f"(alpha={self.alpha}" + format_string += f", sigma={self.sigma}" + format_string += f", interpolation={self.interpolation}" + format_string += f", fill={self.fill})" + return format_string diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/__init__.py b/lib/python3.12/site-packages/torchvision/transforms/v2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2d66917b6eaceb6281b7d809cbb22462a511a91d --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/__init__.py @@ -0,0 +1,60 @@ +from torchvision.transforms import AutoAugmentPolicy, InterpolationMode # usort: skip + +from . import functional # usort: skip + +from ._transform import Transform # usort: skip + +from ._augment import CutMix, JPEG, MixUp, RandomErasing +from ._auto_augment import AugMix, AutoAugment, RandAugment, TrivialAugmentWide +from ._color import ( + ColorJitter, + Grayscale, + RandomAdjustSharpness, + RandomAutocontrast, + RandomChannelPermutation, + RandomEqualize, + RandomGrayscale, + RandomInvert, + RandomPhotometricDistort, + RandomPosterize, + RandomSolarize, + RGB, +) +from ._container import Compose, RandomApply, RandomChoice, RandomOrder +from ._geometry import ( + CenterCrop, + ElasticTransform, + FiveCrop, + Pad, + RandomAffine, + RandomCrop, + RandomHorizontalFlip, + RandomIoUCrop, + RandomPerspective, + RandomResize, + RandomResizedCrop, + RandomRotation, + RandomShortestSize, + RandomVerticalFlip, + RandomZoomOut, + Resize, + ScaleJitter, + TenCrop, +) +from ._meta import ClampBoundingBoxes, ConvertBoundingBoxFormat +from ._misc import ( + ConvertImageDtype, + GaussianBlur, + GaussianNoise, + Identity, + Lambda, + LinearTransformation, + Normalize, + SanitizeBoundingBoxes, + ToDtype, +) +from ._temporal import UniformTemporalSubsample +from ._type_conversion import PILToTensor, ToImage, ToPILImage, ToPureTensor +from ._utils import check_type, get_bounding_boxes, has_all, has_any, query_chw, query_size + +from ._deprecated import ToTensor # usort: skip diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c57bd768412b429962798bb61f8b0fc18ab54bc0 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_augment.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_augment.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..31354781cafb76887ddd584e0a342658e9e4f435 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_augment.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_auto_augment.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_auto_augment.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad465bbafbab412f5bbadfff3ce9e627971ee40b Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_auto_augment.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_color.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_color.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..235accad5fe97e1b708c32260b613a6275cbf2f9 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_color.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_container.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_container.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8fb874fbc7e088df87bca3147d924dbee3aa5034 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_container.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_deprecated.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_deprecated.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd5021b5e1929d9362153216a937d0330c76110e Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_deprecated.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_geometry.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_geometry.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4f37a8f0e56997b722a2880b6ad3075b4315654e Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_geometry.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_meta.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_meta.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6c079889038f0669ac9e66407799f2a09f5cee98 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_meta.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_misc.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_misc.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..66ea2783cebd7e93ec91472be912233584b086dc Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_misc.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_temporal.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_temporal.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c2b62b73fdb612be06bc01468d6ae0706ed5202e Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_temporal.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_transform.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_transform.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1eca346de27dea91c3e86b6e6b3797996d767d5a Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_transform.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_type_conversion.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_type_conversion.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..91fd8fced72271f7b16262cfa0d26fd1e8d1e575 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_type_conversion.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..29d3a1de12ce6f0f1ad42290c58b59c5ae0146e1 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/__pycache__/_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/_augment.py b/lib/python3.12/site-packages/torchvision/transforms/v2/_augment.py new file mode 100644 index 0000000000000000000000000000000000000000..93d4ba45d6593c02bc010f4a57980ece15476183 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/_augment.py @@ -0,0 +1,369 @@ +import math +import numbers +import warnings +from typing import Any, Callable, Dict, List, Optional, Sequence, Union + +import PIL.Image +import torch +from torch.nn.functional import one_hot +from torch.utils._pytree import tree_flatten, tree_unflatten +from torchvision import transforms as _transforms, tv_tensors +from torchvision.transforms.v2 import functional as F + +from ._transform import _RandomApplyTransform, Transform +from ._utils import _check_sequence_input, _parse_labels_getter, has_any, is_pure_tensor, query_chw, query_size + + +class RandomErasing(_RandomApplyTransform): + """Randomly select a rectangle region in the input image or video and erase its pixels. + + This transform does not support PIL Image. + 'Random Erasing Data Augmentation' by Zhong et al. See https://arxiv.org/abs/1708.04896 + + Args: + p (float, optional): probability that the random erasing operation will be performed. + scale (tuple of float, optional): range of proportion of erased area against input image. + ratio (tuple of float, optional): range of aspect ratio of erased area. + value (number or tuple of numbers): erasing value. Default is 0. If a single int, it is used to + erase all pixels. If a tuple of length 3, it is used to erase + R, G, B channels respectively. + If a str of 'random', erasing each pixel with random values. + inplace (bool, optional): boolean to make this transform inplace. Default set to False. + + Returns: + Erased input. + + Example: + >>> from torchvision.transforms import v2 as transforms + >>> + >>> transform = transforms.Compose([ + >>> transforms.RandomHorizontalFlip(), + >>> transforms.PILToTensor(), + >>> transforms.ConvertImageDtype(torch.float), + >>> transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), + >>> transforms.RandomErasing(), + >>> ]) + """ + + _v1_transform_cls = _transforms.RandomErasing + + def _extract_params_for_v1_transform(self) -> Dict[str, Any]: + return dict( + super()._extract_params_for_v1_transform(), + value="random" if self.value is None else self.value, + ) + + def __init__( + self, + p: float = 0.5, + scale: Sequence[float] = (0.02, 0.33), + ratio: Sequence[float] = (0.3, 3.3), + value: float = 0.0, + inplace: bool = False, + ): + super().__init__(p=p) + if not isinstance(value, (numbers.Number, str, tuple, list)): + raise TypeError("Argument value should be either a number or str or a sequence") + if isinstance(value, str) and value != "random": + raise ValueError("If value is str, it should be 'random'") + if not isinstance(scale, Sequence): + raise TypeError("Scale should be a sequence") + if not isinstance(ratio, Sequence): + raise TypeError("Ratio should be a sequence") + if (scale[0] > scale[1]) or (ratio[0] > ratio[1]): + warnings.warn("Scale and ratio should be of kind (min, max)") + if scale[0] < 0 or scale[1] > 1: + raise ValueError("Scale should be between 0 and 1") + self.scale = scale + self.ratio = ratio + if isinstance(value, (int, float)): + self.value = [float(value)] + elif isinstance(value, str): + self.value = None + elif isinstance(value, (list, tuple)): + self.value = [float(v) for v in value] + else: + self.value = value + self.inplace = inplace + + self._log_ratio = torch.log(torch.tensor(self.ratio)) + + def _call_kernel(self, functional: Callable, inpt: Any, *args: Any, **kwargs: Any) -> Any: + if isinstance(inpt, (tv_tensors.BoundingBoxes, tv_tensors.Mask)): + warnings.warn( + f"{type(self).__name__}() is currently passing through inputs of type " + f"tv_tensors.{type(inpt).__name__}. This will likely change in the future." + ) + return super()._call_kernel(functional, inpt, *args, **kwargs) + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + img_c, img_h, img_w = query_chw(flat_inputs) + + if self.value is not None and not (len(self.value) in (1, img_c)): + raise ValueError( + f"If value is a sequence, it should have either a single value or {img_c} (number of inpt channels)" + ) + + area = img_h * img_w + + log_ratio = self._log_ratio + for _ in range(10): + erase_area = area * torch.empty(1).uniform_(self.scale[0], self.scale[1]).item() + aspect_ratio = torch.exp( + torch.empty(1).uniform_( + log_ratio[0], # type: ignore[arg-type] + log_ratio[1], # type: ignore[arg-type] + ) + ).item() + + h = int(round(math.sqrt(erase_area * aspect_ratio))) + w = int(round(math.sqrt(erase_area / aspect_ratio))) + if not (h < img_h and w < img_w): + continue + + if self.value is None: + v = torch.empty([img_c, h, w], dtype=torch.float32).normal_() + else: + v = torch.tensor(self.value)[:, None, None] + + i = torch.randint(0, img_h - h + 1, size=(1,)).item() + j = torch.randint(0, img_w - w + 1, size=(1,)).item() + break + else: + i, j, h, w, v = 0, 0, img_h, img_w, None + + return dict(i=i, j=j, h=h, w=w, v=v) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + if params["v"] is not None: + inpt = self._call_kernel(F.erase, inpt, **params, inplace=self.inplace) + + return inpt + + +class _BaseMixUpCutMix(Transform): + def __init__(self, *, alpha: float = 1.0, num_classes: Optional[int] = None, labels_getter="default") -> None: + super().__init__() + self.alpha = float(alpha) + self._dist = torch.distributions.Beta(torch.tensor([alpha]), torch.tensor([alpha])) + + self.num_classes = num_classes + + self._labels_getter = _parse_labels_getter(labels_getter) + + def forward(self, *inputs): + inputs = inputs if len(inputs) > 1 else inputs[0] + flat_inputs, spec = tree_flatten(inputs) + needs_transform_list = self._needs_transform_list(flat_inputs) + + if has_any(flat_inputs, PIL.Image.Image, tv_tensors.BoundingBoxes, tv_tensors.Mask): + raise ValueError(f"{type(self).__name__}() does not support PIL images, bounding boxes and masks.") + + labels = self._labels_getter(inputs) + if not isinstance(labels, torch.Tensor): + raise ValueError(f"The labels must be a tensor, but got {type(labels)} instead.") + if labels.ndim not in (1, 2): + raise ValueError( + f"labels should be index based with shape (batch_size,) " + f"or probability based with shape (batch_size, num_classes), " + f"but got a tensor of shape {labels.shape} instead." + ) + if labels.ndim == 2 and self.num_classes is not None and labels.shape[-1] != self.num_classes: + raise ValueError( + f"When passing 2D labels, " + f"the number of elements in last dimension must match num_classes: " + f"{labels.shape[-1]} != {self.num_classes}. " + f"You can Leave num_classes to None." + ) + if labels.ndim == 1 and self.num_classes is None: + raise ValueError("num_classes must be passed if the labels are index-based (1D)") + + params = { + "labels": labels, + "batch_size": labels.shape[0], + **self.make_params( + [inpt for (inpt, needs_transform) in zip(flat_inputs, needs_transform_list) if needs_transform] + ), + } + + # By default, the labels will be False inside needs_transform_list, since they are a torch.Tensor coming + # after an image or video. However, we need to handle them in _transform, so we make sure to set them to True + needs_transform_list[next(idx for idx, inpt in enumerate(flat_inputs) if inpt is labels)] = True + flat_outputs = [ + self.transform(inpt, params) if needs_transform else inpt + for (inpt, needs_transform) in zip(flat_inputs, needs_transform_list) + ] + + return tree_unflatten(flat_outputs, spec) + + def _check_image_or_video(self, inpt: torch.Tensor, *, batch_size: int): + expected_num_dims = 5 if isinstance(inpt, tv_tensors.Video) else 4 + if inpt.ndim != expected_num_dims: + raise ValueError( + f"Expected a batched input with {expected_num_dims} dims, but got {inpt.ndim} dimensions instead." + ) + if inpt.shape[0] != batch_size: + raise ValueError( + f"The batch size of the image or video does not match the batch size of the labels: " + f"{inpt.shape[0]} != {batch_size}." + ) + + def _mixup_label(self, label: torch.Tensor, *, lam: float) -> torch.Tensor: + if label.ndim == 1: + label = one_hot(label, num_classes=self.num_classes) # type: ignore[arg-type] + if not label.dtype.is_floating_point: + label = label.float() + return label.roll(1, 0).mul_(1.0 - lam).add_(label.mul(lam)) + + +class MixUp(_BaseMixUpCutMix): + """Apply MixUp to the provided batch of images and labels. + + Paper: `mixup: Beyond Empirical Risk Minimization `_. + + .. note:: + This transform is meant to be used on **batches** of samples, not + individual images. See + :ref:`sphx_glr_auto_examples_transforms_plot_cutmix_mixup.py` for detailed usage + examples. + The sample pairing is deterministic and done by matching consecutive + samples in the batch, so the batch needs to be shuffled (this is an + implementation detail, not a guaranteed convention.) + + In the input, the labels are expected to be a tensor of shape ``(batch_size,)``. They will be transformed + into a tensor of shape ``(batch_size, num_classes)``. + + Args: + alpha (float, optional): hyperparameter of the Beta distribution used for mixup. Default is 1. + num_classes (int, optional): number of classes in the batch. Used for one-hot-encoding. + Can be None only if the labels are already one-hot-encoded. + labels_getter (callable or "default", optional): indicates how to identify the labels in the input. + By default, this will pick the second parameter as the labels if it's a tensor. This covers the most + common scenario where this transform is called as ``MixUp()(imgs_batch, labels_batch)``. + It can also be a callable that takes the same input as the transform, and returns the labels. + """ + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + return dict(lam=float(self._dist.sample(()))) # type: ignore[arg-type] + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + lam = params["lam"] + + if inpt is params["labels"]: + return self._mixup_label(inpt, lam=lam) + elif isinstance(inpt, (tv_tensors.Image, tv_tensors.Video)) or is_pure_tensor(inpt): + self._check_image_or_video(inpt, batch_size=params["batch_size"]) + + output = inpt.roll(1, 0).mul_(1.0 - lam).add_(inpt.mul(lam)) + + if isinstance(inpt, (tv_tensors.Image, tv_tensors.Video)): + output = tv_tensors.wrap(output, like=inpt) + + return output + else: + return inpt + + +class CutMix(_BaseMixUpCutMix): + """Apply CutMix to the provided batch of images and labels. + + Paper: `CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features + `_. + + .. note:: + This transform is meant to be used on **batches** of samples, not + individual images. See + :ref:`sphx_glr_auto_examples_transforms_plot_cutmix_mixup.py` for detailed usage + examples. + The sample pairing is deterministic and done by matching consecutive + samples in the batch, so the batch needs to be shuffled (this is an + implementation detail, not a guaranteed convention.) + + In the input, the labels are expected to be a tensor of shape ``(batch_size,)``. They will be transformed + into a tensor of shape ``(batch_size, num_classes)``. + + Args: + alpha (float, optional): hyperparameter of the Beta distribution used for mixup. Default is 1. + num_classes (int, optional): number of classes in the batch. Used for one-hot-encoding. + Can be None only if the labels are already one-hot-encoded. + labels_getter (callable or "default", optional): indicates how to identify the labels in the input. + By default, this will pick the second parameter as the labels if it's a tensor. This covers the most + common scenario where this transform is called as ``CutMix()(imgs_batch, labels_batch)``. + It can also be a callable that takes the same input as the transform, and returns the labels. + """ + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + lam = float(self._dist.sample(())) # type: ignore[arg-type] + + H, W = query_size(flat_inputs) + + r_x = torch.randint(W, size=(1,)) + r_y = torch.randint(H, size=(1,)) + + r = 0.5 * math.sqrt(1.0 - lam) + r_w_half = int(r * W) + r_h_half = int(r * H) + + x1 = int(torch.clamp(r_x - r_w_half, min=0)) + y1 = int(torch.clamp(r_y - r_h_half, min=0)) + x2 = int(torch.clamp(r_x + r_w_half, max=W)) + y2 = int(torch.clamp(r_y + r_h_half, max=H)) + box = (x1, y1, x2, y2) + + lam_adjusted = float(1.0 - (x2 - x1) * (y2 - y1) / (W * H)) + + return dict(box=box, lam_adjusted=lam_adjusted) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + if inpt is params["labels"]: + return self._mixup_label(inpt, lam=params["lam_adjusted"]) + elif isinstance(inpt, (tv_tensors.Image, tv_tensors.Video)) or is_pure_tensor(inpt): + self._check_image_or_video(inpt, batch_size=params["batch_size"]) + + x1, y1, x2, y2 = params["box"] + rolled = inpt.roll(1, 0) + output = inpt.clone() + output[..., y1:y2, x1:x2] = rolled[..., y1:y2, x1:x2] + + if isinstance(inpt, (tv_tensors.Image, tv_tensors.Video)): + output = tv_tensors.wrap(output, like=inpt) + + return output + else: + return inpt + + +class JPEG(Transform): + """Apply JPEG compression and decompression to the given images. + + If the input is a :class:`torch.Tensor`, it is expected + to be of dtype uint8, on CPU, and have [..., 3 or 1, H, W] shape, + where ... means an arbitrary number of leading dimensions. + + Args: + quality (sequence or number): JPEG quality, from 1 to 100. Lower means more compression. + If quality is a sequence like (min, max), it specifies the range of JPEG quality to + randomly select from (inclusive of both ends). + + Returns: + image with JPEG compression. + """ + + def __init__(self, quality: Union[int, Sequence[int]]): + super().__init__() + if isinstance(quality, int): + quality = [quality, quality] + else: + _check_sequence_input(quality, "quality", req_sizes=(2,)) + + if not (1 <= quality[0] <= quality[1] <= 100 and isinstance(quality[0], int) and isinstance(quality[1], int)): + raise ValueError(f"quality must be an integer from 1 to 100, got {quality =}") + + self.quality = quality + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + quality = torch.randint(self.quality[0], self.quality[1] + 1, ()).item() + return dict(quality=quality) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.jpeg, inpt, quality=params["quality"]) diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/_auto_augment.py b/lib/python3.12/site-packages/torchvision/transforms/v2/_auto_augment.py new file mode 100644 index 0000000000000000000000000000000000000000..4dd7ba343aa360c38192d2c8ac88481093ed8c93 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/_auto_augment.py @@ -0,0 +1,627 @@ +import math +from typing import Any, Callable, cast, Dict, List, Optional, Tuple, Type, Union + +import PIL.Image +import torch + +from torch.utils._pytree import tree_flatten, tree_unflatten, TreeSpec +from torchvision import transforms as _transforms, tv_tensors +from torchvision.transforms import _functional_tensor as _FT +from torchvision.transforms.v2 import AutoAugmentPolicy, functional as F, InterpolationMode, Transform +from torchvision.transforms.v2.functional._geometry import _check_interpolation +from torchvision.transforms.v2.functional._meta import get_size +from torchvision.transforms.v2.functional._utils import _FillType, _FillTypeJIT + +from ._utils import _get_fill, _setup_fill_arg, check_type, is_pure_tensor + + +ImageOrVideo = Union[torch.Tensor, PIL.Image.Image, tv_tensors.Image, tv_tensors.Video] + + +class _AutoAugmentBase(Transform): + def __init__( + self, + *, + interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST, + fill: Union[_FillType, Dict[Union[Type, str], _FillType]] = None, + ) -> None: + super().__init__() + self.interpolation = _check_interpolation(interpolation) + self.fill = fill + self._fill = _setup_fill_arg(fill) + + def _extract_params_for_v1_transform(self) -> Dict[str, Any]: + params = super()._extract_params_for_v1_transform() + + if isinstance(params["fill"], dict): + raise ValueError(f"{type(self).__name__}() can not be scripted for when `fill` is a dictionary.") + + return params + + def _get_random_item(self, dct: Dict[str, Tuple[Callable, bool]]) -> Tuple[str, Tuple[Callable, bool]]: + keys = tuple(dct.keys()) + key = keys[int(torch.randint(len(keys), ()))] + return key, dct[key] + + def _flatten_and_extract_image_or_video( + self, + inputs: Any, + unsupported_types: Tuple[Type, ...] = (tv_tensors.BoundingBoxes, tv_tensors.Mask), + ) -> Tuple[Tuple[List[Any], TreeSpec, int], ImageOrVideo]: + flat_inputs, spec = tree_flatten(inputs if len(inputs) > 1 else inputs[0]) + needs_transform_list = self._needs_transform_list(flat_inputs) + + image_or_videos = [] + for idx, (inpt, needs_transform) in enumerate(zip(flat_inputs, needs_transform_list)): + if needs_transform and check_type( + inpt, + ( + tv_tensors.Image, + PIL.Image.Image, + is_pure_tensor, + tv_tensors.Video, + ), + ): + image_or_videos.append((idx, inpt)) + elif isinstance(inpt, unsupported_types): + raise TypeError(f"Inputs of type {type(inpt).__name__} are not supported by {type(self).__name__}()") + + if not image_or_videos: + raise TypeError("Found no image in the sample.") + if len(image_or_videos) > 1: + raise TypeError( + f"Auto augment transformations are only properly defined for a single image or video, " + f"but found {len(image_or_videos)}." + ) + + idx, image_or_video = image_or_videos[0] + return (flat_inputs, spec, idx), image_or_video + + def _unflatten_and_insert_image_or_video( + self, + flat_inputs_with_spec: Tuple[List[Any], TreeSpec, int], + image_or_video: ImageOrVideo, + ) -> Any: + flat_inputs, spec, idx = flat_inputs_with_spec + flat_inputs[idx] = image_or_video + return tree_unflatten(flat_inputs, spec) + + def _apply_image_or_video_transform( + self, + image: ImageOrVideo, + transform_id: str, + magnitude: float, + interpolation: Union[InterpolationMode, int], + fill: Dict[Union[Type, str], _FillTypeJIT], + ) -> ImageOrVideo: + # Note: this cast is wrong and is only here to make mypy happy (it disagrees with torchscript) + image = cast(torch.Tensor, image) + fill_ = _get_fill(fill, type(image)) + + if transform_id == "Identity": + return image + elif transform_id == "ShearX": + # magnitude should be arctan(magnitude) + # official autoaug: (1, level, 0, 0, 1, 0) + # https://github.com/tensorflow/models/blob/dd02069717128186b88afa8d857ce57d17957f03/research/autoaugment/augmentation_transforms.py#L290 + # compared to + # torchvision: (1, tan(level), 0, 0, 1, 0) + # https://github.com/pytorch/vision/blob/0c2373d0bba3499e95776e7936e207d8a1676e65/torchvision/transforms/functional.py#L976 + return F.affine( + image, + angle=0.0, + translate=[0, 0], + scale=1.0, + shear=[math.degrees(math.atan(magnitude)), 0.0], + interpolation=interpolation, + fill=fill_, + center=[0, 0], + ) + elif transform_id == "ShearY": + # magnitude should be arctan(magnitude) + # See above + return F.affine( + image, + angle=0.0, + translate=[0, 0], + scale=1.0, + shear=[0.0, math.degrees(math.atan(magnitude))], + interpolation=interpolation, + fill=fill_, + center=[0, 0], + ) + elif transform_id == "TranslateX": + return F.affine( + image, + angle=0.0, + translate=[int(magnitude), 0], + scale=1.0, + interpolation=interpolation, + shear=[0.0, 0.0], + fill=fill_, + ) + elif transform_id == "TranslateY": + return F.affine( + image, + angle=0.0, + translate=[0, int(magnitude)], + scale=1.0, + interpolation=interpolation, + shear=[0.0, 0.0], + fill=fill_, + ) + elif transform_id == "Rotate": + return F.rotate(image, angle=magnitude, interpolation=interpolation, fill=fill_) + elif transform_id == "Brightness": + return F.adjust_brightness(image, brightness_factor=1.0 + magnitude) + elif transform_id == "Color": + return F.adjust_saturation(image, saturation_factor=1.0 + magnitude) + elif transform_id == "Contrast": + return F.adjust_contrast(image, contrast_factor=1.0 + magnitude) + elif transform_id == "Sharpness": + return F.adjust_sharpness(image, sharpness_factor=1.0 + magnitude) + elif transform_id == "Posterize": + return F.posterize(image, bits=int(magnitude)) + elif transform_id == "Solarize": + bound = _FT._max_value(image.dtype) if isinstance(image, torch.Tensor) else 255.0 + return F.solarize(image, threshold=bound * magnitude) + elif transform_id == "AutoContrast": + return F.autocontrast(image) + elif transform_id == "Equalize": + return F.equalize(image) + elif transform_id == "Invert": + return F.invert(image) + else: + raise ValueError(f"No transform available for {transform_id}") + + +class AutoAugment(_AutoAugmentBase): + r"""AutoAugment data augmentation method based on + `"AutoAugment: Learning Augmentation Strategies from Data" `_. + + This transformation works on images and videos only. + + If the input is :class:`torch.Tensor`, it should be of type ``torch.uint8``, and it is expected + to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + + Args: + policy (AutoAugmentPolicy, optional): Desired policy enum defined by + :class:`torchvision.transforms.autoaugment.AutoAugmentPolicy`. Default is ``AutoAugmentPolicy.IMAGENET``. + interpolation (InterpolationMode, optional): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. + fill (sequence or number, optional): Pixel fill value for the area outside the transformed + image. If given a number, the value is used for all bands respectively. + """ + _v1_transform_cls = _transforms.AutoAugment + + _AUGMENTATION_SPACE = { + "ShearX": (lambda num_bins, height, width: torch.linspace(0.0, 0.3, num_bins), True), + "ShearY": (lambda num_bins, height, width: torch.linspace(0.0, 0.3, num_bins), True), + "TranslateX": ( + lambda num_bins, height, width: torch.linspace(0.0, 150.0 / 331.0 * width, num_bins), + True, + ), + "TranslateY": ( + lambda num_bins, height, width: torch.linspace(0.0, 150.0 / 331.0 * height, num_bins), + True, + ), + "Rotate": (lambda num_bins, height, width: torch.linspace(0.0, 30.0, num_bins), True), + "Brightness": (lambda num_bins, height, width: torch.linspace(0.0, 0.9, num_bins), True), + "Color": (lambda num_bins, height, width: torch.linspace(0.0, 0.9, num_bins), True), + "Contrast": (lambda num_bins, height, width: torch.linspace(0.0, 0.9, num_bins), True), + "Sharpness": (lambda num_bins, height, width: torch.linspace(0.0, 0.9, num_bins), True), + "Posterize": ( + lambda num_bins, height, width: (8 - (torch.arange(num_bins) / ((num_bins - 1) / 4))).round().int(), + False, + ), + "Solarize": (lambda num_bins, height, width: torch.linspace(1.0, 0.0, num_bins), False), + "AutoContrast": (lambda num_bins, height, width: None, False), + "Equalize": (lambda num_bins, height, width: None, False), + "Invert": (lambda num_bins, height, width: None, False), + } + + def __init__( + self, + policy: AutoAugmentPolicy = AutoAugmentPolicy.IMAGENET, + interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST, + fill: Union[_FillType, Dict[Union[Type, str], _FillType]] = None, + ) -> None: + super().__init__(interpolation=interpolation, fill=fill) + self.policy = policy + self._policies = self._get_policies(policy) + + def _get_policies( + self, policy: AutoAugmentPolicy + ) -> List[Tuple[Tuple[str, float, Optional[int]], Tuple[str, float, Optional[int]]]]: + if policy == AutoAugmentPolicy.IMAGENET: + return [ + (("Posterize", 0.4, 8), ("Rotate", 0.6, 9)), + (("Solarize", 0.6, 5), ("AutoContrast", 0.6, None)), + (("Equalize", 0.8, None), ("Equalize", 0.6, None)), + (("Posterize", 0.6, 7), ("Posterize", 0.6, 6)), + (("Equalize", 0.4, None), ("Solarize", 0.2, 4)), + (("Equalize", 0.4, None), ("Rotate", 0.8, 8)), + (("Solarize", 0.6, 3), ("Equalize", 0.6, None)), + (("Posterize", 0.8, 5), ("Equalize", 1.0, None)), + (("Rotate", 0.2, 3), ("Solarize", 0.6, 8)), + (("Equalize", 0.6, None), ("Posterize", 0.4, 6)), + (("Rotate", 0.8, 8), ("Color", 0.4, 0)), + (("Rotate", 0.4, 9), ("Equalize", 0.6, None)), + (("Equalize", 0.0, None), ("Equalize", 0.8, None)), + (("Invert", 0.6, None), ("Equalize", 1.0, None)), + (("Color", 0.6, 4), ("Contrast", 1.0, 8)), + (("Rotate", 0.8, 8), ("Color", 1.0, 2)), + (("Color", 0.8, 8), ("Solarize", 0.8, 7)), + (("Sharpness", 0.4, 7), ("Invert", 0.6, None)), + (("ShearX", 0.6, 5), ("Equalize", 1.0, None)), + (("Color", 0.4, 0), ("Equalize", 0.6, None)), + (("Equalize", 0.4, None), ("Solarize", 0.2, 4)), + (("Solarize", 0.6, 5), ("AutoContrast", 0.6, None)), + (("Invert", 0.6, None), ("Equalize", 1.0, None)), + (("Color", 0.6, 4), ("Contrast", 1.0, 8)), + (("Equalize", 0.8, None), ("Equalize", 0.6, None)), + ] + elif policy == AutoAugmentPolicy.CIFAR10: + return [ + (("Invert", 0.1, None), ("Contrast", 0.2, 6)), + (("Rotate", 0.7, 2), ("TranslateX", 0.3, 9)), + (("Sharpness", 0.8, 1), ("Sharpness", 0.9, 3)), + (("ShearY", 0.5, 8), ("TranslateY", 0.7, 9)), + (("AutoContrast", 0.5, None), ("Equalize", 0.9, None)), + (("ShearY", 0.2, 7), ("Posterize", 0.3, 7)), + (("Color", 0.4, 3), ("Brightness", 0.6, 7)), + (("Sharpness", 0.3, 9), ("Brightness", 0.7, 9)), + (("Equalize", 0.6, None), ("Equalize", 0.5, None)), + (("Contrast", 0.6, 7), ("Sharpness", 0.6, 5)), + (("Color", 0.7, 7), ("TranslateX", 0.5, 8)), + (("Equalize", 0.3, None), ("AutoContrast", 0.4, None)), + (("TranslateY", 0.4, 3), ("Sharpness", 0.2, 6)), + (("Brightness", 0.9, 6), ("Color", 0.2, 8)), + (("Solarize", 0.5, 2), ("Invert", 0.0, None)), + (("Equalize", 0.2, None), ("AutoContrast", 0.6, None)), + (("Equalize", 0.2, None), ("Equalize", 0.6, None)), + (("Color", 0.9, 9), ("Equalize", 0.6, None)), + (("AutoContrast", 0.8, None), ("Solarize", 0.2, 8)), + (("Brightness", 0.1, 3), ("Color", 0.7, 0)), + (("Solarize", 0.4, 5), ("AutoContrast", 0.9, None)), + (("TranslateY", 0.9, 9), ("TranslateY", 0.7, 9)), + (("AutoContrast", 0.9, None), ("Solarize", 0.8, 3)), + (("Equalize", 0.8, None), ("Invert", 0.1, None)), + (("TranslateY", 0.7, 9), ("AutoContrast", 0.9, None)), + ] + elif policy == AutoAugmentPolicy.SVHN: + return [ + (("ShearX", 0.9, 4), ("Invert", 0.2, None)), + (("ShearY", 0.9, 8), ("Invert", 0.7, None)), + (("Equalize", 0.6, None), ("Solarize", 0.6, 6)), + (("Invert", 0.9, None), ("Equalize", 0.6, None)), + (("Equalize", 0.6, None), ("Rotate", 0.9, 3)), + (("ShearX", 0.9, 4), ("AutoContrast", 0.8, None)), + (("ShearY", 0.9, 8), ("Invert", 0.4, None)), + (("ShearY", 0.9, 5), ("Solarize", 0.2, 6)), + (("Invert", 0.9, None), ("AutoContrast", 0.8, None)), + (("Equalize", 0.6, None), ("Rotate", 0.9, 3)), + (("ShearX", 0.9, 4), ("Solarize", 0.3, 3)), + (("ShearY", 0.8, 8), ("Invert", 0.7, None)), + (("Equalize", 0.9, None), ("TranslateY", 0.6, 6)), + (("Invert", 0.9, None), ("Equalize", 0.6, None)), + (("Contrast", 0.3, 3), ("Rotate", 0.8, 4)), + (("Invert", 0.8, None), ("TranslateY", 0.0, 2)), + (("ShearY", 0.7, 6), ("Solarize", 0.4, 8)), + (("Invert", 0.6, None), ("Rotate", 0.8, 4)), + (("ShearY", 0.3, 7), ("TranslateX", 0.9, 3)), + (("ShearX", 0.1, 6), ("Invert", 0.6, None)), + (("Solarize", 0.7, 2), ("TranslateY", 0.6, 7)), + (("ShearY", 0.8, 4), ("Invert", 0.8, None)), + (("ShearX", 0.7, 9), ("TranslateY", 0.8, 3)), + (("ShearY", 0.8, 5), ("AutoContrast", 0.7, None)), + (("ShearX", 0.7, 2), ("Invert", 0.1, None)), + ] + else: + raise ValueError(f"The provided policy {policy} is not recognized.") + + def forward(self, *inputs: Any) -> Any: + flat_inputs_with_spec, image_or_video = self._flatten_and_extract_image_or_video(inputs) + height, width = get_size(image_or_video) # type: ignore[arg-type] + + policy = self._policies[int(torch.randint(len(self._policies), ()))] + + for transform_id, probability, magnitude_idx in policy: + if not torch.rand(()) <= probability: + continue + + magnitudes_fn, signed = self._AUGMENTATION_SPACE[transform_id] + + magnitudes = magnitudes_fn(10, height, width) + if magnitudes is not None: + magnitude = float(magnitudes[magnitude_idx]) + if signed and torch.rand(()) <= 0.5: + magnitude *= -1 + else: + magnitude = 0.0 + + image_or_video = self._apply_image_or_video_transform( + image_or_video, transform_id, magnitude, interpolation=self.interpolation, fill=self._fill + ) + + return self._unflatten_and_insert_image_or_video(flat_inputs_with_spec, image_or_video) + + +class RandAugment(_AutoAugmentBase): + r"""RandAugment data augmentation method based on + `"RandAugment: Practical automated data augmentation with a reduced search space" + `_. + + This transformation works on images and videos only. + + If the input is :class:`torch.Tensor`, it should be of type ``torch.uint8``, and it is expected + to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + + Args: + num_ops (int, optional): Number of augmentation transformations to apply sequentially. + magnitude (int, optional): Magnitude for all the transformations. + num_magnitude_bins (int, optional): The number of different magnitude values. + interpolation (InterpolationMode, optional): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. + fill (sequence or number, optional): Pixel fill value for the area outside the transformed + image. If given a number, the value is used for all bands respectively. + """ + + _v1_transform_cls = _transforms.RandAugment + _AUGMENTATION_SPACE = { + "Identity": (lambda num_bins, height, width: None, False), + "ShearX": (lambda num_bins, height, width: torch.linspace(0.0, 0.3, num_bins), True), + "ShearY": (lambda num_bins, height, width: torch.linspace(0.0, 0.3, num_bins), True), + "TranslateX": ( + lambda num_bins, height, width: torch.linspace(0.0, 150.0 / 331.0 * width, num_bins), + True, + ), + "TranslateY": ( + lambda num_bins, height, width: torch.linspace(0.0, 150.0 / 331.0 * height, num_bins), + True, + ), + "Rotate": (lambda num_bins, height, width: torch.linspace(0.0, 30.0, num_bins), True), + "Brightness": (lambda num_bins, height, width: torch.linspace(0.0, 0.9, num_bins), True), + "Color": (lambda num_bins, height, width: torch.linspace(0.0, 0.9, num_bins), True), + "Contrast": (lambda num_bins, height, width: torch.linspace(0.0, 0.9, num_bins), True), + "Sharpness": (lambda num_bins, height, width: torch.linspace(0.0, 0.9, num_bins), True), + "Posterize": ( + lambda num_bins, height, width: (8 - (torch.arange(num_bins) / ((num_bins - 1) / 4))).round().int(), + False, + ), + "Solarize": (lambda num_bins, height, width: torch.linspace(1.0, 0.0, num_bins), False), + "AutoContrast": (lambda num_bins, height, width: None, False), + "Equalize": (lambda num_bins, height, width: None, False), + } + + def __init__( + self, + num_ops: int = 2, + magnitude: int = 9, + num_magnitude_bins: int = 31, + interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST, + fill: Union[_FillType, Dict[Union[Type, str], _FillType]] = None, + ) -> None: + super().__init__(interpolation=interpolation, fill=fill) + self.num_ops = num_ops + self.magnitude = magnitude + self.num_magnitude_bins = num_magnitude_bins + + def forward(self, *inputs: Any) -> Any: + flat_inputs_with_spec, image_or_video = self._flatten_and_extract_image_or_video(inputs) + height, width = get_size(image_or_video) # type: ignore[arg-type] + + for _ in range(self.num_ops): + transform_id, (magnitudes_fn, signed) = self._get_random_item(self._AUGMENTATION_SPACE) + magnitudes = magnitudes_fn(self.num_magnitude_bins, height, width) + if magnitudes is not None: + magnitude = float(magnitudes[self.magnitude]) + if signed and torch.rand(()) <= 0.5: + magnitude *= -1 + else: + magnitude = 0.0 + image_or_video = self._apply_image_or_video_transform( + image_or_video, transform_id, magnitude, interpolation=self.interpolation, fill=self._fill + ) + + return self._unflatten_and_insert_image_or_video(flat_inputs_with_spec, image_or_video) + + +class TrivialAugmentWide(_AutoAugmentBase): + r"""Dataset-independent data-augmentation with TrivialAugment Wide, as described in + `"TrivialAugment: Tuning-free Yet State-of-the-Art Data Augmentation" `_. + + This transformation works on images and videos only. + + If the input is :class:`torch.Tensor`, it should be of type ``torch.uint8``, and it is expected + to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + + Args: + num_magnitude_bins (int, optional): The number of different magnitude values. + interpolation (InterpolationMode, optional): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. + fill (sequence or number, optional): Pixel fill value for the area outside the transformed + image. If given a number, the value is used for all bands respectively. + """ + + _v1_transform_cls = _transforms.TrivialAugmentWide + _AUGMENTATION_SPACE = { + "Identity": (lambda num_bins, height, width: None, False), + "ShearX": (lambda num_bins, height, width: torch.linspace(0.0, 0.99, num_bins), True), + "ShearY": (lambda num_bins, height, width: torch.linspace(0.0, 0.99, num_bins), True), + "TranslateX": (lambda num_bins, height, width: torch.linspace(0.0, 32.0, num_bins), True), + "TranslateY": (lambda num_bins, height, width: torch.linspace(0.0, 32.0, num_bins), True), + "Rotate": (lambda num_bins, height, width: torch.linspace(0.0, 135.0, num_bins), True), + "Brightness": (lambda num_bins, height, width: torch.linspace(0.0, 0.99, num_bins), True), + "Color": (lambda num_bins, height, width: torch.linspace(0.0, 0.99, num_bins), True), + "Contrast": (lambda num_bins, height, width: torch.linspace(0.0, 0.99, num_bins), True), + "Sharpness": (lambda num_bins, height, width: torch.linspace(0.0, 0.99, num_bins), True), + "Posterize": ( + lambda num_bins, height, width: (8 - (torch.arange(num_bins) / ((num_bins - 1) / 6))).round().int(), + False, + ), + "Solarize": (lambda num_bins, height, width: torch.linspace(1.0, 0.0, num_bins), False), + "AutoContrast": (lambda num_bins, height, width: None, False), + "Equalize": (lambda num_bins, height, width: None, False), + } + + def __init__( + self, + num_magnitude_bins: int = 31, + interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST, + fill: Union[_FillType, Dict[Union[Type, str], _FillType]] = None, + ): + super().__init__(interpolation=interpolation, fill=fill) + self.num_magnitude_bins = num_magnitude_bins + + def forward(self, *inputs: Any) -> Any: + flat_inputs_with_spec, image_or_video = self._flatten_and_extract_image_or_video(inputs) + height, width = get_size(image_or_video) # type: ignore[arg-type] + + transform_id, (magnitudes_fn, signed) = self._get_random_item(self._AUGMENTATION_SPACE) + + magnitudes = magnitudes_fn(self.num_magnitude_bins, height, width) + if magnitudes is not None: + magnitude = float(magnitudes[int(torch.randint(self.num_magnitude_bins, ()))]) + if signed and torch.rand(()) <= 0.5: + magnitude *= -1 + else: + magnitude = 0.0 + + image_or_video = self._apply_image_or_video_transform( + image_or_video, transform_id, magnitude, interpolation=self.interpolation, fill=self._fill + ) + return self._unflatten_and_insert_image_or_video(flat_inputs_with_spec, image_or_video) + + +class AugMix(_AutoAugmentBase): + r"""AugMix data augmentation method based on + `"AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty" `_. + + This transformation works on images and videos only. + + If the input is :class:`torch.Tensor`, it should be of type ``torch.uint8``, and it is expected + to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + + Args: + severity (int, optional): The severity of base augmentation operators. Default is ``3``. + mixture_width (int, optional): The number of augmentation chains. Default is ``3``. + chain_depth (int, optional): The depth of augmentation chains. A negative value denotes stochastic depth sampled from the interval [1, 3]. + Default is ``-1``. + alpha (float, optional): The hyperparameter for the probability distributions. Default is ``1.0``. + all_ops (bool, optional): Use all operations (including brightness, contrast, color and sharpness). Default is ``True``. + interpolation (InterpolationMode, optional): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. + fill (sequence or number, optional): Pixel fill value for the area outside the transformed + image. If given a number, the value is used for all bands respectively. + """ + + _v1_transform_cls = _transforms.AugMix + + _PARTIAL_AUGMENTATION_SPACE = { + "ShearX": (lambda num_bins, height, width: torch.linspace(0.0, 0.3, num_bins), True), + "ShearY": (lambda num_bins, height, width: torch.linspace(0.0, 0.3, num_bins), True), + "TranslateX": (lambda num_bins, height, width: torch.linspace(0.0, width / 3.0, num_bins), True), + "TranslateY": (lambda num_bins, height, width: torch.linspace(0.0, height / 3.0, num_bins), True), + "Rotate": (lambda num_bins, height, width: torch.linspace(0.0, 30.0, num_bins), True), + "Posterize": ( + lambda num_bins, height, width: (4 - (torch.arange(num_bins) / ((num_bins - 1) / 4))).round().int(), + False, + ), + "Solarize": (lambda num_bins, height, width: torch.linspace(1.0, 0.0, num_bins), False), + "AutoContrast": (lambda num_bins, height, width: None, False), + "Equalize": (lambda num_bins, height, width: None, False), + } + _AUGMENTATION_SPACE: Dict[str, Tuple[Callable[[int, int, int], Optional[torch.Tensor]], bool]] = { + **_PARTIAL_AUGMENTATION_SPACE, + "Brightness": (lambda num_bins, height, width: torch.linspace(0.0, 0.9, num_bins), True), + "Color": (lambda num_bins, height, width: torch.linspace(0.0, 0.9, num_bins), True), + "Contrast": (lambda num_bins, height, width: torch.linspace(0.0, 0.9, num_bins), True), + "Sharpness": (lambda num_bins, height, width: torch.linspace(0.0, 0.9, num_bins), True), + } + + def __init__( + self, + severity: int = 3, + mixture_width: int = 3, + chain_depth: int = -1, + alpha: float = 1.0, + all_ops: bool = True, + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + fill: Union[_FillType, Dict[Union[Type, str], _FillType]] = None, + ) -> None: + super().__init__(interpolation=interpolation, fill=fill) + self._PARAMETER_MAX = 10 + if not (1 <= severity <= self._PARAMETER_MAX): + raise ValueError(f"The severity must be between [1, {self._PARAMETER_MAX}]. Got {severity} instead.") + self.severity = severity + self.mixture_width = mixture_width + self.chain_depth = chain_depth + self.alpha = alpha + self.all_ops = all_ops + + def _sample_dirichlet(self, params: torch.Tensor) -> torch.Tensor: + # Must be on a separate method so that we can overwrite it in tests. + return torch._sample_dirichlet(params) + + def forward(self, *inputs: Any) -> Any: + flat_inputs_with_spec, orig_image_or_video = self._flatten_and_extract_image_or_video(inputs) + height, width = get_size(orig_image_or_video) # type: ignore[arg-type] + + if isinstance(orig_image_or_video, torch.Tensor): + image_or_video = orig_image_or_video + else: # isinstance(inpt, PIL.Image.Image): + image_or_video = F.pil_to_tensor(orig_image_or_video) + + augmentation_space = self._AUGMENTATION_SPACE if self.all_ops else self._PARTIAL_AUGMENTATION_SPACE + + orig_dims = list(image_or_video.shape) + expected_ndim = 5 if isinstance(orig_image_or_video, tv_tensors.Video) else 4 + batch = image_or_video.reshape([1] * max(expected_ndim - image_or_video.ndim, 0) + orig_dims) + batch_dims = [batch.size(0)] + [1] * (batch.ndim - 1) + + # Sample the beta weights for combining the original and augmented image or video. To get Beta, we use a + # Dirichlet with 2 parameters. The 1st column stores the weights of the original and the 2nd the ones of + # augmented image or video. + m = self._sample_dirichlet( + torch.tensor([self.alpha, self.alpha], device=batch.device).expand(batch_dims[0], -1) + ) + + # Sample the mixing weights and combine them with the ones sampled from Beta for the augmented images or videos. + combined_weights = self._sample_dirichlet( + torch.tensor([self.alpha] * self.mixture_width, device=batch.device).expand(batch_dims[0], -1) + ) * m[:, 1].reshape([batch_dims[0], -1]) + + mix = m[:, 0].reshape(batch_dims) * batch + for i in range(self.mixture_width): + aug = batch + depth = self.chain_depth if self.chain_depth > 0 else int(torch.randint(low=1, high=4, size=(1,)).item()) + for _ in range(depth): + transform_id, (magnitudes_fn, signed) = self._get_random_item(augmentation_space) + + magnitudes = magnitudes_fn(self._PARAMETER_MAX, height, width) + if magnitudes is not None: + magnitude = float(magnitudes[int(torch.randint(self.severity, ()))]) + if signed and torch.rand(()) <= 0.5: + magnitude *= -1 + else: + magnitude = 0.0 + + aug = self._apply_image_or_video_transform(aug, transform_id, magnitude, interpolation=self.interpolation, fill=self._fill) # type: ignore[assignment] + mix.add_(combined_weights[:, i].reshape(batch_dims) * aug) + mix = mix.reshape(orig_dims).to(dtype=image_or_video.dtype) + + if isinstance(orig_image_or_video, (tv_tensors.Image, tv_tensors.Video)): + mix = tv_tensors.wrap(mix, like=orig_image_or_video) + elif isinstance(orig_image_or_video, PIL.Image.Image): + mix = F.to_pil_image(mix) + + return self._unflatten_and_insert_image_or_video(flat_inputs_with_spec, mix) diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/_color.py b/lib/python3.12/site-packages/torchvision/transforms/v2/_color.py new file mode 100644 index 0000000000000000000000000000000000000000..7a471e7c1f632b8dd68ca6ad02b92f947b8d5354 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/_color.py @@ -0,0 +1,376 @@ +import collections.abc +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union + +import torch +from torchvision import transforms as _transforms +from torchvision.transforms.v2 import functional as F, Transform + +from ._transform import _RandomApplyTransform +from ._utils import query_chw + + +class Grayscale(Transform): + """Convert images or videos to grayscale. + + If the input is a :class:`torch.Tensor`, it is expected + to have [..., 3 or 1, H, W] shape, where ... means an arbitrary number of leading dimensions + + Args: + num_output_channels (int): (1 or 3) number of channels desired for output image + """ + + _v1_transform_cls = _transforms.Grayscale + + def __init__(self, num_output_channels: int = 1): + super().__init__() + self.num_output_channels = num_output_channels + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.rgb_to_grayscale, inpt, num_output_channels=self.num_output_channels) + + +class RandomGrayscale(_RandomApplyTransform): + """Randomly convert image or videos to grayscale with a probability of p (default 0.1). + + If the input is a :class:`torch.Tensor`, it is expected to have [..., 3 or 1, H, W] shape, + where ... means an arbitrary number of leading dimensions + + The output has the same number of channels as the input. + + Args: + p (float): probability that image should be converted to grayscale. + """ + + _v1_transform_cls = _transforms.RandomGrayscale + + def __init__(self, p: float = 0.1) -> None: + super().__init__(p=p) + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + num_input_channels, *_ = query_chw(flat_inputs) + return dict(num_input_channels=num_input_channels) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.rgb_to_grayscale, inpt, num_output_channels=params["num_input_channels"]) + + +class RGB(Transform): + """Convert images or videos to RGB (if they are already not RGB). + + If the input is a :class:`torch.Tensor`, it is expected + to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions + """ + + def __init__(self): + super().__init__() + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.grayscale_to_rgb, inpt) + + +class ColorJitter(Transform): + """Randomly change the brightness, contrast, saturation and hue of an image or video. + + If the input is a :class:`torch.Tensor`, it is expected + to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. + If img is PIL Image, mode "1", "I", "F" and modes with transparency (alpha channel) are not supported. + + Args: + brightness (float or tuple of float (min, max)): How much to jitter brightness. + brightness_factor is chosen uniformly from [max(0, 1 - brightness), 1 + brightness] + or the given [min, max]. Should be non negative numbers. + contrast (float or tuple of float (min, max)): How much to jitter contrast. + contrast_factor is chosen uniformly from [max(0, 1 - contrast), 1 + contrast] + or the given [min, max]. Should be non-negative numbers. + saturation (float or tuple of float (min, max)): How much to jitter saturation. + saturation_factor is chosen uniformly from [max(0, 1 - saturation), 1 + saturation] + or the given [min, max]. Should be non negative numbers. + hue (float or tuple of float (min, max)): How much to jitter hue. + hue_factor is chosen uniformly from [-hue, hue] or the given [min, max]. + Should have 0<= hue <= 0.5 or -0.5 <= min <= max <= 0.5. + To jitter hue, the pixel values of the input image has to be non-negative for conversion to HSV space; + thus it does not work if you normalize your image to an interval with negative values, + or use an interpolation that generates negative values before using this function. + """ + + _v1_transform_cls = _transforms.ColorJitter + + def _extract_params_for_v1_transform(self) -> Dict[str, Any]: + return {attr: value or 0 for attr, value in super()._extract_params_for_v1_transform().items()} + + def __init__( + self, + brightness: Optional[Union[float, Sequence[float]]] = None, + contrast: Optional[Union[float, Sequence[float]]] = None, + saturation: Optional[Union[float, Sequence[float]]] = None, + hue: Optional[Union[float, Sequence[float]]] = None, + ) -> None: + super().__init__() + self.brightness = self._check_input(brightness, "brightness") + self.contrast = self._check_input(contrast, "contrast") + self.saturation = self._check_input(saturation, "saturation") + self.hue = self._check_input(hue, "hue", center=0, bound=(-0.5, 0.5), clip_first_on_zero=False) + + def _check_input( + self, + value: Optional[Union[float, Sequence[float]]], + name: str, + center: float = 1.0, + bound: Tuple[float, float] = (0, float("inf")), + clip_first_on_zero: bool = True, + ) -> Optional[Tuple[float, float]]: + if value is None: + return None + + if isinstance(value, (int, float)): + if value < 0: + raise ValueError(f"If {name} is a single number, it must be non negative.") + value = [center - value, center + value] + if clip_first_on_zero: + value[0] = max(value[0], 0.0) + elif isinstance(value, collections.abc.Sequence) and len(value) == 2: + value = [float(v) for v in value] + else: + raise TypeError(f"{name}={value} should be a single number or a sequence with length 2.") + + if not bound[0] <= value[0] <= value[1] <= bound[1]: + raise ValueError(f"{name} values should be between {bound}, but got {value}.") + + return None if value[0] == value[1] == center else (float(value[0]), float(value[1])) + + @staticmethod + def _generate_value(left: float, right: float) -> float: + return torch.empty(1).uniform_(left, right).item() + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + fn_idx = torch.randperm(4) + + b = None if self.brightness is None else self._generate_value(self.brightness[0], self.brightness[1]) + c = None if self.contrast is None else self._generate_value(self.contrast[0], self.contrast[1]) + s = None if self.saturation is None else self._generate_value(self.saturation[0], self.saturation[1]) + h = None if self.hue is None else self._generate_value(self.hue[0], self.hue[1]) + + return dict(fn_idx=fn_idx, brightness_factor=b, contrast_factor=c, saturation_factor=s, hue_factor=h) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + output = inpt + brightness_factor = params["brightness_factor"] + contrast_factor = params["contrast_factor"] + saturation_factor = params["saturation_factor"] + hue_factor = params["hue_factor"] + for fn_id in params["fn_idx"]: + if fn_id == 0 and brightness_factor is not None: + output = self._call_kernel(F.adjust_brightness, output, brightness_factor=brightness_factor) + elif fn_id == 1 and contrast_factor is not None: + output = self._call_kernel(F.adjust_contrast, output, contrast_factor=contrast_factor) + elif fn_id == 2 and saturation_factor is not None: + output = self._call_kernel(F.adjust_saturation, output, saturation_factor=saturation_factor) + elif fn_id == 3 and hue_factor is not None: + output = self._call_kernel(F.adjust_hue, output, hue_factor=hue_factor) + return output + + +class RandomChannelPermutation(Transform): + """Randomly permute the channels of an image or video""" + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + num_channels, *_ = query_chw(flat_inputs) + return dict(permutation=torch.randperm(num_channels)) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.permute_channels, inpt, params["permutation"]) + + +class RandomPhotometricDistort(Transform): + """Randomly distorts the image or video as used in `SSD: Single Shot + MultiBox Detector `_. + + This transform relies on :class:`~torchvision.transforms.v2.ColorJitter` + under the hood to adjust the contrast, saturation, hue, brightness, and also + randomly permutes channels. + + Args: + brightness (tuple of float (min, max), optional): How much to jitter brightness. + brightness_factor is chosen uniformly from [min, max]. Should be non negative numbers. + contrast (tuple of float (min, max), optional): How much to jitter contrast. + contrast_factor is chosen uniformly from [min, max]. Should be non-negative numbers. + saturation (tuple of float (min, max), optional): How much to jitter saturation. + saturation_factor is chosen uniformly from [min, max]. Should be non negative numbers. + hue (tuple of float (min, max), optional): How much to jitter hue. + hue_factor is chosen uniformly from [min, max]. Should have -0.5 <= min <= max <= 0.5. + To jitter hue, the pixel values of the input image has to be non-negative for conversion to HSV space; + thus it does not work if you normalize your image to an interval with negative values, + or use an interpolation that generates negative values before using this function. + p (float, optional) probability each distortion operation (contrast, saturation, ...) to be applied. + Default is 0.5. + """ + + def __init__( + self, + brightness: Tuple[float, float] = (0.875, 1.125), + contrast: Tuple[float, float] = (0.5, 1.5), + saturation: Tuple[float, float] = (0.5, 1.5), + hue: Tuple[float, float] = (-0.05, 0.05), + p: float = 0.5, + ): + super().__init__() + self.brightness = brightness + self.contrast = contrast + self.hue = hue + self.saturation = saturation + self.p = p + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + num_channels, *_ = query_chw(flat_inputs) + params: Dict[str, Any] = { + key: ColorJitter._generate_value(range[0], range[1]) if torch.rand(1) < self.p else None + for key, range in [ + ("brightness_factor", self.brightness), + ("contrast_factor", self.contrast), + ("saturation_factor", self.saturation), + ("hue_factor", self.hue), + ] + } + params["contrast_before"] = bool(torch.rand(()) < 0.5) + params["channel_permutation"] = torch.randperm(num_channels) if torch.rand(1) < self.p else None + return params + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + if params["brightness_factor"] is not None: + inpt = self._call_kernel(F.adjust_brightness, inpt, brightness_factor=params["brightness_factor"]) + if params["contrast_factor"] is not None and params["contrast_before"]: + inpt = self._call_kernel(F.adjust_contrast, inpt, contrast_factor=params["contrast_factor"]) + if params["saturation_factor"] is not None: + inpt = self._call_kernel(F.adjust_saturation, inpt, saturation_factor=params["saturation_factor"]) + if params["hue_factor"] is not None: + inpt = self._call_kernel(F.adjust_hue, inpt, hue_factor=params["hue_factor"]) + if params["contrast_factor"] is not None and not params["contrast_before"]: + inpt = self._call_kernel(F.adjust_contrast, inpt, contrast_factor=params["contrast_factor"]) + if params["channel_permutation"] is not None: + inpt = self._call_kernel(F.permute_channels, inpt, permutation=params["channel_permutation"]) + return inpt + + +class RandomEqualize(_RandomApplyTransform): + """Equalize the histogram of the given image or video with a given probability. + + If the input is a :class:`torch.Tensor`, it is expected + to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "P", "L" or "RGB". + + Args: + p (float): probability of the image being equalized. Default value is 0.5 + """ + + _v1_transform_cls = _transforms.RandomEqualize + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.equalize, inpt) + + +class RandomInvert(_RandomApplyTransform): + """Inverts the colors of the given image or video with a given probability. + + If img is a Tensor, it is expected to be in [..., 1 or 3, H, W] format, + where ... means it can have an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + + Args: + p (float): probability of the image being color inverted. Default value is 0.5 + """ + + _v1_transform_cls = _transforms.RandomInvert + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.invert, inpt) + + +class RandomPosterize(_RandomApplyTransform): + """Posterize the image or video with a given probability by reducing the + number of bits for each color channel. + + If the input is a :class:`torch.Tensor`, it should be of type torch.uint8, + and it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + + Args: + bits (int): number of bits to keep for each channel (0-8) + p (float): probability of the image being posterized. Default value is 0.5 + """ + + _v1_transform_cls = _transforms.RandomPosterize + + def __init__(self, bits: int, p: float = 0.5) -> None: + super().__init__(p=p) + self.bits = bits + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.posterize, inpt, bits=self.bits) + + +class RandomSolarize(_RandomApplyTransform): + """Solarize the image or video with a given probability by inverting all pixel + values above a threshold. + + If img is a Tensor, it is expected to be in [..., 1 or 3, H, W] format, + where ... means it can have an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + + Args: + threshold (float): all pixels equal or above this value are inverted. + p (float): probability of the image being solarized. Default value is 0.5 + """ + + _v1_transform_cls = _transforms.RandomSolarize + + def _extract_params_for_v1_transform(self) -> Dict[str, Any]: + params = super()._extract_params_for_v1_transform() + params["threshold"] = float(params["threshold"]) + return params + + def __init__(self, threshold: float, p: float = 0.5) -> None: + super().__init__(p=p) + self.threshold = threshold + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.solarize, inpt, threshold=self.threshold) + + +class RandomAutocontrast(_RandomApplyTransform): + """Autocontrast the pixels of the given image or video with a given probability. + + If the input is a :class:`torch.Tensor`, it is expected + to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. + If img is PIL Image, it is expected to be in mode "L" or "RGB". + + Args: + p (float): probability of the image being autocontrasted. Default value is 0.5 + """ + + _v1_transform_cls = _transforms.RandomAutocontrast + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.autocontrast, inpt) + + +class RandomAdjustSharpness(_RandomApplyTransform): + """Adjust the sharpness of the image or video with a given probability. + + If the input is a :class:`torch.Tensor`, + it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. + + Args: + sharpness_factor (float): How much to adjust the sharpness. Can be + any non-negative number. 0 gives a blurred image, 1 gives the + original image while 2 increases the sharpness by a factor of 2. + p (float): probability of the image being sharpened. Default value is 0.5 + """ + + _v1_transform_cls = _transforms.RandomAdjustSharpness + + def __init__(self, sharpness_factor: float, p: float = 0.5) -> None: + super().__init__(p=p) + self.sharpness_factor = sharpness_factor + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.adjust_sharpness, inpt, sharpness_factor=self.sharpness_factor) diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/_container.py b/lib/python3.12/site-packages/torchvision/transforms/v2/_container.py new file mode 100644 index 0000000000000000000000000000000000000000..54de601c6967001f4f7b5951d84b36b7353dbf23 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/_container.py @@ -0,0 +1,174 @@ +from typing import Any, Callable, Dict, List, Optional, Sequence, Union + +import torch + +from torch import nn +from torchvision import transforms as _transforms +from torchvision.transforms.v2 import Transform + + +class Compose(Transform): + """Composes several transforms together. + + This transform does not support torchscript. + Please, see the note below. + + Args: + transforms (list of ``Transform`` objects): list of transforms to compose. + + Example: + >>> transforms.Compose([ + >>> transforms.CenterCrop(10), + >>> transforms.PILToTensor(), + >>> transforms.ConvertImageDtype(torch.float), + >>> ]) + + .. note:: + In order to script the transformations, please use ``torch.nn.Sequential`` as below. + + >>> transforms = torch.nn.Sequential( + >>> transforms.CenterCrop(10), + >>> transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), + >>> ) + >>> scripted_transforms = torch.jit.script(transforms) + + Make sure to use only scriptable transformations, i.e. that work with ``torch.Tensor``, does not require + `lambda` functions or ``PIL.Image``. + + """ + + def __init__(self, transforms: Sequence[Callable]) -> None: + super().__init__() + if not isinstance(transforms, Sequence): + raise TypeError("Argument transforms should be a sequence of callables") + elif not transforms: + raise ValueError("Pass at least one transform") + self.transforms = transforms + + def forward(self, *inputs: Any) -> Any: + needs_unpacking = len(inputs) > 1 + for transform in self.transforms: + outputs = transform(*inputs) + inputs = outputs if needs_unpacking else (outputs,) + return outputs + + def extra_repr(self) -> str: + format_string = [] + for t in self.transforms: + format_string.append(f" {t}") + return "\n".join(format_string) + + +class RandomApply(Transform): + """Apply randomly a list of transformations with a given probability. + + .. note:: + In order to script the transformation, please use ``torch.nn.ModuleList`` as input instead of list/tuple of + transforms as shown below: + + >>> transforms = transforms.RandomApply(torch.nn.ModuleList([ + >>> transforms.ColorJitter(), + >>> ]), p=0.3) + >>> scripted_transforms = torch.jit.script(transforms) + + Make sure to use only scriptable transformations, i.e. that work with ``torch.Tensor``, does not require + `lambda` functions or ``PIL.Image``. + + Args: + transforms (sequence or torch.nn.Module): list of transformations + p (float): probability of applying the list of transforms + """ + + _v1_transform_cls = _transforms.RandomApply + + def __init__(self, transforms: Union[Sequence[Callable], nn.ModuleList], p: float = 0.5) -> None: + super().__init__() + + if not isinstance(transforms, (Sequence, nn.ModuleList)): + raise TypeError("Argument transforms should be a sequence of callables or a `nn.ModuleList`") + self.transforms = transforms + + if not (0.0 <= p <= 1.0): + raise ValueError("`p` should be a floating point value in the interval [0.0, 1.0].") + self.p = p + + def _extract_params_for_v1_transform(self) -> Dict[str, Any]: + return {"transforms": self.transforms, "p": self.p} + + def forward(self, *inputs: Any) -> Any: + needs_unpacking = len(inputs) > 1 + + if torch.rand(1) >= self.p: + return inputs if needs_unpacking else inputs[0] + + for transform in self.transforms: + outputs = transform(*inputs) + inputs = outputs if needs_unpacking else (outputs,) + return outputs + + def extra_repr(self) -> str: + format_string = [] + for t in self.transforms: + format_string.append(f" {t}") + return "\n".join(format_string) + + +class RandomChoice(Transform): + """Apply single transformation randomly picked from a list. + + This transform does not support torchscript. + + Args: + transforms (sequence or torch.nn.Module): list of transformations + p (list of floats or None, optional): probability of each transform being picked. + If ``p`` doesn't sum to 1, it is automatically normalized. If ``None`` + (default), all transforms have the same probability. + """ + + def __init__( + self, + transforms: Sequence[Callable], + p: Optional[List[float]] = None, + ) -> None: + if not isinstance(transforms, Sequence): + raise TypeError("Argument transforms should be a sequence of callables") + + if p is None: + p = [1] * len(transforms) + elif len(p) != len(transforms): + raise ValueError(f"Length of p doesn't match the number of transforms: {len(p)} != {len(transforms)}") + + super().__init__() + + self.transforms = transforms + total = sum(p) + self.p = [prob / total for prob in p] + + def forward(self, *inputs: Any) -> Any: + idx = int(torch.multinomial(torch.tensor(self.p), 1)) + transform = self.transforms[idx] + return transform(*inputs) + + +class RandomOrder(Transform): + """Apply a list of transformations in a random order. + + This transform does not support torchscript. + + Args: + transforms (sequence or torch.nn.Module): list of transformations + """ + + def __init__(self, transforms: Sequence[Callable]) -> None: + if not isinstance(transforms, Sequence): + raise TypeError("Argument transforms should be a sequence of callables") + super().__init__() + self.transforms = transforms + + def forward(self, *inputs: Any) -> Any: + needs_unpacking = len(inputs) > 1 + for idx in torch.randperm(len(self.transforms)): + transform = self.transforms[idx] + outputs = transform(*inputs) + inputs = outputs if needs_unpacking else (outputs,) + return outputs diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/_deprecated.py b/lib/python3.12/site-packages/torchvision/transforms/v2/_deprecated.py new file mode 100644 index 0000000000000000000000000000000000000000..51a4f076e49d16b8ccde6e9ed04f610ed9c22bff --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/_deprecated.py @@ -0,0 +1,50 @@ +import warnings +from typing import Any, Dict, Union + +import numpy as np +import PIL.Image +import torch +from torchvision.transforms import functional as _F + +from torchvision.transforms.v2 import Transform + + +class ToTensor(Transform): + """[DEPRECATED] Use ``v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])`` instead. + + Convert a PIL Image or ndarray to tensor and scale the values accordingly. + + .. warning:: + :class:`v2.ToTensor` is deprecated and will be removed in a future release. + Please use instead ``v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])``. + Output is equivalent up to float precision. + + This transform does not support torchscript. + + + Converts a PIL Image or numpy.ndarray (H x W x C) in the range + [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0] + if the PIL Image belongs to one of the modes (L, LA, P, I, F, RGB, YCbCr, RGBA, CMYK, 1) + or if the numpy.ndarray has dtype = np.uint8 + + In the other cases, tensors are returned without scaling. + + .. note:: + Because the input image is scaled to [0.0, 1.0], this transformation should not be used when + transforming target image masks. See the `references`_ for implementing the transforms for image masks. + + .. _references: https://github.com/pytorch/vision/tree/main/references/segmentation + """ + + _transformed_types = (PIL.Image.Image, np.ndarray) + + def __init__(self) -> None: + warnings.warn( + "The transform `ToTensor()` is deprecated and will be removed in a future release. " + "Instead, please use `v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])`." + "Output is equivalent up to float precision." + ) + super().__init__() + + def transform(self, inpt: Union[PIL.Image.Image, np.ndarray], params: Dict[str, Any]) -> torch.Tensor: + return _F.to_tensor(inpt) diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/_geometry.py b/lib/python3.12/site-packages/torchvision/transforms/v2/_geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..c2461418a422861159e8b27cb4e8a3f7dfa1dcf8 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/_geometry.py @@ -0,0 +1,1416 @@ +import math +import numbers +import warnings +from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple, Type, Union + +import PIL.Image +import torch + +from torchvision import transforms as _transforms, tv_tensors +from torchvision.ops.boxes import box_iou +from torchvision.transforms.functional import _get_perspective_coeffs +from torchvision.transforms.v2 import functional as F, InterpolationMode, Transform +from torchvision.transforms.v2.functional._utils import _FillType + +from ._transform import _RandomApplyTransform +from ._utils import ( + _check_padding_arg, + _check_padding_mode_arg, + _check_sequence_input, + _get_fill, + _setup_angle, + _setup_fill_arg, + _setup_number_or_seq, + _setup_size, + get_bounding_boxes, + has_all, + has_any, + is_pure_tensor, + query_size, +) + + +class RandomHorizontalFlip(_RandomApplyTransform): + """Horizontally flip the input with a given probability. + + If the input is a :class:`torch.Tensor` or a ``TVTensor`` (e.g. :class:`~torchvision.tv_tensors.Image`, + :class:`~torchvision.tv_tensors.Video`, :class:`~torchvision.tv_tensors.BoundingBoxes` etc.) + it can have arbitrary number of leading batch dimensions. For example, + the image can have ``[..., C, H, W]`` shape. A bounding box can have ``[..., 4]`` shape. + + Args: + p (float, optional): probability of the input being flipped. Default value is 0.5 + """ + + _v1_transform_cls = _transforms.RandomHorizontalFlip + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.horizontal_flip, inpt) + + +class RandomVerticalFlip(_RandomApplyTransform): + """Vertically flip the input with a given probability. + + If the input is a :class:`torch.Tensor` or a ``TVTensor`` (e.g. :class:`~torchvision.tv_tensors.Image`, + :class:`~torchvision.tv_tensors.Video`, :class:`~torchvision.tv_tensors.BoundingBoxes` etc.) + it can have arbitrary number of leading batch dimensions. For example, + the image can have ``[..., C, H, W]`` shape. A bounding box can have ``[..., 4]`` shape. + + Args: + p (float, optional): probability of the input being flipped. Default value is 0.5 + """ + + _v1_transform_cls = _transforms.RandomVerticalFlip + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.vertical_flip, inpt) + + +class Resize(Transform): + """Resize the input to the given size. + + If the input is a :class:`torch.Tensor` or a ``TVTensor`` (e.g. :class:`~torchvision.tv_tensors.Image`, + :class:`~torchvision.tv_tensors.Video`, :class:`~torchvision.tv_tensors.BoundingBoxes` etc.) + it can have arbitrary number of leading batch dimensions. For example, + the image can have ``[..., C, H, W]`` shape. A bounding box can have ``[..., 4]`` shape. + + Args: + size (sequence, int, or None): Desired + output size. + + - If size is a sequence like (h, w), output size will be matched to this. + - If size is an int, smaller edge of the image will be matched to this + number. i.e, if height > width, then image will be rescaled to + (size * height / width, size). + - If size is None, the output shape is determined by the ``max_size`` + parameter. + + .. note:: + In torchscript mode size as single int is not supported, use a sequence of length 1: ``[size, ]``. + interpolation (InterpolationMode, optional): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.NEAREST_EXACT``, + ``InterpolationMode.BILINEAR`` and ``InterpolationMode.BICUBIC`` are supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + max_size (int, optional): The maximum allowed for the longer edge of + the resized image. + + - If ``size`` is an int: if the longer edge of the image is greater + than ``max_size`` after being resized according to ``size``, + ``size`` will be overruled so that the longer edge is equal to + ``max_size``. As a result, the smaller edge may be shorter than + ``size``. This is only supported if ``size`` is an int (or a + sequence of length 1 in torchscript mode). + - If ``size`` is None: the longer edge of the image will be matched + to max_size. i.e, if height > width, then image will be rescaled + to (max_size, max_size * width / height). + + This should be left to ``None`` (default) when ``size`` is a + sequence. + + antialias (bool, optional): Whether to apply antialiasing. + It only affects **tensors** with bilinear or bicubic modes and it is + ignored otherwise: on PIL images, antialiasing is always applied on + bilinear or bicubic modes; on other modes (for PIL images and + tensors), antialiasing makes no sense and this parameter is ignored. + Possible values are: + + - ``True`` (default): will apply antialiasing for bilinear or bicubic modes. + Other mode aren't affected. This is probably what you want to use. + - ``False``: will not apply antialiasing for tensors on any mode. PIL + images are still antialiased on bilinear or bicubic modes, because + PIL doesn't support no antialias. + - ``None``: equivalent to ``False`` for tensors and ``True`` for + PIL images. This value exists for legacy reasons and you probably + don't want to use it unless you really know what you are doing. + + The default value changed from ``None`` to ``True`` in + v0.17, for the PIL and Tensor backends to be consistent. + """ + + _v1_transform_cls = _transforms.Resize + + def __init__( + self, + size: Union[int, Sequence[int], None], + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + max_size: Optional[int] = None, + antialias: Optional[bool] = True, + ) -> None: + super().__init__() + + if isinstance(size, int): + size = [size] + elif isinstance(size, Sequence) and len(size) in {1, 2}: + size = list(size) + elif size is None: + if not isinstance(max_size, int): + raise ValueError(f"max_size must be an integer when size is None, but got {max_size} instead.") + else: + raise ValueError( + f"size can be an integer, a sequence of one or two integers, or None, but got {size} instead." + ) + self.size = size + + self.interpolation = interpolation + self.max_size = max_size + self.antialias = antialias + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel( + F.resize, + inpt, + self.size, + interpolation=self.interpolation, + max_size=self.max_size, + antialias=self.antialias, + ) + + +class CenterCrop(Transform): + """Crop the input at the center. + + If the input is a :class:`torch.Tensor` or a ``TVTensor`` (e.g. :class:`~torchvision.tv_tensors.Image`, + :class:`~torchvision.tv_tensors.Video`, :class:`~torchvision.tv_tensors.BoundingBoxes` etc.) + it can have arbitrary number of leading batch dimensions. For example, + the image can have ``[..., C, H, W]`` shape. A bounding box can have ``[..., 4]`` shape. + + If image size is smaller than output size along any edge, image is padded with 0 and then center cropped. + + Args: + size (sequence or int): Desired output size of the crop. If size is an + int instead of sequence like (h, w), a square crop (size, size) is + made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). + """ + + _v1_transform_cls = _transforms.CenterCrop + + def __init__(self, size: Union[int, Sequence[int]]): + super().__init__() + self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.") + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.center_crop, inpt, output_size=self.size) + + +class RandomResizedCrop(Transform): + """Crop a random portion of the input and resize it to a given size. + + If the input is a :class:`torch.Tensor` or a ``TVTensor`` (e.g. :class:`~torchvision.tv_tensors.Image`, + :class:`~torchvision.tv_tensors.Video`, :class:`~torchvision.tv_tensors.BoundingBoxes` etc.) + it can have arbitrary number of leading batch dimensions. For example, + the image can have ``[..., C, H, W]`` shape. A bounding box can have ``[..., 4]`` shape. + + A crop of the original input is made: the crop has a random area (H * W) + and a random aspect ratio. This crop is finally resized to the given + size. This is popularly used to train the Inception networks. + + Args: + size (int or sequence): expected output size of the crop, for each edge. If size is an + int instead of sequence like (h, w), a square output size ``(size, size)`` is + made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). + + .. note:: + In torchscript mode size as single int is not supported, use a sequence of length 1: ``[size, ]``. + scale (tuple of float, optional): Specifies the lower and upper bounds for the random area of the crop, + before resizing. The scale is defined with respect to the area of the original image. + ratio (tuple of float, optional): lower and upper bounds for the random aspect ratio of the crop, before + resizing. + interpolation (InterpolationMode, optional): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.NEAREST_EXACT``, + ``InterpolationMode.BILINEAR`` and ``InterpolationMode.BICUBIC`` are supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + antialias (bool, optional): Whether to apply antialiasing. + It only affects **tensors** with bilinear or bicubic modes and it is + ignored otherwise: on PIL images, antialiasing is always applied on + bilinear or bicubic modes; on other modes (for PIL images and + tensors), antialiasing makes no sense and this parameter is ignored. + Possible values are: + + - ``True`` (default): will apply antialiasing for bilinear or bicubic modes. + Other mode aren't affected. This is probably what you want to use. + - ``False``: will not apply antialiasing for tensors on any mode. PIL + images are still antialiased on bilinear or bicubic modes, because + PIL doesn't support no antialias. + - ``None``: equivalent to ``False`` for tensors and ``True`` for + PIL images. This value exists for legacy reasons and you probably + don't want to use it unless you really know what you are doing. + + The default value changed from ``None`` to ``True`` in + v0.17, for the PIL and Tensor backends to be consistent. + """ + + _v1_transform_cls = _transforms.RandomResizedCrop + + def __init__( + self, + size: Union[int, Sequence[int]], + scale: Tuple[float, float] = (0.08, 1.0), + ratio: Tuple[float, float] = (3.0 / 4.0, 4.0 / 3.0), + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + antialias: Optional[bool] = True, + ) -> None: + super().__init__() + self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.") + + if not isinstance(scale, Sequence): + raise TypeError("Scale should be a sequence") + if not isinstance(ratio, Sequence): + raise TypeError("Ratio should be a sequence") + if (scale[0] > scale[1]) or (ratio[0] > ratio[1]): + warnings.warn("Scale and ratio should be of kind (min, max)") + + self.scale = scale + self.ratio = ratio + self.interpolation = interpolation + self.antialias = antialias + + self._log_ratio = torch.log(torch.tensor(self.ratio)) + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + height, width = query_size(flat_inputs) + area = height * width + + log_ratio = self._log_ratio + for _ in range(10): + target_area = area * torch.empty(1).uniform_(self.scale[0], self.scale[1]).item() + aspect_ratio = torch.exp( + torch.empty(1).uniform_( + log_ratio[0], # type: ignore[arg-type] + log_ratio[1], # type: ignore[arg-type] + ) + ).item() + + w = int(round(math.sqrt(target_area * aspect_ratio))) + h = int(round(math.sqrt(target_area / aspect_ratio))) + + if 0 < w <= width and 0 < h <= height: + i = torch.randint(0, height - h + 1, size=(1,)).item() + j = torch.randint(0, width - w + 1, size=(1,)).item() + break + else: + # Fallback to central crop + in_ratio = float(width) / float(height) + if in_ratio < min(self.ratio): + w = width + h = int(round(w / min(self.ratio))) + elif in_ratio > max(self.ratio): + h = height + w = int(round(h * max(self.ratio))) + else: # whole image + w = width + h = height + i = (height - h) // 2 + j = (width - w) // 2 + + return dict(top=i, left=j, height=h, width=w) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel( + F.resized_crop, inpt, **params, size=self.size, interpolation=self.interpolation, antialias=self.antialias + ) + + +class FiveCrop(Transform): + """Crop the image or video into four corners and the central crop. + + If the input is a :class:`torch.Tensor` or a :class:`~torchvision.tv_tensors.Image` or a + :class:`~torchvision.tv_tensors.Video` it can have arbitrary number of leading batch dimensions. + For example, the image can have ``[..., C, H, W]`` shape. + + .. Note:: + This transform returns a tuple of images and there may be a mismatch in the number of + inputs and targets your Dataset returns. See below for an example of how to deal with + this. + + Args: + size (sequence or int): Desired output size of the crop. If size is an ``int`` + instead of sequence like (h, w), a square crop of size (size, size) is made. + If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). + + Example: + >>> class BatchMultiCrop(transforms.Transform): + ... def forward(self, sample: Tuple[Tuple[Union[tv_tensors.Image, tv_tensors.Video], ...], int]): + ... images_or_videos, labels = sample + ... batch_size = len(images_or_videos) + ... image_or_video = images_or_videos[0] + ... images_or_videos = tv_tensors.wrap(torch.stack(images_or_videos), like=image_or_video) + ... labels = torch.full((batch_size,), label, device=images_or_videos.device) + ... return images_or_videos, labels + ... + >>> image = tv_tensors.Image(torch.rand(3, 256, 256)) + >>> label = 3 + >>> transform = transforms.Compose([transforms.FiveCrop(224), BatchMultiCrop()]) + >>> images, labels = transform(image, label) + >>> images.shape + torch.Size([5, 3, 224, 224]) + >>> labels + tensor([3, 3, 3, 3, 3]) + """ + + _v1_transform_cls = _transforms.FiveCrop + + def __init__(self, size: Union[int, Sequence[int]]) -> None: + super().__init__() + self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.") + + def _call_kernel(self, functional: Callable, inpt: Any, *args: Any, **kwargs: Any) -> Any: + if isinstance(inpt, (tv_tensors.BoundingBoxes, tv_tensors.Mask)): + warnings.warn( + f"{type(self).__name__}() is currently passing through inputs of type " + f"tv_tensors.{type(inpt).__name__}. This will likely change in the future." + ) + return super()._call_kernel(functional, inpt, *args, **kwargs) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.five_crop, inpt, self.size) + + def check_inputs(self, flat_inputs: List[Any]) -> None: + if has_any(flat_inputs, tv_tensors.BoundingBoxes, tv_tensors.Mask): + raise TypeError(f"BoundingBoxes'es and Mask's are not supported by {type(self).__name__}()") + + +class TenCrop(Transform): + """Crop the image or video into four corners and the central crop plus the flipped version of + these (horizontal flipping is used by default). + + If the input is a :class:`torch.Tensor` or a :class:`~torchvision.tv_tensors.Image` or a + :class:`~torchvision.tv_tensors.Video` it can have arbitrary number of leading batch dimensions. + For example, the image can have ``[..., C, H, W]`` shape. + + See :class:`~torchvision.transforms.v2.FiveCrop` for an example. + + .. Note:: + This transform returns a tuple of images and there may be a mismatch in the number of + inputs and targets your Dataset returns. See below for an example of how to deal with + this. + + Args: + size (sequence or int): Desired output size of the crop. If size is an + int instead of sequence like (h, w), a square crop (size, size) is + made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). + vertical_flip (bool, optional): Use vertical flipping instead of horizontal + """ + + _v1_transform_cls = _transforms.TenCrop + + def __init__(self, size: Union[int, Sequence[int]], vertical_flip: bool = False) -> None: + super().__init__() + self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.") + self.vertical_flip = vertical_flip + + def _call_kernel(self, functional: Callable, inpt: Any, *args: Any, **kwargs: Any) -> Any: + if isinstance(inpt, (tv_tensors.BoundingBoxes, tv_tensors.Mask)): + warnings.warn( + f"{type(self).__name__}() is currently passing through inputs of type " + f"tv_tensors.{type(inpt).__name__}. This will likely change in the future." + ) + return super()._call_kernel(functional, inpt, *args, **kwargs) + + def check_inputs(self, flat_inputs: List[Any]) -> None: + if has_any(flat_inputs, tv_tensors.BoundingBoxes, tv_tensors.Mask): + raise TypeError(f"BoundingBoxes'es and Mask's are not supported by {type(self).__name__}()") + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.ten_crop, inpt, self.size, vertical_flip=self.vertical_flip) + + +class Pad(Transform): + """Pad the input on all sides with the given "pad" value. + + If the input is a :class:`torch.Tensor` or a ``TVTensor`` (e.g. :class:`~torchvision.tv_tensors.Image`, + :class:`~torchvision.tv_tensors.Video`, :class:`~torchvision.tv_tensors.BoundingBoxes` etc.) + it can have arbitrary number of leading batch dimensions. For example, + the image can have ``[..., C, H, W]`` shape. A bounding box can have ``[..., 4]`` shape. + + Args: + padding (int or sequence): Padding on each border. If a single int is provided this + is used to pad all borders. If sequence of length 2 is provided this is the padding + on left/right and top/bottom respectively. If a sequence of length 4 is provided + this is the padding for the left, top, right and bottom borders respectively. + + .. note:: + In torchscript mode padding as single int is not supported, use a sequence of + length 1: ``[padding, ]``. + fill (number or tuple or dict, optional): Pixel fill value used when the ``padding_mode`` is constant. + Default is 0. If a tuple of length 3, it is used to fill R, G, B channels respectively. + Fill value can be also a dictionary mapping data type to the fill value, e.g. + ``fill={tv_tensors.Image: 127, tv_tensors.Mask: 0}`` where ``Image`` will be filled with 127 and + ``Mask`` will be filled with 0. + padding_mode (str, optional): Type of padding. Should be: constant, edge, reflect or symmetric. + Default is "constant". + + - constant: pads with a constant value, this value is specified with fill + + - edge: pads with the last value at the edge of the image. + + - reflect: pads with reflection of image without repeating the last value on the edge. + For example, padding [1, 2, 3, 4] with 2 elements on both sides in reflect mode + will result in [3, 2, 1, 2, 3, 4, 3, 2] + + - symmetric: pads with reflection of image repeating the last value on the edge. + For example, padding [1, 2, 3, 4] with 2 elements on both sides in symmetric mode + will result in [2, 1, 1, 2, 3, 4, 4, 3] + """ + + _v1_transform_cls = _transforms.Pad + + def _extract_params_for_v1_transform(self) -> Dict[str, Any]: + params = super()._extract_params_for_v1_transform() + + if not (params["fill"] is None or isinstance(params["fill"], (int, float))): + raise ValueError(f"{type(self).__name__}() can only be scripted for a scalar `fill`, but got {self.fill}.") + + return params + + def __init__( + self, + padding: Union[int, Sequence[int]], + fill: Union[_FillType, Dict[Union[Type, str], _FillType]] = 0, + padding_mode: Literal["constant", "edge", "reflect", "symmetric"] = "constant", + ) -> None: + super().__init__() + + _check_padding_arg(padding) + _check_padding_mode_arg(padding_mode) + + # This cast does Sequence[int] -> List[int] and is required to make mypy happy + if not isinstance(padding, int): + padding = list(padding) + self.padding = padding + self.fill = fill + self._fill = _setup_fill_arg(fill) + self.padding_mode = padding_mode + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + fill = _get_fill(self._fill, type(inpt)) + return self._call_kernel(F.pad, inpt, padding=self.padding, fill=fill, padding_mode=self.padding_mode) # type: ignore[arg-type] + + +class RandomZoomOut(_RandomApplyTransform): + """ "Zoom out" transformation from + `"SSD: Single Shot MultiBox Detector" `_. + + This transformation randomly pads images, videos, bounding boxes and masks creating a zoom out effect. + Output spatial size is randomly sampled from original size up to a maximum size configured + with ``side_range`` parameter: + + .. code-block:: python + + r = uniform_sample(side_range[0], side_range[1]) + output_width = input_width * r + output_height = input_height * r + + If the input is a :class:`torch.Tensor` or a ``TVTensor`` (e.g. :class:`~torchvision.tv_tensors.Image`, + :class:`~torchvision.tv_tensors.Video`, :class:`~torchvision.tv_tensors.BoundingBoxes` etc.) + it can have arbitrary number of leading batch dimensions. For example, + the image can have ``[..., C, H, W]`` shape. A bounding box can have ``[..., 4]`` shape. + + Args: + fill (number or tuple or dict, optional): Pixel fill value used when the ``padding_mode`` is constant. + Default is 0. If a tuple of length 3, it is used to fill R, G, B channels respectively. + Fill value can be also a dictionary mapping data type to the fill value, e.g. + ``fill={tv_tensors.Image: 127, tv_tensors.Mask: 0}`` where ``Image`` will be filled with 127 and + ``Mask`` will be filled with 0. + side_range (sequence of floats, optional): tuple of two floats defines minimum and maximum factors to + scale the input size. + p (float, optional): probability that the zoom operation will be performed. + """ + + def __init__( + self, + fill: Union[_FillType, Dict[Union[Type, str], _FillType]] = 0, + side_range: Sequence[float] = (1.0, 4.0), + p: float = 0.5, + ) -> None: + super().__init__(p=p) + + self.fill = fill + self._fill = _setup_fill_arg(fill) + + _check_sequence_input(side_range, "side_range", req_sizes=(2,)) + + self.side_range = side_range + if side_range[0] < 1.0 or side_range[0] > side_range[1]: + raise ValueError(f"Invalid side range provided {side_range}.") + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + orig_h, orig_w = query_size(flat_inputs) + + r = self.side_range[0] + torch.rand(1) * (self.side_range[1] - self.side_range[0]) + canvas_width = int(orig_w * r) + canvas_height = int(orig_h * r) + + r = torch.rand(2) + left = int((canvas_width - orig_w) * r[0]) + top = int((canvas_height - orig_h) * r[1]) + right = canvas_width - (left + orig_w) + bottom = canvas_height - (top + orig_h) + padding = [left, top, right, bottom] + + return dict(padding=padding) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + fill = _get_fill(self._fill, type(inpt)) + return self._call_kernel(F.pad, inpt, **params, fill=fill) + + +class RandomRotation(Transform): + """Rotate the input by angle. + + If the input is a :class:`torch.Tensor` or a ``TVTensor`` (e.g. :class:`~torchvision.tv_tensors.Image`, + :class:`~torchvision.tv_tensors.Video`, :class:`~torchvision.tv_tensors.BoundingBoxes` etc.) + it can have arbitrary number of leading batch dimensions. For example, + the image can have ``[..., C, H, W]`` shape. A bounding box can have ``[..., 4]`` shape. + + Args: + degrees (sequence or number): Range of degrees to select from. + If degrees is a number instead of sequence like (min, max), the range of degrees + will be (-degrees, +degrees). + interpolation (InterpolationMode, optional): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + expand (bool, optional): Optional expansion flag. + If true, expands the output to make it large enough to hold the entire rotated image. + If false or omitted, make the output image the same size as the input image. + Note that the expand flag assumes rotation around the center (see note below) and no translation. + center (sequence, optional): Optional center of rotation, (x, y). Origin is the upper left corner. + Default is the center of the image. + + .. note:: + + In theory, setting ``center`` has no effect if ``expand=True``, since the image center will become the + center of rotation. In practice however, due to numerical precision, this can lead to off-by-one + differences of the resulting image size compared to using the image center in the first place. Thus, when + setting ``expand=True``, it's best to leave ``center=None`` (default). + fill (number or tuple or dict, optional): Pixel fill value used when the ``padding_mode`` is constant. + Default is 0. If a tuple of length 3, it is used to fill R, G, B channels respectively. + Fill value can be also a dictionary mapping data type to the fill value, e.g. + ``fill={tv_tensors.Image: 127, tv_tensors.Mask: 0}`` where ``Image`` will be filled with 127 and + ``Mask`` will be filled with 0. + + .. _filters: https://pillow.readthedocs.io/en/latest/handbook/concepts.html#filters + + """ + + _v1_transform_cls = _transforms.RandomRotation + + def __init__( + self, + degrees: Union[numbers.Number, Sequence], + interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST, + expand: bool = False, + center: Optional[List[float]] = None, + fill: Union[_FillType, Dict[Union[Type, str], _FillType]] = 0, + ) -> None: + super().__init__() + self.degrees = _setup_angle(degrees, name="degrees", req_sizes=(2,)) + self.interpolation = interpolation + self.expand = expand + + self.fill = fill + self._fill = _setup_fill_arg(fill) + + if center is not None: + _check_sequence_input(center, "center", req_sizes=(2,)) + + self.center = center + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + angle = torch.empty(1).uniform_(self.degrees[0], self.degrees[1]).item() + return dict(angle=angle) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + fill = _get_fill(self._fill, type(inpt)) + return self._call_kernel( + F.rotate, + inpt, + **params, + interpolation=self.interpolation, + expand=self.expand, + center=self.center, + fill=fill, + ) + + +class RandomAffine(Transform): + """Random affine transformation the input keeping center invariant. + + If the input is a :class:`torch.Tensor` or a ``TVTensor`` (e.g. :class:`~torchvision.tv_tensors.Image`, + :class:`~torchvision.tv_tensors.Video`, :class:`~torchvision.tv_tensors.BoundingBoxes` etc.) + it can have arbitrary number of leading batch dimensions. For example, + the image can have ``[..., C, H, W]`` shape. A bounding box can have ``[..., 4]`` shape. + + Args: + degrees (sequence or number): Range of degrees to select from. + If degrees is a number instead of sequence like (min, max), the range of degrees + will be (-degrees, +degrees). Set to 0 to deactivate rotations. + translate (tuple, optional): tuple of maximum absolute fraction for horizontal + and vertical translations. For example translate=(a, b), then horizontal shift + is randomly sampled in the range -img_width * a < dx < img_width * a and vertical shift is + randomly sampled in the range -img_height * b < dy < img_height * b. Will not translate by default. + scale (tuple, optional): scaling factor interval, e.g (a, b), then scale is + randomly sampled from the range a <= scale <= b. Will keep original scale by default. + shear (sequence or number, optional): Range of degrees to select from. + If shear is a number, a shear parallel to the x-axis in the range (-shear, +shear) + will be applied. Else if shear is a sequence of 2 values a shear parallel to the x-axis in the + range (shear[0], shear[1]) will be applied. Else if shear is a sequence of 4 values, + an x-axis shear in (shear[0], shear[1]) and y-axis shear in (shear[2], shear[3]) will be applied. + Will not apply shear by default. + interpolation (InterpolationMode, optional): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + fill (number or tuple or dict, optional): Pixel fill value used when the ``padding_mode`` is constant. + Default is 0. If a tuple of length 3, it is used to fill R, G, B channels respectively. + Fill value can be also a dictionary mapping data type to the fill value, e.g. + ``fill={tv_tensors.Image: 127, tv_tensors.Mask: 0}`` where ``Image`` will be filled with 127 and + ``Mask`` will be filled with 0. + center (sequence, optional): Optional center of rotation, (x, y). Origin is the upper left corner. + Default is the center of the image. + + .. _filters: https://pillow.readthedocs.io/en/latest/handbook/concepts.html#filters + + """ + + _v1_transform_cls = _transforms.RandomAffine + + def __init__( + self, + degrees: Union[numbers.Number, Sequence], + translate: Optional[Sequence[float]] = None, + scale: Optional[Sequence[float]] = None, + shear: Optional[Union[int, float, Sequence[float]]] = None, + interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST, + fill: Union[_FillType, Dict[Union[Type, str], _FillType]] = 0, + center: Optional[List[float]] = None, + ) -> None: + super().__init__() + self.degrees = _setup_angle(degrees, name="degrees", req_sizes=(2,)) + if translate is not None: + _check_sequence_input(translate, "translate", req_sizes=(2,)) + for t in translate: + if not (0.0 <= t <= 1.0): + raise ValueError("translation values should be between 0 and 1") + self.translate = translate + if scale is not None: + _check_sequence_input(scale, "scale", req_sizes=(2,)) + for s in scale: + if s <= 0: + raise ValueError("scale values should be positive") + self.scale = scale + + if shear is not None: + self.shear = _setup_angle(shear, name="shear", req_sizes=(2, 4)) + else: + self.shear = shear + + self.interpolation = interpolation + self.fill = fill + self._fill = _setup_fill_arg(fill) + + if center is not None: + _check_sequence_input(center, "center", req_sizes=(2,)) + + self.center = center + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + height, width = query_size(flat_inputs) + + angle = torch.empty(1).uniform_(self.degrees[0], self.degrees[1]).item() + if self.translate is not None: + max_dx = float(self.translate[0] * width) + max_dy = float(self.translate[1] * height) + tx = int(round(torch.empty(1).uniform_(-max_dx, max_dx).item())) + ty = int(round(torch.empty(1).uniform_(-max_dy, max_dy).item())) + translate = (tx, ty) + else: + translate = (0, 0) + + if self.scale is not None: + scale = torch.empty(1).uniform_(self.scale[0], self.scale[1]).item() + else: + scale = 1.0 + + shear_x = shear_y = 0.0 + if self.shear is not None: + shear_x = torch.empty(1).uniform_(self.shear[0], self.shear[1]).item() + if len(self.shear) == 4: + shear_y = torch.empty(1).uniform_(self.shear[2], self.shear[3]).item() + + shear = (shear_x, shear_y) + return dict(angle=angle, translate=translate, scale=scale, shear=shear) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + fill = _get_fill(self._fill, type(inpt)) + return self._call_kernel( + F.affine, + inpt, + **params, + interpolation=self.interpolation, + fill=fill, + center=self.center, + ) + + +class RandomCrop(Transform): + """Crop the input at a random location. + + If the input is a :class:`torch.Tensor` or a ``TVTensor`` (e.g. :class:`~torchvision.tv_tensors.Image`, + :class:`~torchvision.tv_tensors.Video`, :class:`~torchvision.tv_tensors.BoundingBoxes` etc.) + it can have arbitrary number of leading batch dimensions. For example, + the image can have ``[..., C, H, W]`` shape. A bounding box can have ``[..., 4]`` shape. + + Args: + size (sequence or int): Desired output size of the crop. If size is an + int instead of sequence like (h, w), a square crop (size, size) is + made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). + padding (int or sequence, optional): Optional padding on each border + of the image. Default is None. If a single int is provided this + is used to pad all borders. If sequence of length 2 is provided this is the padding + on left/right and top/bottom respectively. If a sequence of length 4 is provided + this is the padding for the left, top, right and bottom borders respectively. + + .. note:: + In torchscript mode padding as single int is not supported, use a sequence of + length 1: ``[padding, ]``. + pad_if_needed (boolean, optional): It will pad the image if smaller than the + desired size to avoid raising an exception. Since cropping is done + after padding, the padding seems to be done at a random offset. + fill (number or tuple or dict, optional): Pixel fill value used when the ``padding_mode`` is constant. + Default is 0. If a tuple of length 3, it is used to fill R, G, B channels respectively. + Fill value can be also a dictionary mapping data type to the fill value, e.g. + ``fill={tv_tensors.Image: 127, tv_tensors.Mask: 0}`` where ``Image`` will be filled with 127 and + ``Mask`` will be filled with 0. + padding_mode (str, optional): Type of padding. Should be: constant, edge, reflect or symmetric. + Default is constant. + + - constant: pads with a constant value, this value is specified with fill + + - edge: pads with the last value at the edge of the image. + + - reflect: pads with reflection of image without repeating the last value on the edge. + For example, padding [1, 2, 3, 4] with 2 elements on both sides in reflect mode + will result in [3, 2, 1, 2, 3, 4, 3, 2] + + - symmetric: pads with reflection of image repeating the last value on the edge. + For example, padding [1, 2, 3, 4] with 2 elements on both sides in symmetric mode + will result in [2, 1, 1, 2, 3, 4, 4, 3] + """ + + _v1_transform_cls = _transforms.RandomCrop + + def _extract_params_for_v1_transform(self) -> Dict[str, Any]: + params = super()._extract_params_for_v1_transform() + + if not (params["fill"] is None or isinstance(params["fill"], (int, float))): + raise ValueError(f"{type(self).__name__}() can only be scripted for a scalar `fill`, but got {self.fill}.") + + padding = self.padding + if padding is not None: + pad_left, pad_right, pad_top, pad_bottom = padding + padding = [pad_left, pad_top, pad_right, pad_bottom] + params["padding"] = padding + + return params + + def __init__( + self, + size: Union[int, Sequence[int]], + padding: Optional[Union[int, Sequence[int]]] = None, + pad_if_needed: bool = False, + fill: Union[_FillType, Dict[Union[Type, str], _FillType]] = 0, + padding_mode: Literal["constant", "edge", "reflect", "symmetric"] = "constant", + ) -> None: + super().__init__() + + self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.") + + if pad_if_needed or padding is not None: + if padding is not None: + _check_padding_arg(padding) + _check_padding_mode_arg(padding_mode) + + self.padding = F._geometry._parse_pad_padding(padding) if padding else None # type: ignore[arg-type] + self.pad_if_needed = pad_if_needed + self.fill = fill + self._fill = _setup_fill_arg(fill) + self.padding_mode = padding_mode + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + padded_height, padded_width = query_size(flat_inputs) + + if self.padding is not None: + pad_left, pad_right, pad_top, pad_bottom = self.padding + padded_height += pad_top + pad_bottom + padded_width += pad_left + pad_right + else: + pad_left = pad_right = pad_top = pad_bottom = 0 + + cropped_height, cropped_width = self.size + + if self.pad_if_needed: + if padded_height < cropped_height: + diff = cropped_height - padded_height + + pad_top += diff + pad_bottom += diff + padded_height += 2 * diff + + if padded_width < cropped_width: + diff = cropped_width - padded_width + + pad_left += diff + pad_right += diff + padded_width += 2 * diff + + if padded_height < cropped_height or padded_width < cropped_width: + raise ValueError( + f"Required crop size {(cropped_height, cropped_width)} is larger than " + f"{'padded ' if self.padding is not None else ''}input image size {(padded_height, padded_width)}." + ) + + # We need a different order here than we have in self.padding since this padding will be parsed again in `F.pad` + padding = [pad_left, pad_top, pad_right, pad_bottom] + needs_pad = any(padding) + + needs_vert_crop, top = ( + (True, int(torch.randint(0, padded_height - cropped_height + 1, size=()))) + if padded_height > cropped_height + else (False, 0) + ) + needs_horz_crop, left = ( + (True, int(torch.randint(0, padded_width - cropped_width + 1, size=()))) + if padded_width > cropped_width + else (False, 0) + ) + + return dict( + needs_crop=needs_vert_crop or needs_horz_crop, + top=top, + left=left, + height=cropped_height, + width=cropped_width, + needs_pad=needs_pad, + padding=padding, + ) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + if params["needs_pad"]: + fill = _get_fill(self._fill, type(inpt)) + inpt = self._call_kernel(F.pad, inpt, padding=params["padding"], fill=fill, padding_mode=self.padding_mode) + + if params["needs_crop"]: + inpt = self._call_kernel( + F.crop, inpt, top=params["top"], left=params["left"], height=params["height"], width=params["width"] + ) + + return inpt + + +class RandomPerspective(_RandomApplyTransform): + """Perform a random perspective transformation of the input with a given probability. + + If the input is a :class:`torch.Tensor` or a ``TVTensor`` (e.g. :class:`~torchvision.tv_tensors.Image`, + :class:`~torchvision.tv_tensors.Video`, :class:`~torchvision.tv_tensors.BoundingBoxes` etc.) + it can have arbitrary number of leading batch dimensions. For example, + the image can have ``[..., C, H, W]`` shape. A bounding box can have ``[..., 4]`` shape. + + Args: + distortion_scale (float, optional): argument to control the degree of distortion and ranges from 0 to 1. + Default is 0.5. + p (float, optional): probability of the input being transformed. Default is 0.5. + interpolation (InterpolationMode, optional): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + fill (number or tuple or dict, optional): Pixel fill value used when the ``padding_mode`` is constant. + Default is 0. If a tuple of length 3, it is used to fill R, G, B channels respectively. + Fill value can be also a dictionary mapping data type to the fill value, e.g. + ``fill={tv_tensors.Image: 127, tv_tensors.Mask: 0}`` where ``Image`` will be filled with 127 and + ``Mask`` will be filled with 0. + """ + + _v1_transform_cls = _transforms.RandomPerspective + + def __init__( + self, + distortion_scale: float = 0.5, + p: float = 0.5, + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + fill: Union[_FillType, Dict[Union[Type, str], _FillType]] = 0, + ) -> None: + super().__init__(p=p) + + if not (0 <= distortion_scale <= 1): + raise ValueError("Argument distortion_scale value should be between 0 and 1") + + self.distortion_scale = distortion_scale + self.interpolation = interpolation + self.fill = fill + self._fill = _setup_fill_arg(fill) + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + height, width = query_size(flat_inputs) + + distortion_scale = self.distortion_scale + + half_height = height // 2 + half_width = width // 2 + bound_height = int(distortion_scale * half_height) + 1 + bound_width = int(distortion_scale * half_width) + 1 + topleft = [ + int(torch.randint(0, bound_width, size=(1,))), + int(torch.randint(0, bound_height, size=(1,))), + ] + topright = [ + int(torch.randint(width - bound_width, width, size=(1,))), + int(torch.randint(0, bound_height, size=(1,))), + ] + botright = [ + int(torch.randint(width - bound_width, width, size=(1,))), + int(torch.randint(height - bound_height, height, size=(1,))), + ] + botleft = [ + int(torch.randint(0, bound_width, size=(1,))), + int(torch.randint(height - bound_height, height, size=(1,))), + ] + startpoints = [[0, 0], [width - 1, 0], [width - 1, height - 1], [0, height - 1]] + endpoints = [topleft, topright, botright, botleft] + perspective_coeffs = _get_perspective_coeffs(startpoints, endpoints) + return dict(coefficients=perspective_coeffs) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + fill = _get_fill(self._fill, type(inpt)) + return self._call_kernel( + F.perspective, + inpt, + startpoints=None, + endpoints=None, + fill=fill, + interpolation=self.interpolation, + **params, + ) + + +class ElasticTransform(Transform): + """Transform the input with elastic transformations. + + If the input is a :class:`torch.Tensor` or a ``TVTensor`` (e.g. :class:`~torchvision.tv_tensors.Image`, + :class:`~torchvision.tv_tensors.Video`, :class:`~torchvision.tv_tensors.BoundingBoxes` etc.) + it can have arbitrary number of leading batch dimensions. For example, + the image can have ``[..., C, H, W]`` shape. A bounding box can have ``[..., 4]`` shape. + + Given alpha and sigma, it will generate displacement + vectors for all pixels based on random offsets. Alpha controls the strength + and sigma controls the smoothness of the displacements. + The displacements are added to an identity grid and the resulting grid is + used to transform the input. + + .. note:: + Implementation to transform bounding boxes is approximative (not exact). + We construct an approximation of the inverse grid as ``inverse_grid = identity - displacement``. + This is not an exact inverse of the grid used to transform images, i.e. ``grid = identity + displacement``. + Our assumption is that ``displacement * displacement`` is small and can be ignored. + Large displacements would lead to large errors in the approximation. + + Applications: + Randomly transforms the morphology of objects in images and produces a + see-through-water-like effect. + + Args: + alpha (float or sequence of floats, optional): Magnitude of displacements. Default is 50.0. + sigma (float or sequence of floats, optional): Smoothness of displacements. Default is 5.0. + interpolation (InterpolationMode, optional): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + fill (number or tuple or dict, optional): Pixel fill value used when the ``padding_mode`` is constant. + Default is 0. If a tuple of length 3, it is used to fill R, G, B channels respectively. + Fill value can be also a dictionary mapping data type to the fill value, e.g. + ``fill={tv_tensors.Image: 127, tv_tensors.Mask: 0}`` where ``Image`` will be filled with 127 and + ``Mask`` will be filled with 0. + """ + + _v1_transform_cls = _transforms.ElasticTransform + + def __init__( + self, + alpha: Union[float, Sequence[float]] = 50.0, + sigma: Union[float, Sequence[float]] = 5.0, + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + fill: Union[_FillType, Dict[Union[Type, str], _FillType]] = 0, + ) -> None: + super().__init__() + self.alpha = _setup_number_or_seq(alpha, "alpha") + self.sigma = _setup_number_or_seq(sigma, "sigma") + + self.interpolation = interpolation + self.fill = fill + self._fill = _setup_fill_arg(fill) + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + size = list(query_size(flat_inputs)) + + dx = torch.rand([1, 1] + size) * 2 - 1 + if self.sigma[0] > 0.0: + kx = int(8 * self.sigma[0] + 1) + # if kernel size is even we have to make it odd + if kx % 2 == 0: + kx += 1 + dx = self._call_kernel(F.gaussian_blur, dx, [kx, kx], list(self.sigma)) + dx = dx * self.alpha[0] / size[0] + + dy = torch.rand([1, 1] + size) * 2 - 1 + if self.sigma[1] > 0.0: + ky = int(8 * self.sigma[1] + 1) + # if kernel size is even we have to make it odd + if ky % 2 == 0: + ky += 1 + dy = self._call_kernel(F.gaussian_blur, dy, [ky, ky], list(self.sigma)) + dy = dy * self.alpha[1] / size[1] + displacement = torch.concat([dx, dy], 1).permute([0, 2, 3, 1]) # 1 x H x W x 2 + return dict(displacement=displacement) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + fill = _get_fill(self._fill, type(inpt)) + return self._call_kernel( + F.elastic, + inpt, + **params, + fill=fill, + interpolation=self.interpolation, + ) + + +class RandomIoUCrop(Transform): + """Random IoU crop transformation from + `"SSD: Single Shot MultiBox Detector" `_. + + This transformation requires an image or video data and ``tv_tensors.BoundingBoxes`` in the input. + + .. warning:: + In order to properly remove the bounding boxes below the IoU threshold, `RandomIoUCrop` + must be followed by :class:`~torchvision.transforms.v2.SanitizeBoundingBoxes`, either immediately + after or later in the transforms pipeline. + + If the input is a :class:`torch.Tensor` or a ``TVTensor`` (e.g. :class:`~torchvision.tv_tensors.Image`, + :class:`~torchvision.tv_tensors.Video`, :class:`~torchvision.tv_tensors.BoundingBoxes` etc.) + it can have arbitrary number of leading batch dimensions. For example, + the image can have ``[..., C, H, W]`` shape. A bounding box can have ``[..., 4]`` shape. + + Args: + min_scale (float, optional): Minimum factors to scale the input size. + max_scale (float, optional): Maximum factors to scale the input size. + min_aspect_ratio (float, optional): Minimum aspect ratio for the cropped image or video. + max_aspect_ratio (float, optional): Maximum aspect ratio for the cropped image or video. + sampler_options (list of float, optional): List of minimal IoU (Jaccard) overlap between all the boxes and + a cropped image or video. Default, ``None`` which corresponds to ``[0.0, 0.1, 0.3, 0.5, 0.7, 0.9, 1.0]`` + trials (int, optional): Number of trials to find a crop for a given value of minimal IoU (Jaccard) overlap. + Default, 40. + """ + + def __init__( + self, + min_scale: float = 0.3, + max_scale: float = 1.0, + min_aspect_ratio: float = 0.5, + max_aspect_ratio: float = 2.0, + sampler_options: Optional[List[float]] = None, + trials: int = 40, + ): + super().__init__() + # Configuration similar to https://github.com/weiliu89/caffe/blob/ssd/examples/ssd/ssd_coco.py#L89-L174 + self.min_scale = min_scale + self.max_scale = max_scale + self.min_aspect_ratio = min_aspect_ratio + self.max_aspect_ratio = max_aspect_ratio + if sampler_options is None: + sampler_options = [0.0, 0.1, 0.3, 0.5, 0.7, 0.9, 1.0] + self.options = sampler_options + self.trials = trials + + def check_inputs(self, flat_inputs: List[Any]) -> None: + if not ( + has_all(flat_inputs, tv_tensors.BoundingBoxes) + and has_any(flat_inputs, PIL.Image.Image, tv_tensors.Image, is_pure_tensor) + ): + raise TypeError( + f"{type(self).__name__}() requires input sample to contain tensor or PIL images " + "and bounding boxes. Sample can also contain masks." + ) + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + orig_h, orig_w = query_size(flat_inputs) + bboxes = get_bounding_boxes(flat_inputs) + + while True: + # sample an option + idx = int(torch.randint(low=0, high=len(self.options), size=(1,))) + min_jaccard_overlap = self.options[idx] + if min_jaccard_overlap >= 1.0: # a value larger than 1 encodes the leave as-is option + return dict() + + for _ in range(self.trials): + # check the aspect ratio limitations + r = self.min_scale + (self.max_scale - self.min_scale) * torch.rand(2) + new_w = int(orig_w * r[0]) + new_h = int(orig_h * r[1]) + aspect_ratio = new_w / new_h + if not (self.min_aspect_ratio <= aspect_ratio <= self.max_aspect_ratio): + continue + + # check for 0 area crops + r = torch.rand(2) + left = int((orig_w - new_w) * r[0]) + top = int((orig_h - new_h) * r[1]) + right = left + new_w + bottom = top + new_h + if left == right or top == bottom: + continue + + # check for any valid boxes with centers within the crop area + xyxy_bboxes = F.convert_bounding_box_format( + bboxes.as_subclass(torch.Tensor), + bboxes.format, + tv_tensors.BoundingBoxFormat.XYXY, + ) + cx = 0.5 * (xyxy_bboxes[..., 0] + xyxy_bboxes[..., 2]) + cy = 0.5 * (xyxy_bboxes[..., 1] + xyxy_bboxes[..., 3]) + is_within_crop_area = (left < cx) & (cx < right) & (top < cy) & (cy < bottom) + if not is_within_crop_area.any(): + continue + + # check at least 1 box with jaccard limitations + xyxy_bboxes = xyxy_bboxes[is_within_crop_area] + ious = box_iou( + xyxy_bboxes, + torch.tensor([[left, top, right, bottom]], dtype=xyxy_bboxes.dtype, device=xyxy_bboxes.device), + ) + if ious.max() < min_jaccard_overlap: + continue + + return dict(top=top, left=left, height=new_h, width=new_w, is_within_crop_area=is_within_crop_area) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + + if len(params) < 1: + return inpt + + output = self._call_kernel( + F.crop, inpt, top=params["top"], left=params["left"], height=params["height"], width=params["width"] + ) + + if isinstance(output, tv_tensors.BoundingBoxes): + # We "mark" the invalid boxes as degenreate, and they can be + # removed by a later call to SanitizeBoundingBoxes() + output[~params["is_within_crop_area"]] = 0 + + return output + + +class ScaleJitter(Transform): + """Perform Large Scale Jitter on the input according to + `"Simple Copy-Paste is a Strong Data Augmentation Method for Instance Segmentation" `_. + + If the input is a :class:`torch.Tensor` or a ``TVTensor`` (e.g. :class:`~torchvision.tv_tensors.Image`, + :class:`~torchvision.tv_tensors.Video`, :class:`~torchvision.tv_tensors.BoundingBoxes` etc.) + it can have arbitrary number of leading batch dimensions. For example, + the image can have ``[..., C, H, W]`` shape. A bounding box can have ``[..., 4]`` shape. + + Args: + target_size (tuple of int): Target size. This parameter defines base scale for jittering, + e.g. ``min(target_size[0] / width, target_size[1] / height)``. + scale_range (tuple of float, optional): Minimum and maximum of the scale range. Default, ``(0.1, 2.0)``. + interpolation (InterpolationMode, optional): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.NEAREST_EXACT``, + ``InterpolationMode.BILINEAR`` and ``InterpolationMode.BICUBIC`` are supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + antialias (bool, optional): Whether to apply antialiasing. + It only affects **tensors** with bilinear or bicubic modes and it is + ignored otherwise: on PIL images, antialiasing is always applied on + bilinear or bicubic modes; on other modes (for PIL images and + tensors), antialiasing makes no sense and this parameter is ignored. + Possible values are: + + - ``True`` (default): will apply antialiasing for bilinear or bicubic modes. + Other mode aren't affected. This is probably what you want to use. + - ``False``: will not apply antialiasing for tensors on any mode. PIL + images are still antialiased on bilinear or bicubic modes, because + PIL doesn't support no antialias. + - ``None``: equivalent to ``False`` for tensors and ``True`` for + PIL images. This value exists for legacy reasons and you probably + don't want to use it unless you really know what you are doing. + + The default value changed from ``None`` to ``True`` in + v0.17, for the PIL and Tensor backends to be consistent. + """ + + def __init__( + self, + target_size: Tuple[int, int], + scale_range: Tuple[float, float] = (0.1, 2.0), + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + antialias: Optional[bool] = True, + ): + super().__init__() + self.target_size = target_size + self.scale_range = scale_range + self.interpolation = interpolation + self.antialias = antialias + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + orig_height, orig_width = query_size(flat_inputs) + + scale = self.scale_range[0] + torch.rand(1) * (self.scale_range[1] - self.scale_range[0]) + r = min(self.target_size[1] / orig_height, self.target_size[0] / orig_width) * scale + new_width = int(orig_width * r) + new_height = int(orig_height * r) + + return dict(size=(new_height, new_width)) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel( + F.resize, inpt, size=params["size"], interpolation=self.interpolation, antialias=self.antialias + ) + + +class RandomShortestSize(Transform): + """Randomly resize the input. + + If the input is a :class:`torch.Tensor` or a ``TVTensor`` (e.g. :class:`~torchvision.tv_tensors.Image`, + :class:`~torchvision.tv_tensors.Video`, :class:`~torchvision.tv_tensors.BoundingBoxes` etc.) + it can have arbitrary number of leading batch dimensions. For example, + the image can have ``[..., C, H, W]`` shape. A bounding box can have ``[..., 4]`` shape. + + Args: + min_size (int or sequence of int): Minimum spatial size. Single integer value or a sequence of integer values. + max_size (int, optional): Maximum spatial size. Default, None. + interpolation (InterpolationMode, optional): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.NEAREST_EXACT``, + ``InterpolationMode.BILINEAR`` and ``InterpolationMode.BICUBIC`` are supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + antialias (bool, optional): Whether to apply antialiasing. + It only affects **tensors** with bilinear or bicubic modes and it is + ignored otherwise: on PIL images, antialiasing is always applied on + bilinear or bicubic modes; on other modes (for PIL images and + tensors), antialiasing makes no sense and this parameter is ignored. + Possible values are: + + - ``True`` (default): will apply antialiasing for bilinear or bicubic modes. + Other mode aren't affected. This is probably what you want to use. + - ``False``: will not apply antialiasing for tensors on any mode. PIL + images are still antialiased on bilinear or bicubic modes, because + PIL doesn't support no antialias. + - ``None``: equivalent to ``False`` for tensors and ``True`` for + PIL images. This value exists for legacy reasons and you probably + don't want to use it unless you really know what you are doing. + + The default value changed from ``None`` to ``True`` in + v0.17, for the PIL and Tensor backends to be consistent. + """ + + def __init__( + self, + min_size: Union[List[int], Tuple[int], int], + max_size: Optional[int] = None, + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + antialias: Optional[bool] = True, + ): + super().__init__() + self.min_size = [min_size] if isinstance(min_size, int) else list(min_size) + self.max_size = max_size + self.interpolation = interpolation + self.antialias = antialias + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + orig_height, orig_width = query_size(flat_inputs) + + min_size = self.min_size[int(torch.randint(len(self.min_size), ()))] + r = min_size / min(orig_height, orig_width) + if self.max_size is not None: + r = min(r, self.max_size / max(orig_height, orig_width)) + + new_width = int(orig_width * r) + new_height = int(orig_height * r) + + return dict(size=(new_height, new_width)) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel( + F.resize, inpt, size=params["size"], interpolation=self.interpolation, antialias=self.antialias + ) + + +class RandomResize(Transform): + """Randomly resize the input. + + This transformation can be used together with ``RandomCrop`` as data augmentations to train + models on image segmentation task. + + Output spatial size is randomly sampled from the interval ``[min_size, max_size]``: + + .. code-block:: python + + size = uniform_sample(min_size, max_size) + output_width = size + output_height = size + + If the input is a :class:`torch.Tensor` or a ``TVTensor`` (e.g. :class:`~torchvision.tv_tensors.Image`, + :class:`~torchvision.tv_tensors.Video`, :class:`~torchvision.tv_tensors.BoundingBoxes` etc.) + it can have arbitrary number of leading batch dimensions. For example, + the image can have ``[..., C, H, W]`` shape. A bounding box can have ``[..., 4]`` shape. + + Args: + min_size (int): Minimum output size for random sampling + max_size (int): Maximum output size for random sampling + interpolation (InterpolationMode, optional): Desired interpolation enum defined by + :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``. + If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.NEAREST_EXACT``, + ``InterpolationMode.BILINEAR`` and ``InterpolationMode.BICUBIC`` are supported. + The corresponding Pillow integer constants, e.g. ``PIL.Image.BILINEAR`` are accepted as well. + antialias (bool, optional): Whether to apply antialiasing. + It only affects **tensors** with bilinear or bicubic modes and it is + ignored otherwise: on PIL images, antialiasing is always applied on + bilinear or bicubic modes; on other modes (for PIL images and + tensors), antialiasing makes no sense and this parameter is ignored. + Possible values are: + + - ``True`` (default): will apply antialiasing for bilinear or bicubic modes. + Other mode aren't affected. This is probably what you want to use. + - ``False``: will not apply antialiasing for tensors on any mode. PIL + images are still antialiased on bilinear or bicubic modes, because + PIL doesn't support no antialias. + - ``None``: equivalent to ``False`` for tensors and ``True`` for + PIL images. This value exists for legacy reasons and you probably + don't want to use it unless you really know what you are doing. + + The default value changed from ``None`` to ``True`` in + v0.17, for the PIL and Tensor backends to be consistent. + """ + + def __init__( + self, + min_size: int, + max_size: int, + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + antialias: Optional[bool] = True, + ) -> None: + super().__init__() + self.min_size = min_size + self.max_size = max_size + self.interpolation = interpolation + self.antialias = antialias + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + size = int(torch.randint(self.min_size, self.max_size, ())) + return dict(size=[size]) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel( + F.resize, inpt, params["size"], interpolation=self.interpolation, antialias=self.antialias + ) diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/_meta.py b/lib/python3.12/site-packages/torchvision/transforms/v2/_meta.py new file mode 100644 index 0000000000000000000000000000000000000000..1890b43115afb7eda76e13688daee5121029dc62 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/_meta.py @@ -0,0 +1,36 @@ +from typing import Any, Dict, Union + +from torchvision import tv_tensors +from torchvision.transforms.v2 import functional as F, Transform + + +class ConvertBoundingBoxFormat(Transform): + """Convert bounding box coordinates to the given ``format``, eg from "CXCYWH" to "XYXY". + + Args: + format (str or tv_tensors.BoundingBoxFormat): output bounding box format. + Possible values are defined by :class:`~torchvision.tv_tensors.BoundingBoxFormat` and + string values match the enums, e.g. "XYXY" or "XYWH" etc. + """ + + _transformed_types = (tv_tensors.BoundingBoxes,) + + def __init__(self, format: Union[str, tv_tensors.BoundingBoxFormat]) -> None: + super().__init__() + self.format = format + + def transform(self, inpt: tv_tensors.BoundingBoxes, params: Dict[str, Any]) -> tv_tensors.BoundingBoxes: + return F.convert_bounding_box_format(inpt, new_format=self.format) # type: ignore[return-value, arg-type] + + +class ClampBoundingBoxes(Transform): + """Clamp bounding boxes to their corresponding image dimensions. + + The clamping is done according to the bounding boxes' ``canvas_size`` meta-data. + + """ + + _transformed_types = (tv_tensors.BoundingBoxes,) + + def transform(self, inpt: tv_tensors.BoundingBoxes, params: Dict[str, Any]) -> tv_tensors.BoundingBoxes: + return F.clamp_bounding_boxes(inpt) # type: ignore[return-value] diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/_misc.py b/lib/python3.12/site-packages/torchvision/transforms/v2/_misc.py new file mode 100644 index 0000000000000000000000000000000000000000..d38a6ad8767b7e366564bff4a28c0838abd8379b --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/_misc.py @@ -0,0 +1,451 @@ +import warnings +from typing import Any, Callable, Dict, List, Optional, Sequence, Type, Union + +import PIL.Image + +import torch +from torch.utils._pytree import tree_flatten, tree_unflatten + +from torchvision import transforms as _transforms, tv_tensors +from torchvision.transforms.v2 import functional as F, Transform + +from ._utils import _parse_labels_getter, _setup_number_or_seq, _setup_size, get_bounding_boxes, has_any, is_pure_tensor + + +# TODO: do we want/need to expose this? +class Identity(Transform): + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return inpt + + +class Lambda(Transform): + """Apply a user-defined function as a transform. + + This transform does not support torchscript. + + Args: + lambd (function): Lambda/function to be used for transform. + """ + + _transformed_types = (object,) + + def __init__(self, lambd: Callable[[Any], Any], *types: Type): + super().__init__() + self.lambd = lambd + self.types = types or self._transformed_types + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + if isinstance(inpt, self.types): + return self.lambd(inpt) + else: + return inpt + + def extra_repr(self) -> str: + extras = [] + name = getattr(self.lambd, "__name__", None) + if name: + extras.append(name) + extras.append(f"types={[type.__name__ for type in self.types]}") + return ", ".join(extras) + + +class LinearTransformation(Transform): + """Transform a tensor image or video with a square transformation matrix and a mean_vector computed offline. + + This transform does not support PIL Image. + Given transformation_matrix and mean_vector, will flatten the torch.*Tensor and + subtract mean_vector from it which is then followed by computing the dot + product with the transformation matrix and then reshaping the tensor to its + original shape. + + Applications: + whitening transformation: Suppose X is a column vector zero-centered data. + Then compute the data covariance matrix [D x D] with torch.mm(X.t(), X), + perform SVD on this matrix and pass it as transformation_matrix. + + Args: + transformation_matrix (Tensor): tensor [D x D], D = C x H x W + mean_vector (Tensor): tensor [D], D = C x H x W + """ + + _v1_transform_cls = _transforms.LinearTransformation + + _transformed_types = (is_pure_tensor, tv_tensors.Image, tv_tensors.Video) + + def __init__(self, transformation_matrix: torch.Tensor, mean_vector: torch.Tensor): + super().__init__() + if transformation_matrix.size(0) != transformation_matrix.size(1): + raise ValueError( + "transformation_matrix should be square. Got " + f"{tuple(transformation_matrix.size())} rectangular matrix." + ) + + if mean_vector.size(0) != transformation_matrix.size(0): + raise ValueError( + f"mean_vector should have the same length {mean_vector.size(0)}" + f" as any one of the dimensions of the transformation_matrix [{tuple(transformation_matrix.size())}]" + ) + + if transformation_matrix.device != mean_vector.device: + raise ValueError( + f"Input tensors should be on the same device. Got {transformation_matrix.device} and {mean_vector.device}" + ) + + if transformation_matrix.dtype != mean_vector.dtype: + raise ValueError( + f"Input tensors should have the same dtype. Got {transformation_matrix.dtype} and {mean_vector.dtype}" + ) + + self.transformation_matrix = transformation_matrix + self.mean_vector = mean_vector + + def check_inputs(self, sample: Any) -> Any: + if has_any(sample, PIL.Image.Image): + raise TypeError(f"{type(self).__name__}() does not support PIL images.") + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + shape = inpt.shape + n = shape[-3] * shape[-2] * shape[-1] + if n != self.transformation_matrix.shape[0]: + raise ValueError( + "Input tensor and transformation matrix have incompatible shape." + + f"[{shape[-3]} x {shape[-2]} x {shape[-1]}] != " + + f"{self.transformation_matrix.shape[0]}" + ) + + if inpt.device.type != self.mean_vector.device.type: + raise ValueError( + "Input tensor should be on the same device as transformation matrix and mean vector. " + f"Got {inpt.device} vs {self.mean_vector.device}" + ) + + flat_inpt = inpt.reshape(-1, n) - self.mean_vector + + transformation_matrix = self.transformation_matrix.to(flat_inpt.dtype) + output = torch.mm(flat_inpt, transformation_matrix) + output = output.reshape(shape) + + if isinstance(inpt, (tv_tensors.Image, tv_tensors.Video)): + output = tv_tensors.wrap(output, like=inpt) + return output + + +class Normalize(Transform): + """Normalize a tensor image or video with mean and standard deviation. + + This transform does not support PIL Image. + Given mean: ``(mean[1],...,mean[n])`` and std: ``(std[1],..,std[n])`` for ``n`` + channels, this transform will normalize each channel of the input + ``torch.*Tensor`` i.e., + ``output[channel] = (input[channel] - mean[channel]) / std[channel]`` + + .. note:: + This transform acts out of place, i.e., it does not mutate the input tensor. + + Args: + mean (sequence): Sequence of means for each channel. + std (sequence): Sequence of standard deviations for each channel. + inplace(bool,optional): Bool to make this operation in-place. + + """ + + _v1_transform_cls = _transforms.Normalize + + def __init__(self, mean: Sequence[float], std: Sequence[float], inplace: bool = False): + super().__init__() + self.mean = list(mean) + self.std = list(std) + self.inplace = inplace + + def check_inputs(self, sample: Any) -> Any: + if has_any(sample, PIL.Image.Image): + raise TypeError(f"{type(self).__name__}() does not support PIL images.") + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.normalize, inpt, mean=self.mean, std=self.std, inplace=self.inplace) + + +class GaussianBlur(Transform): + """Blurs image with randomly chosen Gaussian blur kernel. + + The convolution will be using reflection padding corresponding to the kernel size, to maintain the input shape. + + If the input is a Tensor, it is expected + to have [..., C, H, W] shape, where ... means an arbitrary number of leading dimensions. + + Args: + kernel_size (int or sequence): Size of the Gaussian kernel. + sigma (float or tuple of float (min, max)): Standard deviation to be used for + creating kernel to perform blurring. If float, sigma is fixed. If it is tuple + of float (min, max), sigma is chosen uniformly at random to lie in the + given range. + """ + + _v1_transform_cls = _transforms.GaussianBlur + + def __init__( + self, kernel_size: Union[int, Sequence[int]], sigma: Union[int, float, Sequence[float]] = (0.1, 2.0) + ) -> None: + super().__init__() + self.kernel_size = _setup_size(kernel_size, "Kernel size should be a tuple/list of two integers") + for ks in self.kernel_size: + if ks <= 0 or ks % 2 == 0: + raise ValueError("Kernel size value should be an odd and positive number.") + + self.sigma = _setup_number_or_seq(sigma, "sigma") + + if not 0.0 < self.sigma[0] <= self.sigma[1]: + raise ValueError(f"sigma values should be positive and of the form (min, max). Got {self.sigma}") + + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + sigma = torch.empty(1).uniform_(self.sigma[0], self.sigma[1]).item() + return dict(sigma=[sigma, sigma]) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.gaussian_blur, inpt, self.kernel_size, **params) + + +class GaussianNoise(Transform): + """Add gaussian noise to images or videos. + + The input tensor is expected to be in [..., 1 or 3, H, W] format, + where ... means it can have an arbitrary number of leading dimensions. + Each image or frame in a batch will be transformed independently i.e. the + noise added to each image will be different. + + The input tensor is also expected to be of float dtype in ``[0, 1]``. + This transform does not support PIL images. + + Args: + mean (float): Mean of the sampled normal distribution. Default is 0. + sigma (float): Standard deviation of the sampled normal distribution. Default is 0.1. + clip (bool, optional): Whether to clip the values in ``[0, 1]`` after adding noise. Default is True. + """ + + def __init__(self, mean: float = 0.0, sigma: float = 0.1, clip=True) -> None: + super().__init__() + self.mean = mean + self.sigma = sigma + self.clip = clip + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.gaussian_noise, inpt, mean=self.mean, sigma=self.sigma, clip=self.clip) + + +class ToDtype(Transform): + """Converts the input to a specific dtype, optionally scaling the values for images or videos. + + .. note:: + ``ToDtype(dtype, scale=True)`` is the recommended replacement for ``ConvertImageDtype(dtype)``. + + Args: + dtype (``torch.dtype`` or dict of ``TVTensor`` -> ``torch.dtype``): The dtype to convert to. + If a ``torch.dtype`` is passed, e.g. ``torch.float32``, only images and videos will be converted + to that dtype: this is for compatibility with :class:`~torchvision.transforms.v2.ConvertImageDtype`. + A dict can be passed to specify per-tv_tensor conversions, e.g. + ``dtype={tv_tensors.Image: torch.float32, tv_tensors.Mask: torch.int64, "others":None}``. The "others" + key can be used as a catch-all for any other tv_tensor type, and ``None`` means no conversion. + scale (bool, optional): Whether to scale the values for images or videos. See :ref:`range_and_dtype`. + Default: ``False``. + """ + + _transformed_types = (torch.Tensor,) + + def __init__( + self, dtype: Union[torch.dtype, Dict[Union[Type, str], Optional[torch.dtype]]], scale: bool = False + ) -> None: + super().__init__() + + if not isinstance(dtype, (dict, torch.dtype)): + raise ValueError(f"dtype must be a dict or a torch.dtype, got {type(dtype)} instead") + + if ( + isinstance(dtype, dict) + and torch.Tensor in dtype + and any(cls in dtype for cls in [tv_tensors.Image, tv_tensors.Video]) + ): + warnings.warn( + "Got `dtype` values for `torch.Tensor` and either `tv_tensors.Image` or `tv_tensors.Video`. " + "Note that a plain `torch.Tensor` will *not* be transformed by this (or any other transformation) " + "in case a `tv_tensors.Image` or `tv_tensors.Video` is present in the input." + ) + self.dtype = dtype + self.scale = scale + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + if isinstance(self.dtype, torch.dtype): + # For consistency / BC with ConvertImageDtype, we only care about images or videos when dtype + # is a simple torch.dtype + if not is_pure_tensor(inpt) and not isinstance(inpt, (tv_tensors.Image, tv_tensors.Video)): + return inpt + + dtype: Optional[torch.dtype] = self.dtype + elif type(inpt) in self.dtype: + dtype = self.dtype[type(inpt)] + elif "others" in self.dtype: + dtype = self.dtype["others"] + else: + raise ValueError( + f"No dtype was specified for type {type(inpt)}. " + "If you only need to convert the dtype of images or videos, you can just pass e.g. dtype=torch.float32. " + "If you're passing a dict as dtype, " + 'you can use "others" as a catch-all key ' + 'e.g. dtype={tv_tensors.Mask: torch.int64, "others": None} to pass-through the rest of the inputs.' + ) + + supports_scaling = is_pure_tensor(inpt) or isinstance(inpt, (tv_tensors.Image, tv_tensors.Video)) + if dtype is None: + if self.scale and supports_scaling: + warnings.warn( + "scale was set to True but no dtype was specified for images or videos: no scaling will be done." + ) + return inpt + + return self._call_kernel(F.to_dtype, inpt, dtype=dtype, scale=self.scale) + + +class ConvertImageDtype(Transform): + """[DEPRECATED] Use ``v2.ToDtype(dtype, scale=True)`` instead. + + Convert input image to the given ``dtype`` and scale the values accordingly. + + .. warning:: + Consider using ``ToDtype(dtype, scale=True)`` instead. See :class:`~torchvision.transforms.v2.ToDtype`. + + This function does not support PIL Image. + + Args: + dtype (torch.dtype): Desired data type of the output + + .. note:: + + When converting from a smaller to a larger integer ``dtype`` the maximum values are **not** mapped exactly. + If converted back and forth, this mismatch has no effect. + + Raises: + RuntimeError: When trying to cast :class:`torch.float32` to :class:`torch.int32` or :class:`torch.int64` as + well as for trying to cast :class:`torch.float64` to :class:`torch.int64`. These conversions might lead to + overflow errors since the floating point ``dtype`` cannot store consecutive integers over the whole range + of the integer ``dtype``. + """ + + _v1_transform_cls = _transforms.ConvertImageDtype + + def __init__(self, dtype: torch.dtype = torch.float32) -> None: + super().__init__() + self.dtype = dtype + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.to_dtype, inpt, dtype=self.dtype, scale=True) + + +class SanitizeBoundingBoxes(Transform): + """Remove degenerate/invalid bounding boxes and their corresponding labels and masks. + + This transform removes bounding boxes and their associated labels/masks that: + + - are below a given ``min_size`` or ``min_area``: by default this also removes degenerate boxes that have e.g. X2 <= X1. + - have any coordinate outside of their corresponding image. You may want to + call :class:`~torchvision.transforms.v2.ClampBoundingBoxes` first to avoid undesired removals. + + It can also sanitize other tensors like the "iscrowd" or "area" properties from COCO + (see ``labels_getter`` parameter). + + It is recommended to call it at the end of a pipeline, before passing the + input to the models. It is critical to call this transform if + :class:`~torchvision.transforms.v2.RandomIoUCrop` was called. + If you want to be extra careful, you may call it after all transforms that + may modify bounding boxes but once at the end should be enough in most + cases. + + Args: + min_size (float, optional): The size below which bounding boxes are removed. Default is 1. + min_area (float, optional): The area below which bounding boxes are removed. Default is 1. + labels_getter (callable or str or None, optional): indicates how to identify the labels in the input + (or anything else that needs to be sanitized along with the bounding boxes). + By default, this will try to find a "labels" key in the input (case-insensitive), if + the input is a dict or it is a tuple whose second element is a dict. + This heuristic should work well with a lot of datasets, including the built-in torchvision datasets. + + It can also be a callable that takes the same input as the transform, and returns either: + + - A single tensor (the labels) + - A tuple/list of tensors, each of which will be subject to the same sanitization as the bounding boxes. + This is useful to sanitize multiple tensors like the labels, and the "iscrowd" or "area" properties + from COCO. + + If ``labels_getter`` is None then only bounding boxes are sanitized. + """ + + def __init__( + self, + min_size: float = 1.0, + min_area: float = 1.0, + labels_getter: Union[Callable[[Any], Any], str, None] = "default", + ) -> None: + super().__init__() + + if min_size < 1: + raise ValueError(f"min_size must be >= 1, got {min_size}.") + self.min_size = min_size + + if min_area < 1: + raise ValueError(f"min_area must be >= 1, got {min_area}.") + self.min_area = min_area + + self.labels_getter = labels_getter + self._labels_getter = _parse_labels_getter(labels_getter) + + def forward(self, *inputs: Any) -> Any: + inputs = inputs if len(inputs) > 1 else inputs[0] + + labels = self._labels_getter(inputs) + if labels is not None: + msg = "The labels in the input to forward() must be a tensor or None, got {type} instead." + if isinstance(labels, torch.Tensor): + labels = (labels,) + elif isinstance(labels, (tuple, list)): + for entry in labels: + if not isinstance(entry, torch.Tensor): + # TODO: we don't need to enforce tensors, just that entries are indexable as t[bool_mask] + raise ValueError(msg.format(type=type(entry))) + else: + raise ValueError(msg.format(type=type(labels))) + + flat_inputs, spec = tree_flatten(inputs) + boxes = get_bounding_boxes(flat_inputs) + + if labels is not None: + for label in labels: + if boxes.shape[0] != label.shape[0]: + raise ValueError( + f"Number of boxes (shape={boxes.shape}) and must match the number of labels." + f"Found labels with shape={label.shape})." + ) + + valid = F._misc._get_sanitize_bounding_boxes_mask( + boxes, + format=boxes.format, + canvas_size=boxes.canvas_size, + min_size=self.min_size, + min_area=self.min_area, + ) + + params = dict(valid=valid, labels=labels) + flat_outputs = [self.transform(inpt, params) for inpt in flat_inputs] + + return tree_unflatten(flat_outputs, spec) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + is_label = params["labels"] is not None and any(inpt is label for label in params["labels"]) + is_bounding_boxes_or_mask = isinstance(inpt, (tv_tensors.BoundingBoxes, tv_tensors.Mask)) + + if not (is_label or is_bounding_boxes_or_mask): + return inpt + + output = inpt[params["valid"]] + + if is_label: + return output + else: + return tv_tensors.wrap(output, like=inpt) diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/_temporal.py b/lib/python3.12/site-packages/torchvision/transforms/v2/_temporal.py new file mode 100644 index 0000000000000000000000000000000000000000..687b50188a8d2eb7b03c8f04863652a1998974f3 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/_temporal.py @@ -0,0 +1,26 @@ +from typing import Any, Dict + +import torch +from torchvision.transforms.v2 import functional as F, Transform + + +class UniformTemporalSubsample(Transform): + """Uniformly subsample ``num_samples`` indices from the temporal dimension of the video. + + Videos are expected to be of shape ``[..., T, C, H, W]`` where ``T`` denotes the temporal dimension. + + When ``num_samples`` is larger than the size of temporal dimension of the video, it + will sample frames based on nearest neighbor interpolation. + + Args: + num_samples (int): The number of equispaced samples to be selected + """ + + _transformed_types = (torch.Tensor,) + + def __init__(self, num_samples: int): + super().__init__() + self.num_samples = num_samples + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + return self._call_kernel(F.uniform_temporal_subsample, inpt, self.num_samples) diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/_transform.py b/lib/python3.12/site-packages/torchvision/transforms/v2/_transform.py new file mode 100644 index 0000000000000000000000000000000000000000..5f27458970990c86a8ee422de8e194675d7c4eb7 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/_transform.py @@ -0,0 +1,194 @@ +from __future__ import annotations + +import enum +from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union + +import PIL.Image +import torch +from torch import nn +from torch.utils._pytree import tree_flatten, tree_unflatten +from torchvision import tv_tensors +from torchvision.transforms.v2._utils import check_type, has_any, is_pure_tensor +from torchvision.utils import _log_api_usage_once + +from .functional._utils import _get_kernel + + +class Transform(nn.Module): + """Base class to implement your own v2 transforms. + + See :ref:`sphx_glr_auto_examples_transforms_plot_custom_transforms.py` for + more details. + """ + + # Class attribute defining transformed types. Other types are passed-through without any transformation + # We support both Types and callables that are able to do further checks on the type of the input. + _transformed_types: Tuple[Union[Type, Callable[[Any], bool]], ...] = (torch.Tensor, PIL.Image.Image) + + def __init__(self) -> None: + super().__init__() + _log_api_usage_once(self) + + def check_inputs(self, flat_inputs: List[Any]) -> None: + pass + + # When v2 was introduced, this method was private and called + # `_get_params()`. Now it's publicly exposed as `make_params()`. It cannot + # be exposed as `get_params()` because there is already a `get_params()` + # methods for v2 transforms: it's the v1's `get_params()` that we have to + # keep in order to guarantee 100% BC with v1. (It's defined in + # __init_subclass__ below). + def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]: + """Method to override for custom transforms. + + See :ref:`sphx_glr_auto_examples_transforms_plot_custom_transforms.py`""" + return dict() + + def _call_kernel(self, functional: Callable, inpt: Any, *args: Any, **kwargs: Any) -> Any: + kernel = _get_kernel(functional, type(inpt), allow_passthrough=True) + return kernel(inpt, *args, **kwargs) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> Any: + """Method to override for custom transforms. + + See :ref:`sphx_glr_auto_examples_transforms_plot_custom_transforms.py`""" + raise NotImplementedError + + def forward(self, *inputs: Any) -> Any: + """Do not override this! Use ``transform()`` instead.""" + flat_inputs, spec = tree_flatten(inputs if len(inputs) > 1 else inputs[0]) + + self.check_inputs(flat_inputs) + + needs_transform_list = self._needs_transform_list(flat_inputs) + params = self.make_params( + [inpt for (inpt, needs_transform) in zip(flat_inputs, needs_transform_list) if needs_transform] + ) + + flat_outputs = [ + self.transform(inpt, params) if needs_transform else inpt + for (inpt, needs_transform) in zip(flat_inputs, needs_transform_list) + ] + + return tree_unflatten(flat_outputs, spec) + + def _needs_transform_list(self, flat_inputs: List[Any]) -> List[bool]: + # Below is a heuristic on how to deal with pure tensor inputs: + # 1. Pure tensors, i.e. tensors that are not a tv_tensor, are passed through if there is an explicit image + # (`tv_tensors.Image` or `PIL.Image.Image`) or video (`tv_tensors.Video`) in the sample. + # 2. If there is no explicit image or video in the sample, only the first encountered pure tensor is + # transformed as image, while the rest is passed through. The order is defined by the returned `flat_inputs` + # of `tree_flatten`, which recurses depth-first through the input. + # + # This heuristic stems from two requirements: + # 1. We need to keep BC for single input pure tensors and treat them as images. + # 2. We don't want to treat all pure tensors as images, because some datasets like `CelebA` or `Widerface` + # return supplemental numerical data as tensors that cannot be transformed as images. + # + # The heuristic should work well for most people in practice. The only case where it doesn't is if someone + # tries to transform multiple pure tensors at the same time, expecting them all to be treated as images. + # However, this case wasn't supported by transforms v1 either, so there is no BC concern. + + needs_transform_list = [] + transform_pure_tensor = not has_any(flat_inputs, tv_tensors.Image, tv_tensors.Video, PIL.Image.Image) + for inpt in flat_inputs: + needs_transform = True + + if not check_type(inpt, self._transformed_types): + needs_transform = False + elif is_pure_tensor(inpt): + if transform_pure_tensor: + transform_pure_tensor = False + else: + needs_transform = False + needs_transform_list.append(needs_transform) + return needs_transform_list + + def extra_repr(self) -> str: + extra = [] + for name, value in self.__dict__.items(): + if name.startswith("_") or name == "training": + continue + + if not isinstance(value, (bool, int, float, str, tuple, list, enum.Enum)): + continue + + extra.append(f"{name}={value}") + + return ", ".join(extra) + + # This attribute should be set on all transforms that have a v1 equivalent. Doing so enables two things: + # 1. In case the v1 transform has a static `get_params` method, it will also be available under the same name on + # the v2 transform. See `__init_subclass__` for details. + # 2. The v2 transform will be JIT scriptable. See `_extract_params_for_v1_transform` and `__prepare_scriptable__` + # for details. + _v1_transform_cls: Optional[Type[nn.Module]] = None + + def __init_subclass__(cls) -> None: + # Since `get_params` is a `@staticmethod`, we have to bind it to the class itself rather than to an instance. + # This method is called after subclassing has happened, i.e. `cls` is the subclass, e.g. `Resize`. + if cls._v1_transform_cls is not None and hasattr(cls._v1_transform_cls, "get_params"): + cls.get_params = staticmethod(cls._v1_transform_cls.get_params) # type: ignore[attr-defined] + + def _extract_params_for_v1_transform(self) -> Dict[str, Any]: + # This method is called by `__prepare_scriptable__` to instantiate the equivalent v1 transform from the current + # v2 transform instance. It extracts all available public attributes that are specific to that transform and + # not `nn.Module` in general. + # Overwrite this method on the v2 transform class if the above is not sufficient. For example, this might happen + # if the v2 transform introduced new parameters that are not support by the v1 transform. + common_attrs = nn.Module().__dict__.keys() + return { + attr: value + for attr, value in self.__dict__.items() + if not attr.startswith("_") and attr not in common_attrs + } + + def __prepare_scriptable__(self) -> nn.Module: + # This method is called early on when `torch.jit.script`'ing an `nn.Module` instance. If it succeeds, the return + # value is used for scripting over the original object that should have been scripted. Since the v1 transforms + # are JIT scriptable, and we made sure that for single image inputs v1 and v2 are equivalent, we just return the + # equivalent v1 transform here. This of course only makes transforms v2 JIT scriptable as long as transforms v1 + # is around. + if self._v1_transform_cls is None: + raise RuntimeError( + f"Transform {type(self).__name__} cannot be JIT scripted. " + "torchscript is only supported for backward compatibility with transforms " + "which are already in torchvision.transforms. " + "For torchscript support (on tensors only), you can use the functional API instead." + ) + + return self._v1_transform_cls(**self._extract_params_for_v1_transform()) + + +class _RandomApplyTransform(Transform): + def __init__(self, p: float = 0.5) -> None: + if not (0.0 <= p <= 1.0): + raise ValueError("`p` should be a floating point value in the interval [0.0, 1.0].") + + super().__init__() + self.p = p + + def forward(self, *inputs: Any) -> Any: + # We need to almost duplicate `Transform.forward()` here since we always want to check the inputs, but return + # early afterwards in case the random check triggers. The same result could be achieved by calling + # `super().forward()` after the random check, but that would call `self.check_inputs` twice. + + inputs = inputs if len(inputs) > 1 else inputs[0] + flat_inputs, spec = tree_flatten(inputs) + + self.check_inputs(flat_inputs) + + if torch.rand(1) >= self.p: + return inputs + + needs_transform_list = self._needs_transform_list(flat_inputs) + params = self.make_params( + [inpt for (inpt, needs_transform) in zip(flat_inputs, needs_transform_list) if needs_transform] + ) + + flat_outputs = [ + self.transform(inpt, params) if needs_transform else inpt + for (inpt, needs_transform) in zip(flat_inputs, needs_transform_list) + ] + + return tree_unflatten(flat_outputs, spec) diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/_type_conversion.py b/lib/python3.12/site-packages/torchvision/transforms/v2/_type_conversion.py new file mode 100644 index 0000000000000000000000000000000000000000..bf9f71852393efe9c83f3b012a54f22568e341c1 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/_type_conversion.py @@ -0,0 +1,84 @@ +from typing import Any, Dict, Optional, Union + +import numpy as np +import PIL.Image +import torch + +from torchvision import tv_tensors +from torchvision.transforms.v2 import functional as F, Transform + +from torchvision.transforms.v2._utils import is_pure_tensor + + +class PILToTensor(Transform): + """Convert a PIL Image to a tensor of the same type - this does not scale values. + + This transform does not support torchscript. + + Converts a PIL Image (H x W x C) to a Tensor of shape (C x H x W). + """ + + _transformed_types = (PIL.Image.Image,) + + def transform(self, inpt: PIL.Image.Image, params: Dict[str, Any]) -> torch.Tensor: + return F.pil_to_tensor(inpt) + + +class ToImage(Transform): + """Convert a tensor, ndarray, or PIL Image to :class:`~torchvision.tv_tensors.Image` + ; this does not scale values. + + This transform does not support torchscript. + """ + + _transformed_types = (is_pure_tensor, PIL.Image.Image, np.ndarray) + + def transform( + self, inpt: Union[torch.Tensor, PIL.Image.Image, np.ndarray], params: Dict[str, Any] + ) -> tv_tensors.Image: + return F.to_image(inpt) + + +class ToPILImage(Transform): + """Convert a tensor or an ndarray to PIL Image + + This transform does not support torchscript. + + Converts a torch.*Tensor of shape C x H x W or a numpy ndarray of shape + H x W x C to a PIL Image while adjusting the value range depending on the ``mode``. + + Args: + mode (`PIL.Image mode`_): color space and pixel depth of input data (optional). + If ``mode`` is ``None`` (default) there are some assumptions made about the input data: + + - If the input has 4 channels, the ``mode`` is assumed to be ``RGBA``. + - If the input has 3 channels, the ``mode`` is assumed to be ``RGB``. + - If the input has 2 channels, the ``mode`` is assumed to be ``LA``. + - If the input has 1 channel, the ``mode`` is determined by the data type (i.e ``int``, ``float``, + ``short``). + + .. _PIL.Image mode: https://pillow.readthedocs.io/en/latest/handbook/concepts.html#concept-modes + """ + + _transformed_types = (is_pure_tensor, tv_tensors.Image, np.ndarray) + + def __init__(self, mode: Optional[str] = None) -> None: + super().__init__() + self.mode = mode + + def transform( + self, inpt: Union[torch.Tensor, PIL.Image.Image, np.ndarray], params: Dict[str, Any] + ) -> PIL.Image.Image: + return F.to_pil_image(inpt, mode=self.mode) + + +class ToPureTensor(Transform): + """Convert all TVTensors to pure tensors, removing associated metadata (if any). + + This doesn't scale or change the values, only the type. + """ + + _transformed_types = (tv_tensors.TVTensor,) + + def transform(self, inpt: Any, params: Dict[str, Any]) -> torch.Tensor: + return inpt.as_subclass(torch.Tensor) diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/_utils.py b/lib/python3.12/site-packages/torchvision/transforms/v2/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..dd65ca4d9c98b73b8e17eae972e352df08bacbdb --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/_utils.py @@ -0,0 +1,226 @@ +from __future__ import annotations + +import collections.abc +import numbers +from contextlib import suppress + +from typing import Any, Callable, Dict, List, Literal, Sequence, Tuple, Type, Union + +import PIL.Image +import torch + +from torchvision import tv_tensors + +from torchvision._utils import sequence_to_str + +from torchvision.transforms.transforms import _check_sequence_input, _setup_angle, _setup_size # noqa: F401 +from torchvision.transforms.v2.functional import get_dimensions, get_size, is_pure_tensor +from torchvision.transforms.v2.functional._utils import _FillType, _FillTypeJIT + + +def _setup_number_or_seq(arg: Union[int, float, Sequence[Union[int, float]]], name: str) -> Sequence[float]: + if not isinstance(arg, (int, float, Sequence)): + raise TypeError(f"{name} should be a number or a sequence of numbers. Got {type(arg)}") + if isinstance(arg, Sequence) and len(arg) not in (1, 2): + raise ValueError(f"If {name} is a sequence its length should be 1 or 2. Got {len(arg)}") + if isinstance(arg, Sequence): + for element in arg: + if not isinstance(element, (int, float)): + raise ValueError(f"{name} should be a sequence of numbers. Got {type(element)}") + + if isinstance(arg, (int, float)): + arg = [float(arg), float(arg)] + elif isinstance(arg, Sequence): + if len(arg) == 1: + arg = [float(arg[0]), float(arg[0])] + else: + arg = [float(arg[0]), float(arg[1])] + return arg + + +def _check_fill_arg(fill: Union[_FillType, Dict[Union[Type, str], _FillType]]) -> None: + if isinstance(fill, dict): + for value in fill.values(): + _check_fill_arg(value) + else: + if fill is not None and not isinstance(fill, (numbers.Number, tuple, list)): + raise TypeError("Got inappropriate fill arg, only Numbers, tuples, lists and dicts are allowed.") + + +def _convert_fill_arg(fill: _FillType) -> _FillTypeJIT: + # Fill = 0 is not equivalent to None, https://github.com/pytorch/vision/issues/6517 + # So, we can't reassign fill to 0 + # if fill is None: + # fill = 0 + if fill is None: + return fill + + if not isinstance(fill, (int, float)): + fill = [float(v) for v in list(fill)] + return fill # type: ignore[return-value] + + +def _setup_fill_arg(fill: Union[_FillType, Dict[Union[Type, str], _FillType]]) -> Dict[Union[Type, str], _FillTypeJIT]: + _check_fill_arg(fill) + + if isinstance(fill, dict): + for k, v in fill.items(): + fill[k] = _convert_fill_arg(v) + return fill # type: ignore[return-value] + else: + return {"others": _convert_fill_arg(fill)} + + +def _get_fill(fill_dict, inpt_type): + if inpt_type in fill_dict: + return fill_dict[inpt_type] + elif "others" in fill_dict: + return fill_dict["others"] + else: + RuntimeError("This should never happen, please open an issue on the torchvision repo if you hit this.") + + +def _check_padding_arg(padding: Union[int, Sequence[int]]) -> None: + if not isinstance(padding, (numbers.Number, tuple, list)): + raise TypeError("Got inappropriate padding arg") + + if isinstance(padding, (tuple, list)) and len(padding) not in [1, 2, 4]: + raise ValueError(f"Padding must be an int or a 1, 2, or 4 element tuple, not a {len(padding)} element tuple") + + +# TODO: let's use torchvision._utils.StrEnum to have the best of both worlds (strings and enums) +# https://github.com/pytorch/vision/issues/6250 +def _check_padding_mode_arg(padding_mode: Literal["constant", "edge", "reflect", "symmetric"]) -> None: + if padding_mode not in ["constant", "edge", "reflect", "symmetric"]: + raise ValueError("Padding mode should be either constant, edge, reflect or symmetric") + + +def _find_labels_default_heuristic(inputs: Any) -> torch.Tensor: + """ + This heuristic covers three cases: + + 1. The input is tuple or list whose second item is a labels tensor. This happens for already batched + classification inputs for MixUp and CutMix (typically after the Dataloder). + 2. The input is a tuple or list whose second item is a dictionary that contains the labels tensor + under a label-like (see below) key. This happens for the inputs of detection models. + 3. The input is a dictionary that is structured as the one from 2. + + What is "label-like" key? We first search for an case-insensitive match of 'labels' inside the keys of the + dictionary. This is the name our detection models expect. If we can't find that, we look for a case-insensitive + match of the term 'label' anywhere inside the key, i.e. 'FooLaBeLBar'. If we can't find that either, the dictionary + contains no "label-like" key. + """ + + if isinstance(inputs, (tuple, list)): + inputs = inputs[1] + + # MixUp, CutMix + if is_pure_tensor(inputs): + return inputs + + if not isinstance(inputs, collections.abc.Mapping): + raise ValueError( + f"When using the default labels_getter, the input passed to forward must be a dictionary or a two-tuple " + f"whose second item is a dictionary or a tensor, but got {inputs} instead." + ) + + candidate_key = None + with suppress(StopIteration): + candidate_key = next(key for key in inputs.keys() if key.lower() == "labels") + if candidate_key is None: + with suppress(StopIteration): + candidate_key = next(key for key in inputs.keys() if "label" in key.lower()) + if candidate_key is None: + raise ValueError( + "Could not infer where the labels are in the sample. Try passing a callable as the labels_getter parameter?" + "If there are no labels in the sample by design, pass labels_getter=None." + ) + + return inputs[candidate_key] + + +def _parse_labels_getter(labels_getter: Union[str, Callable[[Any], Any], None]) -> Callable[[Any], Any]: + if labels_getter == "default": + return _find_labels_default_heuristic + elif callable(labels_getter): + return labels_getter + elif labels_getter is None: + return lambda _: None + else: + raise ValueError(f"labels_getter should either be 'default', a callable, or None, but got {labels_getter}.") + + +def get_bounding_boxes(flat_inputs: List[Any]) -> tv_tensors.BoundingBoxes: + """Return the Bounding Boxes in the input. + + Assumes only one ``BoundingBoxes`` object is present. + """ + # This assumes there is only one bbox per sample as per the general convention + try: + return next(inpt for inpt in flat_inputs if isinstance(inpt, tv_tensors.BoundingBoxes)) + except StopIteration: + raise ValueError("No bounding boxes were found in the sample") + + +def query_chw(flat_inputs: List[Any]) -> Tuple[int, int, int]: + """Return Channel, Height, and Width.""" + chws = { + tuple(get_dimensions(inpt)) + for inpt in flat_inputs + if check_type(inpt, (is_pure_tensor, tv_tensors.Image, PIL.Image.Image, tv_tensors.Video)) + } + if not chws: + raise TypeError("No image or video was found in the sample") + elif len(chws) > 1: + raise ValueError(f"Found multiple CxHxW dimensions in the sample: {sequence_to_str(sorted(chws))}") + c, h, w = chws.pop() + return c, h, w + + +def query_size(flat_inputs: List[Any]) -> Tuple[int, int]: + """Return Height and Width.""" + sizes = { + tuple(get_size(inpt)) + for inpt in flat_inputs + if check_type( + inpt, + ( + is_pure_tensor, + tv_tensors.Image, + PIL.Image.Image, + tv_tensors.Video, + tv_tensors.Mask, + tv_tensors.BoundingBoxes, + ), + ) + } + if not sizes: + raise TypeError("No image, video, mask or bounding box was found in the sample") + elif len(sizes) > 1: + raise ValueError(f"Found multiple HxW dimensions in the sample: {sequence_to_str(sorted(sizes))}") + h, w = sizes.pop() + return h, w + + +def check_type(obj: Any, types_or_checks: Tuple[Union[Type, Callable[[Any], bool]], ...]) -> bool: + for type_or_check in types_or_checks: + if isinstance(obj, type_or_check) if isinstance(type_or_check, type) else type_or_check(obj): + return True + return False + + +def has_any(flat_inputs: List[Any], *types_or_checks: Union[Type, Callable[[Any], bool]]) -> bool: + for inpt in flat_inputs: + if check_type(inpt, types_or_checks): + return True + return False + + +def has_all(flat_inputs: List[Any], *types_or_checks: Union[Type, Callable[[Any], bool]]) -> bool: + for type_or_check in types_or_checks: + for inpt in flat_inputs: + if isinstance(inpt, type_or_check) if isinstance(type_or_check, type) else type_or_check(inpt): + break + else: + return False + return True diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__init__.py b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d5705d55c4b2d429e450dc0033c99f7fe16a8ddb --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__init__.py @@ -0,0 +1,153 @@ +from torchvision.transforms import InterpolationMode # usort: skip + +from ._utils import is_pure_tensor, register_kernel # usort: skip + +from ._meta import ( + clamp_bounding_boxes, + convert_bounding_box_format, + get_dimensions_image, + get_dimensions_video, + get_dimensions, + get_num_frames_video, + get_num_frames, + get_image_num_channels, + get_num_channels_image, + get_num_channels_video, + get_num_channels, + get_size_bounding_boxes, + get_size_image, + get_size_mask, + get_size_video, + get_size, +) # usort: skip + +from ._augment import erase, erase_image, erase_video, jpeg, jpeg_image, jpeg_video +from ._color import ( + adjust_brightness, + adjust_brightness_image, + adjust_brightness_video, + adjust_contrast, + adjust_contrast_image, + adjust_contrast_video, + adjust_gamma, + adjust_gamma_image, + adjust_gamma_video, + adjust_hue, + adjust_hue_image, + adjust_hue_video, + adjust_saturation, + adjust_saturation_image, + adjust_saturation_video, + adjust_sharpness, + adjust_sharpness_image, + adjust_sharpness_video, + autocontrast, + autocontrast_image, + autocontrast_video, + equalize, + equalize_image, + equalize_video, + grayscale_to_rgb, + grayscale_to_rgb_image, + invert, + invert_image, + invert_video, + permute_channels, + permute_channels_image, + permute_channels_video, + posterize, + posterize_image, + posterize_video, + rgb_to_grayscale, + rgb_to_grayscale_image, + solarize, + solarize_image, + solarize_video, + to_grayscale, +) +from ._geometry import ( + affine, + affine_bounding_boxes, + affine_image, + affine_mask, + affine_video, + center_crop, + center_crop_bounding_boxes, + center_crop_image, + center_crop_mask, + center_crop_video, + crop, + crop_bounding_boxes, + crop_image, + crop_mask, + crop_video, + elastic, + elastic_bounding_boxes, + elastic_image, + elastic_mask, + elastic_transform, + elastic_video, + five_crop, + five_crop_image, + five_crop_video, + hflip, # TODO: Consider moving all pure alias definitions at the bottom of the file + horizontal_flip, + horizontal_flip_bounding_boxes, + horizontal_flip_image, + horizontal_flip_mask, + horizontal_flip_video, + pad, + pad_bounding_boxes, + pad_image, + pad_mask, + pad_video, + perspective, + perspective_bounding_boxes, + perspective_image, + perspective_mask, + perspective_video, + resize, + resize_bounding_boxes, + resize_image, + resize_mask, + resize_video, + resized_crop, + resized_crop_bounding_boxes, + resized_crop_image, + resized_crop_mask, + resized_crop_video, + rotate, + rotate_bounding_boxes, + rotate_image, + rotate_mask, + rotate_video, + ten_crop, + ten_crop_image, + ten_crop_video, + vertical_flip, + vertical_flip_bounding_boxes, + vertical_flip_image, + vertical_flip_mask, + vertical_flip_video, + vflip, +) +from ._misc import ( + convert_image_dtype, + gaussian_blur, + gaussian_blur_image, + gaussian_blur_video, + gaussian_noise, + gaussian_noise_image, + gaussian_noise_video, + normalize, + normalize_image, + normalize_video, + sanitize_bounding_boxes, + to_dtype, + to_dtype_image, + to_dtype_video, +) +from ._temporal import uniform_temporal_subsample, uniform_temporal_subsample_video +from ._type_conversion import pil_to_tensor, to_image, to_pil_image + +from ._deprecated import get_image_size, to_tensor # usort: skip diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..595a675491a928bfb7ef112ff6a0945e9564a06a Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_augment.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_augment.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..573165893a7f7ece96cc39bc8c4087f535df6c30 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_augment.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_color.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_color.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fe45e44ec90a7b7f3a13c8380c10f4cf853098e0 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_color.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_deprecated.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_deprecated.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..72fb36c071fff4060a9c25f940ab74c8f0dd11c7 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_deprecated.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_meta.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_meta.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..85c09bac3135a7fddf8fa7c556627b71da2c041d Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_meta.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_misc.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_misc.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d721095b60b25a01f52edbdcc10422cd8e3e4562 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_misc.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_temporal.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_temporal.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..07557cd573513180100d09001111cc0df537cbdb Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_temporal.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_type_conversion.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_type_conversion.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..71e4333b4c38f61ee41c688c949ef34f91df6922 Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_type_conversion.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1336c3f5949f5181ed68b4e09b3d2404ce9efbdd Binary files /dev/null and b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/__pycache__/_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_augment.py b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_augment.py new file mode 100644 index 0000000000000000000000000000000000000000..a904d8d7cbdfeb78588abbf43c8bca37b3431735 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_augment.py @@ -0,0 +1,106 @@ +import io + +import PIL.Image + +import torch +from torchvision import tv_tensors +from torchvision.io import decode_jpeg, encode_jpeg +from torchvision.transforms.functional import pil_to_tensor, to_pil_image +from torchvision.utils import _log_api_usage_once + +from ._utils import _get_kernel, _register_kernel_internal + + +def erase( + inpt: torch.Tensor, + i: int, + j: int, + h: int, + w: int, + v: torch.Tensor, + inplace: bool = False, +) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.RandomErase` for details.""" + if torch.jit.is_scripting(): + return erase_image(inpt, i=i, j=j, h=h, w=w, v=v, inplace=inplace) + + _log_api_usage_once(erase) + + kernel = _get_kernel(erase, type(inpt)) + return kernel(inpt, i=i, j=j, h=h, w=w, v=v, inplace=inplace) + + +@_register_kernel_internal(erase, torch.Tensor) +@_register_kernel_internal(erase, tv_tensors.Image) +def erase_image( + image: torch.Tensor, i: int, j: int, h: int, w: int, v: torch.Tensor, inplace: bool = False +) -> torch.Tensor: + if not inplace: + image = image.clone() + + image[..., i : i + h, j : j + w] = v + return image + + +@_register_kernel_internal(erase, PIL.Image.Image) +def _erase_image_pil( + image: PIL.Image.Image, i: int, j: int, h: int, w: int, v: torch.Tensor, inplace: bool = False +) -> PIL.Image.Image: + t_img = pil_to_tensor(image) + output = erase_image(t_img, i=i, j=j, h=h, w=w, v=v, inplace=inplace) + return to_pil_image(output, mode=image.mode) + + +@_register_kernel_internal(erase, tv_tensors.Video) +def erase_video( + video: torch.Tensor, i: int, j: int, h: int, w: int, v: torch.Tensor, inplace: bool = False +) -> torch.Tensor: + return erase_image(video, i=i, j=j, h=h, w=w, v=v, inplace=inplace) + + +def jpeg(image: torch.Tensor, quality: int) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.JPEG` for details.""" + if torch.jit.is_scripting(): + return jpeg_image(image, quality=quality) + + _log_api_usage_once(jpeg) + + kernel = _get_kernel(jpeg, type(image)) + return kernel(image, quality=quality) + + +@_register_kernel_internal(jpeg, torch.Tensor) +@_register_kernel_internal(jpeg, tv_tensors.Image) +def jpeg_image(image: torch.Tensor, quality: int) -> torch.Tensor: + original_shape = image.shape + image = image.view((-1,) + image.shape[-3:]) + + if image.shape[0] == 0: # degenerate + return image.reshape(original_shape).clone() + + images = [] + for i in range(image.shape[0]): + # isinstance checks are needed for torchscript. + encoded_image = encode_jpeg(image[i], quality=quality) + assert isinstance(encoded_image, torch.Tensor) + decoded_image = decode_jpeg(encoded_image) + assert isinstance(decoded_image, torch.Tensor) + images.append(decoded_image) + + images = torch.stack(images, dim=0).view(original_shape) + return images + + +@_register_kernel_internal(jpeg, tv_tensors.Video) +def jpeg_video(video: torch.Tensor, quality: int) -> torch.Tensor: + return jpeg_image(video, quality=quality) + + +@_register_kernel_internal(jpeg, PIL.Image.Image) +def _jpeg_image_pil(image: PIL.Image.Image, quality: int) -> PIL.Image.Image: + raw_jpeg = io.BytesIO() + image.save(raw_jpeg, format="JPEG", quality=quality) + + # we need to copy since PIL.Image.open() will return PIL.JpegImagePlugin.JpegImageFile + # which is a sub-class of PIL.Image.Image. this will fail check_transform() test. + return PIL.Image.open(raw_jpeg).copy() diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_color.py b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_color.py new file mode 100644 index 0000000000000000000000000000000000000000..eb75f58cb7ac85a9478c26c4e86d9ab73b93d331 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_color.py @@ -0,0 +1,739 @@ +from typing import List + +import PIL.Image +import torch +from torch.nn.functional import conv2d +from torchvision import tv_tensors +from torchvision.transforms import _functional_pil as _FP +from torchvision.transforms._functional_tensor import _max_value + +from torchvision.utils import _log_api_usage_once + +from ._misc import _num_value_bits, to_dtype_image +from ._type_conversion import pil_to_tensor, to_pil_image +from ._utils import _get_kernel, _register_kernel_internal + + +def rgb_to_grayscale(inpt: torch.Tensor, num_output_channels: int = 1) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.Grayscale` for details.""" + if torch.jit.is_scripting(): + return rgb_to_grayscale_image(inpt, num_output_channels=num_output_channels) + + _log_api_usage_once(rgb_to_grayscale) + + kernel = _get_kernel(rgb_to_grayscale, type(inpt)) + return kernel(inpt, num_output_channels=num_output_channels) + + +# `to_grayscale` actually predates `rgb_to_grayscale` in v1, but only handles PIL images. Since `rgb_to_grayscale` is a +# superset in terms of functionality and has the same signature, we alias here to avoid disruption. +to_grayscale = rgb_to_grayscale + + +def _rgb_to_grayscale_image( + image: torch.Tensor, num_output_channels: int = 1, preserve_dtype: bool = True +) -> torch.Tensor: + # TODO: Maybe move the validation that num_output_channels is 1 or 3 to this function instead of callers. + if image.shape[-3] == 1 and num_output_channels == 1: + return image.clone() + if image.shape[-3] == 1 and num_output_channels == 3: + s = [1] * len(image.shape) + s[-3] = 3 + return image.repeat(s) + r, g, b = image.unbind(dim=-3) + l_img = r.mul(0.2989).add_(g, alpha=0.587).add_(b, alpha=0.114) + l_img = l_img.unsqueeze(dim=-3) + if preserve_dtype: + l_img = l_img.to(image.dtype) + if num_output_channels == 3: + l_img = l_img.expand(image.shape) + return l_img + + +@_register_kernel_internal(rgb_to_grayscale, torch.Tensor) +@_register_kernel_internal(rgb_to_grayscale, tv_tensors.Image) +def rgb_to_grayscale_image(image: torch.Tensor, num_output_channels: int = 1) -> torch.Tensor: + if num_output_channels not in (1, 3): + raise ValueError(f"num_output_channels must be 1 or 3, got {num_output_channels}.") + return _rgb_to_grayscale_image(image, num_output_channels=num_output_channels, preserve_dtype=True) + + +@_register_kernel_internal(rgb_to_grayscale, PIL.Image.Image) +def _rgb_to_grayscale_image_pil(image: PIL.Image.Image, num_output_channels: int = 1) -> PIL.Image.Image: + if num_output_channels not in (1, 3): + raise ValueError(f"num_output_channels must be 1 or 3, got {num_output_channels}.") + return _FP.to_grayscale(image, num_output_channels=num_output_channels) + + +def grayscale_to_rgb(inpt: torch.Tensor) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.RGB` for details.""" + if torch.jit.is_scripting(): + return grayscale_to_rgb_image(inpt) + + _log_api_usage_once(grayscale_to_rgb) + + kernel = _get_kernel(grayscale_to_rgb, type(inpt)) + return kernel(inpt) + + +@_register_kernel_internal(grayscale_to_rgb, torch.Tensor) +@_register_kernel_internal(grayscale_to_rgb, tv_tensors.Image) +def grayscale_to_rgb_image(image: torch.Tensor) -> torch.Tensor: + if image.shape[-3] >= 3: + # Image already has RGB channels. We don't need to do anything. + return image + # rgb_to_grayscale can be used to add channels so we reuse that function. + return _rgb_to_grayscale_image(image, num_output_channels=3, preserve_dtype=True) + + +@_register_kernel_internal(grayscale_to_rgb, PIL.Image.Image) +def grayscale_to_rgb_image_pil(image: PIL.Image.Image) -> PIL.Image.Image: + return image.convert(mode="RGB") + + +def _blend(image1: torch.Tensor, image2: torch.Tensor, ratio: float) -> torch.Tensor: + ratio = float(ratio) + fp = image1.is_floating_point() + bound = _max_value(image1.dtype) + output = image1.mul(ratio).add_(image2, alpha=(1.0 - ratio)).clamp_(0, bound) + return output if fp else output.to(image1.dtype) + + +def adjust_brightness(inpt: torch.Tensor, brightness_factor: float) -> torch.Tensor: + """Adjust brightness.""" + + if torch.jit.is_scripting(): + return adjust_brightness_image(inpt, brightness_factor=brightness_factor) + + _log_api_usage_once(adjust_brightness) + + kernel = _get_kernel(adjust_brightness, type(inpt)) + return kernel(inpt, brightness_factor=brightness_factor) + + +@_register_kernel_internal(adjust_brightness, torch.Tensor) +@_register_kernel_internal(adjust_brightness, tv_tensors.Image) +def adjust_brightness_image(image: torch.Tensor, brightness_factor: float) -> torch.Tensor: + if brightness_factor < 0: + raise ValueError(f"brightness_factor ({brightness_factor}) is not non-negative.") + + c = image.shape[-3] + if c not in [1, 3]: + raise TypeError(f"Input image tensor permitted channel values are 1 or 3, but found {c}") + + fp = image.is_floating_point() + bound = _max_value(image.dtype) + output = image.mul(brightness_factor).clamp_(0, bound) + return output if fp else output.to(image.dtype) + + +@_register_kernel_internal(adjust_brightness, PIL.Image.Image) +def _adjust_brightness_image_pil(image: PIL.Image.Image, brightness_factor: float) -> PIL.Image.Image: + return _FP.adjust_brightness(image, brightness_factor=brightness_factor) + + +@_register_kernel_internal(adjust_brightness, tv_tensors.Video) +def adjust_brightness_video(video: torch.Tensor, brightness_factor: float) -> torch.Tensor: + return adjust_brightness_image(video, brightness_factor=brightness_factor) + + +def adjust_saturation(inpt: torch.Tensor, saturation_factor: float) -> torch.Tensor: + """Adjust saturation.""" + if torch.jit.is_scripting(): + return adjust_saturation_image(inpt, saturation_factor=saturation_factor) + + _log_api_usage_once(adjust_saturation) + + kernel = _get_kernel(adjust_saturation, type(inpt)) + return kernel(inpt, saturation_factor=saturation_factor) + + +@_register_kernel_internal(adjust_saturation, torch.Tensor) +@_register_kernel_internal(adjust_saturation, tv_tensors.Image) +def adjust_saturation_image(image: torch.Tensor, saturation_factor: float) -> torch.Tensor: + if saturation_factor < 0: + raise ValueError(f"saturation_factor ({saturation_factor}) is not non-negative.") + + c = image.shape[-3] + if c not in [1, 3]: + raise TypeError(f"Input image tensor permitted channel values are 1 or 3, but found {c}") + + if c == 1: # Match PIL behaviour + return image + + grayscale_image = _rgb_to_grayscale_image(image, num_output_channels=1, preserve_dtype=False) + if not image.is_floating_point(): + grayscale_image = grayscale_image.floor_() + + return _blend(image, grayscale_image, saturation_factor) + + +_adjust_saturation_image_pil = _register_kernel_internal(adjust_saturation, PIL.Image.Image)(_FP.adjust_saturation) + + +@_register_kernel_internal(adjust_saturation, tv_tensors.Video) +def adjust_saturation_video(video: torch.Tensor, saturation_factor: float) -> torch.Tensor: + return adjust_saturation_image(video, saturation_factor=saturation_factor) + + +def adjust_contrast(inpt: torch.Tensor, contrast_factor: float) -> torch.Tensor: + """See :class:`~torchvision.transforms.RandomAutocontrast`""" + if torch.jit.is_scripting(): + return adjust_contrast_image(inpt, contrast_factor=contrast_factor) + + _log_api_usage_once(adjust_contrast) + + kernel = _get_kernel(adjust_contrast, type(inpt)) + return kernel(inpt, contrast_factor=contrast_factor) + + +@_register_kernel_internal(adjust_contrast, torch.Tensor) +@_register_kernel_internal(adjust_contrast, tv_tensors.Image) +def adjust_contrast_image(image: torch.Tensor, contrast_factor: float) -> torch.Tensor: + if contrast_factor < 0: + raise ValueError(f"contrast_factor ({contrast_factor}) is not non-negative.") + + c = image.shape[-3] + if c not in [1, 3]: + raise TypeError(f"Input image tensor permitted channel values are 1 or 3, but found {c}") + fp = image.is_floating_point() + if c == 3: + grayscale_image = _rgb_to_grayscale_image(image, num_output_channels=1, preserve_dtype=False) + if not fp: + grayscale_image = grayscale_image.floor_() + else: + grayscale_image = image if fp else image.to(torch.float32) + mean = torch.mean(grayscale_image, dim=(-3, -2, -1), keepdim=True) + return _blend(image, mean, contrast_factor) + + +_adjust_contrast_image_pil = _register_kernel_internal(adjust_contrast, PIL.Image.Image)(_FP.adjust_contrast) + + +@_register_kernel_internal(adjust_contrast, tv_tensors.Video) +def adjust_contrast_video(video: torch.Tensor, contrast_factor: float) -> torch.Tensor: + return adjust_contrast_image(video, contrast_factor=contrast_factor) + + +def adjust_sharpness(inpt: torch.Tensor, sharpness_factor: float) -> torch.Tensor: + """See :class:`~torchvision.transforms.RandomAdjustSharpness`""" + if torch.jit.is_scripting(): + return adjust_sharpness_image(inpt, sharpness_factor=sharpness_factor) + + _log_api_usage_once(adjust_sharpness) + + kernel = _get_kernel(adjust_sharpness, type(inpt)) + return kernel(inpt, sharpness_factor=sharpness_factor) + + +@_register_kernel_internal(adjust_sharpness, torch.Tensor) +@_register_kernel_internal(adjust_sharpness, tv_tensors.Image) +def adjust_sharpness_image(image: torch.Tensor, sharpness_factor: float) -> torch.Tensor: + num_channels, height, width = image.shape[-3:] + if num_channels not in (1, 3): + raise TypeError(f"Input image tensor can have 1 or 3 channels, but found {num_channels}") + + if sharpness_factor < 0: + raise ValueError(f"sharpness_factor ({sharpness_factor}) is not non-negative.") + + if image.numel() == 0 or height <= 2 or width <= 2: + return image + + bound = _max_value(image.dtype) + fp = image.is_floating_point() + shape = image.shape + + if image.ndim > 4: + image = image.reshape(-1, num_channels, height, width) + needs_unsquash = True + else: + needs_unsquash = False + + # The following is a normalized 3x3 kernel with 1s in the edges and a 5 in the middle. + kernel_dtype = image.dtype if fp else torch.float32 + a, b = 1.0 / 13.0, 5.0 / 13.0 + kernel = torch.tensor([[a, a, a], [a, b, a], [a, a, a]], dtype=kernel_dtype, device=image.device) + kernel = kernel.expand(num_channels, 1, 3, 3) + + # We copy and cast at the same time to avoid modifications on the original data + output = image.to(dtype=kernel_dtype, copy=True) + blurred_degenerate = conv2d(output, kernel, groups=num_channels) + if not fp: + # it is better to round before cast + blurred_degenerate = blurred_degenerate.round_() + + # Create a view on the underlying output while pointing at the same data. We do this to avoid indexing twice. + view = output[..., 1:-1, 1:-1] + + # We speed up blending by minimizing flops and doing in-place. The 2 blend options are mathematically equivalent: + # x+(1-r)*(y-x) = x + (1-r)*y - (1-r)*x = x*r + y*(1-r) + view.add_(blurred_degenerate.sub_(view), alpha=(1.0 - sharpness_factor)) + + # The actual data of output have been modified by the above. We only need to clamp and cast now. + output = output.clamp_(0, bound) + if not fp: + output = output.to(image.dtype) + + if needs_unsquash: + output = output.reshape(shape) + + return output + + +_adjust_sharpness_image_pil = _register_kernel_internal(adjust_sharpness, PIL.Image.Image)(_FP.adjust_sharpness) + + +@_register_kernel_internal(adjust_sharpness, tv_tensors.Video) +def adjust_sharpness_video(video: torch.Tensor, sharpness_factor: float) -> torch.Tensor: + return adjust_sharpness_image(video, sharpness_factor=sharpness_factor) + + +def adjust_hue(inpt: torch.Tensor, hue_factor: float) -> torch.Tensor: + """Adjust hue""" + if torch.jit.is_scripting(): + return adjust_hue_image(inpt, hue_factor=hue_factor) + + _log_api_usage_once(adjust_hue) + + kernel = _get_kernel(adjust_hue, type(inpt)) + return kernel(inpt, hue_factor=hue_factor) + + +def _rgb_to_hsv(image: torch.Tensor) -> torch.Tensor: + r, g, _ = image.unbind(dim=-3) + + # Implementation is based on + # https://github.com/python-pillow/Pillow/blob/4174d4267616897df3746d315d5a2d0f82c656ee/src/libImaging/Convert.c#L330 + minc, maxc = torch.aminmax(image, dim=-3) + + # The algorithm erases S and H channel where `maxc = minc`. This avoids NaN + # from happening in the results, because + # + S channel has division by `maxc`, which is zero only if `maxc = minc` + # + H channel has division by `(maxc - minc)`. + # + # Instead of overwriting NaN afterwards, we just prevent it from occurring so + # we don't need to deal with it in case we save the NaN in a buffer in + # backprop, if it is ever supported, but it doesn't hurt to do so. + eqc = maxc == minc + + channels_range = maxc - minc + # Since `eqc => channels_range = 0`, replacing denominator with 1 when `eqc` is fine. + ones = torch.ones_like(maxc) + s = channels_range / torch.where(eqc, ones, maxc) + # Note that `eqc => maxc = minc = r = g = b`. So the following calculation + # of `h` would reduce to `bc - gc + 2 + rc - bc + 4 + rc - bc = 6` so it + # would not matter what values `rc`, `gc`, and `bc` have here, and thus + # replacing denominator with 1 when `eqc` is fine. + channels_range_divisor = torch.where(eqc, ones, channels_range).unsqueeze_(dim=-3) + rc, gc, bc = ((maxc.unsqueeze(dim=-3) - image) / channels_range_divisor).unbind(dim=-3) + + mask_maxc_neq_r = maxc != r + mask_maxc_eq_g = maxc == g + + hg = rc.add(2.0).sub_(bc).mul_(mask_maxc_eq_g & mask_maxc_neq_r) + hr = bc.sub_(gc).mul_(~mask_maxc_neq_r) + hb = gc.add_(4.0).sub_(rc).mul_(mask_maxc_neq_r.logical_and_(mask_maxc_eq_g.logical_not_())) + + h = hr.add_(hg).add_(hb) + h = h.mul_(1.0 / 6.0).add_(1.0).fmod_(1.0) + return torch.stack((h, s, maxc), dim=-3) + + +def _hsv_to_rgb(img: torch.Tensor) -> torch.Tensor: + h, s, v = img.unbind(dim=-3) + h6 = h.mul(6) + i = torch.floor(h6) + f = h6.sub_(i) + i = i.to(dtype=torch.int32) + + sxf = s * f + one_minus_s = 1.0 - s + q = (1.0 - sxf).mul_(v).clamp_(0.0, 1.0) + t = sxf.add_(one_minus_s).mul_(v).clamp_(0.0, 1.0) + p = one_minus_s.mul_(v).clamp_(0.0, 1.0) + i.remainder_(6) + + vpqt = torch.stack((v, p, q, t), dim=-3) + + # vpqt -> rgb mapping based on i + select = torch.tensor([[0, 2, 1, 1, 3, 0], [3, 0, 0, 2, 1, 1], [1, 1, 3, 0, 0, 2]], dtype=torch.long) + select = select.to(device=img.device, non_blocking=True) + + select = select[:, i] + if select.ndim > 3: + # if input.shape is (B, ..., C, H, W) then + # select.shape is (C, B, ..., H, W) + # thus we move C axis to get (B, ..., C, H, W) + select = select.moveaxis(0, -3) + + return vpqt.gather(-3, select) + + +@_register_kernel_internal(adjust_hue, torch.Tensor) +@_register_kernel_internal(adjust_hue, tv_tensors.Image) +def adjust_hue_image(image: torch.Tensor, hue_factor: float) -> torch.Tensor: + if not (-0.5 <= hue_factor <= 0.5): + raise ValueError(f"hue_factor ({hue_factor}) is not in [-0.5, 0.5].") + + c = image.shape[-3] + if c not in [1, 3]: + raise TypeError(f"Input image tensor permitted channel values are 1 or 3, but found {c}") + + if c == 1: # Match PIL behaviour + return image + + if image.numel() == 0: + # exit earlier on empty images + return image + + orig_dtype = image.dtype + image = to_dtype_image(image, torch.float32, scale=True) + + image = _rgb_to_hsv(image) + h, s, v = image.unbind(dim=-3) + h.add_(hue_factor).remainder_(1.0) + image = torch.stack((h, s, v), dim=-3) + image_hue_adj = _hsv_to_rgb(image) + + return to_dtype_image(image_hue_adj, orig_dtype, scale=True) + + +_adjust_hue_image_pil = _register_kernel_internal(adjust_hue, PIL.Image.Image)(_FP.adjust_hue) + + +@_register_kernel_internal(adjust_hue, tv_tensors.Video) +def adjust_hue_video(video: torch.Tensor, hue_factor: float) -> torch.Tensor: + return adjust_hue_image(video, hue_factor=hue_factor) + + +def adjust_gamma(inpt: torch.Tensor, gamma: float, gain: float = 1) -> torch.Tensor: + """Adjust gamma.""" + if torch.jit.is_scripting(): + return adjust_gamma_image(inpt, gamma=gamma, gain=gain) + + _log_api_usage_once(adjust_gamma) + + kernel = _get_kernel(adjust_gamma, type(inpt)) + return kernel(inpt, gamma=gamma, gain=gain) + + +@_register_kernel_internal(adjust_gamma, torch.Tensor) +@_register_kernel_internal(adjust_gamma, tv_tensors.Image) +def adjust_gamma_image(image: torch.Tensor, gamma: float, gain: float = 1.0) -> torch.Tensor: + if gamma < 0: + raise ValueError("Gamma should be a non-negative real number") + + # The input image is either assumed to be at [0, 1] scale (if float) or is converted to that scale (if integer). + # Since the gamma is non-negative, the output remains at [0, 1] scale. + if not torch.is_floating_point(image): + output = to_dtype_image(image, torch.float32, scale=True).pow_(gamma) + else: + output = image.pow(gamma) + + if gain != 1.0: + # The clamp operation is needed only if multiplication is performed. It's only when gain != 1, that the scale + # of the output can go beyond [0, 1]. + output = output.mul_(gain).clamp_(0.0, 1.0) + + return to_dtype_image(output, image.dtype, scale=True) + + +_adjust_gamma_image_pil = _register_kernel_internal(adjust_gamma, PIL.Image.Image)(_FP.adjust_gamma) + + +@_register_kernel_internal(adjust_gamma, tv_tensors.Video) +def adjust_gamma_video(video: torch.Tensor, gamma: float, gain: float = 1) -> torch.Tensor: + return adjust_gamma_image(video, gamma=gamma, gain=gain) + + +def posterize(inpt: torch.Tensor, bits: int) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.RandomPosterize` for details.""" + if torch.jit.is_scripting(): + return posterize_image(inpt, bits=bits) + + _log_api_usage_once(posterize) + + kernel = _get_kernel(posterize, type(inpt)) + return kernel(inpt, bits=bits) + + +@_register_kernel_internal(posterize, torch.Tensor) +@_register_kernel_internal(posterize, tv_tensors.Image) +def posterize_image(image: torch.Tensor, bits: int) -> torch.Tensor: + if image.is_floating_point(): + levels = 1 << bits + return image.mul(levels).floor_().clamp_(0, levels - 1).mul_(1.0 / levels) + else: + num_value_bits = _num_value_bits(image.dtype) + if bits >= num_value_bits: + return image + + mask = ((1 << bits) - 1) << (num_value_bits - bits) + return image & mask + + +_posterize_image_pil = _register_kernel_internal(posterize, PIL.Image.Image)(_FP.posterize) + + +@_register_kernel_internal(posterize, tv_tensors.Video) +def posterize_video(video: torch.Tensor, bits: int) -> torch.Tensor: + return posterize_image(video, bits=bits) + + +def solarize(inpt: torch.Tensor, threshold: float) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.RandomSolarize` for details.""" + if torch.jit.is_scripting(): + return solarize_image(inpt, threshold=threshold) + + _log_api_usage_once(solarize) + + kernel = _get_kernel(solarize, type(inpt)) + return kernel(inpt, threshold=threshold) + + +@_register_kernel_internal(solarize, torch.Tensor) +@_register_kernel_internal(solarize, tv_tensors.Image) +def solarize_image(image: torch.Tensor, threshold: float) -> torch.Tensor: + if threshold > _max_value(image.dtype): + raise TypeError(f"Threshold should be less or equal the maximum value of the dtype, but got {threshold}") + + return torch.where(image >= threshold, invert_image(image), image) + + +_solarize_image_pil = _register_kernel_internal(solarize, PIL.Image.Image)(_FP.solarize) + + +@_register_kernel_internal(solarize, tv_tensors.Video) +def solarize_video(video: torch.Tensor, threshold: float) -> torch.Tensor: + return solarize_image(video, threshold=threshold) + + +def autocontrast(inpt: torch.Tensor) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.RandomAutocontrast` for details.""" + if torch.jit.is_scripting(): + return autocontrast_image(inpt) + + _log_api_usage_once(autocontrast) + + kernel = _get_kernel(autocontrast, type(inpt)) + return kernel(inpt) + + +@_register_kernel_internal(autocontrast, torch.Tensor) +@_register_kernel_internal(autocontrast, tv_tensors.Image) +def autocontrast_image(image: torch.Tensor) -> torch.Tensor: + c = image.shape[-3] + if c not in [1, 3]: + raise TypeError(f"Input image tensor permitted channel values are 1 or 3, but found {c}") + + if image.numel() == 0: + # exit earlier on empty images + return image + + bound = _max_value(image.dtype) + fp = image.is_floating_point() + float_image = image if fp else image.to(torch.float32) + + minimum = float_image.amin(dim=(-2, -1), keepdim=True) + maximum = float_image.amax(dim=(-2, -1), keepdim=True) + + eq_idxs = maximum == minimum + inv_scale = maximum.sub_(minimum).mul_(1.0 / bound) + minimum[eq_idxs] = 0.0 + inv_scale[eq_idxs] = 1.0 + + if fp: + diff = float_image.sub(minimum) + else: + diff = float_image.sub_(minimum) + + return diff.div_(inv_scale).clamp_(0, bound).to(image.dtype) + + +_autocontrast_image_pil = _register_kernel_internal(autocontrast, PIL.Image.Image)(_FP.autocontrast) + + +@_register_kernel_internal(autocontrast, tv_tensors.Video) +def autocontrast_video(video: torch.Tensor) -> torch.Tensor: + return autocontrast_image(video) + + +def equalize(inpt: torch.Tensor) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.RandomEqualize` for details.""" + if torch.jit.is_scripting(): + return equalize_image(inpt) + + _log_api_usage_once(equalize) + + kernel = _get_kernel(equalize, type(inpt)) + return kernel(inpt) + + +@_register_kernel_internal(equalize, torch.Tensor) +@_register_kernel_internal(equalize, tv_tensors.Image) +def equalize_image(image: torch.Tensor) -> torch.Tensor: + if image.numel() == 0: + return image + + # 1. The algorithm below can easily be extended to support arbitrary integer dtypes. However, the histogram that + # would be needed to computed will have at least `torch.iinfo(dtype).max + 1` values. That is perfectly fine for + # `torch.int8`, `torch.uint8`, and `torch.int16`, at least questionable for `torch.int32` and completely + # unfeasible for `torch.int64`. + # 2. Floating point inputs need to be binned for this algorithm. Apart from converting them to an integer dtype, we + # could also use PyTorch's builtin histogram functionality. However, that has its own set of issues: in addition + # to being slow in general, PyTorch's implementation also doesn't support batches. In total, that makes it slower + # and more complicated to implement than a simple conversion and a fast histogram implementation for integers. + # Since we need to convert in most cases anyway and out of the acceptable dtypes mentioned in 1. `torch.uint8` is + # by far the most common, we choose it as base. + output_dtype = image.dtype + image = to_dtype_image(image, torch.uint8, scale=True) + + # The histogram is computed by using the flattened image as index. For example, a pixel value of 127 in the image + # corresponds to adding 1 to index 127 in the histogram. + batch_shape = image.shape[:-2] + flat_image = image.flatten(start_dim=-2).to(torch.long) + hist = flat_image.new_zeros(batch_shape + (256,), dtype=torch.int32) + hist.scatter_add_(dim=-1, index=flat_image, src=hist.new_ones(1).expand_as(flat_image)) + cum_hist = hist.cumsum(dim=-1) + + # The simplest form of lookup-table (LUT) that also achieves histogram equalization is + # `lut = cum_hist / flat_image.shape[-1] * 255` + # However, PIL uses a more elaborate scheme: + # https://github.com/python-pillow/Pillow/blob/eb59cb61d5239ee69cbbf12709a0c6fd7314e6d7/src/PIL/ImageOps.py#L368-L385 + # `lut = ((cum_hist + num_non_max_pixels // (2 * 255)) // num_non_max_pixels) * 255` + + # The last non-zero element in the histogram is the first element in the cumulative histogram with the maximum + # value. Thus, the "max" in `num_non_max_pixels` does not refer to 255 as the maximum value of uint8 images, but + # rather the maximum value in the image, which might be or not be 255. + index = cum_hist.argmax(dim=-1) + num_non_max_pixels = flat_image.shape[-1] - hist.gather(dim=-1, index=index.unsqueeze_(-1)) + + # This is performance optimization that saves us one multiplication later. With this, the LUT computation simplifies + # to `lut = (cum_hist + step // 2) // step` and thus saving the final multiplication by 255 while keeping the + # division count the same. PIL uses the variable name `step` for this, so we keep that for easier comparison. + step = num_non_max_pixels.div_(255, rounding_mode="floor") + + # Although it looks like we could return early if we find `step == 0` like PIL does, that is unfortunately not as + # easy due to our support for batched images. We can only return early if `(step == 0).all()` holds. If it doesn't, + # we have to go through the computation below anyway. Since `step == 0` is an edge case anyway, it makes no sense to + # pay the runtime cost for checking it every time. + valid_equalization = step.ne(0).unsqueeze_(-1) + + # `lut[k]` is computed with `cum_hist[k-1]` with `lut[0] == (step // 2) // step == 0`. Thus, we perform the + # computation only for `lut[1:]` with `cum_hist[:-1]` and add `lut[0] == 0` afterwards. + cum_hist = cum_hist[..., :-1] + ( + cum_hist.add_(step // 2) + # We need the `clamp_`(min=1) call here to avoid zero division since they fail for integer dtypes. This has no + # effect on the returned result of this kernel since images inside the batch with `step == 0` are returned as is + # instead of equalized version. + .div_(step.clamp_(min=1), rounding_mode="floor") + # We need the `clamp_` call here since PILs LUT computation scheme can produce values outside the valid value + # range of uint8 images + .clamp_(0, 255) + ) + lut = cum_hist.to(torch.uint8) + lut = torch.cat([lut.new_zeros(1).expand(batch_shape + (1,)), lut], dim=-1) + equalized_image = lut.gather(dim=-1, index=flat_image).view_as(image) + + output = torch.where(valid_equalization, equalized_image, image) + return to_dtype_image(output, output_dtype, scale=True) + + +_equalize_image_pil = _register_kernel_internal(equalize, PIL.Image.Image)(_FP.equalize) + + +@_register_kernel_internal(equalize, tv_tensors.Video) +def equalize_video(video: torch.Tensor) -> torch.Tensor: + return equalize_image(video) + + +def invert(inpt: torch.Tensor) -> torch.Tensor: + """See :func:`~torchvision.transforms.v2.RandomInvert`.""" + if torch.jit.is_scripting(): + return invert_image(inpt) + + _log_api_usage_once(invert) + + kernel = _get_kernel(invert, type(inpt)) + return kernel(inpt) + + +@_register_kernel_internal(invert, torch.Tensor) +@_register_kernel_internal(invert, tv_tensors.Image) +def invert_image(image: torch.Tensor) -> torch.Tensor: + if image.is_floating_point(): + return 1.0 - image + elif image.dtype == torch.uint8: + return image.bitwise_not() + else: # signed integer dtypes + # We can't use `Tensor.bitwise_not` here, since we want to retain the leading zero bit that encodes the sign + return image.bitwise_xor((1 << _num_value_bits(image.dtype)) - 1) + + +_invert_image_pil = _register_kernel_internal(invert, PIL.Image.Image)(_FP.invert) + + +@_register_kernel_internal(invert, tv_tensors.Video) +def invert_video(video: torch.Tensor) -> torch.Tensor: + return invert_image(video) + + +def permute_channels(inpt: torch.Tensor, permutation: List[int]) -> torch.Tensor: + """Permute the channels of the input according to the given permutation. + + This function supports plain :class:`~torch.Tensor`'s, :class:`PIL.Image.Image`'s, and + :class:`torchvision.tv_tensors.Image` and :class:`torchvision.tv_tensors.Video`. + + Example: + >>> rgb_image = torch.rand(3, 256, 256) + >>> bgr_image = F.permute_channels(rgb_image, permutation=[2, 1, 0]) + + Args: + permutation (List[int]): Valid permutation of the input channel indices. The index of the element determines the + channel index in the input and the value determines the channel index in the output. For example, + ``permutation=[2, 0 , 1]`` + + - takes ``ìnpt[..., 0, :, :]`` and puts it at ``output[..., 2, :, :]``, + - takes ``ìnpt[..., 1, :, :]`` and puts it at ``output[..., 0, :, :]``, and + - takes ``ìnpt[..., 2, :, :]`` and puts it at ``output[..., 1, :, :]``. + + Raises: + ValueError: If ``len(permutation)`` doesn't match the number of channels in the input. + """ + if torch.jit.is_scripting(): + return permute_channels_image(inpt, permutation=permutation) + + _log_api_usage_once(permute_channels) + + kernel = _get_kernel(permute_channels, type(inpt)) + return kernel(inpt, permutation=permutation) + + +@_register_kernel_internal(permute_channels, torch.Tensor) +@_register_kernel_internal(permute_channels, tv_tensors.Image) +def permute_channels_image(image: torch.Tensor, permutation: List[int]) -> torch.Tensor: + shape = image.shape + num_channels, height, width = shape[-3:] + + if len(permutation) != num_channels: + raise ValueError( + f"Length of permutation does not match number of channels: " f"{len(permutation)} != {num_channels}" + ) + + if image.numel() == 0: + return image + + image = image.reshape(-1, num_channels, height, width) + image = image[:, permutation, :, :] + return image.reshape(shape) + + +@_register_kernel_internal(permute_channels, PIL.Image.Image) +def _permute_channels_image_pil(image: PIL.Image.Image, permutation: List[int]) -> PIL.Image.Image: + return to_pil_image(permute_channels_image(pil_to_tensor(image), permutation=permutation)) + + +@_register_kernel_internal(permute_channels, tv_tensors.Video) +def permute_channels_video(video: torch.Tensor, permutation: List[int]) -> torch.Tensor: + return permute_channels_image(video, permutation=permutation) diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_deprecated.py b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_deprecated.py new file mode 100644 index 0000000000000000000000000000000000000000..116ea31587a1d7f7172267898152e0167531f303 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_deprecated.py @@ -0,0 +1,24 @@ +import warnings +from typing import Any, List + +import torch + +from torchvision.transforms import functional as _F + + +@torch.jit.unused +def to_tensor(inpt: Any) -> torch.Tensor: + """[DEPREACTED] Use to_image() and to_dtype() instead.""" + warnings.warn( + "The function `to_tensor(...)` is deprecated and will be removed in a future release. " + "Instead, please use `to_image(...)` followed by `to_dtype(..., dtype=torch.float32, scale=True)`." + ) + return _F.to_tensor(inpt) + + +def get_image_size(inpt: torch.Tensor) -> List[int]: + warnings.warn( + "The function `get_image_size(...)` is deprecated and will be removed in a future release. " + "Instead, please use `get_size(...)` which returns `[h, w]` instead of `[w, h]`." + ) + return _F.get_image_size(inpt) diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_geometry.py b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..da080e437c99e0b173c496a9097bbae6ae020d7c --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_geometry.py @@ -0,0 +1,2377 @@ +import math +import numbers +import warnings +from typing import Any, List, Optional, Sequence, Tuple, Union + +import PIL.Image +import torch +from torch.nn.functional import grid_sample, interpolate, pad as torch_pad + +from torchvision import tv_tensors +from torchvision.transforms import _functional_pil as _FP +from torchvision.transforms._functional_tensor import _pad_symmetric +from torchvision.transforms.functional import ( + _compute_resized_output_size as __compute_resized_output_size, + _get_perspective_coeffs, + _interpolation_modes_from_int, + InterpolationMode, + pil_modes_mapping, + pil_to_tensor, + to_pil_image, +) + +from torchvision.utils import _log_api_usage_once + +from ._meta import _get_size_image_pil, clamp_bounding_boxes, convert_bounding_box_format + +from ._utils import _FillTypeJIT, _get_kernel, _register_five_ten_crop_kernel_internal, _register_kernel_internal + + +def _check_interpolation(interpolation: Union[InterpolationMode, int]) -> InterpolationMode: + if isinstance(interpolation, int): + interpolation = _interpolation_modes_from_int(interpolation) + elif not isinstance(interpolation, InterpolationMode): + raise ValueError( + f"Argument interpolation should be an `InterpolationMode` or a corresponding Pillow integer constant, " + f"but got {interpolation}." + ) + return interpolation + + +def horizontal_flip(inpt: torch.Tensor) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.RandomHorizontalFlip` for details.""" + if torch.jit.is_scripting(): + return horizontal_flip_image(inpt) + + _log_api_usage_once(horizontal_flip) + + kernel = _get_kernel(horizontal_flip, type(inpt)) + return kernel(inpt) + + +@_register_kernel_internal(horizontal_flip, torch.Tensor) +@_register_kernel_internal(horizontal_flip, tv_tensors.Image) +def horizontal_flip_image(image: torch.Tensor) -> torch.Tensor: + return image.flip(-1) + + +@_register_kernel_internal(horizontal_flip, PIL.Image.Image) +def _horizontal_flip_image_pil(image: PIL.Image.Image) -> PIL.Image.Image: + return _FP.hflip(image) + + +@_register_kernel_internal(horizontal_flip, tv_tensors.Mask) +def horizontal_flip_mask(mask: torch.Tensor) -> torch.Tensor: + return horizontal_flip_image(mask) + + +def horizontal_flip_bounding_boxes( + bounding_boxes: torch.Tensor, format: tv_tensors.BoundingBoxFormat, canvas_size: Tuple[int, int] +) -> torch.Tensor: + shape = bounding_boxes.shape + + bounding_boxes = bounding_boxes.clone().reshape(-1, 4) + + if format == tv_tensors.BoundingBoxFormat.XYXY: + bounding_boxes[:, [2, 0]] = bounding_boxes[:, [0, 2]].sub_(canvas_size[1]).neg_() + elif format == tv_tensors.BoundingBoxFormat.XYWH: + bounding_boxes[:, 0].add_(bounding_boxes[:, 2]).sub_(canvas_size[1]).neg_() + else: # format == tv_tensors.BoundingBoxFormat.CXCYWH: + bounding_boxes[:, 0].sub_(canvas_size[1]).neg_() + + return bounding_boxes.reshape(shape) + + +@_register_kernel_internal(horizontal_flip, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False) +def _horizontal_flip_bounding_boxes_dispatch(inpt: tv_tensors.BoundingBoxes) -> tv_tensors.BoundingBoxes: + output = horizontal_flip_bounding_boxes( + inpt.as_subclass(torch.Tensor), format=inpt.format, canvas_size=inpt.canvas_size + ) + return tv_tensors.wrap(output, like=inpt) + + +@_register_kernel_internal(horizontal_flip, tv_tensors.Video) +def horizontal_flip_video(video: torch.Tensor) -> torch.Tensor: + return horizontal_flip_image(video) + + +def vertical_flip(inpt: torch.Tensor) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.RandomVerticalFlip` for details.""" + if torch.jit.is_scripting(): + return vertical_flip_image(inpt) + + _log_api_usage_once(vertical_flip) + + kernel = _get_kernel(vertical_flip, type(inpt)) + return kernel(inpt) + + +@_register_kernel_internal(vertical_flip, torch.Tensor) +@_register_kernel_internal(vertical_flip, tv_tensors.Image) +def vertical_flip_image(image: torch.Tensor) -> torch.Tensor: + return image.flip(-2) + + +@_register_kernel_internal(vertical_flip, PIL.Image.Image) +def _vertical_flip_image_pil(image: PIL.Image.Image) -> PIL.Image.Image: + return _FP.vflip(image) + + +@_register_kernel_internal(vertical_flip, tv_tensors.Mask) +def vertical_flip_mask(mask: torch.Tensor) -> torch.Tensor: + return vertical_flip_image(mask) + + +def vertical_flip_bounding_boxes( + bounding_boxes: torch.Tensor, format: tv_tensors.BoundingBoxFormat, canvas_size: Tuple[int, int] +) -> torch.Tensor: + shape = bounding_boxes.shape + + bounding_boxes = bounding_boxes.clone().reshape(-1, 4) + + if format == tv_tensors.BoundingBoxFormat.XYXY: + bounding_boxes[:, [1, 3]] = bounding_boxes[:, [3, 1]].sub_(canvas_size[0]).neg_() + elif format == tv_tensors.BoundingBoxFormat.XYWH: + bounding_boxes[:, 1].add_(bounding_boxes[:, 3]).sub_(canvas_size[0]).neg_() + else: # format == tv_tensors.BoundingBoxFormat.CXCYWH: + bounding_boxes[:, 1].sub_(canvas_size[0]).neg_() + + return bounding_boxes.reshape(shape) + + +@_register_kernel_internal(vertical_flip, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False) +def _vertical_flip_bounding_boxes_dispatch(inpt: tv_tensors.BoundingBoxes) -> tv_tensors.BoundingBoxes: + output = vertical_flip_bounding_boxes( + inpt.as_subclass(torch.Tensor), format=inpt.format, canvas_size=inpt.canvas_size + ) + return tv_tensors.wrap(output, like=inpt) + + +@_register_kernel_internal(vertical_flip, tv_tensors.Video) +def vertical_flip_video(video: torch.Tensor) -> torch.Tensor: + return vertical_flip_image(video) + + +# We changed the names to align them with the transforms, i.e. `RandomHorizontalFlip`. Still, `hflip` and `vflip` are +# prevalent and well understood. Thus, we just alias them without deprecating the old names. +hflip = horizontal_flip +vflip = vertical_flip + + +def _compute_resized_output_size( + canvas_size: Tuple[int, int], size: Optional[List[int]], max_size: Optional[int] = None +) -> List[int]: + if isinstance(size, int): + size = [size] + elif max_size is not None and size is not None and len(size) != 1: + raise ValueError( + "max_size should only be passed if size is None or specifies the length of the smaller edge, " + "i.e. size should be an int or a sequence of length 1 in torchscript mode." + ) + return __compute_resized_output_size(canvas_size, size=size, max_size=max_size, allow_size_none=True) + + +def resize( + inpt: torch.Tensor, + size: Optional[List[int]], + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + max_size: Optional[int] = None, + antialias: Optional[bool] = True, +) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.Resize` for details.""" + if torch.jit.is_scripting(): + return resize_image(inpt, size=size, interpolation=interpolation, max_size=max_size, antialias=antialias) + + _log_api_usage_once(resize) + + kernel = _get_kernel(resize, type(inpt)) + return kernel(inpt, size=size, interpolation=interpolation, max_size=max_size, antialias=antialias) + + +# This is an internal helper method for resize_image. We should put it here instead of keeping it +# inside resize_image due to torchscript. +# uint8 dtype support for bilinear and bicubic is limited to cpu and +# according to our benchmarks on eager, non-AVX CPUs should still prefer u8->f32->interpolate->u8 path for bilinear +def _do_native_uint8_resize_on_cpu(interpolation: InterpolationMode) -> bool: + if interpolation == InterpolationMode.BILINEAR: + if torch.compiler.is_compiling(): + return True + else: + return "AVX2" in torch.backends.cpu.get_cpu_capability() + + return interpolation == InterpolationMode.BICUBIC + + +@_register_kernel_internal(resize, torch.Tensor) +@_register_kernel_internal(resize, tv_tensors.Image) +def resize_image( + image: torch.Tensor, + size: Optional[List[int]], + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + max_size: Optional[int] = None, + antialias: Optional[bool] = True, +) -> torch.Tensor: + interpolation = _check_interpolation(interpolation) + antialias = False if antialias is None else antialias + align_corners: Optional[bool] = None + if interpolation == InterpolationMode.BILINEAR or interpolation == InterpolationMode.BICUBIC: + align_corners = False + else: + # The default of antialias is True from 0.17, so we don't warn or + # error if other interpolation modes are used. This is documented. + antialias = False + + shape = image.shape + numel = image.numel() + num_channels, old_height, old_width = shape[-3:] + new_height, new_width = _compute_resized_output_size((old_height, old_width), size=size, max_size=max_size) + + if (new_height, new_width) == (old_height, old_width): + return image + elif numel > 0: + dtype = image.dtype + acceptable_dtypes = [torch.float32, torch.float64] + if interpolation == InterpolationMode.NEAREST or interpolation == InterpolationMode.NEAREST_EXACT: + # uint8 dtype can be included for cpu and cuda input if nearest mode + acceptable_dtypes.append(torch.uint8) + elif image.device.type == "cpu": + if _do_native_uint8_resize_on_cpu(interpolation): + acceptable_dtypes.append(torch.uint8) + + image = image.reshape(-1, num_channels, old_height, old_width) + strides = image.stride() + if image.is_contiguous(memory_format=torch.channels_last) and image.shape[0] == 1 and numel != strides[0]: + # There is a weird behaviour in torch core where the output tensor of `interpolate()` can be allocated as + # contiguous even though the input is un-ambiguously channels_last (https://github.com/pytorch/pytorch/issues/68430). + # In particular this happens for the typical torchvision use-case of single CHW images where we fake the batch dim + # to become 1CHW. Below, we restride those tensors to trick torch core into properly allocating the output as + # channels_last, thus preserving the memory format of the input. This is not just for format consistency: + # for uint8 bilinear images, this also avoids an extra copy (re-packing) of the output and saves time. + # TODO: when https://github.com/pytorch/pytorch/issues/68430 is fixed (possibly by https://github.com/pytorch/pytorch/pull/100373), + # we should be able to remove this hack. + new_strides = list(strides) + new_strides[0] = numel + image = image.as_strided((1, num_channels, old_height, old_width), new_strides) + + need_cast = dtype not in acceptable_dtypes + if need_cast: + image = image.to(dtype=torch.float32) + + image = interpolate( + image, + size=[new_height, new_width], + mode=interpolation.value, + align_corners=align_corners, + antialias=antialias, + ) + + if need_cast: + if interpolation == InterpolationMode.BICUBIC and dtype == torch.uint8: + # This path is hit on non-AVX archs, or on GPU. + image = image.clamp_(min=0, max=255) + if dtype in (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64): + image = image.round_() + image = image.to(dtype=dtype) + + return image.reshape(shape[:-3] + (num_channels, new_height, new_width)) + + +def _resize_image_pil( + image: PIL.Image.Image, + size: Union[Sequence[int], int], + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + max_size: Optional[int] = None, +) -> PIL.Image.Image: + old_height, old_width = image.height, image.width + new_height, new_width = _compute_resized_output_size( + (old_height, old_width), + size=size, # type: ignore[arg-type] + max_size=max_size, + ) + + interpolation = _check_interpolation(interpolation) + + if (new_height, new_width) == (old_height, old_width): + return image + + return image.resize((new_width, new_height), resample=pil_modes_mapping[interpolation]) + + +@_register_kernel_internal(resize, PIL.Image.Image) +def __resize_image_pil_dispatch( + image: PIL.Image.Image, + size: Union[Sequence[int], int], + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + max_size: Optional[int] = None, + antialias: Optional[bool] = True, +) -> PIL.Image.Image: + if antialias is False: + warnings.warn("Anti-alias option is always applied for PIL Image input. Argument antialias is ignored.") + return _resize_image_pil(image, size=size, interpolation=interpolation, max_size=max_size) + + +def resize_mask(mask: torch.Tensor, size: Optional[List[int]], max_size: Optional[int] = None) -> torch.Tensor: + if mask.ndim < 3: + mask = mask.unsqueeze(0) + needs_squeeze = True + else: + needs_squeeze = False + + output = resize_image(mask, size=size, interpolation=InterpolationMode.NEAREST, max_size=max_size) + + if needs_squeeze: + output = output.squeeze(0) + + return output + + +@_register_kernel_internal(resize, tv_tensors.Mask, tv_tensor_wrapper=False) +def _resize_mask_dispatch( + inpt: tv_tensors.Mask, size: List[int], max_size: Optional[int] = None, **kwargs: Any +) -> tv_tensors.Mask: + output = resize_mask(inpt.as_subclass(torch.Tensor), size, max_size=max_size) + return tv_tensors.wrap(output, like=inpt) + + +def resize_bounding_boxes( + bounding_boxes: torch.Tensor, + canvas_size: Tuple[int, int], + size: Optional[List[int]], + max_size: Optional[int] = None, +) -> Tuple[torch.Tensor, Tuple[int, int]]: + old_height, old_width = canvas_size + new_height, new_width = _compute_resized_output_size(canvas_size, size=size, max_size=max_size) + + if (new_height, new_width) == (old_height, old_width): + return bounding_boxes, canvas_size + + w_ratio = new_width / old_width + h_ratio = new_height / old_height + ratios = torch.tensor([w_ratio, h_ratio, w_ratio, h_ratio], device=bounding_boxes.device) + return ( + bounding_boxes.mul(ratios).to(bounding_boxes.dtype), + (new_height, new_width), + ) + + +@_register_kernel_internal(resize, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False) +def _resize_bounding_boxes_dispatch( + inpt: tv_tensors.BoundingBoxes, size: Optional[List[int]], max_size: Optional[int] = None, **kwargs: Any +) -> tv_tensors.BoundingBoxes: + output, canvas_size = resize_bounding_boxes( + inpt.as_subclass(torch.Tensor), inpt.canvas_size, size, max_size=max_size + ) + return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size) + + +@_register_kernel_internal(resize, tv_tensors.Video) +def resize_video( + video: torch.Tensor, + size: Optional[List[int]], + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + max_size: Optional[int] = None, + antialias: Optional[bool] = True, +) -> torch.Tensor: + return resize_image(video, size=size, interpolation=interpolation, max_size=max_size, antialias=antialias) + + +def affine( + inpt: torch.Tensor, + angle: Union[int, float], + translate: List[float], + scale: float, + shear: List[float], + interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST, + fill: _FillTypeJIT = None, + center: Optional[List[float]] = None, +) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.RandomAffine` for details.""" + if torch.jit.is_scripting(): + return affine_image( + inpt, + angle=angle, + translate=translate, + scale=scale, + shear=shear, + interpolation=interpolation, + fill=fill, + center=center, + ) + + _log_api_usage_once(affine) + + kernel = _get_kernel(affine, type(inpt)) + return kernel( + inpt, + angle=angle, + translate=translate, + scale=scale, + shear=shear, + interpolation=interpolation, + fill=fill, + center=center, + ) + + +def _affine_parse_args( + angle: Union[int, float], + translate: List[float], + scale: float, + shear: List[float], + interpolation: InterpolationMode = InterpolationMode.NEAREST, + center: Optional[List[float]] = None, +) -> Tuple[float, List[float], List[float], Optional[List[float]]]: + if not isinstance(angle, (int, float)): + raise TypeError("Argument angle should be int or float") + + if not isinstance(translate, (list, tuple)): + raise TypeError("Argument translate should be a sequence") + + if len(translate) != 2: + raise ValueError("Argument translate should be a sequence of length 2") + + if scale <= 0.0: + raise ValueError("Argument scale should be positive") + + if not isinstance(shear, (numbers.Number, (list, tuple))): + raise TypeError("Shear should be either a single value or a sequence of two values") + + if not isinstance(interpolation, InterpolationMode): + raise TypeError("Argument interpolation should be a InterpolationMode") + + if isinstance(angle, int): + angle = float(angle) + + if isinstance(translate, tuple): + translate = list(translate) + + if isinstance(shear, numbers.Number): + shear = [shear, 0.0] + + if isinstance(shear, tuple): + shear = list(shear) + + if len(shear) == 1: + shear = [shear[0], shear[0]] + + if len(shear) != 2: + raise ValueError(f"Shear should be a sequence containing two values. Got {shear}") + + if center is not None: + if not isinstance(center, (list, tuple)): + raise TypeError("Argument center should be a sequence") + else: + center = [float(c) for c in center] + + return angle, translate, shear, center + + +def _get_inverse_affine_matrix( + center: List[float], angle: float, translate: List[float], scale: float, shear: List[float], inverted: bool = True +) -> List[float]: + # Helper method to compute inverse matrix for affine transformation + + # Pillow requires inverse affine transformation matrix: + # Affine matrix is : M = T * C * RotateScaleShear * C^-1 + # + # where T is translation matrix: [1, 0, tx | 0, 1, ty | 0, 0, 1] + # C is translation matrix to keep center: [1, 0, cx | 0, 1, cy | 0, 0, 1] + # RotateScaleShear is rotation with scale and shear matrix + # + # RotateScaleShear(a, s, (sx, sy)) = + # = R(a) * S(s) * SHy(sy) * SHx(sx) + # = [ s*cos(a - sy)/cos(sy), s*(-cos(a - sy)*tan(sx)/cos(sy) - sin(a)), 0 ] + # [ s*sin(a - sy)/cos(sy), s*(-sin(a - sy)*tan(sx)/cos(sy) + cos(a)), 0 ] + # [ 0 , 0 , 1 ] + # where R is a rotation matrix, S is a scaling matrix, and SHx and SHy are the shears: + # SHx(s) = [1, -tan(s)] and SHy(s) = [1 , 0] + # [0, 1 ] [-tan(s), 1] + # + # Thus, the inverse is M^-1 = C * RotateScaleShear^-1 * C^-1 * T^-1 + + rot = math.radians(angle) + sx = math.radians(shear[0]) + sy = math.radians(shear[1]) + + cx, cy = center + tx, ty = translate + + # Cached results + cos_sy = math.cos(sy) + tan_sx = math.tan(sx) + rot_minus_sy = rot - sy + cx_plus_tx = cx + tx + cy_plus_ty = cy + ty + + # Rotate Scale Shear (RSS) without scaling + a = math.cos(rot_minus_sy) / cos_sy + b = -(a * tan_sx + math.sin(rot)) + c = math.sin(rot_minus_sy) / cos_sy + d = math.cos(rot) - c * tan_sx + + if inverted: + # Inverted rotation matrix with scale and shear + # det([[a, b], [c, d]]) == 1, since det(rotation) = 1 and det(shear) = 1 + matrix = [d / scale, -b / scale, 0.0, -c / scale, a / scale, 0.0] + # Apply inverse of translation and of center translation: RSS^-1 * C^-1 * T^-1 + # and then apply center translation: C * RSS^-1 * C^-1 * T^-1 + matrix[2] += cx - matrix[0] * cx_plus_tx - matrix[1] * cy_plus_ty + matrix[5] += cy - matrix[3] * cx_plus_tx - matrix[4] * cy_plus_ty + else: + matrix = [a * scale, b * scale, 0.0, c * scale, d * scale, 0.0] + # Apply inverse of center translation: RSS * C^-1 + # and then apply translation and center : T * C * RSS * C^-1 + matrix[2] += cx_plus_tx - matrix[0] * cx - matrix[1] * cy + matrix[5] += cy_plus_ty - matrix[3] * cx - matrix[4] * cy + + return matrix + + +def _compute_affine_output_size(matrix: List[float], w: int, h: int) -> Tuple[int, int]: + if torch.compiler.is_compiling() and not torch.jit.is_scripting(): + return _compute_affine_output_size_python(matrix, w, h) + else: + return _compute_affine_output_size_tensor(matrix, w, h) + + +def _compute_affine_output_size_tensor(matrix: List[float], w: int, h: int) -> Tuple[int, int]: + # Inspired of PIL implementation: + # https://github.com/python-pillow/Pillow/blob/11de3318867e4398057373ee9f12dcb33db7335c/src/PIL/Image.py#L2054 + + # pts are Top-Left, Top-Right, Bottom-Left, Bottom-Right points. + # Points are shifted due to affine matrix torch convention about + # the center point. Center is (0, 0) for image center pivot point (w * 0.5, h * 0.5) + half_w = 0.5 * w + half_h = 0.5 * h + pts = torch.tensor( + [ + [-half_w, -half_h, 1.0], + [-half_w, half_h, 1.0], + [half_w, half_h, 1.0], + [half_w, -half_h, 1.0], + ] + ) + theta = torch.tensor(matrix, dtype=torch.float).view(2, 3) + new_pts = torch.matmul(pts, theta.T) + min_vals, max_vals = new_pts.aminmax(dim=0) + + # shift points to [0, w] and [0, h] interval to match PIL results + halfs = torch.tensor((half_w, half_h)) + min_vals.add_(halfs) + max_vals.add_(halfs) + + # Truncate precision to 1e-4 to avoid ceil of Xe-15 to 1.0 + tol = 1e-4 + inv_tol = 1.0 / tol + cmax = max_vals.mul_(inv_tol).trunc_().mul_(tol).ceil_() + cmin = min_vals.mul_(inv_tol).trunc_().mul_(tol).floor_() + size = cmax.sub_(cmin) + return int(size[0]), int(size[1]) # w, h + + +def _compute_affine_output_size_python(matrix: List[float], w: int, h: int) -> Tuple[int, int]: + # Mostly copied from PIL implementation: + # The only difference is with transformed points as input matrix has zero translation part here and + # PIL has a centered translation part. + # https://github.com/python-pillow/Pillow/blob/11de3318867e4398057373ee9f12dcb33db7335c/src/PIL/Image.py#L2054 + + a, b, c, d, e, f = matrix + xx = [] + yy = [] + + half_w = 0.5 * w + half_h = 0.5 * h + for x, y in ((-half_w, -half_h), (half_w, -half_h), (half_w, half_h), (-half_w, half_h)): + nx = a * x + b * y + c + ny = d * x + e * y + f + xx.append(nx + half_w) + yy.append(ny + half_h) + + nw = math.ceil(max(xx)) - math.floor(min(xx)) + nh = math.ceil(max(yy)) - math.floor(min(yy)) + return int(nw), int(nh) # w, h + + +def _apply_grid_transform(img: torch.Tensor, grid: torch.Tensor, mode: str, fill: _FillTypeJIT) -> torch.Tensor: + input_shape = img.shape + output_height, output_width = grid.shape[1], grid.shape[2] + num_channels, input_height, input_width = input_shape[-3:] + output_shape = input_shape[:-3] + (num_channels, output_height, output_width) + + if img.numel() == 0: + return img.reshape(output_shape) + + img = img.reshape(-1, num_channels, input_height, input_width) + squashed_batch_size = img.shape[0] + + # We are using context knowledge that grid should have float dtype + fp = img.dtype == grid.dtype + float_img = img if fp else img.to(grid.dtype) + + if squashed_batch_size > 1: + # Apply same grid to a batch of images + grid = grid.expand(squashed_batch_size, -1, -1, -1) + + # Append a dummy mask for customized fill colors, should be faster than grid_sample() twice + if fill is not None: + mask = torch.ones( + (squashed_batch_size, 1, input_height, input_width), dtype=float_img.dtype, device=float_img.device + ) + float_img = torch.cat((float_img, mask), dim=1) + + float_img = grid_sample(float_img, grid, mode=mode, padding_mode="zeros", align_corners=False) + + # Fill with required color + if fill is not None: + float_img, mask = torch.tensor_split(float_img, indices=(-1,), dim=-3) + mask = mask.expand_as(float_img) + fill_list = fill if isinstance(fill, (tuple, list)) else [float(fill)] # type: ignore[arg-type] + fill_img = torch.tensor(fill_list, dtype=float_img.dtype, device=float_img.device).view(1, -1, 1, 1) + if mode == "nearest": + float_img = torch.where(mask < 0.5, fill_img.expand_as(float_img), float_img) + else: # 'bilinear' + # The following is mathematically equivalent to: + # img * mask + (1.0 - mask) * fill = img * mask - fill * mask + fill = mask * (img - fill) + fill + float_img = float_img.sub_(fill_img).mul_(mask).add_(fill_img) + + img = float_img.round_().to(img.dtype) if not fp else float_img + + return img.reshape(output_shape) + + +def _assert_grid_transform_inputs( + image: torch.Tensor, + matrix: Optional[List[float]], + interpolation: str, + fill: _FillTypeJIT, + supported_interpolation_modes: List[str], + coeffs: Optional[List[float]] = None, +) -> None: + if matrix is not None: + if not isinstance(matrix, list): + raise TypeError("Argument matrix should be a list") + elif len(matrix) != 6: + raise ValueError("Argument matrix should have 6 float values") + + if coeffs is not None and len(coeffs) != 8: + raise ValueError("Argument coeffs should have 8 float values") + + if fill is not None: + if isinstance(fill, (tuple, list)): + length = len(fill) + num_channels = image.shape[-3] + if length > 1 and length != num_channels: + raise ValueError( + "The number of elements in 'fill' cannot broadcast to match the number of " + f"channels of the image ({length} != {num_channels})" + ) + elif not isinstance(fill, (int, float)): + raise ValueError("Argument fill should be either int, float, tuple or list") + + if interpolation not in supported_interpolation_modes: + raise ValueError(f"Interpolation mode '{interpolation}' is unsupported with Tensor input") + + +def _affine_grid( + theta: torch.Tensor, + w: int, + h: int, + ow: int, + oh: int, +) -> torch.Tensor: + # https://github.com/pytorch/pytorch/blob/74b65c32be68b15dc7c9e8bb62459efbfbde33d8/aten/src/ATen/native/ + # AffineGridGenerator.cpp#L18 + # Difference with AffineGridGenerator is that: + # 1) we normalize grid values after applying theta + # 2) we can normalize by other image size, such that it covers "extend" option like in PIL.Image.rotate + dtype = theta.dtype + device = theta.device + + base_grid = torch.empty(1, oh, ow, 3, dtype=dtype, device=device) + x_grid = torch.linspace((1.0 - ow) * 0.5, (ow - 1.0) * 0.5, steps=ow, device=device) + base_grid[..., 0].copy_(x_grid) + y_grid = torch.linspace((1.0 - oh) * 0.5, (oh - 1.0) * 0.5, steps=oh, device=device).unsqueeze_(-1) + base_grid[..., 1].copy_(y_grid) + base_grid[..., 2].fill_(1) + + rescaled_theta = theta.transpose(1, 2).div_(torch.tensor([0.5 * w, 0.5 * h], dtype=dtype, device=device)) + output_grid = base_grid.view(1, oh * ow, 3).bmm(rescaled_theta) + return output_grid.view(1, oh, ow, 2) + + +@_register_kernel_internal(affine, torch.Tensor) +@_register_kernel_internal(affine, tv_tensors.Image) +def affine_image( + image: torch.Tensor, + angle: Union[int, float], + translate: List[float], + scale: float, + shear: List[float], + interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST, + fill: _FillTypeJIT = None, + center: Optional[List[float]] = None, +) -> torch.Tensor: + interpolation = _check_interpolation(interpolation) + + angle, translate, shear, center = _affine_parse_args(angle, translate, scale, shear, interpolation, center) + + height, width = image.shape[-2:] + + center_f = [0.0, 0.0] + if center is not None: + # Center values should be in pixel coordinates but translated such that (0, 0) corresponds to image center. + center_f = [(c - s * 0.5) for c, s in zip(center, [width, height])] + + translate_f = [float(t) for t in translate] + matrix = _get_inverse_affine_matrix(center_f, angle, translate_f, scale, shear) + + _assert_grid_transform_inputs(image, matrix, interpolation.value, fill, ["nearest", "bilinear"]) + + dtype = image.dtype if torch.is_floating_point(image) else torch.float32 + theta = torch.tensor(matrix, dtype=dtype, device=image.device).reshape(1, 2, 3) + grid = _affine_grid(theta, w=width, h=height, ow=width, oh=height) + return _apply_grid_transform(image, grid, interpolation.value, fill=fill) + + +@_register_kernel_internal(affine, PIL.Image.Image) +def _affine_image_pil( + image: PIL.Image.Image, + angle: Union[int, float], + translate: List[float], + scale: float, + shear: List[float], + interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST, + fill: _FillTypeJIT = None, + center: Optional[List[float]] = None, +) -> PIL.Image.Image: + interpolation = _check_interpolation(interpolation) + angle, translate, shear, center = _affine_parse_args(angle, translate, scale, shear, interpolation, center) + + # center = (img_size[0] * 0.5 + 0.5, img_size[1] * 0.5 + 0.5) + # it is visually better to estimate the center without 0.5 offset + # otherwise image rotated by 90 degrees is shifted vs output image of torch.rot90 or F_t.affine + if center is None: + height, width = _get_size_image_pil(image) + center = [width * 0.5, height * 0.5] + matrix = _get_inverse_affine_matrix(center, angle, translate, scale, shear) + + return _FP.affine(image, matrix, interpolation=pil_modes_mapping[interpolation], fill=fill) + + +def _affine_bounding_boxes_with_expand( + bounding_boxes: torch.Tensor, + format: tv_tensors.BoundingBoxFormat, + canvas_size: Tuple[int, int], + angle: Union[int, float], + translate: List[float], + scale: float, + shear: List[float], + center: Optional[List[float]] = None, + expand: bool = False, +) -> Tuple[torch.Tensor, Tuple[int, int]]: + if bounding_boxes.numel() == 0: + return bounding_boxes, canvas_size + + original_shape = bounding_boxes.shape + original_dtype = bounding_boxes.dtype + bounding_boxes = bounding_boxes.clone() if bounding_boxes.is_floating_point() else bounding_boxes.float() + dtype = bounding_boxes.dtype + device = bounding_boxes.device + bounding_boxes = ( + convert_bounding_box_format( + bounding_boxes, old_format=format, new_format=tv_tensors.BoundingBoxFormat.XYXY, inplace=True + ) + ).reshape(-1, 4) + + angle, translate, shear, center = _affine_parse_args( + angle, translate, scale, shear, InterpolationMode.NEAREST, center + ) + + if center is None: + height, width = canvas_size + center = [width * 0.5, height * 0.5] + + affine_vector = _get_inverse_affine_matrix(center, angle, translate, scale, shear, inverted=False) + transposed_affine_matrix = ( + torch.tensor( + affine_vector, + dtype=dtype, + device=device, + ) + .reshape(2, 3) + .T + ) + # 1) Let's transform bboxes into a tensor of 4 points (top-left, top-right, bottom-left, bottom-right corners). + # Tensor of points has shape (N * 4, 3), where N is the number of bboxes + # Single point structure is similar to + # [(xmin, ymin, 1), (xmax, ymin, 1), (xmax, ymax, 1), (xmin, ymax, 1)] + points = bounding_boxes[:, [[0, 1], [2, 1], [2, 3], [0, 3]]].reshape(-1, 2) + points = torch.cat([points, torch.ones(points.shape[0], 1, device=device, dtype=dtype)], dim=-1) + # 2) Now let's transform the points using affine matrix + transformed_points = torch.matmul(points, transposed_affine_matrix) + # 3) Reshape transformed points to [N boxes, 4 points, x/y coords] + # and compute bounding box from 4 transformed points: + transformed_points = transformed_points.reshape(-1, 4, 2) + out_bbox_mins, out_bbox_maxs = torch.aminmax(transformed_points, dim=1) + out_bboxes = torch.cat([out_bbox_mins, out_bbox_maxs], dim=1) + + if expand: + # Compute minimum point for transformed image frame: + # Points are Top-Left, Top-Right, Bottom-Left, Bottom-Right points. + height, width = canvas_size + points = torch.tensor( + [ + [0.0, 0.0, 1.0], + [0.0, float(height), 1.0], + [float(width), float(height), 1.0], + [float(width), 0.0, 1.0], + ], + dtype=dtype, + device=device, + ) + new_points = torch.matmul(points, transposed_affine_matrix) + tr = torch.amin(new_points, dim=0, keepdim=True) + # Translate bounding boxes + out_bboxes.sub_(tr.repeat((1, 2))) + # Estimate meta-data for image with inverted=True + affine_vector = _get_inverse_affine_matrix(center, angle, translate, scale, shear) + new_width, new_height = _compute_affine_output_size(affine_vector, width, height) + canvas_size = (new_height, new_width) + + out_bboxes = clamp_bounding_boxes(out_bboxes, format=tv_tensors.BoundingBoxFormat.XYXY, canvas_size=canvas_size) + out_bboxes = convert_bounding_box_format( + out_bboxes, old_format=tv_tensors.BoundingBoxFormat.XYXY, new_format=format, inplace=True + ).reshape(original_shape) + + out_bboxes = out_bboxes.to(original_dtype) + return out_bboxes, canvas_size + + +def affine_bounding_boxes( + bounding_boxes: torch.Tensor, + format: tv_tensors.BoundingBoxFormat, + canvas_size: Tuple[int, int], + angle: Union[int, float], + translate: List[float], + scale: float, + shear: List[float], + center: Optional[List[float]] = None, +) -> torch.Tensor: + out_box, _ = _affine_bounding_boxes_with_expand( + bounding_boxes, + format=format, + canvas_size=canvas_size, + angle=angle, + translate=translate, + scale=scale, + shear=shear, + center=center, + expand=False, + ) + return out_box + + +@_register_kernel_internal(affine, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False) +def _affine_bounding_boxes_dispatch( + inpt: tv_tensors.BoundingBoxes, + angle: Union[int, float], + translate: List[float], + scale: float, + shear: List[float], + center: Optional[List[float]] = None, + **kwargs, +) -> tv_tensors.BoundingBoxes: + output = affine_bounding_boxes( + inpt.as_subclass(torch.Tensor), + format=inpt.format, + canvas_size=inpt.canvas_size, + angle=angle, + translate=translate, + scale=scale, + shear=shear, + center=center, + ) + return tv_tensors.wrap(output, like=inpt) + + +def affine_mask( + mask: torch.Tensor, + angle: Union[int, float], + translate: List[float], + scale: float, + shear: List[float], + fill: _FillTypeJIT = None, + center: Optional[List[float]] = None, +) -> torch.Tensor: + if mask.ndim < 3: + mask = mask.unsqueeze(0) + needs_squeeze = True + else: + needs_squeeze = False + + output = affine_image( + mask, + angle=angle, + translate=translate, + scale=scale, + shear=shear, + interpolation=InterpolationMode.NEAREST, + fill=fill, + center=center, + ) + + if needs_squeeze: + output = output.squeeze(0) + + return output + + +@_register_kernel_internal(affine, tv_tensors.Mask, tv_tensor_wrapper=False) +def _affine_mask_dispatch( + inpt: tv_tensors.Mask, + angle: Union[int, float], + translate: List[float], + scale: float, + shear: List[float], + fill: _FillTypeJIT = None, + center: Optional[List[float]] = None, + **kwargs, +) -> tv_tensors.Mask: + output = affine_mask( + inpt.as_subclass(torch.Tensor), + angle=angle, + translate=translate, + scale=scale, + shear=shear, + fill=fill, + center=center, + ) + return tv_tensors.wrap(output, like=inpt) + + +@_register_kernel_internal(affine, tv_tensors.Video) +def affine_video( + video: torch.Tensor, + angle: Union[int, float], + translate: List[float], + scale: float, + shear: List[float], + interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST, + fill: _FillTypeJIT = None, + center: Optional[List[float]] = None, +) -> torch.Tensor: + return affine_image( + video, + angle=angle, + translate=translate, + scale=scale, + shear=shear, + interpolation=interpolation, + fill=fill, + center=center, + ) + + +def rotate( + inpt: torch.Tensor, + angle: float, + interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST, + expand: bool = False, + center: Optional[List[float]] = None, + fill: _FillTypeJIT = None, +) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.RandomRotation` for details.""" + if torch.jit.is_scripting(): + return rotate_image(inpt, angle=angle, interpolation=interpolation, expand=expand, fill=fill, center=center) + + _log_api_usage_once(rotate) + + kernel = _get_kernel(rotate, type(inpt)) + return kernel(inpt, angle=angle, interpolation=interpolation, expand=expand, fill=fill, center=center) + + +@_register_kernel_internal(rotate, torch.Tensor) +@_register_kernel_internal(rotate, tv_tensors.Image) +def rotate_image( + image: torch.Tensor, + angle: float, + interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST, + expand: bool = False, + center: Optional[List[float]] = None, + fill: _FillTypeJIT = None, +) -> torch.Tensor: + angle = angle % 360 # shift angle to [0, 360) range + + # fast path: transpose without affine transform + if center is None: + if angle == 0: + return image.clone() + if angle == 180: + return torch.rot90(image, k=2, dims=(-2, -1)) + + if expand or image.shape[-1] == image.shape[-2]: + if angle == 90: + return torch.rot90(image, k=1, dims=(-2, -1)) + if angle == 270: + return torch.rot90(image, k=3, dims=(-2, -1)) + + interpolation = _check_interpolation(interpolation) + + input_height, input_width = image.shape[-2:] + + center_f = [0.0, 0.0] + if center is not None: + # Center values should be in pixel coordinates but translated such that (0, 0) corresponds to image center. + center_f = [(c - s * 0.5) for c, s in zip(center, [input_width, input_height])] + + # due to current incoherence of rotation angle direction between affine and rotate implementations + # we need to set -angle. + matrix = _get_inverse_affine_matrix(center_f, -angle, [0.0, 0.0], 1.0, [0.0, 0.0]) + + _assert_grid_transform_inputs(image, matrix, interpolation.value, fill, ["nearest", "bilinear"]) + + output_width, output_height = ( + _compute_affine_output_size(matrix, input_width, input_height) if expand else (input_width, input_height) + ) + dtype = image.dtype if torch.is_floating_point(image) else torch.float32 + theta = torch.tensor(matrix, dtype=dtype, device=image.device).reshape(1, 2, 3) + grid = _affine_grid(theta, w=input_width, h=input_height, ow=output_width, oh=output_height) + return _apply_grid_transform(image, grid, interpolation.value, fill=fill) + + +@_register_kernel_internal(rotate, PIL.Image.Image) +def _rotate_image_pil( + image: PIL.Image.Image, + angle: float, + interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST, + expand: bool = False, + center: Optional[List[float]] = None, + fill: _FillTypeJIT = None, +) -> PIL.Image.Image: + interpolation = _check_interpolation(interpolation) + + return _FP.rotate( + image, angle, interpolation=pil_modes_mapping[interpolation], expand=expand, fill=fill, center=center + ) + + +def rotate_bounding_boxes( + bounding_boxes: torch.Tensor, + format: tv_tensors.BoundingBoxFormat, + canvas_size: Tuple[int, int], + angle: float, + expand: bool = False, + center: Optional[List[float]] = None, +) -> Tuple[torch.Tensor, Tuple[int, int]]: + return _affine_bounding_boxes_with_expand( + bounding_boxes, + format=format, + canvas_size=canvas_size, + angle=-angle, + translate=[0.0, 0.0], + scale=1.0, + shear=[0.0, 0.0], + center=center, + expand=expand, + ) + + +@_register_kernel_internal(rotate, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False) +def _rotate_bounding_boxes_dispatch( + inpt: tv_tensors.BoundingBoxes, angle: float, expand: bool = False, center: Optional[List[float]] = None, **kwargs +) -> tv_tensors.BoundingBoxes: + output, canvas_size = rotate_bounding_boxes( + inpt.as_subclass(torch.Tensor), + format=inpt.format, + canvas_size=inpt.canvas_size, + angle=angle, + expand=expand, + center=center, + ) + return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size) + + +def rotate_mask( + mask: torch.Tensor, + angle: float, + expand: bool = False, + center: Optional[List[float]] = None, + fill: _FillTypeJIT = None, +) -> torch.Tensor: + if mask.ndim < 3: + mask = mask.unsqueeze(0) + needs_squeeze = True + else: + needs_squeeze = False + + output = rotate_image( + mask, + angle=angle, + expand=expand, + interpolation=InterpolationMode.NEAREST, + fill=fill, + center=center, + ) + + if needs_squeeze: + output = output.squeeze(0) + + return output + + +@_register_kernel_internal(rotate, tv_tensors.Mask, tv_tensor_wrapper=False) +def _rotate_mask_dispatch( + inpt: tv_tensors.Mask, + angle: float, + expand: bool = False, + center: Optional[List[float]] = None, + fill: _FillTypeJIT = None, + **kwargs, +) -> tv_tensors.Mask: + output = rotate_mask(inpt.as_subclass(torch.Tensor), angle=angle, expand=expand, fill=fill, center=center) + return tv_tensors.wrap(output, like=inpt) + + +@_register_kernel_internal(rotate, tv_tensors.Video) +def rotate_video( + video: torch.Tensor, + angle: float, + interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST, + expand: bool = False, + center: Optional[List[float]] = None, + fill: _FillTypeJIT = None, +) -> torch.Tensor: + return rotate_image(video, angle, interpolation=interpolation, expand=expand, fill=fill, center=center) + + +def pad( + inpt: torch.Tensor, + padding: List[int], + fill: Optional[Union[int, float, List[float]]] = None, + padding_mode: str = "constant", +) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.Pad` for details.""" + if torch.jit.is_scripting(): + return pad_image(inpt, padding=padding, fill=fill, padding_mode=padding_mode) + + _log_api_usage_once(pad) + + kernel = _get_kernel(pad, type(inpt)) + return kernel(inpt, padding=padding, fill=fill, padding_mode=padding_mode) + + +def _parse_pad_padding(padding: Union[int, List[int]]) -> List[int]: + if isinstance(padding, int): + pad_left = pad_right = pad_top = pad_bottom = padding + elif isinstance(padding, (tuple, list)): + if len(padding) == 1: + pad_left = pad_right = pad_top = pad_bottom = padding[0] + elif len(padding) == 2: + pad_left = pad_right = padding[0] + pad_top = pad_bottom = padding[1] + elif len(padding) == 4: + pad_left = padding[0] + pad_top = padding[1] + pad_right = padding[2] + pad_bottom = padding[3] + else: + raise ValueError( + f"Padding must be an int or a 1, 2, or 4 element tuple, not a {len(padding)} element tuple" + ) + else: + raise TypeError(f"`padding` should be an integer or tuple or list of integers, but got {padding}") + + return [pad_left, pad_right, pad_top, pad_bottom] + + +@_register_kernel_internal(pad, torch.Tensor) +@_register_kernel_internal(pad, tv_tensors.Image) +def pad_image( + image: torch.Tensor, + padding: List[int], + fill: Optional[Union[int, float, List[float]]] = None, + padding_mode: str = "constant", +) -> torch.Tensor: + # Be aware that while `padding` has order `[left, top, right, bottom]`, `torch_padding` uses + # `[left, right, top, bottom]`. This stems from the fact that we align our API with PIL, but need to use `torch_pad` + # internally. + torch_padding = _parse_pad_padding(padding) + + if padding_mode not in ("constant", "edge", "reflect", "symmetric"): + raise ValueError( + f"`padding_mode` should be either `'constant'`, `'edge'`, `'reflect'` or `'symmetric'`, " + f"but got `'{padding_mode}'`." + ) + + if fill is None: + fill = 0 + + if isinstance(fill, (int, float)): + return _pad_with_scalar_fill(image, torch_padding, fill=fill, padding_mode=padding_mode) + elif len(fill) == 1: + return _pad_with_scalar_fill(image, torch_padding, fill=fill[0], padding_mode=padding_mode) + else: + return _pad_with_vector_fill(image, torch_padding, fill=fill, padding_mode=padding_mode) + + +def _pad_with_scalar_fill( + image: torch.Tensor, + torch_padding: List[int], + fill: Union[int, float], + padding_mode: str, +) -> torch.Tensor: + shape = image.shape + num_channels, height, width = shape[-3:] + + batch_size = 1 + for s in shape[:-3]: + batch_size *= s + + image = image.reshape(batch_size, num_channels, height, width) + + if padding_mode == "edge": + # Similar to the padding order, `torch_pad`'s PIL's padding modes don't have the same names. Thus, we map + # the PIL name for the padding mode, which we are also using for our API, to the corresponding `torch_pad` + # name. + padding_mode = "replicate" + + if padding_mode == "constant": + image = torch_pad(image, torch_padding, mode=padding_mode, value=float(fill)) + elif padding_mode in ("reflect", "replicate"): + # `torch_pad` only supports `"reflect"` or `"replicate"` padding for floating point inputs. + # TODO: See https://github.com/pytorch/pytorch/issues/40763 + dtype = image.dtype + if not image.is_floating_point(): + needs_cast = True + image = image.to(torch.float32) + else: + needs_cast = False + + image = torch_pad(image, torch_padding, mode=padding_mode) + + if needs_cast: + image = image.to(dtype) + else: # padding_mode == "symmetric" + image = _pad_symmetric(image, torch_padding) + + new_height, new_width = image.shape[-2:] + + return image.reshape(shape[:-3] + (num_channels, new_height, new_width)) + + +# TODO: This should be removed once torch_pad supports non-scalar padding values +def _pad_with_vector_fill( + image: torch.Tensor, + torch_padding: List[int], + fill: List[float], + padding_mode: str, +) -> torch.Tensor: + if padding_mode != "constant": + raise ValueError(f"Padding mode '{padding_mode}' is not supported if fill is not scalar") + + output = _pad_with_scalar_fill(image, torch_padding, fill=0, padding_mode="constant") + left, right, top, bottom = torch_padding + + # We are creating the tensor in the autodetected dtype first and convert to the right one after to avoid an implicit + # float -> int conversion. That happens for example for the valid input of a uint8 image with floating point fill + # value. + fill = torch.tensor(fill, device=image.device).to(dtype=image.dtype).reshape(-1, 1, 1) + + if top > 0: + output[..., :top, :] = fill + if left > 0: + output[..., :, :left] = fill + if bottom > 0: + output[..., -bottom:, :] = fill + if right > 0: + output[..., :, -right:] = fill + return output + + +_pad_image_pil = _register_kernel_internal(pad, PIL.Image.Image)(_FP.pad) + + +@_register_kernel_internal(pad, tv_tensors.Mask) +def pad_mask( + mask: torch.Tensor, + padding: List[int], + fill: Optional[Union[int, float, List[float]]] = None, + padding_mode: str = "constant", +) -> torch.Tensor: + if fill is None: + fill = 0 + + if isinstance(fill, (tuple, list)): + raise ValueError("Non-scalar fill value is not supported") + + if mask.ndim < 3: + mask = mask.unsqueeze(0) + needs_squeeze = True + else: + needs_squeeze = False + + output = pad_image(mask, padding=padding, fill=fill, padding_mode=padding_mode) + + if needs_squeeze: + output = output.squeeze(0) + + return output + + +def pad_bounding_boxes( + bounding_boxes: torch.Tensor, + format: tv_tensors.BoundingBoxFormat, + canvas_size: Tuple[int, int], + padding: List[int], + padding_mode: str = "constant", +) -> Tuple[torch.Tensor, Tuple[int, int]]: + if padding_mode not in ["constant"]: + # TODO: add support of other padding modes + raise ValueError(f"Padding mode '{padding_mode}' is not supported with bounding boxes") + + left, right, top, bottom = _parse_pad_padding(padding) + + if format == tv_tensors.BoundingBoxFormat.XYXY: + pad = [left, top, left, top] + else: + pad = [left, top, 0, 0] + bounding_boxes = bounding_boxes + torch.tensor(pad, dtype=bounding_boxes.dtype, device=bounding_boxes.device) + + height, width = canvas_size + height += top + bottom + width += left + right + canvas_size = (height, width) + + return clamp_bounding_boxes(bounding_boxes, format=format, canvas_size=canvas_size), canvas_size + + +@_register_kernel_internal(pad, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False) +def _pad_bounding_boxes_dispatch( + inpt: tv_tensors.BoundingBoxes, padding: List[int], padding_mode: str = "constant", **kwargs +) -> tv_tensors.BoundingBoxes: + output, canvas_size = pad_bounding_boxes( + inpt.as_subclass(torch.Tensor), + format=inpt.format, + canvas_size=inpt.canvas_size, + padding=padding, + padding_mode=padding_mode, + ) + return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size) + + +@_register_kernel_internal(pad, tv_tensors.Video) +def pad_video( + video: torch.Tensor, + padding: List[int], + fill: Optional[Union[int, float, List[float]]] = None, + padding_mode: str = "constant", +) -> torch.Tensor: + return pad_image(video, padding, fill=fill, padding_mode=padding_mode) + + +def crop(inpt: torch.Tensor, top: int, left: int, height: int, width: int) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.RandomCrop` for details.""" + if torch.jit.is_scripting(): + return crop_image(inpt, top=top, left=left, height=height, width=width) + + _log_api_usage_once(crop) + + kernel = _get_kernel(crop, type(inpt)) + return kernel(inpt, top=top, left=left, height=height, width=width) + + +@_register_kernel_internal(crop, torch.Tensor) +@_register_kernel_internal(crop, tv_tensors.Image) +def crop_image(image: torch.Tensor, top: int, left: int, height: int, width: int) -> torch.Tensor: + h, w = image.shape[-2:] + + right = left + width + bottom = top + height + + if left < 0 or top < 0 or right > w or bottom > h: + image = image[..., max(top, 0) : bottom, max(left, 0) : right] + torch_padding = [ + max(min(right, 0) - left, 0), + max(right - max(w, left), 0), + max(min(bottom, 0) - top, 0), + max(bottom - max(h, top), 0), + ] + return _pad_with_scalar_fill(image, torch_padding, fill=0, padding_mode="constant") + return image[..., top:bottom, left:right] + + +_crop_image_pil = _FP.crop +_register_kernel_internal(crop, PIL.Image.Image)(_crop_image_pil) + + +def crop_bounding_boxes( + bounding_boxes: torch.Tensor, + format: tv_tensors.BoundingBoxFormat, + top: int, + left: int, + height: int, + width: int, +) -> Tuple[torch.Tensor, Tuple[int, int]]: + + # Crop or implicit pad if left and/or top have negative values: + if format == tv_tensors.BoundingBoxFormat.XYXY: + sub = [left, top, left, top] + else: + sub = [left, top, 0, 0] + + bounding_boxes = bounding_boxes - torch.tensor(sub, dtype=bounding_boxes.dtype, device=bounding_boxes.device) + canvas_size = (height, width) + + return clamp_bounding_boxes(bounding_boxes, format=format, canvas_size=canvas_size), canvas_size + + +@_register_kernel_internal(crop, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False) +def _crop_bounding_boxes_dispatch( + inpt: tv_tensors.BoundingBoxes, top: int, left: int, height: int, width: int +) -> tv_tensors.BoundingBoxes: + output, canvas_size = crop_bounding_boxes( + inpt.as_subclass(torch.Tensor), format=inpt.format, top=top, left=left, height=height, width=width + ) + return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size) + + +@_register_kernel_internal(crop, tv_tensors.Mask) +def crop_mask(mask: torch.Tensor, top: int, left: int, height: int, width: int) -> torch.Tensor: + if mask.ndim < 3: + mask = mask.unsqueeze(0) + needs_squeeze = True + else: + needs_squeeze = False + + output = crop_image(mask, top, left, height, width) + + if needs_squeeze: + output = output.squeeze(0) + + return output + + +@_register_kernel_internal(crop, tv_tensors.Video) +def crop_video(video: torch.Tensor, top: int, left: int, height: int, width: int) -> torch.Tensor: + return crop_image(video, top, left, height, width) + + +def perspective( + inpt: torch.Tensor, + startpoints: Optional[List[List[int]]], + endpoints: Optional[List[List[int]]], + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + fill: _FillTypeJIT = None, + coefficients: Optional[List[float]] = None, +) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.RandomPerspective` for details.""" + if torch.jit.is_scripting(): + return perspective_image( + inpt, + startpoints=startpoints, + endpoints=endpoints, + interpolation=interpolation, + fill=fill, + coefficients=coefficients, + ) + + _log_api_usage_once(perspective) + + kernel = _get_kernel(perspective, type(inpt)) + return kernel( + inpt, + startpoints=startpoints, + endpoints=endpoints, + interpolation=interpolation, + fill=fill, + coefficients=coefficients, + ) + + +def _perspective_grid(coeffs: List[float], ow: int, oh: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor: + # https://github.com/python-pillow/Pillow/blob/4634eafe3c695a014267eefdce830b4a825beed7/ + # src/libImaging/Geometry.c#L394 + + # + # x_out = (coeffs[0] * x + coeffs[1] * y + coeffs[2]) / (coeffs[6] * x + coeffs[7] * y + 1) + # y_out = (coeffs[3] * x + coeffs[4] * y + coeffs[5]) / (coeffs[6] * x + coeffs[7] * y + 1) + # + theta1 = torch.tensor( + [[[coeffs[0], coeffs[1], coeffs[2]], [coeffs[3], coeffs[4], coeffs[5]]]], dtype=dtype, device=device + ) + theta2 = torch.tensor([[[coeffs[6], coeffs[7], 1.0], [coeffs[6], coeffs[7], 1.0]]], dtype=dtype, device=device) + + d = 0.5 + base_grid = torch.empty(1, oh, ow, 3, dtype=dtype, device=device) + x_grid = torch.linspace(d, ow + d - 1.0, steps=ow, device=device, dtype=dtype) + base_grid[..., 0].copy_(x_grid) + y_grid = torch.linspace(d, oh + d - 1.0, steps=oh, device=device, dtype=dtype).unsqueeze_(-1) + base_grid[..., 1].copy_(y_grid) + base_grid[..., 2].fill_(1) + + rescaled_theta1 = theta1.transpose(1, 2).div_(torch.tensor([0.5 * ow, 0.5 * oh], dtype=dtype, device=device)) + shape = (1, oh * ow, 3) + output_grid1 = base_grid.view(shape).bmm(rescaled_theta1) + output_grid2 = base_grid.view(shape).bmm(theta2.transpose(1, 2)) + + output_grid = output_grid1.div_(output_grid2).sub_(1.0) + return output_grid.view(1, oh, ow, 2) + + +def _perspective_coefficients( + startpoints: Optional[List[List[int]]], + endpoints: Optional[List[List[int]]], + coefficients: Optional[List[float]], +) -> List[float]: + if coefficients is not None: + if startpoints is not None and endpoints is not None: + raise ValueError("The startpoints/endpoints and the coefficients shouldn't be defined concurrently.") + elif len(coefficients) != 8: + raise ValueError("Argument coefficients should have 8 float values") + return coefficients + elif startpoints is not None and endpoints is not None: + return _get_perspective_coeffs(startpoints, endpoints) + else: + raise ValueError("Either the startpoints/endpoints or the coefficients must have non `None` values.") + + +@_register_kernel_internal(perspective, torch.Tensor) +@_register_kernel_internal(perspective, tv_tensors.Image) +def perspective_image( + image: torch.Tensor, + startpoints: Optional[List[List[int]]], + endpoints: Optional[List[List[int]]], + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + fill: _FillTypeJIT = None, + coefficients: Optional[List[float]] = None, +) -> torch.Tensor: + perspective_coeffs = _perspective_coefficients(startpoints, endpoints, coefficients) + interpolation = _check_interpolation(interpolation) + + _assert_grid_transform_inputs( + image, + matrix=None, + interpolation=interpolation.value, + fill=fill, + supported_interpolation_modes=["nearest", "bilinear"], + coeffs=perspective_coeffs, + ) + + oh, ow = image.shape[-2:] + dtype = image.dtype if torch.is_floating_point(image) else torch.float32 + grid = _perspective_grid(perspective_coeffs, ow=ow, oh=oh, dtype=dtype, device=image.device) + return _apply_grid_transform(image, grid, interpolation.value, fill=fill) + + +@_register_kernel_internal(perspective, PIL.Image.Image) +def _perspective_image_pil( + image: PIL.Image.Image, + startpoints: Optional[List[List[int]]], + endpoints: Optional[List[List[int]]], + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + fill: _FillTypeJIT = None, + coefficients: Optional[List[float]] = None, +) -> PIL.Image.Image: + perspective_coeffs = _perspective_coefficients(startpoints, endpoints, coefficients) + interpolation = _check_interpolation(interpolation) + return _FP.perspective(image, perspective_coeffs, interpolation=pil_modes_mapping[interpolation], fill=fill) + + +def perspective_bounding_boxes( + bounding_boxes: torch.Tensor, + format: tv_tensors.BoundingBoxFormat, + canvas_size: Tuple[int, int], + startpoints: Optional[List[List[int]]], + endpoints: Optional[List[List[int]]], + coefficients: Optional[List[float]] = None, +) -> torch.Tensor: + if bounding_boxes.numel() == 0: + return bounding_boxes + + perspective_coeffs = _perspective_coefficients(startpoints, endpoints, coefficients) + + original_shape = bounding_boxes.shape + # TODO: first cast to float if bbox is int64 before convert_bounding_box_format + bounding_boxes = ( + convert_bounding_box_format(bounding_boxes, old_format=format, new_format=tv_tensors.BoundingBoxFormat.XYXY) + ).reshape(-1, 4) + + dtype = bounding_boxes.dtype if torch.is_floating_point(bounding_boxes) else torch.float32 + device = bounding_boxes.device + + # perspective_coeffs are computed as endpoint -> start point + # We have to invert perspective_coeffs for bboxes: + # (x, y) - end point and (x_out, y_out) - start point + # x_out = (coeffs[0] * x + coeffs[1] * y + coeffs[2]) / (coeffs[6] * x + coeffs[7] * y + 1) + # y_out = (coeffs[3] * x + coeffs[4] * y + coeffs[5]) / (coeffs[6] * x + coeffs[7] * y + 1) + # and we would like to get: + # x = (inv_coeffs[0] * x_out + inv_coeffs[1] * y_out + inv_coeffs[2]) + # / (inv_coeffs[6] * x_out + inv_coeffs[7] * y_out + 1) + # y = (inv_coeffs[3] * x_out + inv_coeffs[4] * y_out + inv_coeffs[5]) + # / (inv_coeffs[6] * x_out + inv_coeffs[7] * y_out + 1) + # and compute inv_coeffs in terms of coeffs + + denom = perspective_coeffs[0] * perspective_coeffs[4] - perspective_coeffs[1] * perspective_coeffs[3] + if denom == 0: + raise RuntimeError( + f"Provided perspective_coeffs {perspective_coeffs} can not be inverted to transform bounding boxes. " + f"Denominator is zero, denom={denom}" + ) + + inv_coeffs = [ + (perspective_coeffs[4] - perspective_coeffs[5] * perspective_coeffs[7]) / denom, + (-perspective_coeffs[1] + perspective_coeffs[2] * perspective_coeffs[7]) / denom, + (perspective_coeffs[1] * perspective_coeffs[5] - perspective_coeffs[2] * perspective_coeffs[4]) / denom, + (-perspective_coeffs[3] + perspective_coeffs[5] * perspective_coeffs[6]) / denom, + (perspective_coeffs[0] - perspective_coeffs[2] * perspective_coeffs[6]) / denom, + (-perspective_coeffs[0] * perspective_coeffs[5] + perspective_coeffs[2] * perspective_coeffs[3]) / denom, + (-perspective_coeffs[4] * perspective_coeffs[6] + perspective_coeffs[3] * perspective_coeffs[7]) / denom, + (-perspective_coeffs[0] * perspective_coeffs[7] + perspective_coeffs[1] * perspective_coeffs[6]) / denom, + ] + + theta1 = torch.tensor( + [[inv_coeffs[0], inv_coeffs[1], inv_coeffs[2]], [inv_coeffs[3], inv_coeffs[4], inv_coeffs[5]]], + dtype=dtype, + device=device, + ) + + theta2 = torch.tensor( + [[inv_coeffs[6], inv_coeffs[7], 1.0], [inv_coeffs[6], inv_coeffs[7], 1.0]], dtype=dtype, device=device + ) + + # 1) Let's transform bboxes into a tensor of 4 points (top-left, top-right, bottom-left, bottom-right corners). + # Tensor of points has shape (N * 4, 3), where N is the number of bboxes + # Single point structure is similar to + # [(xmin, ymin, 1), (xmax, ymin, 1), (xmax, ymax, 1), (xmin, ymax, 1)] + points = bounding_boxes[:, [[0, 1], [2, 1], [2, 3], [0, 3]]].reshape(-1, 2) + points = torch.cat([points, torch.ones(points.shape[0], 1, device=points.device)], dim=-1) + # 2) Now let's transform the points using perspective matrices + # x_out = (coeffs[0] * x + coeffs[1] * y + coeffs[2]) / (coeffs[6] * x + coeffs[7] * y + 1) + # y_out = (coeffs[3] * x + coeffs[4] * y + coeffs[5]) / (coeffs[6] * x + coeffs[7] * y + 1) + + numer_points = torch.matmul(points, theta1.T) + denom_points = torch.matmul(points, theta2.T) + transformed_points = numer_points.div_(denom_points) + + # 3) Reshape transformed points to [N boxes, 4 points, x/y coords] + # and compute bounding box from 4 transformed points: + transformed_points = transformed_points.reshape(-1, 4, 2) + out_bbox_mins, out_bbox_maxs = torch.aminmax(transformed_points, dim=1) + + out_bboxes = clamp_bounding_boxes( + torch.cat([out_bbox_mins, out_bbox_maxs], dim=1).to(bounding_boxes.dtype), + format=tv_tensors.BoundingBoxFormat.XYXY, + canvas_size=canvas_size, + ) + + # out_bboxes should be of shape [N boxes, 4] + + return convert_bounding_box_format( + out_bboxes, old_format=tv_tensors.BoundingBoxFormat.XYXY, new_format=format, inplace=True + ).reshape(original_shape) + + +@_register_kernel_internal(perspective, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False) +def _perspective_bounding_boxes_dispatch( + inpt: tv_tensors.BoundingBoxes, + startpoints: Optional[List[List[int]]], + endpoints: Optional[List[List[int]]], + coefficients: Optional[List[float]] = None, + **kwargs, +) -> tv_tensors.BoundingBoxes: + output = perspective_bounding_boxes( + inpt.as_subclass(torch.Tensor), + format=inpt.format, + canvas_size=inpt.canvas_size, + startpoints=startpoints, + endpoints=endpoints, + coefficients=coefficients, + ) + return tv_tensors.wrap(output, like=inpt) + + +def perspective_mask( + mask: torch.Tensor, + startpoints: Optional[List[List[int]]], + endpoints: Optional[List[List[int]]], + fill: _FillTypeJIT = None, + coefficients: Optional[List[float]] = None, +) -> torch.Tensor: + if mask.ndim < 3: + mask = mask.unsqueeze(0) + needs_squeeze = True + else: + needs_squeeze = False + + output = perspective_image( + mask, startpoints, endpoints, interpolation=InterpolationMode.NEAREST, fill=fill, coefficients=coefficients + ) + + if needs_squeeze: + output = output.squeeze(0) + + return output + + +@_register_kernel_internal(perspective, tv_tensors.Mask, tv_tensor_wrapper=False) +def _perspective_mask_dispatch( + inpt: tv_tensors.Mask, + startpoints: Optional[List[List[int]]], + endpoints: Optional[List[List[int]]], + fill: _FillTypeJIT = None, + coefficients: Optional[List[float]] = None, + **kwargs, +) -> tv_tensors.Mask: + output = perspective_mask( + inpt.as_subclass(torch.Tensor), + startpoints=startpoints, + endpoints=endpoints, + fill=fill, + coefficients=coefficients, + ) + return tv_tensors.wrap(output, like=inpt) + + +@_register_kernel_internal(perspective, tv_tensors.Video) +def perspective_video( + video: torch.Tensor, + startpoints: Optional[List[List[int]]], + endpoints: Optional[List[List[int]]], + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + fill: _FillTypeJIT = None, + coefficients: Optional[List[float]] = None, +) -> torch.Tensor: + return perspective_image( + video, startpoints, endpoints, interpolation=interpolation, fill=fill, coefficients=coefficients + ) + + +def elastic( + inpt: torch.Tensor, + displacement: torch.Tensor, + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + fill: _FillTypeJIT = None, +) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.ElasticTransform` for details.""" + if torch.jit.is_scripting(): + return elastic_image(inpt, displacement=displacement, interpolation=interpolation, fill=fill) + + _log_api_usage_once(elastic) + + kernel = _get_kernel(elastic, type(inpt)) + return kernel(inpt, displacement=displacement, interpolation=interpolation, fill=fill) + + +elastic_transform = elastic + + +@_register_kernel_internal(elastic, torch.Tensor) +@_register_kernel_internal(elastic, tv_tensors.Image) +def elastic_image( + image: torch.Tensor, + displacement: torch.Tensor, + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + fill: _FillTypeJIT = None, +) -> torch.Tensor: + if not isinstance(displacement, torch.Tensor): + raise TypeError("Argument displacement should be a Tensor") + + interpolation = _check_interpolation(interpolation) + + height, width = image.shape[-2:] + device = image.device + dtype = image.dtype if torch.is_floating_point(image) else torch.float32 + + # Patch: elastic transform should support (cpu,f16) input + is_cpu_half = device.type == "cpu" and dtype == torch.float16 + if is_cpu_half: + image = image.to(torch.float32) + dtype = torch.float32 + + # We are aware that if input image dtype is uint8 and displacement is float64 then + # displacement will be cast to float32 and all computations will be done with float32 + # We can fix this later if needed + + expected_shape = (1, height, width, 2) + if expected_shape != displacement.shape: + raise ValueError(f"Argument displacement shape should be {expected_shape}, but given {displacement.shape}") + + grid = _create_identity_grid((height, width), device=device, dtype=dtype).add_( + displacement.to(dtype=dtype, device=device) + ) + output = _apply_grid_transform(image, grid, interpolation.value, fill=fill) + + if is_cpu_half: + output = output.to(torch.float16) + + return output + + +@_register_kernel_internal(elastic, PIL.Image.Image) +def _elastic_image_pil( + image: PIL.Image.Image, + displacement: torch.Tensor, + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + fill: _FillTypeJIT = None, +) -> PIL.Image.Image: + t_img = pil_to_tensor(image) + output = elastic_image(t_img, displacement, interpolation=interpolation, fill=fill) + return to_pil_image(output, mode=image.mode) + + +def _create_identity_grid(size: Tuple[int, int], device: torch.device, dtype: torch.dtype) -> torch.Tensor: + sy, sx = size + base_grid = torch.empty(1, sy, sx, 2, device=device, dtype=dtype) + x_grid = torch.linspace((-sx + 1) / sx, (sx - 1) / sx, sx, device=device, dtype=dtype) + base_grid[..., 0].copy_(x_grid) + + y_grid = torch.linspace((-sy + 1) / sy, (sy - 1) / sy, sy, device=device, dtype=dtype).unsqueeze_(-1) + base_grid[..., 1].copy_(y_grid) + + return base_grid + + +def elastic_bounding_boxes( + bounding_boxes: torch.Tensor, + format: tv_tensors.BoundingBoxFormat, + canvas_size: Tuple[int, int], + displacement: torch.Tensor, +) -> torch.Tensor: + expected_shape = (1, canvas_size[0], canvas_size[1], 2) + if not isinstance(displacement, torch.Tensor): + raise TypeError("Argument displacement should be a Tensor") + elif displacement.shape != expected_shape: + raise ValueError(f"Argument displacement shape should be {expected_shape}, but given {displacement.shape}") + + if bounding_boxes.numel() == 0: + return bounding_boxes + + # TODO: add in docstring about approximation we are doing for grid inversion + device = bounding_boxes.device + dtype = bounding_boxes.dtype if torch.is_floating_point(bounding_boxes) else torch.float32 + + if displacement.dtype != dtype or displacement.device != device: + displacement = displacement.to(dtype=dtype, device=device) + + original_shape = bounding_boxes.shape + # TODO: first cast to float if bbox is int64 before convert_bounding_box_format + bounding_boxes = ( + convert_bounding_box_format(bounding_boxes, old_format=format, new_format=tv_tensors.BoundingBoxFormat.XYXY) + ).reshape(-1, 4) + + id_grid = _create_identity_grid(canvas_size, device=device, dtype=dtype) + # We construct an approximation of inverse grid as inv_grid = id_grid - displacement + # This is not an exact inverse of the grid + inv_grid = id_grid.sub_(displacement) + + # Get points from bboxes + points = bounding_boxes[:, [[0, 1], [2, 1], [2, 3], [0, 3]]].reshape(-1, 2) + if points.is_floating_point(): + points = points.ceil_() + index_xy = points.to(dtype=torch.long) + index_x, index_y = index_xy[:, 0], index_xy[:, 1] + + # Transform points: + t_size = torch.tensor(canvas_size[::-1], device=displacement.device, dtype=displacement.dtype) + transformed_points = inv_grid[0, index_y, index_x, :].add_(1).mul_(0.5 * t_size).sub_(0.5) + + transformed_points = transformed_points.reshape(-1, 4, 2) + out_bbox_mins, out_bbox_maxs = torch.aminmax(transformed_points, dim=1) + out_bboxes = clamp_bounding_boxes( + torch.cat([out_bbox_mins, out_bbox_maxs], dim=1).to(bounding_boxes.dtype), + format=tv_tensors.BoundingBoxFormat.XYXY, + canvas_size=canvas_size, + ) + + return convert_bounding_box_format( + out_bboxes, old_format=tv_tensors.BoundingBoxFormat.XYXY, new_format=format, inplace=True + ).reshape(original_shape) + + +@_register_kernel_internal(elastic, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False) +def _elastic_bounding_boxes_dispatch( + inpt: tv_tensors.BoundingBoxes, displacement: torch.Tensor, **kwargs +) -> tv_tensors.BoundingBoxes: + output = elastic_bounding_boxes( + inpt.as_subclass(torch.Tensor), format=inpt.format, canvas_size=inpt.canvas_size, displacement=displacement + ) + return tv_tensors.wrap(output, like=inpt) + + +def elastic_mask( + mask: torch.Tensor, + displacement: torch.Tensor, + fill: _FillTypeJIT = None, +) -> torch.Tensor: + if mask.ndim < 3: + mask = mask.unsqueeze(0) + needs_squeeze = True + else: + needs_squeeze = False + + output = elastic_image(mask, displacement=displacement, interpolation=InterpolationMode.NEAREST, fill=fill) + + if needs_squeeze: + output = output.squeeze(0) + + return output + + +@_register_kernel_internal(elastic, tv_tensors.Mask, tv_tensor_wrapper=False) +def _elastic_mask_dispatch( + inpt: tv_tensors.Mask, displacement: torch.Tensor, fill: _FillTypeJIT = None, **kwargs +) -> tv_tensors.Mask: + output = elastic_mask(inpt.as_subclass(torch.Tensor), displacement=displacement, fill=fill) + return tv_tensors.wrap(output, like=inpt) + + +@_register_kernel_internal(elastic, tv_tensors.Video) +def elastic_video( + video: torch.Tensor, + displacement: torch.Tensor, + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + fill: _FillTypeJIT = None, +) -> torch.Tensor: + return elastic_image(video, displacement, interpolation=interpolation, fill=fill) + + +def center_crop(inpt: torch.Tensor, output_size: List[int]) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.RandomCrop` for details.""" + if torch.jit.is_scripting(): + return center_crop_image(inpt, output_size=output_size) + + _log_api_usage_once(center_crop) + + kernel = _get_kernel(center_crop, type(inpt)) + return kernel(inpt, output_size=output_size) + + +def _center_crop_parse_output_size(output_size: List[int]) -> List[int]: + if isinstance(output_size, numbers.Number): + s = int(output_size) + return [s, s] + elif isinstance(output_size, (tuple, list)) and len(output_size) == 1: + return [output_size[0], output_size[0]] + else: + return list(output_size) + + +def _center_crop_compute_padding(crop_height: int, crop_width: int, image_height: int, image_width: int) -> List[int]: + return [ + (crop_width - image_width) // 2 if crop_width > image_width else 0, + (crop_height - image_height) // 2 if crop_height > image_height else 0, + (crop_width - image_width + 1) // 2 if crop_width > image_width else 0, + (crop_height - image_height + 1) // 2 if crop_height > image_height else 0, + ] + + +def _center_crop_compute_crop_anchor( + crop_height: int, crop_width: int, image_height: int, image_width: int +) -> Tuple[int, int]: + crop_top = int(round((image_height - crop_height) / 2.0)) + crop_left = int(round((image_width - crop_width) / 2.0)) + return crop_top, crop_left + + +@_register_kernel_internal(center_crop, torch.Tensor) +@_register_kernel_internal(center_crop, tv_tensors.Image) +def center_crop_image(image: torch.Tensor, output_size: List[int]) -> torch.Tensor: + crop_height, crop_width = _center_crop_parse_output_size(output_size) + shape = image.shape + if image.numel() == 0: + return image.reshape(shape[:-2] + (crop_height, crop_width)) + image_height, image_width = shape[-2:] + + if crop_height > image_height or crop_width > image_width: + padding_ltrb = _center_crop_compute_padding(crop_height, crop_width, image_height, image_width) + image = torch_pad(image, _parse_pad_padding(padding_ltrb), value=0.0) + + image_height, image_width = image.shape[-2:] + if crop_width == image_width and crop_height == image_height: + return image + + crop_top, crop_left = _center_crop_compute_crop_anchor(crop_height, crop_width, image_height, image_width) + return image[..., crop_top : (crop_top + crop_height), crop_left : (crop_left + crop_width)] + + +@_register_kernel_internal(center_crop, PIL.Image.Image) +def _center_crop_image_pil(image: PIL.Image.Image, output_size: List[int]) -> PIL.Image.Image: + crop_height, crop_width = _center_crop_parse_output_size(output_size) + image_height, image_width = _get_size_image_pil(image) + + if crop_height > image_height or crop_width > image_width: + padding_ltrb = _center_crop_compute_padding(crop_height, crop_width, image_height, image_width) + image = _pad_image_pil(image, padding_ltrb, fill=0) + + image_height, image_width = _get_size_image_pil(image) + if crop_width == image_width and crop_height == image_height: + return image + + crop_top, crop_left = _center_crop_compute_crop_anchor(crop_height, crop_width, image_height, image_width) + return _crop_image_pil(image, crop_top, crop_left, crop_height, crop_width) + + +def center_crop_bounding_boxes( + bounding_boxes: torch.Tensor, + format: tv_tensors.BoundingBoxFormat, + canvas_size: Tuple[int, int], + output_size: List[int], +) -> Tuple[torch.Tensor, Tuple[int, int]]: + crop_height, crop_width = _center_crop_parse_output_size(output_size) + crop_top, crop_left = _center_crop_compute_crop_anchor(crop_height, crop_width, *canvas_size) + return crop_bounding_boxes( + bounding_boxes, format, top=crop_top, left=crop_left, height=crop_height, width=crop_width + ) + + +@_register_kernel_internal(center_crop, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False) +def _center_crop_bounding_boxes_dispatch( + inpt: tv_tensors.BoundingBoxes, output_size: List[int] +) -> tv_tensors.BoundingBoxes: + output, canvas_size = center_crop_bounding_boxes( + inpt.as_subclass(torch.Tensor), format=inpt.format, canvas_size=inpt.canvas_size, output_size=output_size + ) + return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size) + + +@_register_kernel_internal(center_crop, tv_tensors.Mask) +def center_crop_mask(mask: torch.Tensor, output_size: List[int]) -> torch.Tensor: + if mask.ndim < 3: + mask = mask.unsqueeze(0) + needs_squeeze = True + else: + needs_squeeze = False + + output = center_crop_image(image=mask, output_size=output_size) + + if needs_squeeze: + output = output.squeeze(0) + + return output + + +@_register_kernel_internal(center_crop, tv_tensors.Video) +def center_crop_video(video: torch.Tensor, output_size: List[int]) -> torch.Tensor: + return center_crop_image(video, output_size) + + +def resized_crop( + inpt: torch.Tensor, + top: int, + left: int, + height: int, + width: int, + size: List[int], + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + antialias: Optional[bool] = True, +) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.RandomResizedCrop` for details.""" + if torch.jit.is_scripting(): + return resized_crop_image( + inpt, + top=top, + left=left, + height=height, + width=width, + size=size, + interpolation=interpolation, + antialias=antialias, + ) + + _log_api_usage_once(resized_crop) + + kernel = _get_kernel(resized_crop, type(inpt)) + return kernel( + inpt, + top=top, + left=left, + height=height, + width=width, + size=size, + interpolation=interpolation, + antialias=antialias, + ) + + +@_register_kernel_internal(resized_crop, torch.Tensor) +@_register_kernel_internal(resized_crop, tv_tensors.Image) +def resized_crop_image( + image: torch.Tensor, + top: int, + left: int, + height: int, + width: int, + size: List[int], + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + antialias: Optional[bool] = True, +) -> torch.Tensor: + image = crop_image(image, top, left, height, width) + return resize_image(image, size, interpolation=interpolation, antialias=antialias) + + +def _resized_crop_image_pil( + image: PIL.Image.Image, + top: int, + left: int, + height: int, + width: int, + size: List[int], + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, +) -> PIL.Image.Image: + image = _crop_image_pil(image, top, left, height, width) + return _resize_image_pil(image, size, interpolation=interpolation) + + +@_register_kernel_internal(resized_crop, PIL.Image.Image) +def _resized_crop_image_pil_dispatch( + image: PIL.Image.Image, + top: int, + left: int, + height: int, + width: int, + size: List[int], + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + antialias: Optional[bool] = True, +) -> PIL.Image.Image: + if antialias is False: + warnings.warn("Anti-alias option is always applied for PIL Image input. Argument antialias is ignored.") + return _resized_crop_image_pil( + image, + top=top, + left=left, + height=height, + width=width, + size=size, + interpolation=interpolation, + ) + + +def resized_crop_bounding_boxes( + bounding_boxes: torch.Tensor, + format: tv_tensors.BoundingBoxFormat, + top: int, + left: int, + height: int, + width: int, + size: List[int], +) -> Tuple[torch.Tensor, Tuple[int, int]]: + bounding_boxes, canvas_size = crop_bounding_boxes(bounding_boxes, format, top, left, height, width) + return resize_bounding_boxes(bounding_boxes, canvas_size=canvas_size, size=size) + + +@_register_kernel_internal(resized_crop, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False) +def _resized_crop_bounding_boxes_dispatch( + inpt: tv_tensors.BoundingBoxes, top: int, left: int, height: int, width: int, size: List[int], **kwargs +) -> tv_tensors.BoundingBoxes: + output, canvas_size = resized_crop_bounding_boxes( + inpt.as_subclass(torch.Tensor), format=inpt.format, top=top, left=left, height=height, width=width, size=size + ) + return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size) + + +def resized_crop_mask( + mask: torch.Tensor, + top: int, + left: int, + height: int, + width: int, + size: List[int], +) -> torch.Tensor: + mask = crop_mask(mask, top, left, height, width) + return resize_mask(mask, size) + + +@_register_kernel_internal(resized_crop, tv_tensors.Mask, tv_tensor_wrapper=False) +def _resized_crop_mask_dispatch( + inpt: tv_tensors.Mask, top: int, left: int, height: int, width: int, size: List[int], **kwargs +) -> tv_tensors.Mask: + output = resized_crop_mask( + inpt.as_subclass(torch.Tensor), top=top, left=left, height=height, width=width, size=size + ) + return tv_tensors.wrap(output, like=inpt) + + +@_register_kernel_internal(resized_crop, tv_tensors.Video) +def resized_crop_video( + video: torch.Tensor, + top: int, + left: int, + height: int, + width: int, + size: List[int], + interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR, + antialias: Optional[bool] = True, +) -> torch.Tensor: + return resized_crop_image( + video, top, left, height, width, antialias=antialias, size=size, interpolation=interpolation + ) + + +def five_crop( + inpt: torch.Tensor, size: List[int] +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """See :class:`~torchvision.transforms.v2.FiveCrop` for details.""" + if torch.jit.is_scripting(): + return five_crop_image(inpt, size=size) + + _log_api_usage_once(five_crop) + + kernel = _get_kernel(five_crop, type(inpt)) + return kernel(inpt, size=size) + + +def _parse_five_crop_size(size: List[int]) -> List[int]: + if isinstance(size, numbers.Number): + s = int(size) + size = [s, s] + elif isinstance(size, (tuple, list)) and len(size) == 1: + s = size[0] + size = [s, s] + + if len(size) != 2: + raise ValueError("Please provide only two dimensions (h, w) for size.") + + return size + + +@_register_five_ten_crop_kernel_internal(five_crop, torch.Tensor) +@_register_five_ten_crop_kernel_internal(five_crop, tv_tensors.Image) +def five_crop_image( + image: torch.Tensor, size: List[int] +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + crop_height, crop_width = _parse_five_crop_size(size) + image_height, image_width = image.shape[-2:] + + if crop_width > image_width or crop_height > image_height: + raise ValueError(f"Requested crop size {size} is bigger than input size {(image_height, image_width)}") + + tl = crop_image(image, 0, 0, crop_height, crop_width) + tr = crop_image(image, 0, image_width - crop_width, crop_height, crop_width) + bl = crop_image(image, image_height - crop_height, 0, crop_height, crop_width) + br = crop_image(image, image_height - crop_height, image_width - crop_width, crop_height, crop_width) + center = center_crop_image(image, [crop_height, crop_width]) + + return tl, tr, bl, br, center + + +@_register_five_ten_crop_kernel_internal(five_crop, PIL.Image.Image) +def _five_crop_image_pil( + image: PIL.Image.Image, size: List[int] +) -> Tuple[PIL.Image.Image, PIL.Image.Image, PIL.Image.Image, PIL.Image.Image, PIL.Image.Image]: + crop_height, crop_width = _parse_five_crop_size(size) + image_height, image_width = _get_size_image_pil(image) + + if crop_width > image_width or crop_height > image_height: + raise ValueError(f"Requested crop size {size} is bigger than input size {(image_height, image_width)}") + + tl = _crop_image_pil(image, 0, 0, crop_height, crop_width) + tr = _crop_image_pil(image, 0, image_width - crop_width, crop_height, crop_width) + bl = _crop_image_pil(image, image_height - crop_height, 0, crop_height, crop_width) + br = _crop_image_pil(image, image_height - crop_height, image_width - crop_width, crop_height, crop_width) + center = _center_crop_image_pil(image, [crop_height, crop_width]) + + return tl, tr, bl, br, center + + +@_register_five_ten_crop_kernel_internal(five_crop, tv_tensors.Video) +def five_crop_video( + video: torch.Tensor, size: List[int] +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return five_crop_image(video, size) + + +def ten_crop( + inpt: torch.Tensor, size: List[int], vertical_flip: bool = False +) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + """See :class:`~torchvision.transforms.v2.TenCrop` for details.""" + if torch.jit.is_scripting(): + return ten_crop_image(inpt, size=size, vertical_flip=vertical_flip) + + _log_api_usage_once(ten_crop) + + kernel = _get_kernel(ten_crop, type(inpt)) + return kernel(inpt, size=size, vertical_flip=vertical_flip) + + +@_register_five_ten_crop_kernel_internal(ten_crop, torch.Tensor) +@_register_five_ten_crop_kernel_internal(ten_crop, tv_tensors.Image) +def ten_crop_image( + image: torch.Tensor, size: List[int], vertical_flip: bool = False +) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + non_flipped = five_crop_image(image, size) + + if vertical_flip: + image = vertical_flip_image(image) + else: + image = horizontal_flip_image(image) + + flipped = five_crop_image(image, size) + + return non_flipped + flipped + + +@_register_five_ten_crop_kernel_internal(ten_crop, PIL.Image.Image) +def _ten_crop_image_pil( + image: PIL.Image.Image, size: List[int], vertical_flip: bool = False +) -> Tuple[ + PIL.Image.Image, + PIL.Image.Image, + PIL.Image.Image, + PIL.Image.Image, + PIL.Image.Image, + PIL.Image.Image, + PIL.Image.Image, + PIL.Image.Image, + PIL.Image.Image, + PIL.Image.Image, +]: + non_flipped = _five_crop_image_pil(image, size) + + if vertical_flip: + image = _vertical_flip_image_pil(image) + else: + image = _horizontal_flip_image_pil(image) + + flipped = _five_crop_image_pil(image, size) + + return non_flipped + flipped + + +@_register_five_ten_crop_kernel_internal(ten_crop, tv_tensors.Video) +def ten_crop_video( + video: torch.Tensor, size: List[int], vertical_flip: bool = False +) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + return ten_crop_image(video, size, vertical_flip=vertical_flip) diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_meta.py b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_meta.py new file mode 100644 index 0000000000000000000000000000000000000000..b90e5fb7b5be887d24bb42725d59bb056ac126c2 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_meta.py @@ -0,0 +1,279 @@ +from typing import List, Optional, Tuple + +import PIL.Image +import torch +from torchvision import tv_tensors +from torchvision.transforms import _functional_pil as _FP +from torchvision.tv_tensors import BoundingBoxFormat + +from torchvision.utils import _log_api_usage_once + +from ._utils import _get_kernel, _register_kernel_internal, is_pure_tensor + + +def get_dimensions(inpt: torch.Tensor) -> List[int]: + if torch.jit.is_scripting(): + return get_dimensions_image(inpt) + + _log_api_usage_once(get_dimensions) + + kernel = _get_kernel(get_dimensions, type(inpt)) + return kernel(inpt) + + +@_register_kernel_internal(get_dimensions, torch.Tensor) +@_register_kernel_internal(get_dimensions, tv_tensors.Image, tv_tensor_wrapper=False) +def get_dimensions_image(image: torch.Tensor) -> List[int]: + chw = list(image.shape[-3:]) + ndims = len(chw) + if ndims == 3: + return chw + elif ndims == 2: + chw.insert(0, 1) + return chw + else: + raise TypeError(f"Input tensor should have at least two dimensions, but got {ndims}") + + +_get_dimensions_image_pil = _register_kernel_internal(get_dimensions, PIL.Image.Image)(_FP.get_dimensions) + + +@_register_kernel_internal(get_dimensions, tv_tensors.Video, tv_tensor_wrapper=False) +def get_dimensions_video(video: torch.Tensor) -> List[int]: + return get_dimensions_image(video) + + +def get_num_channels(inpt: torch.Tensor) -> int: + if torch.jit.is_scripting(): + return get_num_channels_image(inpt) + + _log_api_usage_once(get_num_channels) + + kernel = _get_kernel(get_num_channels, type(inpt)) + return kernel(inpt) + + +@_register_kernel_internal(get_num_channels, torch.Tensor) +@_register_kernel_internal(get_num_channels, tv_tensors.Image, tv_tensor_wrapper=False) +def get_num_channels_image(image: torch.Tensor) -> int: + chw = image.shape[-3:] + ndims = len(chw) + if ndims == 3: + return chw[0] + elif ndims == 2: + return 1 + else: + raise TypeError(f"Input tensor should have at least two dimensions, but got {ndims}") + + +_get_num_channels_image_pil = _register_kernel_internal(get_num_channels, PIL.Image.Image)(_FP.get_image_num_channels) + + +@_register_kernel_internal(get_num_channels, tv_tensors.Video, tv_tensor_wrapper=False) +def get_num_channels_video(video: torch.Tensor) -> int: + return get_num_channels_image(video) + + +# We changed the names to ensure it can be used not only for images but also videos. Thus, we just alias it without +# deprecating the old names. +get_image_num_channels = get_num_channels + + +def get_size(inpt: torch.Tensor) -> List[int]: + if torch.jit.is_scripting(): + return get_size_image(inpt) + + _log_api_usage_once(get_size) + + kernel = _get_kernel(get_size, type(inpt)) + return kernel(inpt) + + +@_register_kernel_internal(get_size, torch.Tensor) +@_register_kernel_internal(get_size, tv_tensors.Image, tv_tensor_wrapper=False) +def get_size_image(image: torch.Tensor) -> List[int]: + hw = list(image.shape[-2:]) + ndims = len(hw) + if ndims == 2: + return hw + else: + raise TypeError(f"Input tensor should have at least two dimensions, but got {ndims}") + + +@_register_kernel_internal(get_size, PIL.Image.Image) +def _get_size_image_pil(image: PIL.Image.Image) -> List[int]: + width, height = _FP.get_image_size(image) + return [height, width] + + +@_register_kernel_internal(get_size, tv_tensors.Video, tv_tensor_wrapper=False) +def get_size_video(video: torch.Tensor) -> List[int]: + return get_size_image(video) + + +@_register_kernel_internal(get_size, tv_tensors.Mask, tv_tensor_wrapper=False) +def get_size_mask(mask: torch.Tensor) -> List[int]: + return get_size_image(mask) + + +@_register_kernel_internal(get_size, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False) +def get_size_bounding_boxes(bounding_box: tv_tensors.BoundingBoxes) -> List[int]: + return list(bounding_box.canvas_size) + + +def get_num_frames(inpt: torch.Tensor) -> int: + if torch.jit.is_scripting(): + return get_num_frames_video(inpt) + + _log_api_usage_once(get_num_frames) + + kernel = _get_kernel(get_num_frames, type(inpt)) + return kernel(inpt) + + +@_register_kernel_internal(get_num_frames, torch.Tensor) +@_register_kernel_internal(get_num_frames, tv_tensors.Video, tv_tensor_wrapper=False) +def get_num_frames_video(video: torch.Tensor) -> int: + return video.shape[-4] + + +def _xywh_to_xyxy(xywh: torch.Tensor, inplace: bool) -> torch.Tensor: + xyxy = xywh if inplace else xywh.clone() + xyxy[..., 2:] += xyxy[..., :2] + return xyxy + + +def _xyxy_to_xywh(xyxy: torch.Tensor, inplace: bool) -> torch.Tensor: + xywh = xyxy if inplace else xyxy.clone() + xywh[..., 2:] -= xywh[..., :2] + return xywh + + +def _cxcywh_to_xyxy(cxcywh: torch.Tensor, inplace: bool) -> torch.Tensor: + if not inplace: + cxcywh = cxcywh.clone() + + # Trick to do fast division by 2 and ceil, without casting. It produces the same result as + # `torchvision.ops._box_convert._box_cxcywh_to_xyxy`. + half_wh = cxcywh[..., 2:].div(-2, rounding_mode=None if cxcywh.is_floating_point() else "floor").abs_() + # (cx - width / 2) = x1, same for y1 + cxcywh[..., :2].sub_(half_wh) + # (x1 + width) = x2, same for y2 + cxcywh[..., 2:].add_(cxcywh[..., :2]) + + return cxcywh + + +def _xyxy_to_cxcywh(xyxy: torch.Tensor, inplace: bool) -> torch.Tensor: + if not inplace: + xyxy = xyxy.clone() + + # (x2 - x1) = width, same for height + xyxy[..., 2:].sub_(xyxy[..., :2]) + # (x1 * 2 + width) / 2 = x1 + width / 2 = x1 + (x2-x1)/2 = (x1 + x2)/2 = cx, same for cy + xyxy[..., :2].mul_(2).add_(xyxy[..., 2:]).div_(2, rounding_mode=None if xyxy.is_floating_point() else "floor") + + return xyxy + + +def _convert_bounding_box_format( + bounding_boxes: torch.Tensor, old_format: BoundingBoxFormat, new_format: BoundingBoxFormat, inplace: bool = False +) -> torch.Tensor: + + if new_format == old_format: + return bounding_boxes + + # TODO: Add _xywh_to_cxcywh and _cxcywh_to_xywh to improve performance + if old_format == BoundingBoxFormat.XYWH: + bounding_boxes = _xywh_to_xyxy(bounding_boxes, inplace) + elif old_format == BoundingBoxFormat.CXCYWH: + bounding_boxes = _cxcywh_to_xyxy(bounding_boxes, inplace) + + if new_format == BoundingBoxFormat.XYWH: + bounding_boxes = _xyxy_to_xywh(bounding_boxes, inplace) + elif new_format == BoundingBoxFormat.CXCYWH: + bounding_boxes = _xyxy_to_cxcywh(bounding_boxes, inplace) + + return bounding_boxes + + +def convert_bounding_box_format( + inpt: torch.Tensor, + old_format: Optional[BoundingBoxFormat] = None, + new_format: Optional[BoundingBoxFormat] = None, + inplace: bool = False, +) -> torch.Tensor: + """See :func:`~torchvision.transforms.v2.ConvertBoundingBoxFormat` for details.""" + # This being a kernel / functional hybrid, we need an option to pass `old_format` explicitly for pure tensor + # inputs as well as extract it from `tv_tensors.BoundingBoxes` inputs. However, putting a default value on + # `old_format` means we also need to put one on `new_format` to have syntactically correct Python. Here we mimic the + # default error that would be thrown if `new_format` had no default value. + if new_format is None: + raise TypeError("convert_bounding_box_format() missing 1 required argument: 'new_format'") + + if not torch.jit.is_scripting(): + _log_api_usage_once(convert_bounding_box_format) + + if isinstance(old_format, str): + old_format = BoundingBoxFormat[old_format.upper()] + if isinstance(new_format, str): + new_format = BoundingBoxFormat[new_format.upper()] + + if torch.jit.is_scripting() or is_pure_tensor(inpt): + if old_format is None: + raise ValueError("For pure tensor inputs, `old_format` has to be passed.") + return _convert_bounding_box_format(inpt, old_format=old_format, new_format=new_format, inplace=inplace) + elif isinstance(inpt, tv_tensors.BoundingBoxes): + if old_format is not None: + raise ValueError("For bounding box tv_tensor inputs, `old_format` must not be passed.") + output = _convert_bounding_box_format( + inpt.as_subclass(torch.Tensor), old_format=inpt.format, new_format=new_format, inplace=inplace + ) + return tv_tensors.wrap(output, like=inpt, format=new_format) + else: + raise TypeError( + f"Input can either be a plain tensor or a bounding box tv_tensor, but got {type(inpt)} instead." + ) + + +def _clamp_bounding_boxes( + bounding_boxes: torch.Tensor, format: BoundingBoxFormat, canvas_size: Tuple[int, int] +) -> torch.Tensor: + # TODO: Investigate if it makes sense from a performance perspective to have an implementation for every + # BoundingBoxFormat instead of converting back and forth + in_dtype = bounding_boxes.dtype + bounding_boxes = bounding_boxes.clone() if bounding_boxes.is_floating_point() else bounding_boxes.float() + xyxy_boxes = convert_bounding_box_format( + bounding_boxes, old_format=format, new_format=tv_tensors.BoundingBoxFormat.XYXY, inplace=True + ) + xyxy_boxes[..., 0::2].clamp_(min=0, max=canvas_size[1]) + xyxy_boxes[..., 1::2].clamp_(min=0, max=canvas_size[0]) + out_boxes = convert_bounding_box_format( + xyxy_boxes, old_format=BoundingBoxFormat.XYXY, new_format=format, inplace=True + ) + return out_boxes.to(in_dtype) + + +def clamp_bounding_boxes( + inpt: torch.Tensor, + format: Optional[BoundingBoxFormat] = None, + canvas_size: Optional[Tuple[int, int]] = None, +) -> torch.Tensor: + """See :func:`~torchvision.transforms.v2.ClampBoundingBoxes` for details.""" + if not torch.jit.is_scripting(): + _log_api_usage_once(clamp_bounding_boxes) + + if torch.jit.is_scripting() or is_pure_tensor(inpt): + + if format is None or canvas_size is None: + raise ValueError("For pure tensor inputs, `format` and `canvas_size` have to be passed.") + return _clamp_bounding_boxes(inpt, format=format, canvas_size=canvas_size) + elif isinstance(inpt, tv_tensors.BoundingBoxes): + if format is not None or canvas_size is not None: + raise ValueError("For bounding box tv_tensor inputs, `format` and `canvas_size` must not be passed.") + output = _clamp_bounding_boxes(inpt.as_subclass(torch.Tensor), format=inpt.format, canvas_size=inpt.canvas_size) + return tv_tensors.wrap(output, like=inpt) + else: + raise TypeError( + f"Input can either be a plain tensor or a bounding box tv_tensor, but got {type(inpt)} instead." + ) diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_misc.py b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_misc.py new file mode 100644 index 0000000000000000000000000000000000000000..f40bf117753ef2af57d3e1f0f74df2274a8a05cd --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_misc.py @@ -0,0 +1,420 @@ +import math +from typing import List, Optional, Tuple + +import PIL.Image +import torch +from torch.nn.functional import conv2d, pad as torch_pad + +from torchvision import tv_tensors +from torchvision.transforms._functional_tensor import _max_value +from torchvision.transforms.functional import pil_to_tensor, to_pil_image + +from torchvision.utils import _log_api_usage_once + +from ._meta import _convert_bounding_box_format + +from ._utils import _get_kernel, _register_kernel_internal, is_pure_tensor + + +def normalize( + inpt: torch.Tensor, + mean: List[float], + std: List[float], + inplace: bool = False, +) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.Normalize` for details.""" + if torch.jit.is_scripting(): + return normalize_image(inpt, mean=mean, std=std, inplace=inplace) + + _log_api_usage_once(normalize) + + kernel = _get_kernel(normalize, type(inpt)) + return kernel(inpt, mean=mean, std=std, inplace=inplace) + + +@_register_kernel_internal(normalize, torch.Tensor) +@_register_kernel_internal(normalize, tv_tensors.Image) +def normalize_image(image: torch.Tensor, mean: List[float], std: List[float], inplace: bool = False) -> torch.Tensor: + if not image.is_floating_point(): + raise TypeError(f"Input tensor should be a float tensor. Got {image.dtype}.") + + if image.ndim < 3: + raise ValueError(f"Expected tensor to be a tensor image of size (..., C, H, W). Got {image.shape}.") + + if isinstance(std, (tuple, list)): + divzero = not all(std) + elif isinstance(std, (int, float)): + divzero = std == 0 + else: + divzero = False + if divzero: + raise ValueError("std evaluated to zero, leading to division by zero.") + + dtype = image.dtype + device = image.device + mean = torch.as_tensor(mean, dtype=dtype, device=device) + std = torch.as_tensor(std, dtype=dtype, device=device) + if mean.ndim == 1: + mean = mean.view(-1, 1, 1) + if std.ndim == 1: + std = std.view(-1, 1, 1) + + if inplace: + image = image.sub_(mean) + else: + image = image.sub(mean) + + return image.div_(std) + + +@_register_kernel_internal(normalize, tv_tensors.Video) +def normalize_video(video: torch.Tensor, mean: List[float], std: List[float], inplace: bool = False) -> torch.Tensor: + return normalize_image(video, mean, std, inplace=inplace) + + +def gaussian_blur(inpt: torch.Tensor, kernel_size: List[int], sigma: Optional[List[float]] = None) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.GaussianBlur` for details.""" + if torch.jit.is_scripting(): + return gaussian_blur_image(inpt, kernel_size=kernel_size, sigma=sigma) + + _log_api_usage_once(gaussian_blur) + + kernel = _get_kernel(gaussian_blur, type(inpt)) + return kernel(inpt, kernel_size=kernel_size, sigma=sigma) + + +def _get_gaussian_kernel1d(kernel_size: int, sigma: float, dtype: torch.dtype, device: torch.device) -> torch.Tensor: + lim = (kernel_size - 1) / (2.0 * math.sqrt(2.0)) + x = torch.linspace(-lim, lim, steps=kernel_size, dtype=dtype, device=device) + kernel1d = torch.softmax(x.div(sigma).pow(2).neg(), dim=0) + return kernel1d + + +def _get_gaussian_kernel2d( + kernel_size: List[int], sigma: List[float], dtype: torch.dtype, device: torch.device +) -> torch.Tensor: + kernel1d_x = _get_gaussian_kernel1d(kernel_size[0], sigma[0], dtype, device) + kernel1d_y = _get_gaussian_kernel1d(kernel_size[1], sigma[1], dtype, device) + kernel2d = kernel1d_y.unsqueeze(-1) * kernel1d_x + return kernel2d + + +@_register_kernel_internal(gaussian_blur, torch.Tensor) +@_register_kernel_internal(gaussian_blur, tv_tensors.Image) +def gaussian_blur_image( + image: torch.Tensor, kernel_size: List[int], sigma: Optional[List[float]] = None +) -> torch.Tensor: + # TODO: consider deprecating integers from sigma on the future + if isinstance(kernel_size, int): + kernel_size = [kernel_size, kernel_size] + elif len(kernel_size) != 2: + raise ValueError(f"If kernel_size is a sequence its length should be 2. Got {len(kernel_size)}") + for ksize in kernel_size: + if ksize % 2 == 0 or ksize < 0: + raise ValueError(f"kernel_size should have odd and positive integers. Got {kernel_size}") + + if sigma is None: + sigma = [ksize * 0.15 + 0.35 for ksize in kernel_size] + else: + if isinstance(sigma, (list, tuple)): + length = len(sigma) + if length == 1: + s = sigma[0] + sigma = [s, s] + elif length != 2: + raise ValueError(f"If sigma is a sequence, its length should be 2. Got {length}") + elif isinstance(sigma, (int, float)): + s = float(sigma) + sigma = [s, s] + else: + raise TypeError(f"sigma should be either float or sequence of floats. Got {type(sigma)}") + for s in sigma: + if s <= 0.0: + raise ValueError(f"sigma should have positive values. Got {sigma}") + + if image.numel() == 0: + return image + + dtype = image.dtype + shape = image.shape + ndim = image.ndim + if ndim == 3: + image = image.unsqueeze(dim=0) + elif ndim > 4: + image = image.reshape((-1,) + shape[-3:]) + + fp = torch.is_floating_point(image) + kernel = _get_gaussian_kernel2d(kernel_size, sigma, dtype=dtype if fp else torch.float32, device=image.device) + kernel = kernel.expand(shape[-3], 1, kernel.shape[0], kernel.shape[1]) + + output = image if fp else image.to(dtype=torch.float32) + + # padding = (left, right, top, bottom) + padding = [kernel_size[0] // 2, kernel_size[0] // 2, kernel_size[1] // 2, kernel_size[1] // 2] + output = torch_pad(output, padding, mode="reflect") + output = conv2d(output, kernel, groups=shape[-3]) + + if ndim == 3: + output = output.squeeze(dim=0) + elif ndim > 4: + output = output.reshape(shape) + + if not fp: + output = output.round_().to(dtype=dtype) + + return output + + +@_register_kernel_internal(gaussian_blur, PIL.Image.Image) +def _gaussian_blur_image_pil( + image: PIL.Image.Image, kernel_size: List[int], sigma: Optional[List[float]] = None +) -> PIL.Image.Image: + t_img = pil_to_tensor(image) + output = gaussian_blur_image(t_img, kernel_size=kernel_size, sigma=sigma) + return to_pil_image(output, mode=image.mode) + + +@_register_kernel_internal(gaussian_blur, tv_tensors.Video) +def gaussian_blur_video( + video: torch.Tensor, kernel_size: List[int], sigma: Optional[List[float]] = None +) -> torch.Tensor: + return gaussian_blur_image(video, kernel_size, sigma) + + +def gaussian_noise(inpt: torch.Tensor, mean: float = 0.0, sigma: float = 0.1, clip: bool = True) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.GaussianNoise`""" + if torch.jit.is_scripting(): + return gaussian_noise_image(inpt, mean=mean, sigma=sigma) + + _log_api_usage_once(gaussian_noise) + + kernel = _get_kernel(gaussian_noise, type(inpt)) + return kernel(inpt, mean=mean, sigma=sigma, clip=clip) + + +@_register_kernel_internal(gaussian_noise, torch.Tensor) +@_register_kernel_internal(gaussian_noise, tv_tensors.Image) +def gaussian_noise_image(image: torch.Tensor, mean: float = 0.0, sigma: float = 0.1, clip: bool = True) -> torch.Tensor: + if not image.is_floating_point(): + raise ValueError(f"Input tensor is expected to be in float dtype, got dtype={image.dtype}") + if sigma < 0: + raise ValueError(f"sigma shouldn't be negative. Got {sigma}") + + noise = mean + torch.randn_like(image) * sigma + out = image + noise + if clip: + out = torch.clamp(out, 0, 1) + return out + + +@_register_kernel_internal(gaussian_noise, tv_tensors.Video) +def gaussian_noise_video(video: torch.Tensor, mean: float = 0.0, sigma: float = 0.1, clip: bool = True) -> torch.Tensor: + return gaussian_noise_image(video, mean=mean, sigma=sigma, clip=clip) + + +@_register_kernel_internal(gaussian_noise, PIL.Image.Image) +def _gaussian_noise_pil( + video: torch.Tensor, mean: float = 0.0, sigma: float = 0.1, clip: bool = True +) -> PIL.Image.Image: + raise ValueError("Gaussian Noise is not implemented for PIL images.") + + +def to_dtype(inpt: torch.Tensor, dtype: torch.dtype = torch.float, scale: bool = False) -> torch.Tensor: + """See :func:`~torchvision.transforms.v2.ToDtype` for details.""" + if torch.jit.is_scripting(): + return to_dtype_image(inpt, dtype=dtype, scale=scale) + + _log_api_usage_once(to_dtype) + + kernel = _get_kernel(to_dtype, type(inpt)) + return kernel(inpt, dtype=dtype, scale=scale) + + +def _num_value_bits(dtype: torch.dtype) -> int: + if dtype == torch.uint8: + return 8 + elif dtype == torch.int8: + return 7 + elif dtype == torch.int16: + return 15 + elif dtype == torch.uint16: + return 16 + elif dtype == torch.int32: + return 31 + elif dtype == torch.int64: + return 63 + else: + raise TypeError(f"Number of value bits is only defined for integer dtypes, but got {dtype}.") + + +@_register_kernel_internal(to_dtype, torch.Tensor) +@_register_kernel_internal(to_dtype, tv_tensors.Image) +def to_dtype_image(image: torch.Tensor, dtype: torch.dtype = torch.float, scale: bool = False) -> torch.Tensor: + + if image.dtype == dtype: + return image + elif not scale: + return image.to(dtype) + + float_input = image.is_floating_point() + if torch.jit.is_scripting(): + # TODO: remove this branch as soon as `dtype.is_floating_point` is supported by JIT + float_output = torch.tensor(0, dtype=dtype).is_floating_point() + else: + float_output = dtype.is_floating_point + + if float_input: + # float to float + if float_output: + return image.to(dtype) + + # float to int + if (image.dtype == torch.float32 and dtype in (torch.int32, torch.int64)) or ( + image.dtype == torch.float64 and dtype == torch.int64 + ): + raise RuntimeError(f"The conversion from {image.dtype} to {dtype} cannot be performed safely.") + + # For data in the range `[0.0, 1.0]`, just multiplying by the maximum value of the integer range and converting + # to the integer dtype is not sufficient. For example, `torch.rand(...).mul(255).to(torch.uint8)` will only + # be `255` if the input is exactly `1.0`. See https://github.com/pytorch/vision/pull/2078#issuecomment-612045321 + # for a detailed analysis. + # To mitigate this, we could round before we convert to the integer dtype, but this is an extra operation. + # Instead, we can also multiply by the maximum value plus something close to `1`. See + # https://github.com/pytorch/vision/pull/2078#issuecomment-613524965 for details. + eps = 1e-3 + max_value = float(_max_value(dtype)) + # We need to scale first since the conversion would otherwise turn the input range `[0.0, 1.0]` into the + # discrete set `{0, 1}`. + return image.mul(max_value + 1.0 - eps).to(dtype) + else: + # int to float + if float_output: + return image.to(dtype).mul_(1.0 / _max_value(image.dtype)) + + # int to int + num_value_bits_input = _num_value_bits(image.dtype) + num_value_bits_output = _num_value_bits(dtype) + + # TODO: Remove if/else inner blocks once uint16 dtype supports bitwise shift operations. + shift_by = abs(num_value_bits_input - num_value_bits_output) + if num_value_bits_input > num_value_bits_output: + if image.dtype == torch.uint16: + return (image / 2 ** (shift_by)).to(dtype) + else: + return image.bitwise_right_shift(shift_by).to(dtype) + else: + if dtype == torch.uint16: + return image.to(dtype) * 2 ** (shift_by) + else: + return image.to(dtype).bitwise_left_shift_(shift_by) + + +# We encourage users to use to_dtype() instead but we keep this for BC +def convert_image_dtype(image: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """[DEPRECATED] Use to_dtype() instead.""" + return to_dtype_image(image, dtype=dtype, scale=True) + + +@_register_kernel_internal(to_dtype, tv_tensors.Video) +def to_dtype_video(video: torch.Tensor, dtype: torch.dtype = torch.float, scale: bool = False) -> torch.Tensor: + return to_dtype_image(video, dtype, scale=scale) + + +@_register_kernel_internal(to_dtype, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False) +@_register_kernel_internal(to_dtype, tv_tensors.Mask, tv_tensor_wrapper=False) +def _to_dtype_tensor_dispatch(inpt: torch.Tensor, dtype: torch.dtype, scale: bool = False) -> torch.Tensor: + # We don't need to unwrap and rewrap here, since TVTensor.to() preserves the type + return inpt.to(dtype) + + +def sanitize_bounding_boxes( + bounding_boxes: torch.Tensor, + format: Optional[tv_tensors.BoundingBoxFormat] = None, + canvas_size: Optional[Tuple[int, int]] = None, + min_size: float = 1.0, + min_area: float = 1.0, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Remove degenerate/invalid bounding boxes and return the corresponding indexing mask. + + This removes bounding boxes that: + + - are below a given ``min_size`` or ``min_area``: by default this also removes degenerate boxes that have e.g. X2 <= X1. + - have any coordinate outside of their corresponding image. You may want to + call :func:`~torchvision.transforms.v2.functional.clamp_bounding_boxes` first to avoid undesired removals. + + It is recommended to call it at the end of a pipeline, before passing the + input to the models. It is critical to call this transform if + :class:`~torchvision.transforms.v2.RandomIoUCrop` was called. + If you want to be extra careful, you may call it after all transforms that + may modify bounding boxes but once at the end should be enough in most + cases. + + Args: + bounding_boxes (Tensor or :class:`~torchvision.tv_tensors.BoundingBoxes`): The bounding boxes to be sanitized. + format (str or :class:`~torchvision.tv_tensors.BoundingBoxFormat`, optional): The format of the bounding boxes. + Must be left to none if ``bounding_boxes`` is a :class:`~torchvision.tv_tensors.BoundingBoxes` object. + canvas_size (tuple of int, optional): The canvas_size of the bounding boxes + (size of the corresponding image/video). + Must be left to none if ``bounding_boxes`` is a :class:`~torchvision.tv_tensors.BoundingBoxes` object. + min_size (float, optional) The size below which bounding boxes are removed. Default is 1. + min_area (float, optional) The area below which bounding boxes are removed. Default is 1. + + Returns: + out (tuple of Tensors): The subset of valid bounding boxes, and the corresponding indexing mask. + The mask can then be used to subset other tensors (e.g. labels) that are associated with the bounding boxes. + """ + if torch.jit.is_scripting() or is_pure_tensor(bounding_boxes): + if format is None or canvas_size is None: + raise ValueError( + "format and canvas_size cannot be None if bounding_boxes is a pure tensor. " + f"Got format={format} and canvas_size={canvas_size}." + "Set those to appropriate values or pass bounding_boxes as a tv_tensors.BoundingBoxes object." + ) + if isinstance(format, str): + format = tv_tensors.BoundingBoxFormat[format.upper()] + valid = _get_sanitize_bounding_boxes_mask( + bounding_boxes, format=format, canvas_size=canvas_size, min_size=min_size, min_area=min_area + ) + bounding_boxes = bounding_boxes[valid] + else: + if not isinstance(bounding_boxes, tv_tensors.BoundingBoxes): + raise ValueError("bounding_boxes must be a tv_tensors.BoundingBoxes instance or a pure tensor.") + if format is not None or canvas_size is not None: + raise ValueError( + "format and canvas_size must be None when bounding_boxes is a tv_tensors.BoundingBoxes instance. " + f"Got format={format} and canvas_size={canvas_size}. " + "Leave those to None or pass bounding_boxes as a pure tensor." + ) + valid = _get_sanitize_bounding_boxes_mask( + bounding_boxes, + format=bounding_boxes.format, + canvas_size=bounding_boxes.canvas_size, + min_size=min_size, + min_area=min_area, + ) + bounding_boxes = tv_tensors.wrap(bounding_boxes[valid], like=bounding_boxes) + + return bounding_boxes, valid + + +def _get_sanitize_bounding_boxes_mask( + bounding_boxes: torch.Tensor, + format: tv_tensors.BoundingBoxFormat, + canvas_size: Tuple[int, int], + min_size: float = 1.0, + min_area: float = 1.0, +) -> torch.Tensor: + + bounding_boxes = _convert_bounding_box_format( + bounding_boxes, new_format=tv_tensors.BoundingBoxFormat.XYXY, old_format=format + ) + + image_h, image_w = canvas_size + ws, hs = bounding_boxes[:, 2] - bounding_boxes[:, 0], bounding_boxes[:, 3] - bounding_boxes[:, 1] + valid = (ws >= min_size) & (hs >= min_size) & (bounding_boxes >= 0).all(dim=-1) & (ws * hs >= min_area) + # TODO: Do we really need to check for out of bounds here? All + # transforms should be clamping anyway, so this should never happen? + image_h, image_w = canvas_size + valid &= (bounding_boxes[:, 0] <= image_w) & (bounding_boxes[:, 2] <= image_w) + valid &= (bounding_boxes[:, 1] <= image_h) & (bounding_boxes[:, 3] <= image_h) + return valid diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_temporal.py b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_temporal.py new file mode 100644 index 0000000000000000000000000000000000000000..f932b06a295fd10316fba3e796ec4649053e92db --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_temporal.py @@ -0,0 +1,27 @@ +import torch + +from torchvision import tv_tensors + +from torchvision.utils import _log_api_usage_once + +from ._utils import _get_kernel, _register_kernel_internal + + +def uniform_temporal_subsample(inpt: torch.Tensor, num_samples: int) -> torch.Tensor: + """See :class:`~torchvision.transforms.v2.UniformTemporalSubsample` for details.""" + if torch.jit.is_scripting(): + return uniform_temporal_subsample_video(inpt, num_samples=num_samples) + + _log_api_usage_once(uniform_temporal_subsample) + + kernel = _get_kernel(uniform_temporal_subsample, type(inpt)) + return kernel(inpt, num_samples=num_samples) + + +@_register_kernel_internal(uniform_temporal_subsample, torch.Tensor) +@_register_kernel_internal(uniform_temporal_subsample, tv_tensors.Video) +def uniform_temporal_subsample_video(video: torch.Tensor, num_samples: int) -> torch.Tensor: + # Reference: https://github.com/facebookresearch/pytorchvideo/blob/a0a131e/pytorchvideo/transforms/functional.py#L19 + t_max = video.shape[-4] - 1 + indices = torch.linspace(0, t_max, num_samples, device=video.device).long() + return torch.index_select(video, -4, indices) diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_type_conversion.py b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_type_conversion.py new file mode 100644 index 0000000000000000000000000000000000000000..c5a731fe143c365400d5905db8370c538097583a --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_type_conversion.py @@ -0,0 +1,27 @@ +from typing import Union + +import numpy as np +import PIL.Image +import torch +from torchvision import tv_tensors +from torchvision.transforms import functional as _F + + +@torch.jit.unused +def to_image(inpt: Union[torch.Tensor, PIL.Image.Image, np.ndarray]) -> tv_tensors.Image: + """See :class:`~torchvision.transforms.v2.ToImage` for details.""" + if isinstance(inpt, np.ndarray): + output = torch.from_numpy(np.atleast_3d(inpt)).permute((2, 0, 1)).contiguous() + elif isinstance(inpt, PIL.Image.Image): + output = pil_to_tensor(inpt) + elif isinstance(inpt, torch.Tensor): + output = inpt + else: + raise TypeError( + f"Input can either be a pure Tensor, a numpy array, or a PIL image, but got {type(inpt)} instead." + ) + return tv_tensors.Image(output) + + +to_pil_image = _F.to_pil_image +pil_to_tensor = _F.pil_to_tensor diff --git a/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_utils.py b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..fe0faeddc1b9b56596783e4e4872b604edcd1555 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/transforms/v2/functional/_utils.py @@ -0,0 +1,141 @@ +import functools +from typing import Any, Callable, Dict, List, Optional, Sequence, Type, Union + +import torch +from torchvision import tv_tensors + +_FillType = Union[int, float, Sequence[int], Sequence[float], None] +_FillTypeJIT = Optional[List[float]] + + +def is_pure_tensor(inpt: Any) -> bool: + return isinstance(inpt, torch.Tensor) and not isinstance(inpt, tv_tensors.TVTensor) + + +# {functional: {input_type: type_specific_kernel}} +_KERNEL_REGISTRY: Dict[Callable, Dict[Type, Callable]] = {} + + +def _kernel_tv_tensor_wrapper(kernel): + @functools.wraps(kernel) + def wrapper(inpt, *args, **kwargs): + # If you're wondering whether we could / should get rid of this wrapper, + # the answer is no: we want to pass pure Tensors to avoid the overhead + # of the __torch_function__ machinery. Note that this is always valid, + # regardless of whether we override __torch_function__ in our base class + # or not. + # Also, even if we didn't call `as_subclass` here, we would still need + # this wrapper to call wrap(), because the TVTensor type would be + # lost after the first operation due to our own __torch_function__ + # logic. + output = kernel(inpt.as_subclass(torch.Tensor), *args, **kwargs) + return tv_tensors.wrap(output, like=inpt) + + return wrapper + + +def _register_kernel_internal(functional, input_type, *, tv_tensor_wrapper=True): + registry = _KERNEL_REGISTRY.setdefault(functional, {}) + if input_type in registry: + raise ValueError(f"Functional {functional} already has a kernel registered for type {input_type}.") + + def decorator(kernel): + registry[input_type] = ( + _kernel_tv_tensor_wrapper(kernel) + if issubclass(input_type, tv_tensors.TVTensor) and tv_tensor_wrapper + else kernel + ) + return kernel + + return decorator + + +def _name_to_functional(name): + import torchvision.transforms.v2.functional # noqa + + try: + return getattr(torchvision.transforms.v2.functional, name) + except AttributeError: + raise ValueError( + f"Could not find functional with name '{name}' in torchvision.transforms.v2.functional." + ) from None + + +_BUILTIN_DATAPOINT_TYPES = { + obj for obj in tv_tensors.__dict__.values() if isinstance(obj, type) and issubclass(obj, tv_tensors.TVTensor) +} + + +def register_kernel(functional, tv_tensor_cls): + """Decorate a kernel to register it for a functional and a (custom) tv_tensor type. + + See :ref:`sphx_glr_auto_examples_transforms_plot_custom_tv_tensors.py` for usage + details. + """ + if isinstance(functional, str): + functional = _name_to_functional(name=functional) + elif not ( + callable(functional) + and getattr(functional, "__module__", "").startswith("torchvision.transforms.v2.functional") + ): + raise ValueError( + f"Kernels can only be registered on functionals from the torchvision.transforms.v2.functional namespace, " + f"but got {functional}." + ) + + if not (isinstance(tv_tensor_cls, type) and issubclass(tv_tensor_cls, tv_tensors.TVTensor)): + raise ValueError( + f"Kernels can only be registered for subclasses of torchvision.tv_tensors.TVTensor, " + f"but got {tv_tensor_cls}." + ) + + if tv_tensor_cls in _BUILTIN_DATAPOINT_TYPES: + raise ValueError(f"Kernels cannot be registered for the builtin tv_tensor classes, but got {tv_tensor_cls}") + + return _register_kernel_internal(functional, tv_tensor_cls, tv_tensor_wrapper=False) + + +def _get_kernel(functional, input_type, *, allow_passthrough=False): + registry = _KERNEL_REGISTRY.get(functional) + if not registry: + raise ValueError(f"No kernel registered for functional {functional.__name__}.") + + for cls in input_type.__mro__: + if cls in registry: + return registry[cls] + elif cls is tv_tensors.TVTensor: + # We don't want user-defined tv_tensors to dispatch to the pure Tensor kernels, so we explicit stop the + # MRO traversal before hitting torch.Tensor. We can even stop at tv_tensors.TVTensor, since we don't + # allow kernels to be registered for tv_tensors.TVTensor anyway. + break + + if allow_passthrough: + return lambda inpt, *args, **kwargs: inpt + + raise TypeError( + f"Functional F.{functional.__name__} supports inputs of type {registry.keys()}, " + f"but got {input_type} instead." + ) + + +# This basically replicates _register_kernel_internal, but with a specialized wrapper for five_crop / ten_crop +# We could get rid of this by letting _register_kernel_internal take arbitrary functionals rather than wrap_kernel: bool +def _register_five_ten_crop_kernel_internal(functional, input_type): + registry = _KERNEL_REGISTRY.setdefault(functional, {}) + if input_type in registry: + raise TypeError(f"Functional '{functional}' already has a kernel registered for type '{input_type}'.") + + def wrap(kernel): + @functools.wraps(kernel) + def wrapper(inpt, *args, **kwargs): + output = kernel(inpt, *args, **kwargs) + container_type = type(output) + return container_type(tv_tensors.wrap(o, like=inpt) for o in output) + + return wrapper + + def decorator(kernel): + registry[input_type] = wrap(kernel) if issubclass(input_type, tv_tensors.TVTensor) else kernel + return kernel + + return decorator diff --git a/lib/python3.12/site-packages/torchvision/utils.py b/lib/python3.12/site-packages/torchvision/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..eb1ec2ccf65bf18dacec7c828e527fb760205786 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/utils.py @@ -0,0 +1,666 @@ +import collections +import math +import pathlib +import warnings +from itertools import repeat +from types import FunctionType +from typing import Any, BinaryIO, List, Optional, Tuple, Union + +import numpy as np +import torch +from PIL import Image, ImageColor, ImageDraw, ImageFont + + +__all__ = [ + "make_grid", + "save_image", + "draw_bounding_boxes", + "draw_segmentation_masks", + "draw_keypoints", + "flow_to_image", +] + + +@torch.no_grad() +def make_grid( + tensor: Union[torch.Tensor, List[torch.Tensor]], + nrow: int = 8, + padding: int = 2, + normalize: bool = False, + value_range: Optional[Tuple[int, int]] = None, + scale_each: bool = False, + pad_value: float = 0.0, +) -> torch.Tensor: + """ + Make a grid of images. + + Args: + tensor (Tensor or list): 4D mini-batch Tensor of shape (B x C x H x W) + or a list of images all of the same size. + nrow (int, optional): Number of images displayed in each row of the grid. + The final grid size is ``(B / nrow, nrow)``. Default: ``8``. + padding (int, optional): amount of padding. Default: ``2``. + normalize (bool, optional): If True, shift the image to the range (0, 1), + by the min and max values specified by ``value_range``. Default: ``False``. + value_range (tuple, optional): tuple (min, max) where min and max are numbers, + then these numbers are used to normalize the image. By default, min and max + are computed from the tensor. + scale_each (bool, optional): If ``True``, scale each image in the batch of + images separately rather than the (min, max) over all images. Default: ``False``. + pad_value (float, optional): Value for the padded pixels. Default: ``0``. + + Returns: + grid (Tensor): the tensor containing grid of images. + """ + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(make_grid) + if not torch.is_tensor(tensor): + if isinstance(tensor, list): + for t in tensor: + if not torch.is_tensor(t): + raise TypeError(f"tensor or list of tensors expected, got a list containing {type(t)}") + else: + raise TypeError(f"tensor or list of tensors expected, got {type(tensor)}") + + # if list of tensors, convert to a 4D mini-batch Tensor + if isinstance(tensor, list): + tensor = torch.stack(tensor, dim=0) + + if tensor.dim() == 2: # single image H x W + tensor = tensor.unsqueeze(0) + if tensor.dim() == 3: # single image + if tensor.size(0) == 1: # if single-channel, convert to 3-channel + tensor = torch.cat((tensor, tensor, tensor), 0) + tensor = tensor.unsqueeze(0) + + if tensor.dim() == 4 and tensor.size(1) == 1: # single-channel images + tensor = torch.cat((tensor, tensor, tensor), 1) + + if normalize is True: + tensor = tensor.clone() # avoid modifying tensor in-place + if value_range is not None and not isinstance(value_range, tuple): + raise TypeError("value_range has to be a tuple (min, max) if specified. min and max are numbers") + + def norm_ip(img, low, high): + img.clamp_(min=low, max=high) + img.sub_(low).div_(max(high - low, 1e-5)) + + def norm_range(t, value_range): + if value_range is not None: + norm_ip(t, value_range[0], value_range[1]) + else: + norm_ip(t, float(t.min()), float(t.max())) + + if scale_each is True: + for t in tensor: # loop over mini-batch dimension + norm_range(t, value_range) + else: + norm_range(tensor, value_range) + + if not isinstance(tensor, torch.Tensor): + raise TypeError("tensor should be of type torch.Tensor") + if tensor.size(0) == 1: + return tensor.squeeze(0) + + # make the mini-batch of images into a grid + nmaps = tensor.size(0) + xmaps = min(nrow, nmaps) + ymaps = int(math.ceil(float(nmaps) / xmaps)) + height, width = int(tensor.size(2) + padding), int(tensor.size(3) + padding) + num_channels = tensor.size(1) + grid = tensor.new_full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value) + k = 0 + for y in range(ymaps): + for x in range(xmaps): + if k >= nmaps: + break + # Tensor.copy_() is a valid method but seems to be missing from the stubs + # https://pytorch.org/docs/stable/tensors.html#torch.Tensor.copy_ + grid.narrow(1, y * height + padding, height - padding).narrow( # type: ignore[attr-defined] + 2, x * width + padding, width - padding + ).copy_(tensor[k]) + k = k + 1 + return grid + + +@torch.no_grad() +def save_image( + tensor: Union[torch.Tensor, List[torch.Tensor]], + fp: Union[str, pathlib.Path, BinaryIO], + format: Optional[str] = None, + **kwargs, +) -> None: + """ + Save a given Tensor into an image file. + + Args: + tensor (Tensor or list): Image to be saved. If given a mini-batch tensor, + saves the tensor as a grid of images by calling ``make_grid``. + fp (string or file object): A filename or a file object + format(Optional): If omitted, the format to use is determined from the filename extension. + If a file object was used instead of a filename, this parameter should always be used. + **kwargs: Other arguments are documented in ``make_grid``. + """ + + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(save_image) + grid = make_grid(tensor, **kwargs) + # Add 0.5 after unnormalizing to [0, 255] to round to the nearest integer + ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy() + im = Image.fromarray(ndarr) + im.save(fp, format=format) + + +@torch.no_grad() +def draw_bounding_boxes( + image: torch.Tensor, + boxes: torch.Tensor, + labels: Optional[List[str]] = None, + colors: Optional[Union[List[Union[str, Tuple[int, int, int]]], str, Tuple[int, int, int]]] = None, + fill: Optional[bool] = False, + width: int = 1, + font: Optional[str] = None, + font_size: Optional[int] = None, + label_colors: Optional[Union[List[Union[str, Tuple[int, int, int]]], str, Tuple[int, int, int]]] = None, +) -> torch.Tensor: + + """ + Draws bounding boxes on given RGB image. + The image values should be uint8 in [0, 255] or float in [0, 1]. + If fill is True, Resulting Tensor should be saved as PNG image. + + Args: + image (Tensor): Tensor of shape (C, H, W) and dtype uint8 or float. + boxes (Tensor): Tensor of size (N, 4) containing bounding boxes in (xmin, ymin, xmax, ymax) format. Note that + the boxes are absolute coordinates with respect to the image. In other words: `0 <= xmin < xmax < W` and + `0 <= ymin < ymax < H`. + labels (List[str]): List containing the labels of bounding boxes. + colors (color or list of colors, optional): List containing the colors + of the boxes or single color for all boxes. The color can be represented as + PIL strings e.g. "red" or "#FF00FF", or as RGB tuples e.g. ``(240, 10, 157)``. + By default, random colors are generated for boxes. + fill (bool): If `True` fills the bounding box with specified color. + width (int): Width of bounding box. + font (str): A filename containing a TrueType font. If the file is not found in this filename, the loader may + also search in other directories, such as the `fonts/` directory on Windows or `/Library/Fonts/`, + `/System/Library/Fonts/` and `~/Library/Fonts/` on macOS. + font_size (int): The requested font size in points. + label_colors (color or list of colors, optional): Colors for the label text. See the description of the + `colors` argument for details. Defaults to the same colors used for the boxes. + + Returns: + img (Tensor[C, H, W]): Image Tensor of dtype uint8 with bounding boxes plotted. + + """ + import torchvision.transforms.v2.functional as F # noqa + + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(draw_bounding_boxes) + if not isinstance(image, torch.Tensor): + raise TypeError(f"Tensor expected, got {type(image)}") + elif not (image.dtype == torch.uint8 or image.is_floating_point()): + raise ValueError(f"The image dtype must be uint8 or float, got {image.dtype}") + elif image.dim() != 3: + raise ValueError("Pass individual images, not batches") + elif image.size(0) not in {1, 3}: + raise ValueError("Only grayscale and RGB images are supported") + elif (boxes[:, 0] > boxes[:, 2]).any() or (boxes[:, 1] > boxes[:, 3]).any(): + raise ValueError( + "Boxes need to be in (xmin, ymin, xmax, ymax) format. Use torchvision.ops.box_convert to convert them" + ) + + num_boxes = boxes.shape[0] + + if num_boxes == 0: + warnings.warn("boxes doesn't contain any box. No box was drawn") + return image + + if labels is None: + labels: Union[List[str], List[None]] = [None] * num_boxes # type: ignore[no-redef] + elif len(labels) != num_boxes: + raise ValueError( + f"Number of boxes ({num_boxes}) and labels ({len(labels)}) mismatch. Please specify labels for each box." + ) + + colors = _parse_colors(colors, num_objects=num_boxes) + if label_colors: + label_colors = _parse_colors(label_colors, num_objects=num_boxes) # type: ignore[assignment] + else: + label_colors = colors.copy() # type: ignore[assignment] + + if font is None: + if font_size is not None: + warnings.warn("Argument 'font_size' will be ignored since 'font' is not set.") + txt_font = ImageFont.load_default() + else: + txt_font = ImageFont.truetype(font=font, size=font_size or 10) + + # Handle Grayscale images + if image.size(0) == 1: + image = torch.tile(image, (3, 1, 1)) + + original_dtype = image.dtype + if original_dtype.is_floating_point: + image = F.to_dtype(image, dtype=torch.uint8, scale=True) + + img_to_draw = F.to_pil_image(image) + img_boxes = boxes.to(torch.int64).tolist() + + if fill: + draw = ImageDraw.Draw(img_to_draw, "RGBA") + else: + draw = ImageDraw.Draw(img_to_draw) + + for bbox, color, label, label_color in zip(img_boxes, colors, labels, label_colors): # type: ignore[arg-type] + if fill: + fill_color = color + (100,) + draw.rectangle(bbox, width=width, outline=color, fill=fill_color) + else: + draw.rectangle(bbox, width=width, outline=color) + + if label is not None: + margin = width + 1 + draw.text((bbox[0] + margin, bbox[1] + margin), label, fill=label_color, font=txt_font) # type: ignore[arg-type] + + out = F.pil_to_tensor(img_to_draw) + if original_dtype.is_floating_point: + out = F.to_dtype(out, dtype=original_dtype, scale=True) + return out + + +@torch.no_grad() +def draw_segmentation_masks( + image: torch.Tensor, + masks: torch.Tensor, + alpha: float = 0.8, + colors: Optional[Union[List[Union[str, Tuple[int, int, int]]], str, Tuple[int, int, int]]] = None, +) -> torch.Tensor: + + """ + Draws segmentation masks on given RGB image. + The image values should be uint8 in [0, 255] or float in [0, 1]. + + Args: + image (Tensor): Tensor of shape (3, H, W) and dtype uint8 or float. + masks (Tensor): Tensor of shape (num_masks, H, W) or (H, W) and dtype bool. + alpha (float): Float number between 0 and 1 denoting the transparency of the masks. + 0 means full transparency, 1 means no transparency. + colors (color or list of colors, optional): List containing the colors + of the masks or single color for all masks. The color can be represented as + PIL strings e.g. "red" or "#FF00FF", or as RGB tuples e.g. ``(240, 10, 157)``. + By default, random colors are generated for each mask. + + Returns: + img (Tensor[C, H, W]): Image Tensor, with segmentation masks drawn on top. + """ + + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(draw_segmentation_masks) + if not isinstance(image, torch.Tensor): + raise TypeError(f"The image must be a tensor, got {type(image)}") + elif not (image.dtype == torch.uint8 or image.is_floating_point()): + raise ValueError(f"The image dtype must be uint8 or float, got {image.dtype}") + elif image.dim() != 3: + raise ValueError("Pass individual images, not batches") + elif image.size()[0] != 3: + raise ValueError("Pass an RGB image. Other Image formats are not supported") + if masks.ndim == 2: + masks = masks[None, :, :] + if masks.ndim != 3: + raise ValueError("masks must be of shape (H, W) or (batch_size, H, W)") + if masks.dtype != torch.bool: + raise ValueError(f"The masks must be of dtype bool. Got {masks.dtype}") + if masks.shape[-2:] != image.shape[-2:]: + raise ValueError("The image and the masks must have the same height and width") + + num_masks = masks.size()[0] + overlapping_masks = masks.sum(dim=0) > 1 + + if num_masks == 0: + warnings.warn("masks doesn't contain any mask. No mask was drawn") + return image + + original_dtype = image.dtype + colors = [ + torch.tensor(color, dtype=original_dtype, device=image.device) + for color in _parse_colors(colors, num_objects=num_masks, dtype=original_dtype) + ] + + img_to_draw = image.detach().clone() + # TODO: There might be a way to vectorize this + for mask, color in zip(masks, colors): + img_to_draw[:, mask] = color[:, None] + + img_to_draw[:, overlapping_masks] = 0 + + out = image * (1 - alpha) + img_to_draw * alpha + # Note: at this point, out is a float tensor in [0, 1] or [0, 255] depending on original_dtype + return out.to(original_dtype) + + +@torch.no_grad() +def draw_keypoints( + image: torch.Tensor, + keypoints: torch.Tensor, + connectivity: Optional[List[Tuple[int, int]]] = None, + colors: Optional[Union[str, Tuple[int, int, int]]] = None, + radius: int = 2, + width: int = 3, + visibility: Optional[torch.Tensor] = None, +) -> torch.Tensor: + + """ + Draws Keypoints on given RGB image. + The image values should be uint8 in [0, 255] or float in [0, 1]. + Keypoints can be drawn for multiple instances at a time. + + This method allows that keypoints and their connectivity are drawn based on the visibility of this keypoint. + + Args: + image (Tensor): Tensor of shape (3, H, W) and dtype uint8 or float. + keypoints (Tensor): Tensor of shape (num_instances, K, 2) the K keypoint locations for each of the N instances, + in the format [x, y]. + connectivity (List[Tuple[int, int]]]): A List of tuple where each tuple contains a pair of keypoints + to be connected. + If at least one of the two connected keypoints has a ``visibility`` of False, + this specific connection is not drawn. + Exclusions due to invisibility are computed per-instance. + colors (str, Tuple): The color can be represented as + PIL strings e.g. "red" or "#FF00FF", or as RGB tuples e.g. ``(240, 10, 157)``. + radius (int): Integer denoting radius of keypoint. + width (int): Integer denoting width of line connecting keypoints. + visibility (Tensor): Tensor of shape (num_instances, K) specifying the visibility of the K + keypoints for each of the N instances. + True means that the respective keypoint is visible and should be drawn. + False means invisible, so neither the point nor possible connections containing it are drawn. + The input tensor will be cast to bool. + Default ``None`` means that all the keypoints are visible. + For more details, see :ref:`draw_keypoints_with_visibility`. + + Returns: + img (Tensor[C, H, W]): Image Tensor with keypoints drawn. + """ + + if not torch.jit.is_scripting() and not torch.jit.is_tracing(): + _log_api_usage_once(draw_keypoints) + # validate image + if not isinstance(image, torch.Tensor): + raise TypeError(f"The image must be a tensor, got {type(image)}") + elif not (image.dtype == torch.uint8 or image.is_floating_point()): + raise ValueError(f"The image dtype must be uint8 or float, got {image.dtype}") + elif image.dim() != 3: + raise ValueError("Pass individual images, not batches") + elif image.size()[0] != 3: + raise ValueError("Pass an RGB image. Other Image formats are not supported") + + # validate keypoints + if keypoints.ndim != 3: + raise ValueError("keypoints must be of shape (num_instances, K, 2)") + + # validate visibility + if visibility is None: # set default + visibility = torch.ones(keypoints.shape[:-1], dtype=torch.bool) + if visibility.ndim == 3: + # If visibility was passed as pred.split([2, 1], dim=-1), it will be of shape (num_instances, K, 1). + # We make sure it is of shape (num_instances, K). This isn't documented, we're just being nice. + visibility = visibility.squeeze(-1) + if visibility.ndim != 2: + raise ValueError(f"visibility must be of shape (num_instances, K). Got ndim={visibility.ndim}") + if visibility.shape != keypoints.shape[:-1]: + raise ValueError( + "keypoints and visibility must have the same dimensionality for num_instances and K. " + f"Got {visibility.shape = } and {keypoints.shape = }" + ) + + original_dtype = image.dtype + if original_dtype.is_floating_point: + from torchvision.transforms.v2.functional import to_dtype # noqa + + image = to_dtype(image, dtype=torch.uint8, scale=True) + + ndarr = image.permute(1, 2, 0).cpu().numpy() + img_to_draw = Image.fromarray(ndarr) + draw = ImageDraw.Draw(img_to_draw) + img_kpts = keypoints.to(torch.int64).tolist() + img_vis = visibility.cpu().bool().tolist() + + for kpt_inst, vis_inst in zip(img_kpts, img_vis): + for kpt_coord, kp_vis in zip(kpt_inst, vis_inst): + if not kp_vis: + continue + x1 = kpt_coord[0] - radius + x2 = kpt_coord[0] + radius + y1 = kpt_coord[1] - radius + y2 = kpt_coord[1] + radius + draw.ellipse([x1, y1, x2, y2], fill=colors, outline=None, width=0) + + if connectivity: + for connection in connectivity: + if (not vis_inst[connection[0]]) or (not vis_inst[connection[1]]): + continue + start_pt_x = kpt_inst[connection[0]][0] + start_pt_y = kpt_inst[connection[0]][1] + + end_pt_x = kpt_inst[connection[1]][0] + end_pt_y = kpt_inst[connection[1]][1] + + draw.line( + ((start_pt_x, start_pt_y), (end_pt_x, end_pt_y)), + width=width, + ) + + out = torch.from_numpy(np.array(img_to_draw)).permute(2, 0, 1) + if original_dtype.is_floating_point: + out = to_dtype(out, dtype=original_dtype, scale=True) + return out + + +# Flow visualization code adapted from https://github.com/tomrunia/OpticalFlow_Visualization +@torch.no_grad() +def flow_to_image(flow: torch.Tensor) -> torch.Tensor: + + """ + Converts a flow to an RGB image. + + Args: + flow (Tensor): Flow of shape (N, 2, H, W) or (2, H, W) and dtype torch.float. + + Returns: + img (Tensor): Image Tensor of dtype uint8 where each color corresponds + to a given flow direction. Shape is (N, 3, H, W) or (3, H, W) depending on the input. + """ + + if flow.dtype != torch.float: + raise ValueError(f"Flow should be of dtype torch.float, got {flow.dtype}.") + + orig_shape = flow.shape + if flow.ndim == 3: + flow = flow[None] # Add batch dim + + if flow.ndim != 4 or flow.shape[1] != 2: + raise ValueError(f"Input flow should have shape (2, H, W) or (N, 2, H, W), got {orig_shape}.") + + max_norm = torch.sum(flow**2, dim=1).sqrt().max() + epsilon = torch.finfo((flow).dtype).eps + normalized_flow = flow / (max_norm + epsilon) + img = _normalized_flow_to_image(normalized_flow) + + if len(orig_shape) == 3: + img = img[0] # Remove batch dim + return img + + +@torch.no_grad() +def _normalized_flow_to_image(normalized_flow: torch.Tensor) -> torch.Tensor: + + """ + Converts a batch of normalized flow to an RGB image. + + Args: + normalized_flow (torch.Tensor): Normalized flow tensor of shape (N, 2, H, W) + Returns: + img (Tensor(N, 3, H, W)): Flow visualization image of dtype uint8. + """ + + N, _, H, W = normalized_flow.shape + device = normalized_flow.device + flow_image = torch.zeros((N, 3, H, W), dtype=torch.uint8, device=device) + colorwheel = _make_colorwheel().to(device) # shape [55x3] + num_cols = colorwheel.shape[0] + norm = torch.sum(normalized_flow**2, dim=1).sqrt() + a = torch.atan2(-normalized_flow[:, 1, :, :], -normalized_flow[:, 0, :, :]) / torch.pi + fk = (a + 1) / 2 * (num_cols - 1) + k0 = torch.floor(fk).to(torch.long) + k1 = k0 + 1 + k1[k1 == num_cols] = 0 + f = fk - k0 + + for c in range(colorwheel.shape[1]): + tmp = colorwheel[:, c] + col0 = tmp[k0] / 255.0 + col1 = tmp[k1] / 255.0 + col = (1 - f) * col0 + f * col1 + col = 1 - norm * (1 - col) + flow_image[:, c, :, :] = torch.floor(255 * col) + return flow_image + + +def _make_colorwheel() -> torch.Tensor: + """ + Generates a color wheel for optical flow visualization as presented in: + Baker et al. "A Database and Evaluation Methodology for Optical Flow" (ICCV, 2007) + URL: http://vision.middlebury.edu/flow/flowEval-iccv07.pdf. + + Returns: + colorwheel (Tensor[55, 3]): Colorwheel Tensor. + """ + + RY = 15 + YG = 6 + GC = 4 + CB = 11 + BM = 13 + MR = 6 + + ncols = RY + YG + GC + CB + BM + MR + colorwheel = torch.zeros((ncols, 3)) + col = 0 + + # RY + colorwheel[0:RY, 0] = 255 + colorwheel[0:RY, 1] = torch.floor(255 * torch.arange(0, RY) / RY) + col = col + RY + # YG + colorwheel[col : col + YG, 0] = 255 - torch.floor(255 * torch.arange(0, YG) / YG) + colorwheel[col : col + YG, 1] = 255 + col = col + YG + # GC + colorwheel[col : col + GC, 1] = 255 + colorwheel[col : col + GC, 2] = torch.floor(255 * torch.arange(0, GC) / GC) + col = col + GC + # CB + colorwheel[col : col + CB, 1] = 255 - torch.floor(255 * torch.arange(CB) / CB) + colorwheel[col : col + CB, 2] = 255 + col = col + CB + # BM + colorwheel[col : col + BM, 2] = 255 + colorwheel[col : col + BM, 0] = torch.floor(255 * torch.arange(0, BM) / BM) + col = col + BM + # MR + colorwheel[col : col + MR, 2] = 255 - torch.floor(255 * torch.arange(MR) / MR) + colorwheel[col : col + MR, 0] = 255 + return colorwheel + + +def _generate_color_palette(num_objects: int): + palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1]) + return [tuple((i * palette) % 255) for i in range(num_objects)] + + +def _parse_colors( + colors: Union[None, str, Tuple[int, int, int], List[Union[str, Tuple[int, int, int]]]], + *, + num_objects: int, + dtype: torch.dtype = torch.uint8, +) -> List[Tuple[int, int, int]]: + """ + Parses a specification of colors for a set of objects. + + Args: + colors: A specification of colors for the objects. This can be one of the following: + - None: to generate a color palette automatically. + - A list of colors: where each color is either a string (specifying a named color) or an RGB tuple. + - A string or an RGB tuple: to use the same color for all objects. + + If `colors` is a tuple, it should be a 3-tuple specifying the RGB values of the color. + If `colors` is a list, it should have at least as many elements as the number of objects to color. + + num_objects (int): The number of objects to color. + + Returns: + A list of 3-tuples, specifying the RGB values of the colors. + + Raises: + ValueError: If the number of colors in the list is less than the number of objects to color. + If `colors` is not a list, tuple, string or None. + """ + if colors is None: + colors = _generate_color_palette(num_objects) + elif isinstance(colors, list): + if len(colors) < num_objects: + raise ValueError( + f"Number of colors must be equal or larger than the number of objects, but got {len(colors)} < {num_objects}." + ) + elif not isinstance(colors, (tuple, str)): + raise ValueError(f"`colors` must be a tuple or a string, or a list thereof, but got {colors}.") + elif isinstance(colors, tuple) and len(colors) != 3: + raise ValueError(f"If passed as tuple, colors should be an RGB triplet, but got {colors}.") + else: # colors specifies a single color for all objects + colors = [colors] * num_objects + + colors = [ImageColor.getrgb(color) if isinstance(color, str) else color for color in colors] + if dtype.is_floating_point: # [0, 255] -> [0, 1] + colors = [tuple(v / 255 for v in color) for color in colors] # type: ignore[union-attr] + return colors # type: ignore[return-value] + + +def _log_api_usage_once(obj: Any) -> None: + + """ + Logs API usage(module and name) within an organization. + In a large ecosystem, it's often useful to track the PyTorch and + TorchVision APIs usage. This API provides the similar functionality to the + logging module in the Python stdlib. It can be used for debugging purpose + to log which methods are used and by default it is inactive, unless the user + manually subscribes a logger via the `SetAPIUsageLogger method `_. + Please note it is triggered only once for the same API call within a process. + It does not collect any data from open-source users since it is no-op by default. + For more information, please refer to + * PyTorch note: https://pytorch.org/docs/stable/notes/large_scale_deployments.html#api-usage-logging; + * Logging policy: https://github.com/pytorch/vision/issues/5052; + + Args: + obj (class instance or method): an object to extract info from. + """ + module = obj.__module__ + if not module.startswith("torchvision"): + module = f"torchvision.internal.{module}" + name = obj.__class__.__name__ + if isinstance(obj, FunctionType): + name = obj.__name__ + torch._C._log_api_usage_once(f"{module}.{name}") + + +def _make_ntuple(x: Any, n: int) -> Tuple[Any, ...]: + """ + Make n-tuple from input x. If x is an iterable, then we just convert it to tuple. + Otherwise, we will make a tuple of length n, all with value of x. + reference: https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/utils.py#L8 + + Args: + x (Any): input value + n (int): length of the resulting tuple + """ + if isinstance(x, collections.abc.Iterable): + return tuple(x) + return tuple(repeat(x, n)) diff --git a/lib/python3.12/site-packages/torchvision/version.py b/lib/python3.12/site-packages/torchvision/version.py new file mode 100644 index 0000000000000000000000000000000000000000..3e8c33bd46e080cbacfe14e88d1d44c0b49395d3 --- /dev/null +++ b/lib/python3.12/site-packages/torchvision/version.py @@ -0,0 +1,5 @@ +__version__ = '0.21.0+cu124' +git_version = '7af698794eded568735f9519593603c1ec889eba' +from torchvision.extension import _check_cuda_version +if _check_cuda_version() > 0: + cuda = _check_cuda_version() diff --git a/lib/python3.12/site-packages/transformers/kernels/__init__.py b/lib/python3.12/site-packages/transformers/kernels/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/transformers/kernels/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/kernels/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fc3a2053a25fdfdfa18ea1cac9ed0ab134af53c2 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/kernels/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/kernels/deta/cpu/ms_deform_attn_cpu.cpp b/lib/python3.12/site-packages/transformers/kernels/deta/cpu/ms_deform_attn_cpu.cpp new file mode 100644 index 0000000000000000000000000000000000000000..388a73d22d4c9b561e2a887b50a1897b8cf2def9 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/deta/cpu/ms_deform_attn_cpu.cpp @@ -0,0 +1,40 @@ +/*! +************************************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************************************** +* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 +************************************************************************************************** +*/ + +#include + +#include +#include + + +at::Tensor +ms_deform_attn_cpu_forward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const int im2col_step) +{ + AT_ERROR("Not implement on cpu"); +} + +std::vector +ms_deform_attn_cpu_backward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const at::Tensor &grad_output, + const int im2col_step) +{ + AT_ERROR("Not implement on cpu"); +} diff --git a/lib/python3.12/site-packages/transformers/kernels/deta/cpu/ms_deform_attn_cpu.h b/lib/python3.12/site-packages/transformers/kernels/deta/cpu/ms_deform_attn_cpu.h new file mode 100644 index 0000000000000000000000000000000000000000..7eac8c8bcd1bf529bb9c13d54d2d4215c9e4c89f --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/deta/cpu/ms_deform_attn_cpu.h @@ -0,0 +1,32 @@ +/*! +************************************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************************************** +* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 +************************************************************************************************** +*/ + +#pragma once +#include + +at::Tensor +ms_deform_attn_cpu_forward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const int im2col_step); + +std::vector +ms_deform_attn_cpu_backward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const at::Tensor &grad_output, + const int im2col_step); + diff --git a/lib/python3.12/site-packages/transformers/kernels/deta/cuda/ms_deform_attn_cuda.cu b/lib/python3.12/site-packages/transformers/kernels/deta/cuda/ms_deform_attn_cuda.cu new file mode 100644 index 0000000000000000000000000000000000000000..8ea1d7fabe2684dbb85f00fae2c47b469687cb2c --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/deta/cuda/ms_deform_attn_cuda.cu @@ -0,0 +1,156 @@ +/*! +************************************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************************************** +* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 +************************************************************************************************** +*/ + +#include +#include "cuda/ms_deform_im2col_cuda.cuh" + +#include +#include +#include +#include + +#pragma once +#include + + +at::Tensor ms_deform_attn_cuda_forward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const int im2col_step) +{ + AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous"); + AT_ASSERTM(spatial_shapes.is_contiguous(), "spatial_shapes tensor has to be contiguous"); + AT_ASSERTM(level_start_index.is_contiguous(), "level_start_index tensor has to be contiguous"); + AT_ASSERTM(sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous"); + AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous"); + + AT_ASSERTM(value.type().is_cuda(), "value must be a CUDA tensor"); + AT_ASSERTM(spatial_shapes.type().is_cuda(), "spatial_shapes must be a CUDA tensor"); + AT_ASSERTM(level_start_index.type().is_cuda(), "level_start_index must be a CUDA tensor"); + AT_ASSERTM(sampling_loc.type().is_cuda(), "sampling_loc must be a CUDA tensor"); + AT_ASSERTM(attn_weight.type().is_cuda(), "attn_weight must be a CUDA tensor"); + + const int batch = value.size(0); + const int spatial_size = value.size(1); + const int num_heads = value.size(2); + const int channels = value.size(3); + + const int num_levels = spatial_shapes.size(0); + + const int num_query = sampling_loc.size(1); + const int num_point = sampling_loc.size(4); + + const int im2col_step_ = std::min(batch, im2col_step); + + AT_ASSERTM(batch % im2col_step_ == 0, "batch(%d) must divide im2col_step(%d)", batch, im2col_step_); + + auto output = at::zeros({batch, num_query, num_heads, channels}, value.options()); + + const int batch_n = im2col_step_; + auto output_n = output.view({batch/im2col_step_, batch_n, num_query, num_heads, channels}); + auto per_value_size = spatial_size * num_heads * channels; + auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2; + auto per_attn_weight_size = num_query * num_heads * num_levels * num_point; + for (int n = 0; n < batch/im2col_step_; ++n) + { + auto columns = output_n.select(0, n); + AT_DISPATCH_FLOATING_TYPES(value.type(), "ms_deform_attn_forward_cuda", ([&] { + ms_deformable_im2col_cuda(at::cuda::getCurrentCUDAStream(), + value.data() + n * im2col_step_ * per_value_size, + spatial_shapes.data(), + level_start_index.data(), + sampling_loc.data() + n * im2col_step_ * per_sample_loc_size, + attn_weight.data() + n * im2col_step_ * per_attn_weight_size, + batch_n, spatial_size, num_heads, channels, num_levels, num_query, num_point, + columns.data()); + + })); + } + + output = output.view({batch, num_query, num_heads*channels}); + + return output; +} + + +std::vector ms_deform_attn_cuda_backward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const at::Tensor &grad_output, + const int im2col_step) +{ + + AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous"); + AT_ASSERTM(spatial_shapes.is_contiguous(), "spatial_shapes tensor has to be contiguous"); + AT_ASSERTM(level_start_index.is_contiguous(), "level_start_index tensor has to be contiguous"); + AT_ASSERTM(sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous"); + AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous"); + AT_ASSERTM(grad_output.is_contiguous(), "grad_output tensor has to be contiguous"); + + AT_ASSERTM(value.type().is_cuda(), "value must be a CUDA tensor"); + AT_ASSERTM(spatial_shapes.type().is_cuda(), "spatial_shapes must be a CUDA tensor"); + AT_ASSERTM(level_start_index.type().is_cuda(), "level_start_index must be a CUDA tensor"); + AT_ASSERTM(sampling_loc.type().is_cuda(), "sampling_loc must be a CUDA tensor"); + AT_ASSERTM(attn_weight.type().is_cuda(), "attn_weight must be a CUDA tensor"); + AT_ASSERTM(grad_output.type().is_cuda(), "grad_output must be a CUDA tensor"); + + const int batch = value.size(0); + const int spatial_size = value.size(1); + const int num_heads = value.size(2); + const int channels = value.size(3); + + const int num_levels = spatial_shapes.size(0); + + const int num_query = sampling_loc.size(1); + const int num_point = sampling_loc.size(4); + + const int im2col_step_ = std::min(batch, im2col_step); + + AT_ASSERTM(batch % im2col_step_ == 0, "batch(%d) must divide im2col_step(%d)", batch, im2col_step_); + + auto grad_value = at::zeros_like(value); + auto grad_sampling_loc = at::zeros_like(sampling_loc); + auto grad_attn_weight = at::zeros_like(attn_weight); + + const int batch_n = im2col_step_; + auto per_value_size = spatial_size * num_heads * channels; + auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2; + auto per_attn_weight_size = num_query * num_heads * num_levels * num_point; + auto grad_output_n = grad_output.view({batch/im2col_step_, batch_n, num_query, num_heads, channels}); + + for (int n = 0; n < batch/im2col_step_; ++n) + { + auto grad_output_g = grad_output_n.select(0, n); + AT_DISPATCH_FLOATING_TYPES(value.type(), "ms_deform_attn_backward_cuda", ([&] { + ms_deformable_col2im_cuda(at::cuda::getCurrentCUDAStream(), + grad_output_g.data(), + value.data() + n * im2col_step_ * per_value_size, + spatial_shapes.data(), + level_start_index.data(), + sampling_loc.data() + n * im2col_step_ * per_sample_loc_size, + attn_weight.data() + n * im2col_step_ * per_attn_weight_size, + batch_n, spatial_size, num_heads, channels, num_levels, num_query, num_point, + grad_value.data() + n * im2col_step_ * per_value_size, + grad_sampling_loc.data() + n * im2col_step_ * per_sample_loc_size, + grad_attn_weight.data() + n * im2col_step_ * per_attn_weight_size); + + })); + } + + return { + grad_value, grad_sampling_loc, grad_attn_weight + }; +} diff --git a/lib/python3.12/site-packages/transformers/kernels/deta/cuda/ms_deform_attn_cuda.cuh b/lib/python3.12/site-packages/transformers/kernels/deta/cuda/ms_deform_attn_cuda.cuh new file mode 100644 index 0000000000000000000000000000000000000000..34f8ae9cb77bbaa8cb4dd25e0cb86632db9ad05d --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/deta/cuda/ms_deform_attn_cuda.cuh @@ -0,0 +1,1467 @@ +/*! +************************************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************************************** +* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 +************************************************************************************************** +*/ + +#include + +#include +#include + +#include +#include +#include + +#include +#include + +#include + +#define CUDA_KERNEL_LOOP(i, n) \ + for (int i = blockIdx.x * blockDim.x + threadIdx.x; \ + i < (n); \ + i += blockDim.x * gridDim.x) + + +at::Tensor ms_deform_attn_cuda_forward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const int im2col_step) +{ + AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous"); + AT_ASSERTM(spatial_shapes.is_contiguous(), "spatial_shapes tensor has to be contiguous"); + AT_ASSERTM(level_start_index.is_contiguous(), "level_start_index tensor has to be contiguous"); + AT_ASSERTM(sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous"); + AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous"); + + AT_ASSERTM(value.type().is_cuda(), "value must be a CUDA tensor"); + AT_ASSERTM(spatial_shapes.type().is_cuda(), "spatial_shapes must be a CUDA tensor"); + AT_ASSERTM(level_start_index.type().is_cuda(), "level_start_index must be a CUDA tensor"); + AT_ASSERTM(sampling_loc.type().is_cuda(), "sampling_loc must be a CUDA tensor"); + AT_ASSERTM(attn_weight.type().is_cuda(), "attn_weight must be a CUDA tensor"); + + const int batch = value.size(0); + const int spatial_size = value.size(1); + const int num_heads = value.size(2); + const int channels = value.size(3); + + const int num_levels = spatial_shapes.size(0); + + const int num_query = sampling_loc.size(1); + const int num_point = sampling_loc.size(4); + + const int im2col_step_ = std::min(batch, im2col_step); + + AT_ASSERTM(batch % im2col_step_ == 0, "batch(%d) must divide im2col_step(%d)", batch, im2col_step_); + + auto output = at::zeros({batch, num_query, num_heads, channels}, value.options()); + + const int batch_n = im2col_step_; + auto output_n = output.view({batch/im2col_step_, batch_n, num_query, num_heads, channels}); + auto per_value_size = spatial_size * num_heads * channels; + auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2; + auto per_attn_weight_size = num_query * num_heads * num_levels * num_point; + for (int n = 0; n < batch/im2col_step_; ++n) + { + auto columns = output_n.select(0, n); + AT_DISPATCH_FLOATING_TYPES(value.type(), "ms_deform_attn_forward_cuda", ([&] { + ms_deformable_im2col_cuda(at::cuda::getCurrentCUDAStream(), + value.data() + n * im2col_step_ * per_value_size, + spatial_shapes.data(), + level_start_index.data(), + sampling_loc.data() + n * im2col_step_ * per_sample_loc_size, + attn_weight.data() + n * im2col_step_ * per_attn_weight_size, + batch_n, spatial_size, num_heads, channels, num_levels, num_query, num_point, + columns.data()); + + })); + } + + output = output.view({batch, num_query, num_heads*channels}); + + return output; +} + + +std::vector ms_deform_attn_cuda_backward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const at::Tensor &grad_output, + const int im2col_step) +{ + + AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous"); + AT_ASSERTM(spatial_shapes.is_contiguous(), "spatial_shapes tensor has to be contiguous"); + AT_ASSERTM(level_start_index.is_contiguous(), "level_start_index tensor has to be contiguous"); + AT_ASSERTM(sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous"); + AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous"); + AT_ASSERTM(grad_output.is_contiguous(), "grad_output tensor has to be contiguous"); + + AT_ASSERTM(value.type().is_cuda(), "value must be a CUDA tensor"); + AT_ASSERTM(spatial_shapes.type().is_cuda(), "spatial_shapes must be a CUDA tensor"); + AT_ASSERTM(level_start_index.type().is_cuda(), "level_start_index must be a CUDA tensor"); + AT_ASSERTM(sampling_loc.type().is_cuda(), "sampling_loc must be a CUDA tensor"); + AT_ASSERTM(attn_weight.type().is_cuda(), "attn_weight must be a CUDA tensor"); + AT_ASSERTM(grad_output.type().is_cuda(), "grad_output must be a CUDA tensor"); + + const int batch = value.size(0); + const int spatial_size = value.size(1); + const int num_heads = value.size(2); + const int channels = value.size(3); + + const int num_levels = spatial_shapes.size(0); + + const int num_query = sampling_loc.size(1); + const int num_point = sampling_loc.size(4); + + const int im2col_step_ = std::min(batch, im2col_step); + + AT_ASSERTM(batch % im2col_step_ == 0, "batch(%d) must divide im2col_step(%d)", batch, im2col_step_); + + auto grad_value = at::zeros_like(value); + auto grad_sampling_loc = at::zeros_like(sampling_loc); + auto grad_attn_weight = at::zeros_like(attn_weight); + + const int batch_n = im2col_step_; + auto per_value_size = spatial_size * num_heads * channels; + auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2; + auto per_attn_weight_size = num_query * num_heads * num_levels * num_point; + auto grad_output_n = grad_output.view({batch/im2col_step_, batch_n, num_query, num_heads, channels}); + + for (int n = 0; n < batch/im2col_step_; ++n) + { + auto grad_output_g = grad_output_n.select(0, n); + AT_DISPATCH_FLOATING_TYPES(value.type(), "ms_deform_attn_backward_cuda", ([&] { + ms_deformable_col2im_cuda(at::cuda::getCurrentCUDAStream(), + grad_output_g.data(), + value.data() + n * im2col_step_ * per_value_size, + spatial_shapes.data(), + level_start_index.data(), + sampling_loc.data() + n * im2col_step_ * per_sample_loc_size, + attn_weight.data() + n * im2col_step_ * per_attn_weight_size, + batch_n, spatial_size, num_heads, channels, num_levels, num_query, num_point, + grad_value.data() + n * im2col_step_ * per_value_size, + grad_sampling_loc.data() + n * im2col_step_ * per_sample_loc_size, + grad_attn_weight.data() + n * im2col_step_ * per_attn_weight_size); + + })); + } + + return { + grad_value, grad_sampling_loc, grad_attn_weight + }; +} + +const int CUDA_NUM_THREADS = 1024; +inline int GET_BLOCKS(const int N, const int num_threads) +{ + return (N + num_threads - 1) / num_threads; +} + + +template +__device__ scalar_t ms_deform_attn_im2col_bilinear(const scalar_t* &bottom_data, + const int &height, const int &width, const int &nheads, const int &channels, + const scalar_t &h, const scalar_t &w, const int &m, const int &c) +{ + const int h_low = floor(h); + const int w_low = floor(w); + const int h_high = h_low + 1; + const int w_high = w_low + 1; + + const scalar_t lh = h - h_low; + const scalar_t lw = w - w_low; + const scalar_t hh = 1 - lh, hw = 1 - lw; + + const int w_stride = nheads * channels; + const int h_stride = width * w_stride; + const int h_low_ptr_offset = h_low * h_stride; + const int h_high_ptr_offset = h_low_ptr_offset + h_stride; + const int w_low_ptr_offset = w_low * w_stride; + const int w_high_ptr_offset = w_low_ptr_offset + w_stride; + const int base_ptr = m * channels + c; + + scalar_t v1 = 0; + if (h_low >= 0 && w_low >= 0) + { + const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr; + v1 = bottom_data[ptr1]; + } + scalar_t v2 = 0; + if (h_low >= 0 && w_high <= width - 1) + { + const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr; + v2 = bottom_data[ptr2]; + } + scalar_t v3 = 0; + if (h_high <= height - 1 && w_low >= 0) + { + const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr; + v3 = bottom_data[ptr3]; + } + scalar_t v4 = 0; + if (h_high <= height - 1 && w_high <= width - 1) + { + const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr; + v4 = bottom_data[ptr4]; + } + + const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw; + + const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4); + return val; +} + + +template +__device__ void ms_deform_attn_col2im_bilinear(const scalar_t* &bottom_data, + const int &height, const int &width, const int &nheads, const int &channels, + const scalar_t &h, const scalar_t &w, const int &m, const int &c, + const scalar_t &top_grad, + const scalar_t &attn_weight, + scalar_t* &grad_value, + scalar_t* grad_sampling_loc, + scalar_t* grad_attn_weight) +{ + const int h_low = floor(h); + const int w_low = floor(w); + const int h_high = h_low + 1; + const int w_high = w_low + 1; + + const scalar_t lh = h - h_low; + const scalar_t lw = w - w_low; + const scalar_t hh = 1 - lh, hw = 1 - lw; + + const int w_stride = nheads * channels; + const int h_stride = width * w_stride; + const int h_low_ptr_offset = h_low * h_stride; + const int h_high_ptr_offset = h_low_ptr_offset + h_stride; + const int w_low_ptr_offset = w_low * w_stride; + const int w_high_ptr_offset = w_low_ptr_offset + w_stride; + const int base_ptr = m * channels + c; + + const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw; + const scalar_t top_grad_value = top_grad * attn_weight; + scalar_t grad_h_weight = 0, grad_w_weight = 0; + + scalar_t v1 = 0; + if (h_low >= 0 && w_low >= 0) + { + const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr; + v1 = bottom_data[ptr1]; + grad_h_weight -= hw * v1; + grad_w_weight -= hh * v1; + atomicAdd(grad_value+ptr1, w1*top_grad_value); + } + scalar_t v2 = 0; + if (h_low >= 0 && w_high <= width - 1) + { + const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr; + v2 = bottom_data[ptr2]; + grad_h_weight -= lw * v2; + grad_w_weight += hh * v2; + atomicAdd(grad_value+ptr2, w2*top_grad_value); + } + scalar_t v3 = 0; + if (h_high <= height - 1 && w_low >= 0) + { + const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr; + v3 = bottom_data[ptr3]; + grad_h_weight += hw * v3; + grad_w_weight -= lh * v3; + atomicAdd(grad_value+ptr3, w3*top_grad_value); + } + scalar_t v4 = 0; + if (h_high <= height - 1 && w_high <= width - 1) + { + const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr; + v4 = bottom_data[ptr4]; + grad_h_weight += lw * v4; + grad_w_weight += lh * v4; + atomicAdd(grad_value+ptr4, w4*top_grad_value); + } + + const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4); + *grad_attn_weight = top_grad * val; + *grad_sampling_loc = width * grad_w_weight * top_grad_value; + *(grad_sampling_loc + 1) = height * grad_h_weight * top_grad_value; +} + + +template +__device__ void ms_deform_attn_col2im_bilinear_gm(const scalar_t* &bottom_data, + const int &height, const int &width, const int &nheads, const int &channels, + const scalar_t &h, const scalar_t &w, const int &m, const int &c, + const scalar_t &top_grad, + const scalar_t &attn_weight, + scalar_t* &grad_value, + scalar_t* grad_sampling_loc, + scalar_t* grad_attn_weight) +{ + const int h_low = floor(h); + const int w_low = floor(w); + const int h_high = h_low + 1; + const int w_high = w_low + 1; + + const scalar_t lh = h - h_low; + const scalar_t lw = w - w_low; + const scalar_t hh = 1 - lh, hw = 1 - lw; + + const int w_stride = nheads * channels; + const int h_stride = width * w_stride; + const int h_low_ptr_offset = h_low * h_stride; + const int h_high_ptr_offset = h_low_ptr_offset + h_stride; + const int w_low_ptr_offset = w_low * w_stride; + const int w_high_ptr_offset = w_low_ptr_offset + w_stride; + const int base_ptr = m * channels + c; + + const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw; + const scalar_t top_grad_value = top_grad * attn_weight; + scalar_t grad_h_weight = 0, grad_w_weight = 0; + + scalar_t v1 = 0; + if (h_low >= 0 && w_low >= 0) + { + const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr; + v1 = bottom_data[ptr1]; + grad_h_weight -= hw * v1; + grad_w_weight -= hh * v1; + atomicAdd(grad_value+ptr1, w1*top_grad_value); + } + scalar_t v2 = 0; + if (h_low >= 0 && w_high <= width - 1) + { + const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr; + v2 = bottom_data[ptr2]; + grad_h_weight -= lw * v2; + grad_w_weight += hh * v2; + atomicAdd(grad_value+ptr2, w2*top_grad_value); + } + scalar_t v3 = 0; + if (h_high <= height - 1 && w_low >= 0) + { + const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr; + v3 = bottom_data[ptr3]; + grad_h_weight += hw * v3; + grad_w_weight -= lh * v3; + atomicAdd(grad_value+ptr3, w3*top_grad_value); + } + scalar_t v4 = 0; + if (h_high <= height - 1 && w_high <= width - 1) + { + const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr; + v4 = bottom_data[ptr4]; + grad_h_weight += lw * v4; + grad_w_weight += lh * v4; + atomicAdd(grad_value+ptr4, w4*top_grad_value); + } + + const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4); + atomicAdd(grad_attn_weight, top_grad * val); + atomicAdd(grad_sampling_loc, width * grad_w_weight * top_grad_value); + atomicAdd(grad_sampling_loc + 1, height * grad_h_weight * top_grad_value); +} + + +template +__global__ void ms_deformable_im2col_gpu_kernel(const int n, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *data_col) +{ + CUDA_KERNEL_LOOP(index, n) + { + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + scalar_t *data_col_ptr = data_col + index; + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + scalar_t col = 0; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const scalar_t *data_value_ptr = data_value + (data_value_ptr_init_offset + level_start_id * qid_stride); + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + col += ms_deform_attn_im2col_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col) * weight; + } + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + } + } + *data_col_ptr = col; + } +} + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2]; + __shared__ scalar_t cache_grad_attn_weight[blockSize]; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + if (tid == 0) + { + scalar_t _grad_w=cache_grad_sampling_loc[0], _grad_h=cache_grad_sampling_loc[1], _grad_a=cache_grad_attn_weight[0]; + int sid=2; + for (unsigned int tid = 1; tid < blockSize; ++tid) + { + _grad_w += cache_grad_sampling_loc[sid]; + _grad_h += cache_grad_sampling_loc[sid + 1]; + _grad_a += cache_grad_attn_weight[tid]; + sid += 2; + } + + + *grad_sampling_loc = _grad_w; + *(grad_sampling_loc + 1) = _grad_h; + *grad_attn_weight = _grad_a; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2]; + __shared__ scalar_t cache_grad_attn_weight[blockSize]; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + + for (unsigned int s=blockSize/2; s>0; s>>=1) + { + if (tid < s) { + const unsigned int xid1 = tid << 1; + const unsigned int xid2 = (tid + s) << 1; + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1]; + } + __syncthreads(); + } + + if (tid == 0) + { + *grad_sampling_loc = cache_grad_sampling_loc[0]; + *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1]; + *grad_attn_weight = cache_grad_attn_weight[0]; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v1(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + extern __shared__ int _s[]; + scalar_t* cache_grad_sampling_loc = (scalar_t*)_s; + scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + if (tid == 0) + { + scalar_t _grad_w=cache_grad_sampling_loc[0], _grad_h=cache_grad_sampling_loc[1], _grad_a=cache_grad_attn_weight[0]; + int sid=2; + for (unsigned int tid = 1; tid < blockDim.x; ++tid) + { + _grad_w += cache_grad_sampling_loc[sid]; + _grad_h += cache_grad_sampling_loc[sid + 1]; + _grad_a += cache_grad_attn_weight[tid]; + sid += 2; + } + + + *grad_sampling_loc = _grad_w; + *(grad_sampling_loc + 1) = _grad_h; + *grad_attn_weight = _grad_a; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + extern __shared__ int _s[]; + scalar_t* cache_grad_sampling_loc = (scalar_t*)_s; + scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + + for (unsigned int s=blockDim.x/2, spre=blockDim.x; s>0; s>>=1, spre>>=1) + { + if (tid < s) { + const unsigned int xid1 = tid << 1; + const unsigned int xid2 = (tid + s) << 1; + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1]; + if (tid + (s << 1) < spre) + { + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)]; + } + } + __syncthreads(); + } + + if (tid == 0) + { + *grad_sampling_loc = cache_grad_sampling_loc[0]; + *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1]; + *grad_attn_weight = cache_grad_attn_weight[0]; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + extern __shared__ int _s[]; + scalar_t* cache_grad_sampling_loc = (scalar_t*)_s; + scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + + for (unsigned int s=blockDim.x/2, spre=blockDim.x; s>0; s>>=1, spre>>=1) + { + if (tid < s) { + const unsigned int xid1 = tid << 1; + const unsigned int xid2 = (tid + s) << 1; + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1]; + if (tid + (s << 1) < spre) + { + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)]; + } + } + __syncthreads(); + } + + if (tid == 0) + { + atomicAdd(grad_sampling_loc, cache_grad_sampling_loc[0]); + atomicAdd(grad_sampling_loc + 1, cache_grad_sampling_loc[1]); + atomicAdd(grad_attn_weight, cache_grad_attn_weight[0]); + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +__global__ void ms_deformable_col2im_gpu_kernel_gm(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear_gm( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + grad_sampling_loc, grad_attn_weight); + } + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +void ms_deformable_im2col_cuda(cudaStream_t stream, + const scalar_t* data_value, + const int64_t* data_spatial_shapes, + const int64_t* data_level_start_index, + const scalar_t* data_sampling_loc, + const scalar_t* data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t* data_col) +{ + const int num_kernels = batch_size * num_query * num_heads * channels; + const int num_actual_kernels = batch_size * num_query * num_heads * channels; + const int num_threads = CUDA_NUM_THREADS; + ms_deformable_im2col_gpu_kernel + <<>>( + num_kernels, data_value, data_spatial_shapes, data_level_start_index, data_sampling_loc, data_attn_weight, + batch_size, spatial_size, num_heads, channels, num_levels, num_query, num_point, data_col); + + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + { + printf("error in ms_deformable_im2col_cuda: %s\n", cudaGetErrorString(err)); + } + +} + +template +void ms_deformable_col2im_cuda(cudaStream_t stream, + const scalar_t* grad_col, + const scalar_t* data_value, + const int64_t * data_spatial_shapes, + const int64_t * data_level_start_index, + const scalar_t * data_sampling_loc, + const scalar_t * data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t* grad_value, + scalar_t* grad_sampling_loc, + scalar_t* grad_attn_weight) +{ + const int num_threads = (channels > CUDA_NUM_THREADS)?CUDA_NUM_THREADS:channels; + const int num_kernels = batch_size * num_query * num_heads * channels; + const int num_actual_kernels = batch_size * num_query * num_heads * channels; + if (channels > 1024) + { + if ((channels & 1023) == 0) + { + ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + else + { + ms_deformable_col2im_gpu_kernel_gm + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + } + else{ + switch(channels) + { + case 1: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 2: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 4: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 8: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 16: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 32: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 64: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 128: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 256: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 512: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 1024: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + default: + if (channels < 64) + { + ms_deformable_col2im_gpu_kernel_shm_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + else + { + ms_deformable_col2im_gpu_kernel_shm_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + } + } + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + { + printf("error in ms_deformable_col2im_cuda: %s\n", cudaGetErrorString(err)); + } + +} diff --git a/lib/python3.12/site-packages/transformers/kernels/deta/cuda/ms_deform_attn_cuda.h b/lib/python3.12/site-packages/transformers/kernels/deta/cuda/ms_deform_attn_cuda.h new file mode 100644 index 0000000000000000000000000000000000000000..fbcf4543e66bb1162f42ce2ae57e1bac92243cb4 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/deta/cuda/ms_deform_attn_cuda.h @@ -0,0 +1,29 @@ +/*! +************************************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************************************** +* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 +************************************************************************************************** +*/ + +#pragma once +#include + +at::Tensor ms_deform_attn_cuda_forward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const int im2col_step); + +std::vector ms_deform_attn_cuda_backward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const at::Tensor &grad_output, + const int im2col_step); diff --git a/lib/python3.12/site-packages/transformers/kernels/deta/cuda/ms_deform_im2col_cuda.cuh b/lib/python3.12/site-packages/transformers/kernels/deta/cuda/ms_deform_im2col_cuda.cuh new file mode 100644 index 0000000000000000000000000000000000000000..c0db0c88c9db2c09d7f601937ea0f6ac480913bf --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/deta/cuda/ms_deform_im2col_cuda.cuh @@ -0,0 +1,1327 @@ +/*! +************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************** +* Modified from DCN (https://github.com/msracver/Deformable-ConvNets) +* Copyright (c) 2018 Microsoft +************************************************************************** +*/ + +#include +#include +#include + +#include +#include + +#include + +#define CUDA_KERNEL_LOOP(i, n) \ + for (int i = blockIdx.x * blockDim.x + threadIdx.x; \ + i < (n); \ + i += blockDim.x * gridDim.x) + +const int CUDA_NUM_THREADS = 1024; +inline int GET_BLOCKS(const int N, const int num_threads) +{ + return (N + num_threads - 1) / num_threads; +} + + +template +__device__ scalar_t ms_deform_attn_im2col_bilinear(const scalar_t* &bottom_data, + const int &height, const int &width, const int &nheads, const int &channels, + const scalar_t &h, const scalar_t &w, const int &m, const int &c) +{ + const int h_low = floor(h); + const int w_low = floor(w); + const int h_high = h_low + 1; + const int w_high = w_low + 1; + + const scalar_t lh = h - h_low; + const scalar_t lw = w - w_low; + const scalar_t hh = 1 - lh, hw = 1 - lw; + + const int w_stride = nheads * channels; + const int h_stride = width * w_stride; + const int h_low_ptr_offset = h_low * h_stride; + const int h_high_ptr_offset = h_low_ptr_offset + h_stride; + const int w_low_ptr_offset = w_low * w_stride; + const int w_high_ptr_offset = w_low_ptr_offset + w_stride; + const int base_ptr = m * channels + c; + + scalar_t v1 = 0; + if (h_low >= 0 && w_low >= 0) + { + const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr; + v1 = bottom_data[ptr1]; + } + scalar_t v2 = 0; + if (h_low >= 0 && w_high <= width - 1) + { + const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr; + v2 = bottom_data[ptr2]; + } + scalar_t v3 = 0; + if (h_high <= height - 1 && w_low >= 0) + { + const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr; + v3 = bottom_data[ptr3]; + } + scalar_t v4 = 0; + if (h_high <= height - 1 && w_high <= width - 1) + { + const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr; + v4 = bottom_data[ptr4]; + } + + const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw; + + const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4); + return val; +} + + +template +__device__ void ms_deform_attn_col2im_bilinear(const scalar_t* &bottom_data, + const int &height, const int &width, const int &nheads, const int &channels, + const scalar_t &h, const scalar_t &w, const int &m, const int &c, + const scalar_t &top_grad, + const scalar_t &attn_weight, + scalar_t* &grad_value, + scalar_t* grad_sampling_loc, + scalar_t* grad_attn_weight) +{ + const int h_low = floor(h); + const int w_low = floor(w); + const int h_high = h_low + 1; + const int w_high = w_low + 1; + + const scalar_t lh = h - h_low; + const scalar_t lw = w - w_low; + const scalar_t hh = 1 - lh, hw = 1 - lw; + + const int w_stride = nheads * channels; + const int h_stride = width * w_stride; + const int h_low_ptr_offset = h_low * h_stride; + const int h_high_ptr_offset = h_low_ptr_offset + h_stride; + const int w_low_ptr_offset = w_low * w_stride; + const int w_high_ptr_offset = w_low_ptr_offset + w_stride; + const int base_ptr = m * channels + c; + + const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw; + const scalar_t top_grad_value = top_grad * attn_weight; + scalar_t grad_h_weight = 0, grad_w_weight = 0; + + scalar_t v1 = 0; + if (h_low >= 0 && w_low >= 0) + { + const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr; + v1 = bottom_data[ptr1]; + grad_h_weight -= hw * v1; + grad_w_weight -= hh * v1; + atomicAdd(grad_value+ptr1, w1*top_grad_value); + } + scalar_t v2 = 0; + if (h_low >= 0 && w_high <= width - 1) + { + const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr; + v2 = bottom_data[ptr2]; + grad_h_weight -= lw * v2; + grad_w_weight += hh * v2; + atomicAdd(grad_value+ptr2, w2*top_grad_value); + } + scalar_t v3 = 0; + if (h_high <= height - 1 && w_low >= 0) + { + const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr; + v3 = bottom_data[ptr3]; + grad_h_weight += hw * v3; + grad_w_weight -= lh * v3; + atomicAdd(grad_value+ptr3, w3*top_grad_value); + } + scalar_t v4 = 0; + if (h_high <= height - 1 && w_high <= width - 1) + { + const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr; + v4 = bottom_data[ptr4]; + grad_h_weight += lw * v4; + grad_w_weight += lh * v4; + atomicAdd(grad_value+ptr4, w4*top_grad_value); + } + + const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4); + *grad_attn_weight = top_grad * val; + *grad_sampling_loc = width * grad_w_weight * top_grad_value; + *(grad_sampling_loc + 1) = height * grad_h_weight * top_grad_value; +} + + +template +__device__ void ms_deform_attn_col2im_bilinear_gm(const scalar_t* &bottom_data, + const int &height, const int &width, const int &nheads, const int &channels, + const scalar_t &h, const scalar_t &w, const int &m, const int &c, + const scalar_t &top_grad, + const scalar_t &attn_weight, + scalar_t* &grad_value, + scalar_t* grad_sampling_loc, + scalar_t* grad_attn_weight) +{ + const int h_low = floor(h); + const int w_low = floor(w); + const int h_high = h_low + 1; + const int w_high = w_low + 1; + + const scalar_t lh = h - h_low; + const scalar_t lw = w - w_low; + const scalar_t hh = 1 - lh, hw = 1 - lw; + + const int w_stride = nheads * channels; + const int h_stride = width * w_stride; + const int h_low_ptr_offset = h_low * h_stride; + const int h_high_ptr_offset = h_low_ptr_offset + h_stride; + const int w_low_ptr_offset = w_low * w_stride; + const int w_high_ptr_offset = w_low_ptr_offset + w_stride; + const int base_ptr = m * channels + c; + + const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw; + const scalar_t top_grad_value = top_grad * attn_weight; + scalar_t grad_h_weight = 0, grad_w_weight = 0; + + scalar_t v1 = 0; + if (h_low >= 0 && w_low >= 0) + { + const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr; + v1 = bottom_data[ptr1]; + grad_h_weight -= hw * v1; + grad_w_weight -= hh * v1; + atomicAdd(grad_value+ptr1, w1*top_grad_value); + } + scalar_t v2 = 0; + if (h_low >= 0 && w_high <= width - 1) + { + const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr; + v2 = bottom_data[ptr2]; + grad_h_weight -= lw * v2; + grad_w_weight += hh * v2; + atomicAdd(grad_value+ptr2, w2*top_grad_value); + } + scalar_t v3 = 0; + if (h_high <= height - 1 && w_low >= 0) + { + const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr; + v3 = bottom_data[ptr3]; + grad_h_weight += hw * v3; + grad_w_weight -= lh * v3; + atomicAdd(grad_value+ptr3, w3*top_grad_value); + } + scalar_t v4 = 0; + if (h_high <= height - 1 && w_high <= width - 1) + { + const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr; + v4 = bottom_data[ptr4]; + grad_h_weight += lw * v4; + grad_w_weight += lh * v4; + atomicAdd(grad_value+ptr4, w4*top_grad_value); + } + + const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4); + atomicAdd(grad_attn_weight, top_grad * val); + atomicAdd(grad_sampling_loc, width * grad_w_weight * top_grad_value); + atomicAdd(grad_sampling_loc + 1, height * grad_h_weight * top_grad_value); +} + + +template +__global__ void ms_deformable_im2col_gpu_kernel(const int n, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *data_col) +{ + CUDA_KERNEL_LOOP(index, n) + { + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + scalar_t *data_col_ptr = data_col + index; + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + scalar_t col = 0; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const scalar_t *data_value_ptr = data_value + (data_value_ptr_init_offset + level_start_id * qid_stride); + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + col += ms_deform_attn_im2col_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col) * weight; + } + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + } + } + *data_col_ptr = col; + } +} + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2]; + __shared__ scalar_t cache_grad_attn_weight[blockSize]; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + if (tid == 0) + { + scalar_t _grad_w=cache_grad_sampling_loc[0], _grad_h=cache_grad_sampling_loc[1], _grad_a=cache_grad_attn_weight[0]; + int sid=2; + for (unsigned int tid = 1; tid < blockSize; ++tid) + { + _grad_w += cache_grad_sampling_loc[sid]; + _grad_h += cache_grad_sampling_loc[sid + 1]; + _grad_a += cache_grad_attn_weight[tid]; + sid += 2; + } + + + *grad_sampling_loc = _grad_w; + *(grad_sampling_loc + 1) = _grad_h; + *grad_attn_weight = _grad_a; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2]; + __shared__ scalar_t cache_grad_attn_weight[blockSize]; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + + for (unsigned int s=blockSize/2; s>0; s>>=1) + { + if (tid < s) { + const unsigned int xid1 = tid << 1; + const unsigned int xid2 = (tid + s) << 1; + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1]; + } + __syncthreads(); + } + + if (tid == 0) + { + *grad_sampling_loc = cache_grad_sampling_loc[0]; + *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1]; + *grad_attn_weight = cache_grad_attn_weight[0]; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v1(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + extern __shared__ int _s[]; + scalar_t* cache_grad_sampling_loc = (scalar_t*)_s; + scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + if (tid == 0) + { + scalar_t _grad_w=cache_grad_sampling_loc[0], _grad_h=cache_grad_sampling_loc[1], _grad_a=cache_grad_attn_weight[0]; + int sid=2; + for (unsigned int tid = 1; tid < blockDim.x; ++tid) + { + _grad_w += cache_grad_sampling_loc[sid]; + _grad_h += cache_grad_sampling_loc[sid + 1]; + _grad_a += cache_grad_attn_weight[tid]; + sid += 2; + } + + + *grad_sampling_loc = _grad_w; + *(grad_sampling_loc + 1) = _grad_h; + *grad_attn_weight = _grad_a; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + extern __shared__ int _s[]; + scalar_t* cache_grad_sampling_loc = (scalar_t*)_s; + scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + + for (unsigned int s=blockDim.x/2, spre=blockDim.x; s>0; s>>=1, spre>>=1) + { + if (tid < s) { + const unsigned int xid1 = tid << 1; + const unsigned int xid2 = (tid + s) << 1; + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1]; + if (tid + (s << 1) < spre) + { + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)]; + } + } + __syncthreads(); + } + + if (tid == 0) + { + *grad_sampling_loc = cache_grad_sampling_loc[0]; + *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1]; + *grad_attn_weight = cache_grad_attn_weight[0]; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + extern __shared__ int _s[]; + scalar_t* cache_grad_sampling_loc = (scalar_t*)_s; + scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + + for (unsigned int s=blockDim.x/2, spre=blockDim.x; s>0; s>>=1, spre>>=1) + { + if (tid < s) { + const unsigned int xid1 = tid << 1; + const unsigned int xid2 = (tid + s) << 1; + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1]; + if (tid + (s << 1) < spre) + { + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)]; + } + } + __syncthreads(); + } + + if (tid == 0) + { + atomicAdd(grad_sampling_loc, cache_grad_sampling_loc[0]); + atomicAdd(grad_sampling_loc + 1, cache_grad_sampling_loc[1]); + atomicAdd(grad_attn_weight, cache_grad_attn_weight[0]); + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +__global__ void ms_deformable_col2im_gpu_kernel_gm(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear_gm( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + grad_sampling_loc, grad_attn_weight); + } + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +void ms_deformable_im2col_cuda(cudaStream_t stream, + const scalar_t* data_value, + const int64_t* data_spatial_shapes, + const int64_t* data_level_start_index, + const scalar_t* data_sampling_loc, + const scalar_t* data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t* data_col) +{ + const int num_kernels = batch_size * num_query * num_heads * channels; + const int num_actual_kernels = batch_size * num_query * num_heads * channels; + const int num_threads = CUDA_NUM_THREADS; + ms_deformable_im2col_gpu_kernel + <<>>( + num_kernels, data_value, data_spatial_shapes, data_level_start_index, data_sampling_loc, data_attn_weight, + batch_size, spatial_size, num_heads, channels, num_levels, num_query, num_point, data_col); + + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + { + printf("error in ms_deformable_im2col_cuda: %s\n", cudaGetErrorString(err)); + } + +} + +template +void ms_deformable_col2im_cuda(cudaStream_t stream, + const scalar_t* grad_col, + const scalar_t* data_value, + const int64_t * data_spatial_shapes, + const int64_t * data_level_start_index, + const scalar_t * data_sampling_loc, + const scalar_t * data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t* grad_value, + scalar_t* grad_sampling_loc, + scalar_t* grad_attn_weight) +{ + const int num_threads = (channels > CUDA_NUM_THREADS)?CUDA_NUM_THREADS:channels; + const int num_kernels = batch_size * num_query * num_heads * channels; + const int num_actual_kernels = batch_size * num_query * num_heads * channels; + if (channels > 1024) + { + if ((channels & 1023) == 0) + { + ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + else + { + ms_deformable_col2im_gpu_kernel_gm + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + } + else{ + switch(channels) + { + case 1: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 2: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 4: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 8: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 16: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 32: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 64: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 128: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 256: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 512: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 1024: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + default: + if (channels < 64) + { + ms_deformable_col2im_gpu_kernel_shm_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + else + { + ms_deformable_col2im_gpu_kernel_shm_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + } + } + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + { + printf("error in ms_deformable_col2im_cuda: %s\n", cudaGetErrorString(err)); + } + +} diff --git a/lib/python3.12/site-packages/transformers/kernels/deta/ms_deform_attn.h b/lib/python3.12/site-packages/transformers/kernels/deta/ms_deform_attn.h new file mode 100644 index 0000000000000000000000000000000000000000..119b1fa317d1e5fcfb61a4837e560e9248db05f3 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/deta/ms_deform_attn.h @@ -0,0 +1,61 @@ +/*! +************************************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************************************** +* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 +************************************************************************************************** +*/ + +#pragma once + +#include "cpu/ms_deform_attn_cpu.h" + +#ifdef WITH_CUDA +#include "cuda/ms_deform_attn_cuda.h" +#endif + + +at::Tensor +ms_deform_attn_forward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const int im2col_step) +{ + if (value.type().is_cuda()) + { +#ifdef WITH_CUDA + return ms_deform_attn_cuda_forward( + value, spatial_shapes, level_start_index, sampling_loc, attn_weight, im2col_step); +#else + AT_ERROR("Not compiled with GPU support"); +#endif + } + AT_ERROR("Not implemented on the CPU"); +} + +std::vector +ms_deform_attn_backward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const at::Tensor &grad_output, + const int im2col_step) +{ + if (value.type().is_cuda()) + { +#ifdef WITH_CUDA + return ms_deform_attn_cuda_backward( + value, spatial_shapes, level_start_index, sampling_loc, attn_weight, grad_output, im2col_step); +#else + AT_ERROR("Not compiled with GPU support"); +#endif + } + AT_ERROR("Not implemented on the CPU"); +} diff --git a/lib/python3.12/site-packages/transformers/kernels/deta/vision.cpp b/lib/python3.12/site-packages/transformers/kernels/deta/vision.cpp new file mode 100644 index 0000000000000000000000000000000000000000..6ce3875568b9ba8d660c90acc805077cca98f891 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/deta/vision.cpp @@ -0,0 +1,16 @@ +/*! +************************************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************************************** +* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 +************************************************************************************************** +*/ + +#include "ms_deform_attn.h" + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("ms_deform_attn_forward", &ms_deform_attn_forward, "ms_deform_attn_forward"); + m.def("ms_deform_attn_backward", &ms_deform_attn_backward, "ms_deform_attn_backward"); +} \ No newline at end of file diff --git a/lib/python3.12/site-packages/transformers/kernels/falcon_mamba/__init__.py b/lib/python3.12/site-packages/transformers/kernels/falcon_mamba/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..da88e3394f653369a7443245c67dcbe57f2ed23e --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/falcon_mamba/__init__.py @@ -0,0 +1,15 @@ +# coding=utf-8 +# Copyright 2024 Tri Dao, Albert Gu, Technological Innovation Institute and HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from .selective_scan_with_ln_interface import mamba_inner_fn diff --git a/lib/python3.12/site-packages/transformers/kernels/falcon_mamba/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/kernels/falcon_mamba/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5a89279907f95133faf27581a87b8b72a4eb3239 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/kernels/falcon_mamba/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/kernels/falcon_mamba/__pycache__/selective_scan_with_ln_interface.cpython-312.pyc b/lib/python3.12/site-packages/transformers/kernels/falcon_mamba/__pycache__/selective_scan_with_ln_interface.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ca5243ae54ce0acfdfabb5bde4698902e1863039 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/kernels/falcon_mamba/__pycache__/selective_scan_with_ln_interface.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/kernels/falcon_mamba/selective_scan_with_ln_interface.py b/lib/python3.12/site-packages/transformers/kernels/falcon_mamba/selective_scan_with_ln_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..4a74986a81a13f9428eab353de5b61a4d101972d --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/falcon_mamba/selective_scan_with_ln_interface.py @@ -0,0 +1,525 @@ +# coding=utf-8 +# Copyright 2024 Tri Dao, Albert Gu, Technological Innovation Institute and HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Original code from: https://github.com/state-spaces/mamba/blob/main/mamba_ssm/ops/selective_scan_interface.py + +import torch +import torch.nn.functional as F +from einops import rearrange, repeat +from torch.cuda.amp import custom_bwd, custom_fwd + + +try: + import causal_conv1d_cuda +except ImportError: + causal_conv1d_cuda = None + +import mamba_ssm +import selective_scan_cuda + + +# For BC for old mamba-ssm versions: https://github.com/huggingface/transformers/pull/33195#discussion_r1736401127 +if hasattr(mamba_ssm.ops.triton, "layernorm"): + from mamba_ssm.ops.triton.layernorm import _layer_norm_fwd +else: + from mamba_ssm.ops.triton.layer_norm import _layer_norm_fwd + + +class SelectiveScanFn(torch.autograd.Function): + @staticmethod + def forward( + ctx, u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False, return_last_state=False + ): + if u.stride(-1) != 1: + u = u.contiguous() + if delta.stride(-1) != 1: + delta = delta.contiguous() + if D is not None: + D = D.contiguous() + if B.stride(-1) != 1: + B = B.contiguous() + if C.stride(-1) != 1: + C = C.contiguous() + if z is not None and z.stride(-1) != 1: + z = z.contiguous() + if B.dim() == 3: + B = rearrange(B, "b dstate l -> b 1 dstate l") + ctx.squeeze_B = True + if C.dim() == 3: + C = rearrange(C, "b dstate l -> b 1 dstate l") + ctx.squeeze_C = True + out, x, *rest = selective_scan_cuda.fwd(u, delta, A, B, C, D, z, delta_bias, delta_softplus) + ctx.delta_softplus = delta_softplus + ctx.has_z = z is not None + last_state = x[:, :, -1, 1::2] # (batch, dim, dstate) + if not ctx.has_z: + ctx.save_for_backward(u, delta, A, B, C, D, delta_bias, x) + return out if not return_last_state else (out, last_state) + else: + ctx.save_for_backward(u, delta, A, B, C, D, z, delta_bias, x, out) + out_z = rest[0] + return out_z if not return_last_state else (out_z, last_state) + + @staticmethod + def backward(ctx, dout, *args): + if not ctx.has_z: + u, delta, A, B, C, D, delta_bias, x = ctx.saved_tensors + z = None + out = None + else: + u, delta, A, B, C, D, z, delta_bias, x, out = ctx.saved_tensors + if dout.stride(-1) != 1: + dout = dout.contiguous() + # The kernel supports passing in a pre-allocated dz (e.g., in case we want to fuse the + # backward of selective_scan_cuda with the backward of chunk). + # Here we just pass in None and dz will be allocated in the C++ code. + du, ddelta, dA, dB, dC, dD, ddelta_bias, *rest = selective_scan_cuda.bwd( + u, + delta, + A, + B, + C, + D, + z, + delta_bias, + dout, + x, + out, + None, + ctx.delta_softplus, + False, # option to recompute out_z, not used here + ) + dz = rest[0] if ctx.has_z else None + dB = dB.squeeze(1) if getattr(ctx, "squeeze_B", False) else dB + dC = dC.squeeze(1) if getattr(ctx, "squeeze_C", False) else dC + return ( + du, + ddelta, + dA, + dB, + dC, + dD if D is not None else None, + dz, + ddelta_bias if delta_bias is not None else None, + None, + None, + ) + + +def rms_norm_forward( + x, + weight, + bias, + eps=1e-6, + is_rms_norm=True, +): + # x (b l) d + if x.stride(-1) != 1: + x = x.contiguous() + weight = weight.contiguous() + if bias is not None: + bias = bias.contiguous() + y = _layer_norm_fwd(x, weight, bias, eps, None, residual_dtype=None, is_rms_norm=is_rms_norm)[0] + # y (b l) d + return y + + +def selective_scan_fn( + u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False, return_last_state=False +): + """if return_last_state is True, returns (out, last_state) + last_state has shape (batch, dim, dstate). Note that the gradient of the last state is + not considered in the backward pass. + """ + return SelectiveScanFn.apply(u, delta, A, B, C, D, z, delta_bias, delta_softplus, return_last_state) + + +def selective_scan_ref( + u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False, return_last_state=False +): + """ + u: r(B D L) + delta: r(B D L) + A: c(D N) or r(D N) + B: c(D N) or r(B N L) or r(B N 2L) or r(B G N L) or (B G N L) + C: c(D N) or r(B N L) or r(B N 2L) or r(B G N L) or (B G N L) + D: r(D) + z: r(B D L) + delta_bias: r(D), fp32 + + out: r(B D L) + last_state (optional): r(B D dstate) or c(B D dstate) + """ + dtype_in = u.dtype + u = u.float() + delta = delta.float() + if delta_bias is not None: + delta = delta + delta_bias[..., None].float() + if delta_softplus: + delta = F.softplus(delta) + batch, dim, dstate = u.shape[0], A.shape[0], A.shape[1] + is_variable_B = B.dim() >= 3 + is_variable_C = C.dim() >= 3 + if A.is_complex(): + if is_variable_B: + B = torch.view_as_complex(rearrange(B.float(), "... (L two) -> ... L two", two=2)) + if is_variable_C: + C = torch.view_as_complex(rearrange(C.float(), "... (L two) -> ... L two", two=2)) + else: + B = B.float() + C = C.float() + x = A.new_zeros((batch, dim, dstate)) + ys = [] + deltaA = torch.exp(torch.einsum("bdl,dn->bdln", delta, A)) + if not is_variable_B: + deltaB_u = torch.einsum("bdl,dn,bdl->bdln", delta, B, u) + else: + if B.dim() == 3: + deltaB_u = torch.einsum("bdl,bnl,bdl->bdln", delta, B, u) + else: + B = repeat(B, "B G N L -> B (G H) N L", H=dim // B.shape[1]) + deltaB_u = torch.einsum("bdl,bdnl,bdl->bdln", delta, B, u) + if is_variable_C and C.dim() == 4: + C = repeat(C, "B G N L -> B (G H) N L", H=dim // C.shape[1]) + last_state = None + for i in range(u.shape[2]): + x = deltaA[:, :, i] * x + deltaB_u[:, :, i] + if not is_variable_C: + y = torch.einsum("bdn,dn->bd", x, C) + else: + if C.dim() == 3: + y = torch.einsum("bdn,bn->bd", x, C[:, :, i]) + else: + y = torch.einsum("bdn,bdn->bd", x, C[:, :, :, i]) + if i == u.shape[2] - 1: + last_state = x + if y.is_complex(): + y = y.real * 2 + ys.append(y) + y = torch.stack(ys, dim=2) # (batch dim L) + out = y if D is None else y + u * rearrange(D, "d -> d 1") + if z is not None: + out = out * F.silu(z) + out = out.to(dtype=dtype_in) + return out if not return_last_state else (out, last_state) + + +class MambaInnerFn(torch.autograd.Function): + @staticmethod + @custom_fwd + def forward( + ctx, + xz, + conv1d_weight, + conv1d_bias, + x_proj_weight, + delta_proj_weight, + out_proj_weight, + out_proj_bias, + A, + B=None, + C=None, + D=None, + delta_bias=None, + B_proj_bias=None, + C_proj_bias=None, + delta_softplus=True, + checkpoint_lvl=1, + b_rms_weight=None, + c_rms_weight=None, + dt_rms_weight=None, + b_c_dt_rms_eps=1e-6, + ): + """ + xz: (batch, dim, seqlen) + """ + assert causal_conv1d_cuda is not None, "causal_conv1d_cuda is not available. Please install causal-conv1d." + assert checkpoint_lvl in [0, 1] + L = xz.shape[-1] + delta_rank = delta_proj_weight.shape[1] + d_state = A.shape[-1] * (1 if not A.is_complex() else 2) + if torch.is_autocast_enabled(): + x_proj_weight = x_proj_weight.to(dtype=torch.get_autocast_gpu_dtype()) + delta_proj_weight = delta_proj_weight.to(dtype=torch.get_autocast_gpu_dtype()) + out_proj_weight = out_proj_weight.to(dtype=torch.get_autocast_gpu_dtype()) + out_proj_bias = ( + out_proj_bias.to(dtype=torch.get_autocast_gpu_dtype()) if out_proj_bias is not None else None + ) + if xz.stride(-1) != 1: + xz = xz.contiguous() + conv1d_weight = rearrange(conv1d_weight, "d 1 w -> d w") + x, z = xz.chunk(2, dim=1) + conv1d_bias = conv1d_bias.contiguous() if conv1d_bias is not None else None + conv1d_out = causal_conv1d_cuda.causal_conv1d_fwd(x, conv1d_weight, conv1d_bias, None, None, None, True) + # We're being very careful here about the layout, to avoid extra transposes. + # We want delta to have d as the slowest moving dimension + # and L as the fastest moving dimension, since those are what the ssm_scan kernel expects. + x_dbl = F.linear(rearrange(conv1d_out, "b d l -> (b l) d"), x_proj_weight) # (bl d) + delta = rearrange(delta_proj_weight @ x_dbl[:, :delta_rank].t(), "d (b l) -> b d l", l=L) + ctx.is_variable_B = B is None + ctx.is_variable_C = C is None + ctx.B_proj_bias_is_None = B_proj_bias is None + ctx.C_proj_bias_is_None = C_proj_bias is None + if B is None: # variable B + B = x_dbl[:, delta_rank : delta_rank + d_state] # (bl dstate) + if B_proj_bias is not None: + B = B + B_proj_bias.to(dtype=B.dtype) + if not A.is_complex(): + # B = rearrange(B, "(b l) dstate -> b dstate l", l=L).contiguous() + B = rearrange(B, "(b l) dstate -> b 1 dstate l", l=L).contiguous() + else: + B = rearrange(B, "(b l) (dstate two) -> b 1 dstate (l two)", l=L, two=2).contiguous() + else: + if B.stride(-1) != 1: + B = B.contiguous() + if C is None: # variable C + C = x_dbl[:, -d_state:] # (bl dstate) + if C_proj_bias is not None: + C = C + C_proj_bias.to(dtype=C.dtype) + if not A.is_complex(): + # C = rearrange(C, "(b l) dstate -> b dstate l", l=L).contiguous() + C = rearrange(C, "(b l) dstate -> b 1 dstate l", l=L).contiguous() + else: + C = rearrange(C, "(b l) (dstate two) -> b 1 dstate (l two)", l=L, two=2).contiguous() + else: + if C.stride(-1) != 1: + C = C.contiguous() + if D is not None: + D = D.contiguous() + + if b_rms_weight is not None: + B = rearrange(B, "b 1 dstate l -> (b l) dstate", l=L).contiguous() + B = rms_norm_forward(B, b_rms_weight, bias=None, eps=b_c_dt_rms_eps) + B = rearrange(B, "(b l) dstate -> b 1 dstate l", l=L).contiguous() + if c_rms_weight is not None: + C = rearrange(C, "b 1 dstate l -> (b l) dstate", l=L).contiguous() + C = rms_norm_forward(C, c_rms_weight, bias=None, eps=b_c_dt_rms_eps) + C = rearrange(C, "(b l) dstate -> b 1 dstate l", l=L).contiguous() + if dt_rms_weight is not None: + delta = rearrange(delta, "b d l -> (b l) d", l=L).contiguous() + delta = rms_norm_forward(delta, dt_rms_weight, bias=None, eps=b_c_dt_rms_eps) + delta = rearrange(delta, "(b l) d -> b d l", l=L).contiguous() + + out, scan_intermediates, out_z = selective_scan_cuda.fwd( + conv1d_out, delta, A, B, C, D, z, delta_bias, delta_softplus + ) + ctx.delta_softplus = delta_softplus + ctx.out_proj_bias_is_None = out_proj_bias is None + ctx.checkpoint_lvl = checkpoint_lvl + ctx.b_rms_weight = b_rms_weight + ctx.c_rms_weight = c_rms_weight + ctx.dt_rms_weight = dt_rms_weight + ctx.b_c_dt_rms_eps = b_c_dt_rms_eps + if checkpoint_lvl >= 1: # Will recompute conv1d_out and delta in the backward pass + conv1d_out, delta = None, None + ctx.save_for_backward( + xz, + conv1d_weight, + conv1d_bias, + x_dbl, + x_proj_weight, + delta_proj_weight, + out_proj_weight, + conv1d_out, + delta, + A, + B, + C, + D, + delta_bias, + scan_intermediates, + b_rms_weight, + c_rms_weight, + dt_rms_weight, + out, + ) + return F.linear(rearrange(out_z, "b d l -> b l d"), out_proj_weight, out_proj_bias) + + @staticmethod + @custom_bwd + def backward(ctx, dout): + # dout: (batch, seqlen, dim) + assert causal_conv1d_cuda is not None, "causal_conv1d_cuda is not available. Please install causal-conv1d." + ( + xz, + conv1d_weight, + conv1d_bias, + x_dbl, + x_proj_weight, + delta_proj_weight, + out_proj_weight, + conv1d_out, + delta, + A, + B, + C, + D, + delta_bias, + scan_intermediates, + b_rms_weight, + c_rms_weight, + dt_rms_weight, + out, + ) = ctx.saved_tensors + L = xz.shape[-1] + delta_rank = delta_proj_weight.shape[1] + d_state = A.shape[-1] * (1 if not A.is_complex() else 2) + x, z = xz.chunk(2, dim=1) + if dout.stride(-1) != 1: + dout = dout.contiguous() + if ctx.checkpoint_lvl == 1: + conv1d_out = causal_conv1d_cuda.causal_conv1d_fwd(x, conv1d_weight, conv1d_bias, None, None, None, True) + delta = rearrange(delta_proj_weight @ x_dbl[:, :delta_rank].t(), "d (b l) -> b d l", l=L) + if dt_rms_weight is not None: + delta = rearrange(delta, "b d l -> (b l) d", l=L).contiguous() + delta = rms_norm_forward(delta, ctx.dt_rms_weight, None, ctx.b_c_dt_rms_eps) + delta = rearrange(delta, "(b l) d -> b d l", l=L).contiguous() + if b_rms_weight is not None: + # Recompute & RMSNorm B + B = rearrange(B, "b 1 dstate l -> (b l) dstate", l=L).contiguous() + B = rms_norm_forward(B, ctx.b_rms_weight, None, ctx.b_c_dt_rms_eps) + B = rearrange(B, "(b l) dstate -> b 1 dstate l", l=L).contiguous() + if c_rms_weight is not None: + # Recompute & RMSNorm C + C = rearrange(C, "b 1 dstate l -> (b l) dstate", l=L).contiguous() + C = rms_norm_forward(C, ctx.c_rms_weight, None, ctx.b_c_dt_rms_eps) + C = rearrange(C, "(b l) dstate -> b 1 dstate l", l=L).contiguous() + + # The kernel supports passing in a pre-allocated dz (e.g., in case we want to fuse the + # backward of selective_scan_cuda with the backward of chunk). + dxz = torch.empty_like(xz) # (batch, dim, seqlen) + dx, dz = dxz.chunk(2, dim=1) + dout = rearrange(dout, "b l e -> e (b l)") + dout_y = rearrange(out_proj_weight.t() @ dout, "d (b l) -> b d l", l=L) + dconv1d_out, ddelta, dA, dB, dC, dD, ddelta_bias, dz, out_z = selective_scan_cuda.bwd( + conv1d_out, + delta, + A, + B, + C, + D, + z, + delta_bias, + dout_y, + scan_intermediates, + out, + dz, + ctx.delta_softplus, + True, # option to recompute out_z + ) + dout_proj_weight = torch.einsum("eB,dB->ed", dout, rearrange(out_z, "b d l -> d (b l)")) + dout_proj_bias = dout.sum(dim=(0, 1)) if not ctx.out_proj_bias_is_None else None + dD = dD if D is not None else None + dx_dbl = torch.empty_like(x_dbl) + dB_proj_bias = None + if ctx.is_variable_B: + if not A.is_complex(): + dB = rearrange(dB, "b 1 dstate l -> (b l) dstate").contiguous() + else: + dB = rearrange(dB, "b 1 dstate (l two) -> (b l) (dstate two)", two=2).contiguous() + dB_proj_bias = dB.sum(0) if not ctx.B_proj_bias_is_None else None + dx_dbl[:, delta_rank : delta_rank + d_state] = dB # (bl d) + dB = None + dC_proj_bias = None + if ctx.is_variable_C: + if not A.is_complex(): + dC = rearrange(dC, "b 1 dstate l -> (b l) dstate").contiguous() + else: + dC = rearrange(dC, "b 1 dstate (l two) -> (b l) (dstate two)", two=2).contiguous() + dC_proj_bias = dC.sum(0) if not ctx.C_proj_bias_is_None else None + dx_dbl[:, -d_state:] = dC # (bl d) + dC = None + ddelta = rearrange(ddelta, "b d l -> d (b l)") + ddelta_proj_weight = torch.einsum("dB,Br->dr", ddelta, x_dbl[:, :delta_rank]) + dx_dbl[:, :delta_rank] = torch.einsum("dB,dr->Br", ddelta, delta_proj_weight) + dconv1d_out = rearrange(dconv1d_out, "b d l -> d (b l)") + dx_proj_weight = torch.einsum("Br,Bd->rd", dx_dbl, rearrange(conv1d_out, "b d l -> (b l) d")) + dconv1d_out = torch.addmm(dconv1d_out, x_proj_weight.t(), dx_dbl.t(), out=dconv1d_out) + dconv1d_out = rearrange(dconv1d_out, "d (b l) -> b d l", b=x.shape[0], l=x.shape[-1]) + # The kernel supports passing in a pre-allocated dx (e.g., in case we want to fuse the + # backward of conv1d with the backward of chunk). + dx, dconv1d_weight, dconv1d_bias, *_ = causal_conv1d_cuda.causal_conv1d_bwd( + x, conv1d_weight, conv1d_bias, dconv1d_out, None, None, None, dx, False, True + ) + dconv1d_bias = dconv1d_bias if conv1d_bias is not None else None + dconv1d_weight = rearrange(dconv1d_weight, "d w -> d 1 w") + return ( + dxz, + dconv1d_weight, + dconv1d_bias, + dx_proj_weight, + ddelta_proj_weight, + dout_proj_weight, + dout_proj_bias, + dA, + dB, + dC, + dD, + ddelta_bias if delta_bias is not None else None, + # 6-None are delta_softplus, checkpoint_lvl, b_rms_weight, c_rms_weight, dt_rms_weight, b_c_dt_rms_eps + dB_proj_bias, + dC_proj_bias, + None, + None, + None, + None, + None, + None, + ) + + +def mamba_inner_fn( + xz, + conv1d_weight, + conv1d_bias, + x_proj_weight, + delta_proj_weight, + out_proj_weight, + out_proj_bias, + A, + B=None, + C=None, + D=None, + delta_bias=None, + B_proj_bias=None, + C_proj_bias=None, + delta_softplus=True, + checkpoint_lvl=1, + b_rms_weight=None, + c_rms_weight=None, + dt_rms_weight=None, + b_c_dt_rms_eps=1e-6, +): + return MambaInnerFn.apply( + xz, + conv1d_weight, + conv1d_bias, + x_proj_weight, + delta_proj_weight, + out_proj_weight, + out_proj_bias, + A, + B, + C, + D, + delta_bias, + B_proj_bias, + C_proj_bias, + delta_softplus, + checkpoint_lvl, + b_rms_weight, + c_rms_weight, + dt_rms_weight, + b_c_dt_rms_eps, + ) diff --git a/lib/python3.12/site-packages/transformers/kernels/mra/cuda_kernel.cu b/lib/python3.12/site-packages/transformers/kernels/mra/cuda_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..87ed89052873813153786bd416a981d3e5279af9 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/mra/cuda_kernel.cu @@ -0,0 +1,383 @@ +#include "cuda_kernel.h" + +////////////////////////////////////////////////////////////////////////////////////////////////// +////////////////////////////////////////////////////////////////////////////////////////////////// + +__global__ void index_max_cuda_kernel( + float *index_vals, // [batch_size, 32, num_block] + int *indices, // [batch_size, num_block] + float *max_vals, // [batch_size, A_num_block * 32] + float *max_vals_scatter, // [batch_size, 32, num_block] + long batch_size, + long A_num_block, + long B_num_block, + long num_block +) { + + long batch_idx = blockIdx.x; + + long thread_idx = threadIdx.x; + long num_thread = blockDim.x; + + extern __shared__ float buffer[]; + int *max_buffer = (int*)buffer; + + for (int i = 0; i < A_num_block * 32; i = i + num_thread) { + int idx = i + thread_idx; + if (idx < A_num_block * 32) { + max_buffer[idx] = -1e8; + } + } + __syncthreads(); + + int *indices_pt = &indices[batch_idx * num_block]; + float *index_vals_pt = &index_vals[batch_idx * num_block * 32]; + + for (int idx_start = 0; idx_start < 32 * num_block; idx_start = idx_start + num_thread) { + int idx = idx_start + thread_idx; + int A_block_idx = indices_pt[idx % num_block] / B_num_block; + atomicMax(&max_buffer[A_block_idx * 32 + idx / num_block], (int)(index_vals_pt[idx] * 1000)); + } + __syncthreads(); + + float *max_vals_pt = &max_vals[batch_idx * A_num_block * 32]; + for (int i = 0; i < A_num_block * 32; i = i + num_thread) { + int idx = i + thread_idx; + if (idx < A_num_block * 32) { + max_vals_pt[idx] = (float)max_buffer[idx] / 1000.; + } + } + + float *max_vals_scatter_pt = &max_vals_scatter[batch_idx * num_block * 32]; + for (int idx_start = 0; idx_start < 32 * num_block; idx_start = idx_start + num_thread) { + int idx = idx_start + thread_idx; + int A_block_idx = indices_pt[idx % num_block] / B_num_block; + max_vals_scatter_pt[idx] = (float)max_buffer[A_block_idx * 32 + idx / num_block] / 1000.; + } + +} + +__global__ void mm_to_sparse_cuda_kernel( + float *dense_A, // [batch_size, A_num_block, dim, 32] + float *dense_B, // [batch_size, B_num_block, dim, 32] + int *indices, // [batch_size, num_block] + float *sparse_C, // [batch_size, num_block, 32, 32] + long batch_size, + long A_num_block, + long B_num_block, + long dim, + long num_block +) { + + long batch_idx = blockIdx.y; + long block_idx = blockIdx.x * blockDim.y + threadIdx.y; + + long thread_idx = threadIdx.x; + + __shared__ float buffer[4096]; + float *A_buffer = &buffer[threadIdx.y * 1024]; // [2, 8, 32] + float *B_buffer = &buffer[threadIdx.y * 1024 + 512]; // [2, 8, 32] + + long batch_idx__block_idx = batch_idx * num_block + block_idx; + + long AB_block_idx = indices[batch_idx__block_idx]; + float *dense_A_pt = &dense_A[(batch_idx * A_num_block + AB_block_idx / B_num_block) * dim * 32]; + float *dense_B_pt = &dense_B[(batch_idx * B_num_block + AB_block_idx % B_num_block) * dim * 32]; + + int reg_1_idx = thread_idx / 8; // [0000000011111111222222223333333344444444555555556666666677777777] + int reg_2_idx = thread_idx % 8; // [0123456701234567012345670123456701234567012345670123456701234567] + + float reg_1[8]; + float reg_2[8]; + + float reg_array[16] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; + + #pragma unroll + for (int i = 0; i < 4; i++) { + A_buffer[i * 64 + thread_idx] = dense_A_pt[i * 64 + thread_idx]; + B_buffer[i * 64 + thread_idx] = dense_B_pt[i * 64 + thread_idx]; + } + + __syncthreads(); + + #pragma unroll + for (int i = 0; i < 4; i++) { + reg_1[i] = A_buffer[reg_1_idx * 4 + i]; + reg_2[i] = B_buffer[reg_2_idx * 4 + i]; + } + + for (int dim_stride = 1; dim_stride < (dim / 8); dim_stride++) { + + #pragma unroll + for (int i = 0; i < 4; i++) { + A_buffer[(dim_stride % 2) * 256 + i * 64 + thread_idx] = dense_A_pt[dim_stride * 256 + i * 64 + thread_idx]; + B_buffer[(dim_stride % 2) * 256 + i * 64 + thread_idx] = dense_B_pt[dim_stride * 256 + i * 64 + thread_idx]; + } + + #pragma unroll + for (int mini_dim_idx = 1; mini_dim_idx < 8; mini_dim_idx++) { + #pragma unroll + for (int i = 0; i < 4; i++) { + reg_1[(mini_dim_idx % 2) * 4 + i] = A_buffer[((dim_stride - 1) % 2) * 256 + mini_dim_idx * 32 + reg_1_idx * 4 + i]; + reg_2[(mini_dim_idx % 2) * 4 + i] = B_buffer[((dim_stride - 1) % 2) * 256 + mini_dim_idx * 32 + reg_2_idx * 4 + i]; + } + #pragma unroll + for (int i = 0; i < 4; i++) { + #pragma unroll + for (int j = 0; j < 4; j++) { + reg_array[i * 4 + j] += reg_1[((mini_dim_idx - 1) % 2) * 4 + i] * reg_2[((mini_dim_idx - 1) % 2) * 4 + j]; + } + } + } + + __syncthreads(); + + #pragma unroll + for (int i = 0; i < 4; i++) { + reg_1[i] = A_buffer[(dim_stride % 2) * 256 + reg_1_idx * 4 + i]; + reg_2[i] = B_buffer[(dim_stride % 2) * 256 + reg_2_idx * 4 + i]; + } + + #pragma unroll + for (int i = 0; i < 4; i++) { + #pragma unroll + for (int j = 0; j < 4; j++) { + reg_array[i * 4 + j] += reg_1[4 + i] * reg_2[4 + j]; + } + } + + } + + #pragma unroll + for (int mini_dim_idx = 1; mini_dim_idx < 8; mini_dim_idx++) { + #pragma unroll + for (int i = 0; i < 4; i++) { + reg_1[(mini_dim_idx % 2) * 4 + i] = A_buffer[256 + mini_dim_idx * 32 + reg_1_idx * 4 + i]; + reg_2[(mini_dim_idx % 2) * 4 + i] = B_buffer[256 + mini_dim_idx * 32 + reg_2_idx * 4 + i]; + } + #pragma unroll + for (int i = 0; i < 4; i++) { + #pragma unroll + for (int j = 0; j < 4; j++) { + reg_array[i * 4 + j] += reg_1[((mini_dim_idx - 1) % 2) * 4 + i] * reg_2[((mini_dim_idx - 1) % 2) * 4 + j]; + } + } + } + #pragma unroll + for (int i = 0; i < 4; i++) { + #pragma unroll + for (int j = 0; j < 4; j++) { + reg_array[i * 4 + j] += reg_1[4 + i] * reg_2[4 + j]; + } + } + __syncthreads(); + + float *C_buffer = &buffer[threadIdx.y * 1024]; // [32, 32] + + #pragma unroll + for (int i = 0; i < 4; i++) { + #pragma unroll + for (int j = 0; j < 4; j++) { + C_buffer[(reg_2_idx * 4 + j) * 32 + reg_1_idx * 4 + i] = reg_array[i * 4 + j]; + } + } + __syncthreads(); + + float *sparse_C_pt = &sparse_C[batch_idx__block_idx * 1024]; + + #pragma unroll + for (int i = 0; i < 16; i++) { + sparse_C_pt[i * 64 + thread_idx] = C_buffer[i * 64 + thread_idx]; + } + +} + +__global__ void sparse_dense_mm_cuda_kernel( + float *sparse_A, // [batch_size, num_block, 32, 32] + int *indices, // [batch_size, num_block] + float *dense_B, // [batch_size, B_num_block, dim, 32] + float *dense_C, // [batch_size, A_num_block, dim, 32] + long batch_size, + long A_num_block, + long B_num_block, + long dim, + long num_block +) { + + long batch_idx = blockIdx.y; + long block_idx = blockIdx.x * blockDim.y + threadIdx.y; + + long thread_idx = threadIdx.x; + + __shared__ float buffer[6144]; + float *A_buffer = &buffer[threadIdx.y * 3072]; // [32, 32] + float *B_buffer = &buffer[threadIdx.y * 3072 + 1024]; // [32, 64] + + long batch_idx__block_idx = batch_idx * num_block + block_idx; + + float *sparse_A_pt = &sparse_A[batch_idx__block_idx * 1024]; + #pragma unroll + for (int i = 0; i < 8; i++) { + A_buffer[i * 128 + thread_idx] = sparse_A_pt[i * 128 + thread_idx]; + } + + long AB_block_idx = indices[batch_idx__block_idx]; + float *dense_B_pt = &dense_B[(batch_idx * B_num_block + AB_block_idx % B_num_block) * 32 * dim]; + float *dense_C_pt = &dense_C[(batch_idx * A_num_block + AB_block_idx / B_num_block) * 32 * dim]; + + // [0000000011111111222222223333333344444444555555556666666677777777] + // [0123456701234567012345670123456701234567012345670123456701234567] + int reg_1_idx = thread_idx / 8; + int reg_2_idx = thread_idx % 8; + + float reg_1[8]; + float reg_2[8]; + + float reg_array[16]; + + for (int dim_stride = 0; dim_stride < dim; dim_stride = dim_stride + 64) { + + #pragma unroll + for (int i = 0; i < 16; i++) { + B_buffer[i * 128 + thread_idx] = dense_B_pt[dim_stride * 32 + i * 128 + thread_idx]; + } + + #pragma unroll + for (int i = 0; i < 16; i++) { + reg_array[i] = 0; + } + + __syncthreads(); + + #pragma unroll + for (int i = 0; i < 4; i++) { + reg_1[i] = B_buffer[(reg_1_idx * 4 + i) * 32]; + reg_2[i] = A_buffer[reg_2_idx * 4 + i]; + } + + #pragma unroll + for (int mini_dim_idx = 1; mini_dim_idx < 32; mini_dim_idx++) { + #pragma unroll + for (int i = 0; i < 4; i++) { + reg_1[(mini_dim_idx % 2) * 4 + i] = B_buffer[(reg_1_idx * 4 + i) * 32 + mini_dim_idx]; + reg_2[(mini_dim_idx % 2) * 4 + i] = A_buffer[mini_dim_idx * 32 + reg_2_idx * 4 + i]; + } + #pragma unroll + for (int i = 0; i < 4; i++) { + #pragma unroll + for (int j = 0; j < 4; j++) { + reg_array[i * 4 + j] += reg_1[((mini_dim_idx - 1) % 2) * 4 + i] * reg_2[((mini_dim_idx - 1) % 2) * 4 + j]; + } + } + } + + #pragma unroll + for (int i = 0; i < 4; i++) { + #pragma unroll + for (int j = 0; j < 4; j++) { + reg_array[i * 4 + j] += reg_1[4 + i] * reg_2[4 + j]; + } + } + + __syncthreads(); + + float *C_buffer = &buffer[threadIdx.y * 3072 + 1024]; // [64, 32] + + #pragma unroll + for (int i = 0; i < 4; i++) { + #pragma unroll + for (int j = 0; j < 4; j++) { + C_buffer[(reg_1_idx * 4 + i) * 32 + reg_2_idx * 4 + j] = reg_array[i * 4 + j]; + } + } + __syncthreads(); + + #pragma unroll + for (int i = 0; i < 16; i++) { + atomicAdd(&dense_C_pt[dim_stride * 32 + i * 128 + thread_idx], C_buffer[i * 128 + thread_idx]); + } + __syncthreads(); + + } + +} + + +__global__ void reduce_sum_cuda_kernel( + float *sparse_A, // [batch_size, num_block, 32, 32] + int *indices, // [batch_size, num_block] + float *dense_C, // [batch_size, A_num_block, 32] + long batch_size, + long A_num_block, + long B_num_block, + long num_block +) { + + long batch_idx = blockIdx.y; + long block_idx = blockIdx.x * blockDim.y + threadIdx.y; + + long thread_idx = threadIdx.x; + + long batch_idx__block_idx = batch_idx * num_block + block_idx; + + long AB_block_idx = indices[batch_idx__block_idx]; + float *sparse_A_pt = &sparse_A[batch_idx__block_idx * 1024]; + + float reg_array[16]; + float value = 0; + + #pragma unroll + for (int i = 0; i < 8; i++) { + reg_array[i] = sparse_A_pt[i * 32 + thread_idx]; + } + #pragma unroll + for (int stride = 8; stride < 32; stride = stride + 8) { + #pragma unroll + for (int i = 0; i < 8; i++) { + reg_array[(stride + i) % 16] = sparse_A_pt[(stride + i) * 32 + thread_idx]; + } + #pragma unroll + for (int i = 0; i < 8; i++) { + value = value + reg_array[(stride - 8 + i) % 16]; + } + } + #pragma unroll + for (int i = 0; i < 8; i++) { + value = value + reg_array[8 + i]; + } + + float *dense_C_pt = &dense_C[(batch_idx * A_num_block + AB_block_idx / B_num_block) * 32]; + + atomicAdd(&dense_C_pt[thread_idx], value); + +} + +__global__ void scatter_cuda_kernel( + float *dense_A, // [batch_size, A_num_block, 32] + int *indices, // [batch_size, num_block] + float *sparse_C, // [batch_size, num_block, 32, 32] + long batch_size, + long A_num_block, + long B_num_block, + long num_block +) { + + long batch_idx = blockIdx.y; + long block_idx = blockIdx.x * blockDim.y + threadIdx.y; + + long thread_idx = threadIdx.x; + + long batch_idx__block_idx = batch_idx * num_block + block_idx; + + long AB_block_idx = indices[batch_idx__block_idx]; + float *dense_A_pt = &dense_A[(batch_idx * A_num_block + AB_block_idx / B_num_block) * 32]; + float *sparse_C_pt = &sparse_C[(batch_idx * num_block + block_idx) * 1024]; + + float value = dense_A_pt[thread_idx]; + + #pragma unroll + for (int i = 0; i < 32; i++) { + sparse_C_pt[i * 32 + thread_idx] = value; + } + +} diff --git a/lib/python3.12/site-packages/transformers/kernels/mra/cuda_kernel.h b/lib/python3.12/site-packages/transformers/kernels/mra/cuda_kernel.h new file mode 100644 index 0000000000000000000000000000000000000000..a95b46f7d159b11851143710034cf80c20aa6bf8 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/mra/cuda_kernel.h @@ -0,0 +1,59 @@ + +#define WARP_SIZE 32 +#define FULL_MASK 0xffffffff +#define OPTIMAL_THREADS 256 + +__global__ void index_max_cuda_kernel( + float *index_vals, // [batch_size, 32, num_block] + int *indices, // [batch_size, num_block] + float *max_vals, // [batch_size, A_num_block * 32] + float *max_vals_scatter, // [batch_size, 32, num_block] + long batch_size, + long A_num_block, + long B_num_block, + long num_block +); + +__global__ void mm_to_sparse_cuda_kernel( + float *dense_A, // [batch_size, A_num_block, dim, 32] + float *dense_B, // [batch_size, B_num_block, dim, 32] + int *indices, // [batch_size, num_block] + float *sparse_C, // [batch_size, num_block, 32, 32] + long batch_size, + long A_num_block, + long B_num_block, + long dim, + long num_block +); + +__global__ void sparse_dense_mm_cuda_kernel( + float *sparse_A, // [batch_size, num_block, 32, 32] + int *indices, // [batch_size, num_block] + float *dense_B, // [batch_size, B_num_block, dim, 32] + float *dense_C, // [batch_size, A_num_block, dim, 32] + long batch_size, + long A_num_block, + long B_num_block, + long dim, + long num_block +); + +__global__ void reduce_sum_cuda_kernel( + float *sparse_A, // [batch_size, num_block, 32, 32] + int *indices, // [batch_size, num_block] + float *dense_C, // [batch_size, A_num_block, 32] + long batch_size, + long A_num_block, + long B_num_block, + long num_block +); + +__global__ void scatter_cuda_kernel( + float *dense_A, // [batch_size, A_num_block, 32] + int *indices, // [batch_size, num_block] + float *sparse_C, // [batch_size, num_block, 32, 32] + long batch_size, + long A_num_block, + long B_num_block, + long num_block +); diff --git a/lib/python3.12/site-packages/transformers/kernels/mra/cuda_launch.cu b/lib/python3.12/site-packages/transformers/kernels/mra/cuda_launch.cu new file mode 100644 index 0000000000000000000000000000000000000000..ba2a0cacfe614e75e06d2dde80dc77a6e8a4ec1a --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/mra/cuda_launch.cu @@ -0,0 +1,154 @@ +#include +#include +#include "cuda_launch.h" +#include "cuda_kernel.h" +#include + +////////////////////////////////////////////////////////////////////////////////////////////////// +////////////////////////////////////////////////////////////////////////////////////////////////// + +std::vector index_max_kernel( + at::Tensor index_vals, // [batch_size, 32, num_block] + at::Tensor indices, // [batch_size, num_block], + int A_num_block, + int B_num_block +) { + int batch_size = indices.size(0); + int num_block = indices.size(1); + + at::Tensor max_vals = at::zeros({batch_size, A_num_block * 32}, index_vals.options()); + at::Tensor max_vals_scatter = at::zeros({batch_size, 32, num_block}, index_vals.options()); + + dim3 threads(256); + dim3 blocks(batch_size); + int shared_mem = A_num_block * 32 * sizeof(float); + + index_max_cuda_kernel<<>>( + index_vals.data_ptr(), + indices.data_ptr(), + max_vals.data_ptr(), + max_vals_scatter.data_ptr(), + batch_size, + A_num_block, + B_num_block, + num_block + ); + + return {max_vals, max_vals_scatter}; +} + +at::Tensor mm_to_sparse_kernel( + at::Tensor dense_A, // [batch_size, A_num_block, dim, 32] + at::Tensor dense_B, // [batch_size, B_num_block, dim, 32] + at::Tensor indices // [batch_size, num_block] +) { + int batch_size = dense_A.size(0); + int A_num_block = dense_A.size(1); + int B_num_block = dense_B.size(1); + int dim = dense_A.size(2); + int num_block = indices.size(1); + + at::Tensor sparse_C = at::zeros({batch_size, num_block, 32, 32}, dense_A.options()); + + dim3 threads(64, 4); + dim3 blocks(num_block / 4, batch_size); + + mm_to_sparse_cuda_kernel<<>>( + dense_A.data_ptr(), + dense_B.data_ptr(), + indices.data_ptr(), + sparse_C.data_ptr(), + batch_size, + A_num_block, + B_num_block, + dim, + num_block + ); + + return sparse_C; +} + +at::Tensor sparse_dense_mm_kernel( + at::Tensor sparse_A, // [batch_size, num_block, 32, 32] + at::Tensor indices, // [batch_size, num_block] + at::Tensor dense_B, // [batch_size, B_num_block, dim, 32] + int A_num_block +) { + int batch_size = sparse_A.size(0); + int num_block = sparse_A.size(1); + int B_num_block = dense_B.size(1); + int dim = dense_B.size(2); + + at::Tensor dense_C = at::zeros({batch_size, A_num_block, dim, 32}, dense_B.options()); + + dim3 threads(128, 2); + dim3 blocks(num_block / 2, batch_size); + + sparse_dense_mm_cuda_kernel<<>>( + sparse_A.data_ptr(), + indices.data_ptr(), + dense_B.data_ptr(), + dense_C.data_ptr(), + batch_size, + A_num_block, + B_num_block, + dim, + num_block + ); + + return dense_C; +} + +at::Tensor reduce_sum_kernel( + at::Tensor sparse_A, // [batch_size, num_block, 32, 32] + at::Tensor indices, // [batch_size, num_block] + int A_num_block, + int B_num_block +) { + int batch_size = sparse_A.size(0); + int num_block = sparse_A.size(1); + + at::Tensor dense_C = at::zeros({batch_size, A_num_block, 32}, sparse_A.options()); + + dim3 threads(32, 4); + dim3 blocks(num_block / 4, batch_size); + + reduce_sum_cuda_kernel<<>>( + sparse_A.data_ptr(), + indices.data_ptr(), + dense_C.data_ptr(), + batch_size, + A_num_block, + B_num_block, + num_block + ); + + return dense_C; +} + +at::Tensor scatter_kernel( + at::Tensor dense_A, // [batch_size, A_num_block, 32] + at::Tensor indices, // [batch_size, num_block] + int B_num_block +) { + int batch_size = dense_A.size(0); + int A_num_block = dense_A.size(1); + int num_block = indices.size(1); + + at::Tensor sparse_C = at::zeros({batch_size, num_block, 32, 32}, dense_A.options()); + + dim3 threads(32, 4); + dim3 blocks(num_block / 4, batch_size); + + scatter_cuda_kernel<<>>( + dense_A.data_ptr(), + indices.data_ptr(), + sparse_C.data_ptr(), + batch_size, + A_num_block, + B_num_block, + num_block + ); + + return sparse_C; +} diff --git a/lib/python3.12/site-packages/transformers/kernels/mra/cuda_launch.h b/lib/python3.12/site-packages/transformers/kernels/mra/cuda_launch.h new file mode 100644 index 0000000000000000000000000000000000000000..0200140ee337b8c5d9583767bbad1e842e9d4677 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/mra/cuda_launch.h @@ -0,0 +1,39 @@ +#include +#include +#include + +#define min(a, b) ((a)<(b)?(a):(b)) +#define max(a, b) ((a)>(b)?(a):(b)) + +std::vector index_max_kernel( + at::Tensor index_vals, + at::Tensor indices, + int A_num_block, + int B_num_block +); + +at::Tensor mm_to_sparse_kernel( + at::Tensor dense_A, + at::Tensor dense_B, + at::Tensor indices +); + +at::Tensor sparse_dense_mm_kernel( + at::Tensor sparse_A, + at::Tensor indices, + at::Tensor dense_B, + int A_num_block +); + +at::Tensor reduce_sum_kernel( + at::Tensor sparse_A, + at::Tensor indices, + int A_num_block, + int B_num_block +); + +at::Tensor scatter_kernel( + at::Tensor dense_A, + at::Tensor indices, + int B_num_block +); diff --git a/lib/python3.12/site-packages/transformers/kernels/mra/torch_extension.cpp b/lib/python3.12/site-packages/transformers/kernels/mra/torch_extension.cpp new file mode 100644 index 0000000000000000000000000000000000000000..60c9262b779270a6e95ae54f53a67daa6d740a9e --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/mra/torch_extension.cpp @@ -0,0 +1,78 @@ +#include +#include +#include "cuda_launch.h" +#include + +std::vector index_max( + at::Tensor index_vals, + at::Tensor indices, + int A_num_block, + int B_num_block +) { + return index_max_kernel( + index_vals, + indices, + A_num_block, + B_num_block + ); +} + +at::Tensor mm_to_sparse( + at::Tensor dense_A, + at::Tensor dense_B, + at::Tensor indices +) { + return mm_to_sparse_kernel( + dense_A, + dense_B, + indices + ); +} + +at::Tensor sparse_dense_mm( + at::Tensor sparse_A, + at::Tensor indices, + at::Tensor dense_B, + int A_num_block +) { + return sparse_dense_mm_kernel( + sparse_A, + indices, + dense_B, + A_num_block + ); +} + +at::Tensor reduce_sum( + at::Tensor sparse_A, + at::Tensor indices, + int A_num_block, + int B_num_block +) { + return reduce_sum_kernel( + sparse_A, + indices, + A_num_block, + B_num_block + ); +} + +at::Tensor scatter( + at::Tensor dense_A, + at::Tensor indices, + int B_num_block +) { + return scatter_kernel( + dense_A, + indices, + B_num_block + ); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("index_max", &index_max, "index_max (CUDA)"); + m.def("mm_to_sparse", &mm_to_sparse, "mm_to_sparse (CUDA)"); + m.def("sparse_dense_mm", &sparse_dense_mm, "sparse_dense_mm (CUDA)"); + m.def("reduce_sum", &reduce_sum, "reduce_sum (CUDA)"); + m.def("scatter", &scatter, "scatter (CUDA)"); +} diff --git a/lib/python3.12/site-packages/transformers/kernels/rwkv/wkv_cuda.cu b/lib/python3.12/site-packages/transformers/kernels/rwkv/wkv_cuda.cu new file mode 100644 index 0000000000000000000000000000000000000000..571d5a8a8307e95aac689eb3c9333d1ad350c7de --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/rwkv/wkv_cuda.cu @@ -0,0 +1,187 @@ +#include +#include + +#define MIN_VALUE (-1e38) + +template +__global__ void kernel_forward( + const int B, const int T, const int C, const F *__restrict__ const _w, const F *__restrict__ const _u, + const F *__restrict__ const _k, const F *__restrict__ const _v, F *__restrict__ const _y +) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const int _b = idx / C; + const int _c = idx % C; + const int _offset = _b * T * C + _c; + + F u = _u[_c]; + F w = _w[_c]; + const F *__restrict__ const k = _k + _offset; + const F *__restrict__ const v = _v + _offset; + F *__restrict__ const y = _y + _offset; + + // aa and bb are running sums divided by exp(pp) (to avoid overflow) + F aa = 0, bb = 0, pp = MIN_VALUE; + for (int i = 0; i < T; i++) { + const int ii = i * C; + const F kk = k[ii]; + const F vv = v[ii]; + + F ww = u + kk; + F p = max(pp, ww); + F e1 = exp(pp - p); + F e2 = exp(ww - p); + y[ii] = (e1 * aa + e2 * vv) / (e1 * bb + e2); + + ww = w + pp; + p = max(ww, kk); + e1 = exp(ww - p); + e2 = exp(kk - p); + aa = e1 * aa + e2 * vv; + bb = e1 * bb + e2; + pp = p; + } +} + +template +__global__ void kernel_forward_with_state( + const int B, const int T, const int C, const F *__restrict__ const _w, const F *__restrict__ const _u, + const F *__restrict__ const _k, const F *__restrict__ const _v, F *__restrict__ const _y, F *__restrict__ const _s +) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const int _b = idx / C; + const int _c = idx % C; + const int _offset_s = _b * C * 3 + _c * 3; + const int _offset = _b * T * C + _c; + + F u = _u[_c]; + F w = _w[_c]; + const F *__restrict__ const k = _k + _offset; + const F *__restrict__ const v = _v + _offset; + F *__restrict__ const y = _y + _offset; + F *__restrict__ const s = _s + _offset_s; + + // aa and bb are running sums divided by exp(pp) (to avoid overflow) + F aa = s[0], bb = s[1], pp = s[2]; + for (int i = 0; i < T; i++) { + const int ii = i * C; + const F kk = k[ii]; + const F vv = v[ii]; + + F ww = u + kk; + F p = max(pp, ww); + F e1 = exp(pp - p); + F e2 = exp(ww - p); + y[ii] = (e1 * aa + e2 * vv) / (e1 * bb + e2); + + ww = w + pp; + p = max(ww, kk); + e1 = exp(ww - p); + e2 = exp(kk - p); + aa = e1 * aa + e2 * vv; + bb = e1 * bb + e2; + pp = p; + } + s[0] = aa; + s[1] = bb; + s[2] = pp; +} + +template +__global__ void kernel_backward( + const int B, const int T, const int C, const F *__restrict__ const _w, const F *__restrict__ const _u, + const F *__restrict__ const _k, const F *__restrict__ const _v, const F *__restrict__ const _y, + const F *__restrict__ const _gy, F *__restrict__ const _gw, F *__restrict__ const _gu, F *__restrict__ const _gk, + F *__restrict__ const _gv +) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const int _b = idx / C; + const int _c = idx % C; + const int _offset = _b * T * C + _c; + + F u = _u[_c]; + F w = _w[_c]; + const F *__restrict__ const k = _k + _offset; + const F *__restrict__ const v = _v + _offset; + const F *__restrict__ const y = _y + _offset; + const F *__restrict__ const gy = _gy + _offset; + F *__restrict__ const gk = _gk + _offset; + F *__restrict__ const gv = _gv + _offset; + + F q[Tmax], r[Tmax]; + + F gw = 0, gu = 0, aa = 0, bb = 0, ga = 0, gb = 0, pp = MIN_VALUE; + for (int i = 0; i < T; i++) { + const int ii = i * C; + const F kk = k[ii]; + const F vv = v[ii]; + const F yy = y[ii]; + + F ww = u + kk; + F p = max(pp, ww); + F e1 = exp(pp - p); + F e2 = exp(ww - p); + const F qq = gy[ii] / (e1 * bb + e2); + gw += (ga - gb * yy) * e1 * qq; + gu += (vv - yy) * e2 * qq; + q[i] = qq; + r[i] = ww - p; + + ww = w + pp; + p = max(ww, kk); + e1 = exp(ww - p); + e2 = exp(kk - p); + ga = e1 * (aa + ga); + gb = e1 * (bb + gb); + aa = e1 * aa + e2 * vv; + bb = e1 * bb + e2; + pp = p; + } + const int _offsetBC = _b * C + _c; + _gw[_offsetBC] = gw * _w[_c]; // multiply by w because of w -> -exp(w) in python forward() + _gu[_offsetBC] = gu; + + aa = 0, bb = 0, pp = MIN_VALUE; + for (int i = T - 1; i >= 0; i--) { + const int ii = i * C; + const F kk = k[ii]; + const F vv = v[ii]; + const F yy = y[ii]; + const F qq = q[i]; + const F rr = r[i]; + + F e1 = qq * exp(rr); + F e2 = exp(kk + pp); + gk[ii] = e1 * (vv - yy) + e2 * (aa * vv + bb); + gv[ii] = e1 + e2 * aa; + + const F ww = w + pp; + const F www = rr - u - kk; + const F p = max(ww, www); + e1 = exp(ww - p); + e2 = qq * exp(www - p); + aa = e1 * aa + e2; + bb = e1 * bb - e2 * yy; + pp = p; + } +} + +void cuda_forward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y) { + dim3 threadsPerBlock( min(C, 32) ); // requires --maxrregcount 60 for optimal performance + assert(B * C % threadsPerBlock.x == 0); + dim3 numBlocks(B * C / threadsPerBlock.x); + kernel_forward<<>>(B, T, C, w, u, k, v, y); +} + +void cuda_forward_with_state(int B, int T, int C, float *w, float *u, float *k, float *v, float *y, float *s) { + dim3 threadsPerBlock( min(C, 32) ); // requires --maxrregcount 60 for optimal performance + assert(B * C % threadsPerBlock.x == 0); + dim3 numBlocks(B * C / threadsPerBlock.x); + kernel_forward_with_state<<>>(B, T, C, w, u, k, v, y, s); +} + +void cuda_backward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y, float *gy, float *gw, float *gu, float *gk, float *gv) { + dim3 threadsPerBlock( min(C, 32) ); // requires --maxrregcount 60 for optimal performance + assert(B * C % threadsPerBlock.x == 0); + dim3 numBlocks(B * C / threadsPerBlock.x); + kernel_backward<<>>(B, T, C, w, u, k, v, y, gy, gw, gu, gk, gv); +} diff --git a/lib/python3.12/site-packages/transformers/kernels/rwkv/wkv_cuda_bf16.cu b/lib/python3.12/site-packages/transformers/kernels/rwkv/wkv_cuda_bf16.cu new file mode 100644 index 0000000000000000000000000000000000000000..042cb4aba1db98be5916aea1de86a7fed0b6510d --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/rwkv/wkv_cuda_bf16.cu @@ -0,0 +1,186 @@ +#include +#include +#include "ATen/ATen.h" +#define MIN_VALUE (-1e38) +typedef at::BFloat16 bf16; + +__global__ void kernel_forward_bf16( + const int B, const int T, const int C, const float *__restrict__ const _w, const bf16 *__restrict__ const _u, + const bf16 *__restrict__ const _k, const bf16 *__restrict__ const _v, bf16 *__restrict__ const _y +) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const int _b = idx / C; + const int _c = idx % C; + const int _offset = _b * T * C + _c; + + float u = float(_u[_c]); + float w = _w[_c]; + const bf16 *__restrict__ const k = _k + _offset; + const bf16 *__restrict__ const v = _v + _offset; + bf16 *__restrict__ const y = _y + _offset; + + // aa and bb are running sums divided by exp(pp) (to avoid overflow) + float aa = 0, bb = 0, pp = MIN_VALUE; + for (int i = 0; i < T; i++) { + const int ii = i * C; + const float kk = float(k[ii]); + const float vv = float(v[ii]); + + float ww = u + kk; + float p = max(pp, ww); + float e1 = exp(pp - p); + float e2 = exp(ww - p); + y[ii] = bf16((e1 * aa + e2 * vv) / (e1 * bb + e2)); + + ww = w + pp; + p = max(ww, kk); + e1 = exp(ww - p); + e2 = exp(kk - p); + aa = e1 * aa + e2 * vv; + bb = e1 * bb + e2; + pp = p; + } +} + +__global__ void kernel_forward_with_state_bf16( + const int B, const int T, const int C, const float *__restrict__ const _w, const bf16 *__restrict__ const _u, + const bf16 *__restrict__ const _k, const bf16 *__restrict__ const _v, bf16 *__restrict__ const _y, + float *__restrict__ const _s +) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const int _b = idx / C; + const int _c = idx % C; + const int _offset_s = _b * C * 3 + _c * 3; + const int _offset = _b * T * C + _c; + + float u = float(_u[_c]); + float w = _w[_c]; + const bf16 *__restrict__ const k = _k + _offset; + const bf16 *__restrict__ const v = _v + _offset; + bf16 *__restrict__ const y = _y + _offset; + float *__restrict__ const s = _s + _offset_s; + + // aa and bb are running sums divided by exp(pp) (to avoid overflow) + float aa = s[0], bb = s[1], pp = s[2]; + for (int i = 0; i < T; i++) { + const int ii = i * C; + const float kk = float(k[ii]); + const float vv = float(v[ii]); + + float ww = u + kk; + float p = max(pp, ww); + float e1 = exp(pp - p); + float e2 = exp(ww - p); + y[ii] = bf16(e1 * aa + e2 * vv) / (e1 * bb + e2); + + ww = w + pp; + p = max(ww, kk); + e1 = exp(ww - p); + e2 = exp(kk - p); + aa = e1 * aa + e2 * vv; + bb = e1 * bb + e2; + pp = p; + } + s[0] = aa; + s[1] = bb; + s[2] = pp; +} + +__global__ void kernel_backward_bf16( + const int B, const int T, const int C, const float *__restrict__ const _w, const bf16 *__restrict__ const _u, + const bf16 *__restrict__ const _k, const bf16 *__restrict__ const _v, const bf16 *__restrict__ const _y, + const bf16 *__restrict__ const _gy, bf16 *__restrict__ const _gw, bf16 *__restrict__ const _gu, + bf16 *__restrict__ const _gk, bf16 *__restrict__ const _gv +) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const int _b = idx / C; + const int _c = idx % C; + const int _offset = _b * T * C + _c; + + float u = float(_u[_c]); + float w = _w[_c]; + const bf16 *__restrict__ const k = _k + _offset; + const bf16 *__restrict__ const v = _v + _offset; + const bf16 *__restrict__ const y = _y + _offset; + const bf16 *__restrict__ const gy = _gy + _offset; + bf16 *__restrict__ const gk = _gk + _offset; + bf16 *__restrict__ const gv = _gv + _offset; + + float q[Tmax], r[Tmax]; + + float gw = 0, gu = 0, aa = 0, bb = 0, ga = 0, gb = 0, pp = MIN_VALUE; + for (int i = 0; i < T; i++) { + const int ii = i * C; + const float kk = float(k[ii]); + const float vv = float(v[ii]); + const float yy = float(y[ii]); + + float ww = u + kk; + float p = max(pp, ww); + float e1 = exp(pp - p); + float e2 = exp(ww - p); + const float qq = float(gy[ii]) / (e1 * bb + e2); + gw += (ga - gb * yy) * e1 * qq; + gu += (vv - yy) * e2 * qq; + q[i] = qq; + r[i] = ww - p; + + ww = w + pp; + p = max(ww, kk); + e1 = exp(ww - p); + e2 = exp(kk - p); + ga = e1 * (aa + ga); + gb = e1 * (bb + gb); + aa = e1 * aa + e2 * vv; + bb = e1 * bb + e2; + pp = p; + } + const int _offsetBC = _b * C + _c; + _gw[_offsetBC] = bf16(gw * _w[_c]); // multiply by w because of w -> -exp(w) in python forward() + _gu[_offsetBC] = bf16(gu); + + aa = 0, bb = 0, pp = MIN_VALUE; + for (int i = T - 1; i >= 0; i--) { + const int ii = i * C; + const float kk = float(k[ii]); + const float vv = float(v[ii]); + const float yy = float(y[ii]); + const float qq = q[i]; + const float rr = r[i]; + + float e1 = qq * exp(rr); + float e2 = exp(kk + pp); + gk[ii] = bf16(e1 * (vv - yy) + e2 * (aa * vv + bb)); + gv[ii] = bf16(e1 + e2 * aa); + + const float ww = w + pp; + const float www = rr - u - kk; + const float p = max(ww, www); + e1 = exp(ww - p); + e2 = qq * exp(www - p); + aa = e1 * aa + e2; + bb = e1 * bb - e2 * yy; + pp = p; + } +} + +void cuda_forward_bf16(int B, int T, int C, float *w, bf16 *u, bf16 *k, bf16 *v, bf16 *y) { + dim3 threadsPerBlock( min(C, 32) ); // requires --maxrregcount 60 for optimal performance + assert(B * C % threadsPerBlock.x == 0); + dim3 numBlocks(B * C / threadsPerBlock.x); + kernel_forward_bf16<<>>(B, T, C, w, u, k, v, y); +} + +void cuda_forward_with_state_bf16(int B, int T, int C, float *w, bf16 *u, bf16 *k, bf16 *v, bf16 *y, float *s) { + dim3 threadsPerBlock( min(C, 32) ); // requires --maxrregcount 60 for optimal performance + assert(B * C % threadsPerBlock.x == 0); + dim3 numBlocks(B * C / threadsPerBlock.x); + kernel_forward_with_state_bf16<<>>(B, T, C, w, u, k, v, y, s); +} + +void cuda_backward_bf16(int B, int T, int C, float *w, bf16 *u, bf16 *k, bf16 *v, bf16 *y, bf16 *gy, bf16 *gw, bf16 *gu, bf16 *gk, bf16 *gv) { + dim3 threadsPerBlock( min(C, 32) ); // requires --maxrregcount 60 for optimal performance + assert(B * C % threadsPerBlock.x == 0); + dim3 numBlocks(B * C / threadsPerBlock.x); + kernel_backward_bf16<<>>(B, T, C, w, u, k, v, y, gy, gw, gu, gk, gv); +} diff --git a/lib/python3.12/site-packages/transformers/kernels/rwkv/wkv_op.cpp b/lib/python3.12/site-packages/transformers/kernels/rwkv/wkv_op.cpp new file mode 100644 index 0000000000000000000000000000000000000000..55e7280665927b523a88021d5111daf28a63c905 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/rwkv/wkv_op.cpp @@ -0,0 +1,66 @@ +#include +#include "ATen/ATen.h" +typedef at::BFloat16 bf16; + +void cuda_forward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y); +void cuda_forward_bf16(int B, int T, int C, float *w, bf16 *u, bf16 *k, bf16 *v, bf16 *y); +void cuda_forward_with_state(int B, int T, int C, float *w, float *u, float *k, float *v, float *y, float *s); +void cuda_forward_with_state_bf16(int B, int T, int C, float *w, bf16 *u, bf16 *k, bf16 *v, bf16 *y, float *s); +void cuda_backward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y, float *gy, float *gw, float *gu, float *gk, float *gv); +void cuda_backward_bf16(int B, int T, int C, float *w, bf16 *u, bf16 *k, bf16 *v, bf16 *y, bf16 *gy, bf16 *gw, bf16 *gu, bf16 *gk, bf16 *gv); + +void forward(torch::Tensor &w, torch::Tensor &u, torch::Tensor &k, torch::Tensor &v, torch::Tensor &y) { + const int B = k.size(0); + const int T = k.size(1); + const int C = k.size(2); + cuda_forward(B, T, C, w.data_ptr(), u.data_ptr(), k.data_ptr(), v.data_ptr(), y.data_ptr()); +} +void forward_bf16(torch::Tensor &w, torch::Tensor &u, torch::Tensor &k, torch::Tensor &v, torch::Tensor &y) { + const int B = k.size(0); + const int T = k.size(1); + const int C = k.size(2); + cuda_forward_bf16(B, T, C, w.data_ptr(), u.data_ptr(), k.data_ptr(), v.data_ptr(), y.data_ptr()); +} +void forward_with_state(torch::Tensor &w, torch::Tensor &u, torch::Tensor &k, torch::Tensor &v, torch::Tensor &y, torch::Tensor &s) { + const int B = k.size(0); + const int T = k.size(1); + const int C = k.size(2); + cuda_forward_with_state(B, T, C, w.data_ptr(), u.data_ptr(), k.data_ptr(), v.data_ptr(), y.data_ptr(), s.data_ptr()); +} +void forward_with_state_bf16(torch::Tensor &w, torch::Tensor &u, torch::Tensor &k, torch::Tensor &v, torch::Tensor &y, torch::Tensor &s) { + const int B = k.size(0); + const int T = k.size(1); + const int C = k.size(2); + cuda_forward_with_state_bf16(B, T, C, w.data_ptr(), u.data_ptr(), k.data_ptr(), v.data_ptr(), y.data_ptr(), s.data_ptr()); +} +void backward(torch::Tensor &w, torch::Tensor &u, torch::Tensor &k, torch::Tensor &v, torch::Tensor &y, torch::Tensor &gy, torch::Tensor &gw, torch::Tensor &gu, torch::Tensor &gk, torch::Tensor &gv) { + const int B = k.size(0); + const int T = k.size(1); + const int C = k.size(2); + cuda_backward(B, T, C, w.data_ptr(), u.data_ptr(), k.data_ptr(), v.data_ptr(), y.data_ptr(), gy.data_ptr(), gw.data_ptr(), gu.data_ptr(), gk.data_ptr(), gv.data_ptr()); +} +void backward_bf16(torch::Tensor &w, torch::Tensor &u, torch::Tensor &k, torch::Tensor &v, torch::Tensor &y, torch::Tensor &gy, torch::Tensor &gw, torch::Tensor &gu, torch::Tensor &gk, torch::Tensor &gv) { + const int B = k.size(0); + const int T = k.size(1); + const int C = k.size(2); + cuda_backward_bf16(B, T, C, w.data_ptr(), u.data_ptr(), k.data_ptr(), v.data_ptr(), y.data_ptr(), + gy.data_ptr(), gw.data_ptr(), gu.data_ptr(), gk.data_ptr(), gv.data_ptr()); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("forward", &forward, "wkv forward"); + m.def("forward_bf16", &forward_bf16, "wkv forward bf16"); + m.def("forward_with_state", &forward_with_state, "wkv forward with state"); + m.def("forward_with_state_bf16", &forward_with_state_bf16, "wkv forward with state bf16"); + m.def("backward", &backward, "wkv backward"); + m.def("backward_bf16", &backward_bf16, "wkv backward bf16"); +} + +TORCH_LIBRARY(wkv, m) { + m.def("forward", forward); + m.def("forward_bf16", forward_bf16); + m.def("forward_with_state", forward_with_state); + m.def("forward_with_state_bf16", forward_with_state_bf16); + m.def("backward", backward); + m.def("backward_bf16", backward_bf16); +} diff --git a/lib/python3.12/site-packages/transformers/kernels/yoso/common.h b/lib/python3.12/site-packages/transformers/kernels/yoso/common.h new file mode 100644 index 0000000000000000000000000000000000000000..e5085c88dd3ea9a12eec264a8c48946bf2b80b23 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/yoso/common.h @@ -0,0 +1,10 @@ + +#define min(a, b) ((a)<(b)?(a):(b)) +#define max(a, b) ((a)>(b)?(a):(b)) +#define ceil_divide(a, b) ((a)/(b)+((a)%(b)!=0)) +#define select(cond, a, b) ((cond)?(a):(b)) +#define PI 3.141592 +#define EPSILON 1e-8 +#define MAX_VAL 1e12 +#define MIN_VAL -1e12 +#define EMPTY_VALUE -1 diff --git a/lib/python3.12/site-packages/transformers/kernels/yoso/common_cuda.h b/lib/python3.12/site-packages/transformers/kernels/yoso/common_cuda.h new file mode 100644 index 0000000000000000000000000000000000000000..97030870649a2fdac58cb26cf966e8f5c8cc7909 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/yoso/common_cuda.h @@ -0,0 +1,9 @@ + +#define MAX_THREADS_PER_BLOCK 1024 +#define OPTIMAL_THREADS_PER_BLOCK 256 +#define WARP_SIZE 32 +#define MAX_NUM_BLOCK_X 2147483647 +#define MAX_NUM_BLOCK_Y 65535 +#define MAX_NUM_BLOCK_Z 65535 +#define MAX_SHARED_MEM_PER_BLOCK 48000 +#define FULL_MASK 0xffffffff diff --git a/lib/python3.12/site-packages/transformers/kernels/yoso/common_cuda_device.h b/lib/python3.12/site-packages/transformers/kernels/yoso/common_cuda_device.h new file mode 100644 index 0000000000000000000000000000000000000000..6674f93afdc25ab35c5d83881d00028bcf2989fc --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/yoso/common_cuda_device.h @@ -0,0 +1,79 @@ + +#include "common.h" + +template +__device__ int set_insert(T *set, int set_size, T value) { + int slot = value % set_size; + int start_slot = slot; + while (true) { + T prev = atomicCAS(&set[slot], EMPTY_VALUE, value); + if (prev == EMPTY_VALUE || prev == value) { + return slot; + } + slot = (slot + 1) % set_size; + if (slot == start_slot) { + return -1; + } + } + return -1; +} + +template +__device__ int set_lookup(T *set, int set_size, T value) { + int slot = value % set_size; + int start_slot = slot; + while (true) { + if (set[slot] == value) { + return slot; + } + slot = (slot + 1) % set_size; + if (slot == start_slot) { + return -1; + } + } + return -1; +} + +template +__device__ void init_buffer(T init_value, T *buffer, int buffer_size, int num_threads, int thread_id) { + __syncthreads(); + for (int i = 0; i < buffer_size; i = i + num_threads) { + int offset_idx = i + thread_id; + if (offset_idx < buffer_size) { + buffer[offset_idx] = init_value; + } + } + __syncthreads(); +} + +template +__device__ void copy_data(T *src_pt, T *dist_pt, int data_length, int num_threads, int thread_id) { + __syncthreads(); + for (int i = 0; i < data_length; i = i + num_threads) { + int offset_idx = i + thread_id; + if (offset_idx < data_length) { + dist_pt[offset_idx] = src_pt[offset_idx]; + } + } + __syncthreads(); +} + +template +__device__ void init_buffer_nonblocking(T init_value, T *buffer, int buffer_size, int num_threads, int thread_id) { + for (int i = 0; i < buffer_size; i = i + num_threads) { + int offset_idx = i + thread_id; + if (offset_idx < buffer_size) { + buffer[offset_idx] = init_value; + } + } +} + +template +__device__ void copy_data_nonblocking(T *src_pt, T *dist_pt, int data_length, int num_threads, int thread_id) { + for (int i = 0; i < data_length; i = i + num_threads) { + int offset_idx = i + thread_id; + if (offset_idx < data_length) { + dist_pt[offset_idx] = src_pt[offset_idx]; + } + } +} diff --git a/lib/python3.12/site-packages/transformers/kernels/yoso/fast_lsh_cumulation.cu b/lib/python3.12/site-packages/transformers/kernels/yoso/fast_lsh_cumulation.cu new file mode 100644 index 0000000000000000000000000000000000000000..c6b13e6cb5f53c9c62e51d2c399a14d14dab7037 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/yoso/fast_lsh_cumulation.cu @@ -0,0 +1,588 @@ +// File from https://github.com/mlpen/YOSO/blob/main/encoders/backbones/efficient_attentions/yoso/yoso_v1/cuda/fast_lsh_cumulation.cu + +#include +#include +#include "fast_lsh_cumulation.h" +#include "fast_lsh_cumulation_cuda.h" +#include "common_cuda.h" +#include "common.h" +#include +////////////////////////////////////////////////////////////////////////////////////////////////// +////////////////////////////////////////////////////////////////////////////////////////////////// + +std::vector fast_hash_ver1_kernel( + at::Tensor query_mask, + at::Tensor query_vector, + at::Tensor key_mask, + at::Tensor key_vector, + int num_hash_f, + int hash_code_len, + bool use_cuda +) { + + int batch_size = query_vector.size(0); + int num_query = query_vector.size(1); + int num_key = key_vector.size(1); + int vector_dim = query_vector.size(2); + + int num_hash_per_part = vector_dim / hash_code_len; + int num_part = max(1, ceil_divide(num_hash_f, num_hash_per_part)); + + at::Tensor Dmat = 2 * at::randint(0, 2, {batch_size, 3, num_part, vector_dim}, query_mask.options()) - 1; + at::Tensor query_hash_code = at::zeros({batch_size, num_query, num_hash_f}, query_mask.options()); + at::Tensor key_hash_code = at::zeros({batch_size, num_key, num_hash_f}, key_mask.options()); + + int *query_mask_ptr = query_mask.data_ptr(); + float *query_vector_ptr = query_vector.data_ptr(); + int *key_mask_ptr = key_mask.data_ptr(); + float *key_vector_ptr = key_vector.data_ptr(); + + int *Dmat_ptr = Dmat.data_ptr(); + + int *query_hash_code_ptr = query_hash_code.data_ptr(); + int *key_hash_code_ptr = key_hash_code.data_ptr(); + + if (use_cuda) { + { + dim3 threads(vector_dim); + dim3 blocks(num_part, num_query, batch_size); + int shared_mem = vector_dim * sizeof(float); + fast_hash_ver1_cuda_kernel<<>>( + query_mask_ptr, + query_vector_ptr, + Dmat_ptr, + query_hash_code_ptr, + batch_size, + num_query, + vector_dim, + num_part, + num_hash_f, + hash_code_len + ); + } + { + dim3 threads(vector_dim); + dim3 blocks(num_part, num_key, batch_size); + int shared_mem = vector_dim * sizeof(float); + fast_hash_ver1_cuda_kernel<<>>( + key_mask_ptr, + key_vector_ptr, + Dmat_ptr, + key_hash_code_ptr, + batch_size, + num_key, + vector_dim, + num_part, + num_hash_f, + hash_code_len + ); + } + } + + return {query_hash_code, key_hash_code}; + +} + +at::Tensor lsh_cumulation_ver1_kernel( + at::Tensor query_mask, + at::Tensor query_hash_code, + at::Tensor key_mask, + at::Tensor key_hash_code, + at::Tensor value, + int hashtable_capacity, + bool use_cuda +) { + + int batch_size = query_hash_code.size(0); + int num_hash_f = query_hash_code.size(2); + + int num_query = query_hash_code.size(1); + int num_key = key_hash_code.size(1); + int value_dim = value.size(2); + + at::Tensor hashtable_value = at::empty({batch_size, num_hash_f, hashtable_capacity, WARP_SIZE}, value.options()); + at::Tensor cumulation_value = at::zeros({batch_size, num_query, value_dim}, value.options()); + + if (use_cuda) { + int threads_x = WARP_SIZE; + int threads_y = OPTIMAL_THREADS_PER_BLOCK / WARP_SIZE; + int block_x_step1 = num_key / threads_y; + int block_x_step2 = num_query / threads_y; + int block_y = batch_size; + + dim3 threads(threads_x, threads_y); + dim3 blocks_step1(block_x_step1, block_y); + dim3 blocks_step2(block_x_step2, block_y); + + int *query_mask_ptr = query_mask.data_ptr(); + int *query_hash_code_ptr = query_hash_code.data_ptr(); + int *key_mask_ptr = key_mask.data_ptr(); + int *key_hash_code_ptr = key_hash_code.data_ptr(); + float *value_ptr = value.data_ptr(); + float *hashtable_value_ptr = hashtable_value.data_ptr(); + float *cumulation_value_ptr = cumulation_value.data_ptr(); + + for (int value_offset = 0; value_offset < value_dim; value_offset = value_offset + WARP_SIZE) { + + cudaMemset(hashtable_value_ptr, 0, (batch_size * num_hash_f * hashtable_capacity * WARP_SIZE) * sizeof(float)); + + lsh_cumulation_ver1_step1_cuda_kernel<<>>( + key_mask_ptr, + key_hash_code_ptr, + value_ptr, + hashtable_value_ptr, + batch_size, + num_hash_f, + hashtable_capacity, + num_key, + value_dim, + value_offset + ); + + lsh_cumulation_ver1_step2_cuda_kernel<<>>( + query_mask_ptr, + query_hash_code_ptr, + hashtable_value_ptr, + cumulation_value_ptr, + batch_size, + num_hash_f, + hashtable_capacity, + num_query, + value_dim, + value_offset + ); + } + + } + + return cumulation_value; + +} + +at::Tensor lsh_weighted_cumulation_ver1_kernel( + at::Tensor query_mask, + at::Tensor query_hash_code, + at::Tensor query_weight, + at::Tensor key_mask, + at::Tensor key_hash_code, + at::Tensor key_weight, + at::Tensor value, + int hashtable_capacity, + bool use_cuda +) { + + int batch_size = query_hash_code.size(0); + int num_hash_f = query_hash_code.size(2); + + int num_query = query_hash_code.size(1); + int num_key = key_hash_code.size(1); + int value_dim = value.size(2); + int weight_dim = query_weight.size(2); + + at::Tensor hashtable_value = at::zeros({batch_size, num_hash_f, hashtable_capacity, WARP_SIZE}, value.options()); + at::Tensor cumulation_value = at::zeros({batch_size, num_query, value_dim}, value.options()); + + if (use_cuda) { + int threads_x = WARP_SIZE; + int threads_y = OPTIMAL_THREADS_PER_BLOCK / WARP_SIZE; + int block_x_step1 = num_key / threads_y; + int block_x_step2 = num_query / threads_y; + int block_y = batch_size; + + dim3 threads(threads_x, threads_y); + dim3 blocks_step1(block_x_step1, block_y); + dim3 blocks_step2(block_x_step2, block_y); + + int *query_mask_ptr = query_mask.data_ptr(); + int *query_hash_code_ptr = query_hash_code.data_ptr(); + float *query_weight_ptr = query_weight.data_ptr(); + int *key_mask_ptr = key_mask.data_ptr(); + int *key_hash_code_ptr = key_hash_code.data_ptr(); + float *key_weight_ptr = key_weight.data_ptr(); + float *value_ptr = value.data_ptr(); + float *hashtable_value_ptr = hashtable_value.data_ptr(); + float *cumulation_value_ptr = cumulation_value.data_ptr(); + + for (int value_offset = 0; value_offset < value_dim; value_offset = value_offset + WARP_SIZE) { + for (int weight_idx = 0; weight_idx < weight_dim; weight_idx++) { + + cudaMemset(hashtable_value_ptr, 0, (batch_size * num_hash_f * hashtable_capacity * WARP_SIZE) * sizeof(float)); + + lsh_weighted_cumulation_ver1_step1_cuda_kernel<<>>( + key_mask_ptr, + key_hash_code_ptr, + key_weight_ptr, + value_ptr, + hashtable_value_ptr, + batch_size, + num_hash_f, + hashtable_capacity, + num_key, + value_dim, + weight_dim, + value_offset, + weight_idx + ); + + lsh_weighted_cumulation_ver1_step2_cuda_kernel<<>>( + query_mask_ptr, + query_hash_code_ptr, + query_weight_ptr, + hashtable_value_ptr, + cumulation_value_ptr, + batch_size, + num_hash_f, + hashtable_capacity, + num_query, + value_dim, + weight_dim, + value_offset, + weight_idx + ); + } + } + + } + + return cumulation_value; + +} + +at::Tensor lsh_weighted_cumulation_ver2_kernel( + at::Tensor query_mask, + at::Tensor query_hash_code, + at::Tensor query_weight, + at::Tensor key_mask, + at::Tensor key_hash_code, + at::Tensor key_weight, + at::Tensor value, + int hashtable_capacity, + bool use_cuda +) { + + int batch_size = query_hash_code.size(0); + int num_hash_f = query_hash_code.size(2); + + int num_query = query_hash_code.size(1); + int num_key = key_hash_code.size(1); + int value_dim = value.size(2); + int weight_dim = query_weight.size(2); + + at::Tensor count_sort_table = at::zeros({batch_size, num_hash_f, hashtable_capacity}, query_hash_code.options()); + at::Tensor key_sorted_idxes = at::zeros({batch_size, num_hash_f, num_key}, query_hash_code.options()); + at::Tensor query_info = at::zeros({batch_size, num_query, 2, num_hash_f}, query_hash_code.options()); + at::Tensor cumulation_value = at::zeros({batch_size, num_query, value_dim}, value.options()); + + if (use_cuda) { + + int *query_mask_ptr = query_mask.data_ptr(); + int *query_hash_code_ptr = query_hash_code.data_ptr(); + float *query_weight_ptr = query_weight.data_ptr(); + int *key_mask_ptr = key_mask.data_ptr(); + int *key_hash_code_ptr = key_hash_code.data_ptr(); + float *key_weight_ptr = key_weight.data_ptr(); + float *value_ptr = value.data_ptr(); + + int *count_sort_table_ptr = count_sort_table.data_ptr(); + int *key_sorted_idxes_ptr = key_sorted_idxes.data_ptr(); + int *query_info_ptr = query_info.data_ptr(); + + float *cumulation_value_ptr = cumulation_value.data_ptr(); + + { + dim3 threads_step13(num_hash_f, max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f)); + dim3 blocks_step13(num_key / max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f), batch_size); + dim3 threads_step2(min(hashtable_capacity, OPTIMAL_THREADS_PER_BLOCK)); + dim3 blocks_step2(num_hash_f, batch_size); + int shared_mem = hashtable_capacity * sizeof(float); + count_sort_step1_cuda_kernel<<>>( + key_mask_ptr, + key_hash_code_ptr, + count_sort_table_ptr, + batch_size, + num_hash_f, + hashtable_capacity, + num_key + ); + count_sort_step2_cuda_kernel<<>>( + count_sort_table_ptr, + batch_size, + num_hash_f, + hashtable_capacity + ); + count_sort_step3_cuda_kernel<<>>( + key_mask_ptr, + key_hash_code_ptr, + count_sort_table_ptr, + key_sorted_idxes_ptr, + batch_size, + num_hash_f, + hashtable_capacity, + num_key + ); + } + { + dim3 threads(num_hash_f, max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f)); + dim3 blocks(num_query / max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f), batch_size); + extract_query_info_cuda_kernel<<>>( + query_mask_ptr, + query_hash_code_ptr, + count_sort_table_ptr, + query_info_ptr, + batch_size, + num_hash_f, + hashtable_capacity, + num_query + ); + } + { + dim3 threads(WARP_SIZE, OPTIMAL_THREADS_PER_BLOCK / WARP_SIZE); + dim3 blocks(num_query, num_hash_f, batch_size); + int shared_mem = (weight_dim + WARP_SIZE) * sizeof(float); + lsh_weighted_cumulation_ver2_step2_cuda_kernel<<>>( + query_mask_ptr, + query_info_ptr, + key_sorted_idxes_ptr, + query_weight_ptr, + key_weight_ptr, + value_ptr, + cumulation_value_ptr, + batch_size, + num_hash_f, + num_query, + num_key, + value_dim, + weight_dim + ); + } + } + + return cumulation_value; + +} + +at::Tensor lsh_weighted_cumulation_ver3_kernel( + at::Tensor query_mask, + at::Tensor query_hash_code, + at::Tensor query_weight, + at::Tensor key_mask, + at::Tensor key_hash_code, + at::Tensor key_weight, + at::Tensor value, + int hashtable_capacity, + bool use_cuda +) { + + int batch_size = query_hash_code.size(0); + int num_hash_f = query_hash_code.size(2); + + int num_query = query_hash_code.size(1); + int num_key = key_hash_code.size(1); + int value_dim = value.size(2); + int weight_dim = query_weight.size(2); + + at::Tensor count_sort_table = at::zeros({batch_size, num_hash_f, hashtable_capacity}, query_hash_code.options()); + at::Tensor query_sorted_idxes = at::zeros({batch_size, num_hash_f, num_query}, query_hash_code.options()); + at::Tensor key_info = at::zeros({batch_size, num_key, 2, num_hash_f}, query_hash_code.options()); + at::Tensor cumulation_value = at::zeros({batch_size, num_query, value_dim}, value.options()); + + if (use_cuda) { + + int *query_mask_ptr = query_mask.data_ptr(); + int *query_hash_code_ptr = query_hash_code.data_ptr(); + float *query_weight_ptr = query_weight.data_ptr(); + int *key_mask_ptr = key_mask.data_ptr(); + int *key_hash_code_ptr = key_hash_code.data_ptr(); + float *key_weight_ptr = key_weight.data_ptr(); + float *value_ptr = value.data_ptr(); + + int *count_sort_table_ptr = count_sort_table.data_ptr(); + int *query_sorted_idxes_ptr = query_sorted_idxes.data_ptr(); + int *key_info_ptr = key_info.data_ptr(); + + float *cumulation_value_ptr = cumulation_value.data_ptr(); + + { + dim3 threads_step13(num_hash_f, max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f)); + dim3 blocks_step13(num_query / max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f), batch_size); + dim3 threads_step2(min(hashtable_capacity, OPTIMAL_THREADS_PER_BLOCK)); + dim3 blocks_step2(num_hash_f, batch_size); + int shared_mem = hashtable_capacity * sizeof(float); + count_sort_step1_cuda_kernel<<>>( + query_mask_ptr, + query_hash_code_ptr, + count_sort_table_ptr, + batch_size, + num_hash_f, + hashtable_capacity, + num_query + ); + count_sort_step2_cuda_kernel<<>>( + count_sort_table_ptr, + batch_size, + num_hash_f, + hashtable_capacity + ); + count_sort_step3_cuda_kernel<<>>( + query_mask_ptr, + query_hash_code_ptr, + count_sort_table_ptr, + query_sorted_idxes_ptr, + batch_size, + num_hash_f, + hashtable_capacity, + num_query + ); + } + { + dim3 threads(num_hash_f, max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f)); + dim3 blocks(num_key / max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f), batch_size); + extract_query_info_cuda_kernel<<>>( + key_mask_ptr, + key_hash_code_ptr, + count_sort_table_ptr, + key_info_ptr, + batch_size, + num_hash_f, + hashtable_capacity, + num_key + ); + } + { + dim3 threads(WARP_SIZE, OPTIMAL_THREADS_PER_BLOCK / WARP_SIZE); + dim3 blocks(num_key, num_hash_f, batch_size); + int shared_mem = (weight_dim + value_dim + WARP_SIZE) * sizeof(float); + lsh_weighted_cumulation_ver3_step2_cuda_kernel<<>>( + query_sorted_idxes_ptr, + key_mask_ptr, + key_info_ptr, + query_weight_ptr, + key_weight_ptr, + value_ptr, + cumulation_value_ptr, + batch_size, + num_hash_f, + num_query, + num_key, + value_dim, + weight_dim + ); + } + } + + return cumulation_value; + +} + +at::Tensor lsh_weighted_cumulation_ver4_kernel( + at::Tensor query_mask, + at::Tensor query_hash_code, + at::Tensor query_weight, + at::Tensor key_mask, + at::Tensor key_hash_code, + at::Tensor key_weight, + at::Tensor value, + int hashtable_capacity, + bool use_cuda +) { + + int batch_size = query_hash_code.size(0); + int num_hash_f = query_hash_code.size(2); + + int num_query = query_hash_code.size(1); + int num_key = key_hash_code.size(1); + int value_dim = value.size(2); + int weight_dim = query_weight.size(2); + + at::Tensor count_sort_table = at::zeros({batch_size, num_hash_f, hashtable_capacity}, query_hash_code.options()); + at::Tensor query_sorted_idxes = at::zeros({batch_size, num_hash_f, num_query}, query_hash_code.options()); + at::Tensor key_info = at::zeros({batch_size, num_key, 2, num_hash_f}, query_hash_code.options()); + at::Tensor cumulation_value = at::zeros({batch_size, num_query, value_dim}, value.options()); + + if (use_cuda) { + + int *query_mask_ptr = query_mask.data_ptr(); + int *query_hash_code_ptr = query_hash_code.data_ptr(); + float *query_weight_ptr = query_weight.data_ptr(); + int *key_mask_ptr = key_mask.data_ptr(); + int *key_hash_code_ptr = key_hash_code.data_ptr(); + float *key_weight_ptr = key_weight.data_ptr(); + float *value_ptr = value.data_ptr(); + + int *count_sort_table_ptr = count_sort_table.data_ptr(); + int *query_sorted_idxes_ptr = query_sorted_idxes.data_ptr(); + int *key_info_ptr = key_info.data_ptr(); + + float *cumulation_value_ptr = cumulation_value.data_ptr(); + + { + dim3 threads_step13(num_hash_f, max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f)); + dim3 blocks_step13(num_query / max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f), batch_size); + dim3 threads_step2(min(hashtable_capacity, OPTIMAL_THREADS_PER_BLOCK)); + dim3 blocks_step2(num_hash_f, batch_size); + int shared_mem = hashtable_capacity * sizeof(float); + count_sort_step1_cuda_kernel<<>>( + query_mask_ptr, + query_hash_code_ptr, + count_sort_table_ptr, + batch_size, + num_hash_f, + hashtable_capacity, + num_query + ); + count_sort_step2_cuda_kernel<<>>( + count_sort_table_ptr, + batch_size, + num_hash_f, + hashtable_capacity + ); + count_sort_step3_cuda_kernel<<>>( + query_mask_ptr, + query_hash_code_ptr, + count_sort_table_ptr, + query_sorted_idxes_ptr, + batch_size, + num_hash_f, + hashtable_capacity, + num_query + ); + } + { + dim3 threads(num_hash_f, max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f)); + dim3 blocks(num_key / max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f), batch_size); + extract_query_info_cuda_kernel<<>>( + key_mask_ptr, + key_hash_code_ptr, + count_sort_table_ptr, + key_info_ptr, + batch_size, + num_hash_f, + hashtable_capacity, + num_key + ); + } + { + dim3 threads(WARP_SIZE, OPTIMAL_THREADS_PER_BLOCK / WARP_SIZE); + dim3 blocks(num_key, batch_size); + int shared_mem = (weight_dim + value_dim + 2 * num_hash_f) * sizeof(float); + lsh_weighted_cumulation_ver4_step2_cuda_kernel<<>>( + query_sorted_idxes_ptr, + key_mask_ptr, + key_info_ptr, + query_weight_ptr, + key_weight_ptr, + value_ptr, + cumulation_value_ptr, + batch_size, + num_hash_f, + num_query, + num_key, + value_dim, + weight_dim + ); + } + } + + return cumulation_value; + +} diff --git a/lib/python3.12/site-packages/transformers/kernels/yoso/fast_lsh_cumulation.h b/lib/python3.12/site-packages/transformers/kernels/yoso/fast_lsh_cumulation.h new file mode 100644 index 0000000000000000000000000000000000000000..dd48de0ed159f49ee3afe93b12aaae719fe87688 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/yoso/fast_lsh_cumulation.h @@ -0,0 +1,71 @@ +#include +#include +#include + +std::vector fast_hash_ver1_kernel( + at::Tensor query_mask, + at::Tensor query_vector, + at::Tensor key_mask, + at::Tensor key_vector, + int num_hash_f, + int hash_code_len, + bool use_cuda +); + +at::Tensor lsh_cumulation_ver1_kernel( + at::Tensor query_mask, + at::Tensor query_hash_code, + at::Tensor key_mask, + at::Tensor key_hash_code, + at::Tensor value, + int hashtable_capacity, + bool use_cuda +); + +at::Tensor lsh_weighted_cumulation_ver1_kernel( + at::Tensor query_mask, + at::Tensor query_hash_code, + at::Tensor query_weight, + at::Tensor key_mask, + at::Tensor key_hash_code, + at::Tensor key_weight, + at::Tensor value, + int hashtable_capacity, + bool use_cuda +); + +at::Tensor lsh_weighted_cumulation_ver2_kernel( + at::Tensor query_mask, + at::Tensor query_hash_code, + at::Tensor query_weight, + at::Tensor key_mask, + at::Tensor key_hash_code, + at::Tensor key_weight, + at::Tensor value, + int hashtable_capacity, + bool use_cuda +); + +at::Tensor lsh_weighted_cumulation_ver3_kernel( + at::Tensor query_mask, + at::Tensor query_hash_code, + at::Tensor query_weight, + at::Tensor key_mask, + at::Tensor key_hash_code, + at::Tensor key_weight, + at::Tensor value, + int hashtable_capacity, + bool use_cuda +); + +at::Tensor lsh_weighted_cumulation_ver4_kernel( + at::Tensor query_mask, + at::Tensor query_hash_code, + at::Tensor query_weight, + at::Tensor key_mask, + at::Tensor key_hash_code, + at::Tensor key_weight, + at::Tensor value, + int hashtable_capacity, + bool use_cuda +); diff --git a/lib/python3.12/site-packages/transformers/kernels/yoso/fast_lsh_cumulation_cuda.cu b/lib/python3.12/site-packages/transformers/kernels/yoso/fast_lsh_cumulation_cuda.cu new file mode 100644 index 0000000000000000000000000000000000000000..ebc6260dd6db3ecaf8cb7b35c3c1a6e1ab3851dc --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/yoso/fast_lsh_cumulation_cuda.cu @@ -0,0 +1,825 @@ +// File from https://github.com/mlpen/YOSO/blob/main/encoders/backbones/efficient_attentions/yoso/yoso_v1/cuda/fast_lsh_cumulation_cuda.cu + +#include "fast_lsh_cumulation_cuda.h" +#include "common_cuda_device.h" +#include "common_cuda.h" +#include "common.h" +#include +////////////////////////////////////////////////////////////////////////////////////////////////// +////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void fast_hadamard_transform(float *vector_buffer, int vector_dim, int dim_idx) { + int stride = vector_dim / 2; + while (stride > (WARP_SIZE / 2)) { + __syncthreads(); + int sign = 1 - ((dim_idx / stride) % 2) * 2; + float val1 = vector_buffer[dim_idx]; + float val2 = vector_buffer[dim_idx + sign * stride]; + __syncthreads(); + vector_buffer[dim_idx] = float(sign) * val1 + val2; + stride = stride / 2; + } + + float val = vector_buffer[dim_idx]; + #pragma unroll + for (stride = (WARP_SIZE / 2); stride > 0; stride = stride / 2) { + int sign = 1 - ((dim_idx / stride) % 2) * 2; + val = float(sign) * val + __shfl_xor_sync(FULL_MASK, val, stride); + } + vector_buffer[dim_idx] = val; +} + +__global__ void fast_hash_ver1_cuda_kernel( + int *mask, // [batch_size, num_vector] + float *vector, // [batch_size, num_vector, vector_dim] + int *Dmat, // [batch_size, 3, num_part, vector_dim] + int *hash_code, // [batch_size, num_vector, num_hash_f] + int batch_size, + int num_vector, + int vector_dim, + int num_part, + int num_hash_f, + int hash_code_len +) { + + int batch_idx = blockIdx.z; + int vector_idx = blockIdx.y; + int part_idx = blockIdx.x; + + int dim_idx = threadIdx.x; + + int batch_idx__vector_idx = batch_idx * num_vector + vector_idx; + if (mask[batch_idx__vector_idx] == 0) { + return; + } + + extern __shared__ float buffer[]; + float *vector_buffer = buffer; + + vector_buffer[dim_idx] = vector[batch_idx__vector_idx * vector_dim + dim_idx]; + + vector_buffer[dim_idx] = vector_buffer[dim_idx] * (float)Dmat[((batch_idx * 3 + 0) * num_part + part_idx) * vector_dim + dim_idx]; + fast_hadamard_transform(vector_buffer, vector_dim, dim_idx); + vector_buffer[dim_idx] = vector_buffer[dim_idx] * (float)Dmat[((batch_idx * 3 + 1) * num_part + part_idx) * vector_dim + dim_idx]; + fast_hadamard_transform(vector_buffer, vector_dim, dim_idx); + vector_buffer[dim_idx] = vector_buffer[dim_idx] * (float)Dmat[((batch_idx * 3 + 2) * num_part + part_idx) * vector_dim + dim_idx]; + fast_hadamard_transform(vector_buffer, vector_dim, dim_idx); + + int num_hash_per_part = vector_dim / hash_code_len; + if (hash_code_len == 8 || hash_code_len == 16) { + int code = select(vector_buffer[dim_idx] > 0, 1 << (dim_idx % hash_code_len), 0); + for (int offset = 1; offset < hash_code_len; offset = offset * 2) { + code += __shfl_xor_sync(FULL_MASK, code, offset); + } + if (dim_idx % hash_code_len == 0) { + int hash_f_idx = part_idx * num_hash_per_part + dim_idx / hash_code_len; + if (hash_f_idx < num_hash_f) { + hash_code[batch_idx__vector_idx * num_hash_f + hash_f_idx] = code; + } + } + } else { + vector_buffer[dim_idx] = select(vector_buffer[dim_idx] > 0, 1 << (dim_idx % hash_code_len), 0); + __syncthreads(); + if (dim_idx < num_hash_per_part) { + int code = 0; + for (int i = 0; i < hash_code_len; i++) { + code += vector_buffer[dim_idx * hash_code_len + i]; + } + int hash_f_idx = part_idx * num_hash_per_part + dim_idx; + if (hash_f_idx < num_hash_f) { + hash_code[batch_idx__vector_idx * num_hash_f + hash_f_idx] = code; + } + } + } +} + +__global__ void lsh_cumulation_ver1_step1_cuda_kernel( + int *key_mask, // [batch_size, num_key] + int *key_hash_code, // [batch_size, num_key, num_hash_f] + float *value, // [batch_size, num_key, value_dim] + float *hashtable_value, // [batch_size, num_hash_f, hashtable_capacity, WARP_SIZE] + int batch_size, + int num_hash_f, + int hashtable_capacity, + int num_key, + int value_dim, + int offset_warp +) { + + int warp_thread_idx = threadIdx.x; + + int batch_idx = blockIdx.y; + int key_idx = blockIdx.x * blockDim.y + threadIdx.y; + + int batch_idx__key_idx = batch_idx * num_key + key_idx; + if (key_mask[batch_idx__key_idx] == 0) { + return; + } + + if (num_hash_f > WARP_SIZE) { + float warp_value = value[batch_idx__key_idx * value_dim + offset_warp + warp_thread_idx]; + for (int hash_f_start = 0; hash_f_start < num_hash_f; hash_f_start = hash_f_start + WARP_SIZE) { + int warp_hashcode = key_hash_code[batch_idx__key_idx * num_hash_f + hash_f_start + warp_thread_idx]; + #pragma unroll + for (int hash_f_offset = 0; hash_f_offset < WARP_SIZE; hash_f_offset++) { + int current_hashcode = warp_hashcode; + current_hashcode = __shfl_sync(FULL_MASK, current_hashcode, hash_f_offset); + int hashtable_idx = (batch_idx * num_hash_f + (hash_f_start + hash_f_offset)) * hashtable_capacity + current_hashcode; + atomicAdd(&hashtable_value[hashtable_idx * WARP_SIZE + warp_thread_idx], warp_value); + } + } + } else { + float warp_value = value[batch_idx__key_idx * value_dim + offset_warp + warp_thread_idx]; + int warp_hashcode = 0; + if (warp_thread_idx < num_hash_f) { + warp_hashcode = key_hash_code[batch_idx__key_idx * num_hash_f + warp_thread_idx]; + } + for (int hash_f_idx = 0; hash_f_idx < num_hash_f; hash_f_idx++) { + int current_hashcode = warp_hashcode; + current_hashcode = __shfl_sync(FULL_MASK, current_hashcode, hash_f_idx); + int hashtable_idx = (batch_idx * num_hash_f + hash_f_idx) * hashtable_capacity + current_hashcode; + atomicAdd(&hashtable_value[hashtable_idx * WARP_SIZE + warp_thread_idx], warp_value); + } + } + +} + +__global__ void lsh_cumulation_ver1_step2_cuda_kernel( + int *query_mask, // [batch_size, num_query] + int *query_hash_code, // [batch_size, num_query, num_hash_f] + float *hashtable_value, // [batch_size, num_hash_f, hashtable_capacity, WARP_SIZE] + float *cumulation_value, // [batch_size, num_query, value_dim] + int batch_size, + int num_hash_f, + int hashtable_capacity, + int num_query, + int value_dim, + int offset_warp +) { + + int warp_thread_idx = threadIdx.x; + + int batch_idx = blockIdx.y; + int query_idx = blockIdx.x * blockDim.y + threadIdx.y; + + int batch_idx__query_idx = batch_idx * num_query + query_idx; + if (query_mask[batch_idx__query_idx] == 0) { + return; + } + + if (num_hash_f > WARP_SIZE) { + float warp_value = 0; + for (int hash_f_start = 0; hash_f_start < num_hash_f; hash_f_start = hash_f_start + WARP_SIZE) { + int warp_hashcode = query_hash_code[batch_idx__query_idx * num_hash_f + hash_f_start + warp_thread_idx]; + #pragma unroll + for (int hash_f_offset = 0; hash_f_offset < WARP_SIZE; hash_f_offset++) { + int current_hashcode = warp_hashcode; + current_hashcode = __shfl_sync(FULL_MASK, current_hashcode, hash_f_offset); + int hashtable_idx = (batch_idx * num_hash_f + (hash_f_start + hash_f_offset)) * hashtable_capacity + current_hashcode; + warp_value = warp_value + hashtable_value[hashtable_idx * WARP_SIZE + warp_thread_idx]; + } + } + cumulation_value[batch_idx__query_idx * value_dim + offset_warp + warp_thread_idx] = warp_value / float(num_hash_f); + } else { + float warp_value = 0; + int warp_hashcode = 0; + if (warp_thread_idx < num_hash_f) { + warp_hashcode = query_hash_code[batch_idx__query_idx * num_hash_f + warp_thread_idx]; + } + for (int hash_f_idx = 0; hash_f_idx < num_hash_f; hash_f_idx++) { + int current_hashcode = warp_hashcode; + current_hashcode = __shfl_sync(FULL_MASK, current_hashcode, hash_f_idx); + int hashtable_idx = (batch_idx * num_hash_f + hash_f_idx) * hashtable_capacity + current_hashcode; + warp_value = warp_value + hashtable_value[hashtable_idx * WARP_SIZE + warp_thread_idx]; + } + cumulation_value[batch_idx__query_idx * value_dim + offset_warp + warp_thread_idx] = warp_value / float(num_hash_f); + } + +} + +__global__ void lsh_weighted_cumulation_ver1_step1_cuda_kernel( + int *key_mask, // [batch_size, num_key] + int *key_hash_code, // [batch_size, num_key, num_hash_f] + float *key_weight, // [batch_size, num_key, weight_dim] + float *value, // [batch_size, num_key, value_dim] + float *hashtable_value, // [batch_size, num_hash_f, hashtable_capacity, WARP_SIZE] + int batch_size, + int num_hash_f, + int hashtable_capacity, + int num_key, + int value_dim, + int weight_dim, + int offset_warp, + int weight_idx +) { + + int warp_thread_idx = threadIdx.x; + + int batch_idx = blockIdx.y; + int key_idx = blockIdx.x * blockDim.y + threadIdx.y; + + int batch_idx__key_idx = batch_idx * num_key + key_idx; + if (key_mask[batch_idx__key_idx] == 0) { + return; + } + + if (num_hash_f > WARP_SIZE) { + float warp_value = key_weight[batch_idx__key_idx * weight_dim + weight_idx] * value[batch_idx__key_idx * value_dim + offset_warp + warp_thread_idx]; + for (int hash_f_start = 0; hash_f_start < num_hash_f; hash_f_start = hash_f_start + WARP_SIZE) { + int warp_hashcode = key_hash_code[batch_idx__key_idx * num_hash_f + hash_f_start + warp_thread_idx]; + #pragma unroll + for (int hash_f_offset = 0; hash_f_offset < WARP_SIZE; hash_f_offset++) { + int current_hashcode = warp_hashcode; + current_hashcode = __shfl_sync(FULL_MASK, current_hashcode, hash_f_offset); + int hashtable_idx = (batch_idx * num_hash_f + (hash_f_start + hash_f_offset)) * hashtable_capacity + current_hashcode; + atomicAdd(&hashtable_value[hashtable_idx * WARP_SIZE + warp_thread_idx], warp_value); + } + } + } else { + float warp_value = key_weight[batch_idx__key_idx * weight_dim + weight_idx] * value[batch_idx__key_idx * value_dim + offset_warp + warp_thread_idx]; + int warp_hashcode = 0; + if (warp_thread_idx < num_hash_f) { + warp_hashcode = key_hash_code[batch_idx__key_idx * num_hash_f + warp_thread_idx]; + } + for (int hash_f_idx = 0; hash_f_idx < num_hash_f; hash_f_idx++) { + int current_hashcode = warp_hashcode; + current_hashcode = __shfl_sync(FULL_MASK, current_hashcode, hash_f_idx); + int hashtable_idx = (batch_idx * num_hash_f + hash_f_idx) * hashtable_capacity + current_hashcode; + atomicAdd(&hashtable_value[hashtable_idx * WARP_SIZE + warp_thread_idx], warp_value); + } + } + +} + +__global__ void lsh_weighted_cumulation_ver1_step2_cuda_kernel( + int *query_mask, // [batch_size, num_query] + int *query_hash_code, // [batch_size, num_query, num_hash_f] + float *query_weight, // [batch_size, num_query, weight_dim] + float *hashtable_value, // [batch_size, num_hash_f, hashtable_capacity, WARP_SIZE] + float *cumulation_value, // [batch_size, num_query, value_dim] + int batch_size, + int num_hash_f, + int hashtable_capacity, + int num_query, + int value_dim, + int weight_dim, + int offset_warp, + int weight_idx +) { + + int warp_thread_idx = threadIdx.x; + + int batch_idx = blockIdx.y; + int query_idx = blockIdx.x * blockDim.y + threadIdx.y; + + int batch_idx__query_idx = batch_idx * num_query + query_idx; + if (query_mask[batch_idx__query_idx] == 0) { + return; + } + + if (num_hash_f > WARP_SIZE) { + float warp_value = 0; + for (int hash_f_start = 0; hash_f_start < num_hash_f; hash_f_start = hash_f_start + WARP_SIZE) { + int warp_hashcode = query_hash_code[batch_idx__query_idx * num_hash_f + hash_f_start + warp_thread_idx]; + #pragma unroll + for (int hash_f_offset = 0; hash_f_offset < WARP_SIZE; hash_f_offset++) { + int current_hashcode = warp_hashcode; + current_hashcode = __shfl_sync(FULL_MASK, current_hashcode, hash_f_offset); + int hashtable_idx = (batch_idx * num_hash_f + (hash_f_start + hash_f_offset)) * hashtable_capacity + current_hashcode; + warp_value = warp_value + hashtable_value[hashtable_idx * WARP_SIZE + warp_thread_idx]; + } + } + float warp_weight = query_weight[batch_idx__query_idx * weight_dim + weight_idx]; + cumulation_value[batch_idx__query_idx * value_dim + offset_warp + warp_thread_idx] += warp_weight * warp_value / float(num_hash_f); + } else { + float warp_value = 0; + int warp_hashcode = 0; + if (warp_thread_idx < num_hash_f) { + warp_hashcode = query_hash_code[batch_idx__query_idx * num_hash_f + warp_thread_idx]; + } + for (int hash_f_idx = 0; hash_f_idx < num_hash_f; hash_f_idx++) { + int current_hashcode = warp_hashcode; + current_hashcode = __shfl_sync(FULL_MASK, current_hashcode, hash_f_idx); + int hashtable_idx = (batch_idx * num_hash_f + hash_f_idx) * hashtable_capacity + current_hashcode; + warp_value = warp_value + hashtable_value[hashtable_idx * WARP_SIZE + warp_thread_idx]; + } + float warp_weight = query_weight[batch_idx__query_idx * weight_dim + weight_idx]; + cumulation_value[batch_idx__query_idx * value_dim + offset_warp + warp_thread_idx] += warp_weight * warp_value / float(num_hash_f); + } + +} + +__global__ void count_sort_step1_cuda_kernel( + int *key_mask, // [batch_size, num_key] + int *key_hash_code, // [batch_size, num_key, num_hash_f] + int *count_sort_table, // [batch_size, num_hash_f, hashtable_capacity] + int batch_size, + int num_hash_f, + int hashtable_capacity, + int num_key +) { + + int batch_idx = blockIdx.y; + int key_idx = blockIdx.x * blockDim.y + threadIdx.y; + int hash_f_idx = threadIdx.x; + + int batch_idx__key_idx = batch_idx * num_key + key_idx; + if (key_mask[batch_idx__key_idx] == 0) { + return; + } + + int hash_code = key_hash_code[batch_idx__key_idx * num_hash_f + hash_f_idx]; + atomicAdd(&count_sort_table[(batch_idx * num_hash_f + hash_f_idx) * hashtable_capacity + hash_code], 1); + +} + +__global__ void count_sort_step2_cuda_kernel( + int *count_sort_table, // [batch_size, num_hash_f, hashtable_capacity] + int batch_size, + int num_hash_f, + int hashtable_capacity +) { + + int batch_idx = blockIdx.y; + int hash_f_idx = blockIdx.x; + + int num_threads = blockDim.x; + int thread_id = threadIdx.x; + + int batch_idx__hash_f_idx = batch_idx * num_hash_f + hash_f_idx; + + extern __shared__ float buffer[]; + int *table_buffer = (int*)buffer; + + if (thread_id == 0) { + table_buffer[0] = 0; + } + copy_data(&count_sort_table[batch_idx__hash_f_idx * hashtable_capacity], &table_buffer[1], hashtable_capacity - 1, num_threads, thread_id); + + for (int table_idx_start = 0; table_idx_start < hashtable_capacity; table_idx_start = table_idx_start + num_threads) { + int thread_value = table_buffer[table_idx_start + thread_id]; + int next_thread_value = 0; + for (int offset = 1; offset < WARP_SIZE; offset = offset << 1) { + next_thread_value = __shfl_up_sync(FULL_MASK, thread_value, offset); + if (thread_id % WARP_SIZE >= offset) { + thread_value = thread_value + next_thread_value; + } + } + table_buffer[table_idx_start + thread_id] = thread_value; + } + __syncthreads(); + + if (hashtable_capacity > WARP_SIZE) { + if (thread_id < WARP_SIZE) { + for (int table_idx_start = WARP_SIZE; table_idx_start < hashtable_capacity; table_idx_start = table_idx_start + WARP_SIZE) { + table_buffer[table_idx_start + thread_id] += table_buffer[table_idx_start - 1]; + } + } + } + + copy_data(table_buffer, &count_sort_table[batch_idx__hash_f_idx * hashtable_capacity], hashtable_capacity, num_threads, thread_id); + +} + + +__global__ void count_sort_step3_cuda_kernel( + int *key_mask, // [batch_size, num_key] + int *key_hash_code, // [batch_size, num_key, num_hash_f] + int *count_sort_table, // [batch_size, num_hash_f, hashtable_capacity] + int *key_sorted_idxes, // [batch_size, num_hash_f, num_key] + int batch_size, + int num_hash_f, + int hashtable_capacity, + int num_key +) { + + int batch_idx = blockIdx.y; + int key_idx = blockIdx.x * blockDim.y + threadIdx.y; + int hash_f_idx = threadIdx.x; + + int batch_idx__key_idx = batch_idx * num_key + key_idx; + if (key_mask[batch_idx__key_idx] == 0) { + return; + } + + int batch_idx__hash_f_idx = batch_idx * num_hash_f + hash_f_idx; + + int hash_code = key_hash_code[batch_idx__key_idx * num_hash_f + hash_f_idx]; + int sort_idx = atomicAdd(&count_sort_table[batch_idx__hash_f_idx * hashtable_capacity + hash_code], 1); + key_sorted_idxes[batch_idx__hash_f_idx * num_key + sort_idx] = key_idx; + +} + +__global__ void extract_query_info_cuda_kernel( + int *query_mask, // [batch_size, num_query] + int *query_hash_code, // [batch_size, num_query, num_hash_f] + int *count_sort_table, // [batch_size, num_hash_f, hashtable_capacity] + int *query_info, // [batch_size, num_query, 2, num_hash_f] + int batch_size, + int num_hash_f, + int hashtable_capacity, + int num_query +) { + + int batch_idx = blockIdx.y; + int query_idx = blockIdx.x * blockDim.y + threadIdx.y; + int hash_f_idx = threadIdx.x; + + int batch_idx__query_idx = batch_idx * num_query + query_idx; + if (query_mask[batch_idx__query_idx] == 0) { + return; + } + + int hash_code = query_hash_code[batch_idx__query_idx * num_hash_f + hash_f_idx]; + int batch_idx__hash_f_idx__hash_code = (batch_idx * num_hash_f + hash_f_idx) * hashtable_capacity + hash_code; + + int key_offset = select(hash_code == 0, 0, count_sort_table[batch_idx__hash_f_idx__hash_code - 1]); + int key_count = count_sort_table[batch_idx__hash_f_idx__hash_code] - key_offset; + + query_info[batch_idx__query_idx * 2 * num_hash_f + hash_f_idx] = key_offset; + query_info[(batch_idx__query_idx * 2 + 1) * num_hash_f + hash_f_idx] = key_count; + +} + +__global__ void lsh_weighted_cumulation_ver2_step2_cuda_kernel( + int *query_mask, // [batch_size, num_query] + int *query_info, // [batch_size, num_query, 2, num_hash_f] + int *key_sorted_idxes, // [batch_size, num_hash_f, num_key] + float *query_weight, // [batch_size, num_query, weight_dim] + float *key_weight, // [batch_size, num_key, weight_dim] + float *value, // [batch_size, num_key, value_dim] + float *cumulation_value, // [batch_size, num_query, value_dim] + int batch_size, + int num_hash_f, + int num_query, + int num_key, + int value_dim, + int weight_dim +) { + + int batch_idx = blockIdx.z; + int hash_f_idx = blockIdx.y; + int query_idx = blockIdx.x; + + int num_threads = blockDim.y * blockDim.x; + int thread_id = threadIdx.y * blockDim.x + threadIdx.x; + + int num_warps = blockDim.y; + int warp_idx = threadIdx.y; + int warp_thread_idx = threadIdx.x; + + int batch_idx__query_idx = batch_idx * num_query + query_idx; + if (query_mask[batch_idx__query_idx] == 0) { + return; + } + + int key_offset = query_info[batch_idx__query_idx * 2 * num_hash_f + hash_f_idx]; + int key_count = query_info[(batch_idx__query_idx * 2 + 1) * num_hash_f + hash_f_idx]; + + if (key_count == 0) { + return; + } + + extern __shared__ float buffer[]; + + if (key_count == 1) { + if (warp_idx == 0) { + int key_idx = key_sorted_idxes[(batch_idx * num_hash_f + hash_f_idx) * num_key + key_offset]; + int batch_idx__key_idx = batch_idx * num_key + key_idx; + float weight = 0; + for (int weight_offset = 0; weight_offset < weight_dim; weight_offset = weight_offset + WARP_SIZE) { + int weight_dim_idx = weight_offset + warp_thread_idx; + float val = query_weight[batch_idx__query_idx * weight_dim + weight_dim_idx] * key_weight[batch_idx__key_idx * weight_dim + weight_dim_idx]; + #pragma unroll + for (int offset = 1; offset < WARP_SIZE; offset = offset << 1) { + val += __shfl_xor_sync(FULL_MASK, val, offset); + } + weight = weight + val; + } + weight = weight / float(num_hash_f); + for (int value_offset = 0; value_offset < value_dim; value_offset = value_offset + WARP_SIZE) { + int value_dim_idx = value_offset + warp_thread_idx; + float val = value[batch_idx__key_idx * value_dim + value_dim_idx]; + atomicAdd(&cumulation_value[batch_idx__query_idx * value_dim + value_dim_idx], weight * val); + } + } + } else { + float *weight_buffer = buffer; + int *key_idxes_buffer = (int*)&buffer[weight_dim]; + + copy_data_nonblocking(&query_weight[batch_idx__query_idx * weight_dim], weight_buffer, weight_dim, num_threads, thread_id); + + while (key_count > 0) { + int work_size = min(WARP_SIZE, key_count); + copy_data_nonblocking(&key_sorted_idxes[(batch_idx * num_hash_f + hash_f_idx) * num_key + key_offset], key_idxes_buffer, work_size, num_threads, thread_id); + __syncthreads(); + for (int work_offset = 0; work_offset < WARP_SIZE; work_offset = work_offset + num_warps) { + int work_idx = work_offset + warp_idx; + if (work_idx < key_count) { + int key_idx = key_idxes_buffer[work_idx]; + int batch_idx__key_idx = batch_idx * num_key + key_idx; + float weight = 0; + for (int weight_offset = 0; weight_offset < weight_dim; weight_offset = weight_offset + WARP_SIZE) { + int weight_dim_idx = weight_offset + warp_thread_idx; + float val = weight_buffer[weight_dim_idx] * key_weight[batch_idx__key_idx * weight_dim + weight_dim_idx]; + #pragma unroll + for (int offset = 1; offset < WARP_SIZE; offset = offset << 1) { + val += __shfl_xor_sync(FULL_MASK, val, offset); + } + weight = weight + val; + } + weight = weight / float(num_hash_f); + for (int value_offset = 0; value_offset < value_dim; value_offset = value_offset + WARP_SIZE) { + int value_dim_idx = value_offset + warp_thread_idx; + float val = value[batch_idx__key_idx * value_dim + value_dim_idx]; + atomicAdd(&cumulation_value[batch_idx__query_idx * value_dim + value_dim_idx], weight * val); + } + } + } + key_count = key_count - work_size; + key_offset = key_offset + work_size; + } + } + +} + +__global__ void lsh_weighted_cumulation_ver3_step2_cuda_kernel( + int *query_sorted_idxes, // [batch_size, num_hash_f, num_query] + int *key_mask, // [batch_size, num_key] + int *key_info, // [batch_size, num_key, 2, num_hash_f] + float *query_weight, // [batch_size, num_query, weight_dim] + float *key_weight, // [batch_size, num_key, weight_dim] + float *value, // [batch_size, num_key, value_dim] + float *cumulation_value, // [batch_size, num_query, value_dim] + int batch_size, + int num_hash_f, + int num_query, + int num_key, + int value_dim, + int weight_dim +) { + + int batch_idx = blockIdx.z; + int hash_f_idx = blockIdx.y; + int key_idx = blockIdx.x; + + int num_threads = blockDim.y * blockDim.x; + int thread_id = threadIdx.y * blockDim.x + threadIdx.x; + + int num_warps = blockDim.y; + int warp_idx = threadIdx.y; + int warp_thread_idx = threadIdx.x; + + int batch_idx__key_idx = batch_idx * num_key + key_idx; + if (key_mask[batch_idx__key_idx] == 0) { + return; + } + + int query_offset = key_info[batch_idx__key_idx * 2 * num_hash_f + hash_f_idx]; + int query_count = key_info[(batch_idx__key_idx * 2 + 1) * num_hash_f + hash_f_idx]; + + if (query_count == 0) { + return; + } + + extern __shared__ float buffer[]; + + if (query_count == 1) { + if (warp_idx == 0) { + int query_idx = query_sorted_idxes[(batch_idx * num_hash_f + hash_f_idx) * num_query + query_offset]; + int batch_idx__query_idx = batch_idx * num_query + query_idx; + float weight = 0; + for (int weight_offset = 0; weight_offset < weight_dim; weight_offset = weight_offset + WARP_SIZE) { + int weight_dim_idx = weight_offset + warp_thread_idx; + float val = key_weight[batch_idx__key_idx * weight_dim + weight_dim_idx] * query_weight[batch_idx__query_idx * weight_dim + weight_dim_idx]; + #pragma unroll + for (int offset = 1; offset < WARP_SIZE; offset = offset << 1) { + val += __shfl_xor_sync(FULL_MASK, val, offset); + } + weight = weight + val; + } + weight = weight / float(num_hash_f); + for (int value_offset = 0; value_offset < value_dim; value_offset = value_offset + WARP_SIZE) { + int value_dim_idx = value_offset + warp_thread_idx; + float val = value[batch_idx__key_idx * value_dim + value_dim_idx]; + atomicAdd(&cumulation_value[batch_idx__query_idx * value_dim + value_dim_idx], weight * val); + } + } + } else { + float *weight_buffer = buffer; + float *value_buffer = &buffer[weight_dim]; + int *query_idxes_buffer = (int*)&buffer[weight_dim + value_dim]; + + copy_data_nonblocking(&key_weight[batch_idx__key_idx * weight_dim], weight_buffer, weight_dim, num_threads, thread_id); + copy_data_nonblocking(&value[batch_idx__key_idx * value_dim], value_buffer, value_dim, num_threads, thread_id); + + while (query_count > 0) { + int work_size = min(WARP_SIZE, query_count); + copy_data_nonblocking(&query_sorted_idxes[(batch_idx * num_hash_f + hash_f_idx) * num_query + query_offset], query_idxes_buffer, work_size, num_threads, thread_id); + __syncthreads(); + for (int work_offset = 0; work_offset < WARP_SIZE; work_offset = work_offset + num_warps) { + int work_idx = work_offset + warp_idx; + if (work_idx < query_count) { + int query_idx = query_idxes_buffer[work_idx]; + int batch_idx__query_idx = batch_idx * num_query + query_idx; + float weight = 0; + for (int weight_offset = 0; weight_offset < weight_dim; weight_offset = weight_offset + WARP_SIZE) { + int weight_dim_idx = weight_offset + warp_thread_idx; + float val = weight_buffer[weight_dim_idx] * query_weight[batch_idx__query_idx * weight_dim + weight_dim_idx]; + #pragma unroll + for (int offset = 1; offset < WARP_SIZE; offset = offset << 1) { + val += __shfl_xor_sync(FULL_MASK, val, offset); + } + weight = weight + val; + } + weight = weight / float(num_hash_f); + for (int value_offset = 0; value_offset < value_dim; value_offset = value_offset + WARP_SIZE) { + int value_dim_idx = value_offset + warp_thread_idx; + float val = value_buffer[value_dim_idx]; + atomicAdd(&cumulation_value[batch_idx__query_idx * value_dim + value_dim_idx], weight * val); + } + } + } + query_count = query_count - work_size; + query_offset = query_offset + work_size; + } + } + +} + +__global__ void lsh_weighted_cumulation_ver4_step2_cuda_kernel( + int *query_sorted_idxes, // [batch_size, num_hash_f, num_query] + int *key_mask, // [batch_size, num_key] + int *key_info, // [batch_size, num_key, 2, num_hash_f] + float *query_weight, // [batch_size, num_query, weight_dim] + float *key_weight, // [batch_size, num_key, weight_dim] + float *value, // [batch_size, num_key, value_dim] + float *cumulation_value, // [batch_size, num_query, value_dim] + int batch_size, + int num_hash_f, + int num_query, + int num_key, + int value_dim, + int weight_dim +) { + + int batch_idx = blockIdx.y; + int key_idx = blockIdx.x; + + int num_threads = blockDim.y * blockDim.x; + int thread_id = threadIdx.y * blockDim.x + threadIdx.x; + + int num_warps = blockDim.y; + int warp_idx = threadIdx.y; + int warp_thread_idx = threadIdx.x; + + int batch_idx__key_idx = batch_idx * num_key + key_idx; + if (key_mask[batch_idx__key_idx] == 0) { + return; + } + + extern __shared__ float buffer[]; + float *weight_buffer = buffer; + float *value_buffer = &buffer[weight_dim]; + int *key_info_buffer = (int*)&buffer[weight_dim + value_dim]; + + copy_data_nonblocking(&key_weight[batch_idx__key_idx * weight_dim], weight_buffer, weight_dim, num_threads, thread_id); + copy_data_nonblocking(&value[batch_idx__key_idx * value_dim], value_buffer, value_dim, num_threads, thread_id); + copy_data_nonblocking(&key_info[batch_idx__key_idx * 2 * num_hash_f], key_info_buffer, 2 * num_hash_f, num_threads, thread_id); + + int *query_offset_buffer = key_info_buffer; + int *query_count_buffer = &key_info_buffer[num_hash_f]; + + const int hashtable_size = 1024 + OPTIMAL_THREADS_PER_BLOCK; + __shared__ int hashtable_query[hashtable_size]; + __shared__ int hashtable_count[hashtable_size]; + __shared__ int inserted_query[hashtable_size]; + __shared__ int query_counter[1]; + + int hash_f_idx_base = 0; + + while (true) { + + init_buffer_nonblocking(EMPTY_VALUE, hashtable_query, hashtable_size, num_threads, thread_id); + init_buffer_nonblocking(0, hashtable_count, hashtable_size, num_threads, thread_id); + init_buffer_nonblocking(EMPTY_VALUE, inserted_query, hashtable_size, num_threads, thread_id); + init_buffer_nonblocking(0, query_counter, 1, num_threads, thread_id); + __syncthreads(); + + while (hash_f_idx_base < num_hash_f) { + + int hash_f_idx = hash_f_idx_base + warp_idx; + int batch_idx__hash_f_idx = batch_idx * num_hash_f + hash_f_idx; + + int stop_flag = 0; + + int query_offset = query_offset_buffer[hash_f_idx]; + int query_count = query_count_buffer[hash_f_idx]; + + while (query_count > 0) { + + int work_size = min(query_count, WARP_SIZE); + + // try inserting query to set and check whether the query is new + int found_new_query = 0; + int query_idx = -1; + if (warp_thread_idx < work_size) { + query_idx = query_sorted_idxes[batch_idx__hash_f_idx * num_query + query_offset + warp_thread_idx]; + int slot = set_insert(hashtable_query, hashtable_size, query_idx); + if (slot >= 0) { + found_new_query = atomicAdd(&hashtable_count[slot], 1) == 0; + } + } + + // compute cumulative offset + int position_offset = found_new_query; + int next_position_offset = 0; + #pragma unroll + for (int offset = 1; offset < WARP_SIZE; offset = offset << 1) { + next_position_offset = __shfl_up_sync(FULL_MASK, position_offset, offset); + if (thread_id % WARP_SIZE >= offset) { + position_offset = position_offset + next_position_offset; + } + } + + // get the inserted query list end index + int inserted_query_base = 0; + if (thread_id % WARP_SIZE == WARP_SIZE - 1) { + inserted_query_base = atomicAdd(query_counter, position_offset); + } + inserted_query_base = __shfl_sync(FULL_MASK, inserted_query_base, WARP_SIZE - 1); + + // insert new queries to list + int insert_idx = inserted_query_base + position_offset - 1; + if (found_new_query) { + inserted_query[insert_idx] = query_idx; + } + + // remove inserted queries from list + query_offset_buffer[hash_f_idx] += work_size; + query_count_buffer[hash_f_idx] -= work_size; + query_offset += work_size; + query_count -= work_size; + + // if list is almost full, stop inserting + if (inserted_query_base + OPTIMAL_THREADS_PER_BLOCK > hashtable_size) { + stop_flag = 1; + break; + } + + } + + if (stop_flag) { + break; + } + + hash_f_idx_base = hash_f_idx_base + num_warps; + + } + + __syncthreads(); + + int num_distint_query = query_counter[0]; + + if (num_distint_query > 0) { + for (int idx_base = 0; idx_base < num_distint_query; idx_base = idx_base + num_warps) { + int idx = idx_base + warp_idx; + if (idx < num_distint_query) { + int query_idx = inserted_query[idx]; + int batch_idx__query_idx = batch_idx * num_query + query_idx; + + int slot = set_lookup(hashtable_query, hashtable_size, query_idx); + int duplicate_count = hashtable_count[slot]; + + float weight = 0; + for (int weight_idx_base = 0; weight_idx_base < weight_dim; weight_idx_base = weight_idx_base + WARP_SIZE) { + int weight_dim_idx = weight_idx_base + warp_thread_idx; + float val = weight_buffer[weight_dim_idx] * query_weight[batch_idx__query_idx * weight_dim + weight_dim_idx]; + #pragma unroll + for (int offset = 1; offset < WARP_SIZE; offset = offset << 1) { + val += __shfl_xor_sync(FULL_MASK, val, offset); + } + weight = weight + val; + } + + weight = (float)duplicate_count * weight / float(num_hash_f); + + for (int value_idx_base = 0; value_idx_base < value_dim; value_idx_base = value_idx_base + WARP_SIZE) { + int value_dim_idx = value_idx_base + warp_thread_idx; + float val = value_buffer[value_dim_idx]; + atomicAdd(&cumulation_value[batch_idx__query_idx * value_dim + value_dim_idx], weight * val); + } + } + } + } else { + + // all computation is completed if num_distint_query == 0 + break; + + } + + __syncthreads(); + + } + +} diff --git a/lib/python3.12/site-packages/transformers/kernels/yoso/fast_lsh_cumulation_cuda.h b/lib/python3.12/site-packages/transformers/kernels/yoso/fast_lsh_cumulation_cuda.h new file mode 100644 index 0000000000000000000000000000000000000000..b2adc0f735358d0fcb6a056e7d19ba745977e129 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/yoso/fast_lsh_cumulation_cuda.h @@ -0,0 +1,157 @@ +__global__ void fast_hash_ver1_cuda_kernel( + int *mask, // [batch_size, num_vector] + float *vector, // [batch_size, num_vector, vector_dim] + int *Dmat, // [3, num_part, vector_dim] + int *hash_code, // [batch_size, num_vector, num_hash_f] + int batch_size, + int num_vector, + int vector_dim, + int num_part, + int num_hash_f, + int hash_code_len +); + +__global__ void lsh_cumulation_ver1_step1_cuda_kernel( + int *key_mask, // [batch_size, num_key] + int *key_hash_code, // [batch_size, num_key, num_hash_f] + float *value, // [batch_size, num_key, value_dim] + float *hashtable_value, // [batch_size, num_hash_f, hashtable_capacity, value_dim] + int batch_size, + int num_hash_f, + int hashtable_capacity, + int num_key, + int value_dim, + int offset_warp +); + +__global__ void lsh_cumulation_ver1_step2_cuda_kernel( + int *query_mask, // [batch_size, num_query] + int *query_hash_code, // [batch_size, num_query, num_hash_f] + float *hashtable_value, // [batch_size, num_hash_f, hashtable_capacity, value_dim] + float *cumulation_value, // [batch_size, num_query, value_dim] + int batch_size, + int num_hash_f, + int hashtable_capacity, + int num_query, + int value_dim, + int offset_warp +); + +__global__ void lsh_weighted_cumulation_ver1_step1_cuda_kernel( + int *key_mask, // [batch_size, num_key] + int *key_hash_code, // [batch_size, num_key, num_hash_f] + float *key_weight, // [batch_size, num_key, weight_dim] + float *value, // [batch_size, num_key, value_dim] + float *hashtable_value, // [batch_size, num_hash_f, hashtable_capacity, WARP_SIZE] + int batch_size, + int num_hash_f, + int hashtable_capacity, + int num_key, + int value_dim, + int weight_dim, + int offset_warp, + int weight_idx +); + +__global__ void lsh_weighted_cumulation_ver1_step2_cuda_kernel( + int *query_mask, // [batch_size, num_query] + int *query_hash_code, // [batch_size, num_query, num_hash_f] + float *query_weight, // [batch_size, num_query, weight_dim] + float *hashtable_value, // [batch_size, num_hash_f, hashtable_capacity, WARP_SIZE] + float *cumulation_value, // [batch_size, num_query, value_dim] + int batch_size, + int num_hash_f, + int hashtable_capacity, + int num_query, + int value_dim, + int weight_dim, + int offset_warp, + int weight_idx +); + +__global__ void count_sort_step1_cuda_kernel( + int *key_mask, // [batch_size, num_key] + int *key_hash_code, // [batch_size, num_key, num_hash_f] + int *count_sort_table, // [batch_size, num_hash_f, hashtable_capacity] + int batch_size, + int num_hash_f, + int hashtable_capacity, + int num_key +); + +__global__ void count_sort_step2_cuda_kernel( + int *count_sort_table, // [batch_size, num_hash_f, hashtable_capacity] + int batch_size, + int num_hash_f, + int hashtable_capacity +); + +__global__ void count_sort_step3_cuda_kernel( + int *key_mask, // [batch_size, num_key] + int *key_hash_code, // [batch_size, num_key, num_hash_f] + int *count_sort_table, // [batch_size, num_hash_f, hashtable_capacity] + int *key_sorted_idxes, // [batch_size, num_hash_f, num_key] + int batch_size, + int num_hash_f, + int hashtable_capacity, + int num_key +); + +__global__ void extract_query_info_cuda_kernel( + int *query_mask, // [batch_size, num_query] + int *query_hash_code, // [batch_size, num_query, num_hash_f] + int *count_sort_table, // [batch_size, num_hash_f, hashtable_capacity] + int *query_info, // [batch_size, num_query, 2, num_hash_f] + int batch_size, + int num_hash_f, + int hashtable_capacity, + int num_query +); + +__global__ void lsh_weighted_cumulation_ver2_step2_cuda_kernel( + int *query_mask, // [batch_size, num_query] + int *query_info, // [batch_size, num_query, 2, num_hash_f] + int *key_sorted_idxes, // [batch_size, num_hash_f, num_key] + float *query_weight, // [batch_size, num_query, weight_dim] + float *key_weight, // [batch_size, num_key, weight_dim] + float *value, // [batch_size, num_key, value_dim] + float *cumulation_value, // [batch_size, num_query, value_dim] + int batch_size, + int num_hash_f, + int num_query, + int num_key, + int value_dim, + int weight_dim +); + +__global__ void lsh_weighted_cumulation_ver3_step2_cuda_kernel( + int *query_sorted_idxes, // [batch_size, num_hash_f, num_query] + int *key_mask, // [batch_size, num_key] + int *key_info, // [batch_size, num_key, 2, num_hash_f] + float *query_weight, // [batch_size, num_query, weight_dim] + float *key_weight, // [batch_size, num_key, weight_dim] + float *value, // [batch_size, num_key, value_dim] + float *cumulation_value, // [batch_size, num_query, value_dim] + int batch_size, + int num_hash_f, + int num_query, + int num_key, + int value_dim, + int weight_dim +); + +__global__ void lsh_weighted_cumulation_ver4_step2_cuda_kernel( + int *query_sorted_idxes, // [batch_size, num_hash_f, num_query] + int *key_mask, // [batch_size, num_key] + int *key_info, // [batch_size, num_key, 2, num_hash_f] + float *query_weight, // [batch_size, num_query, weight_dim] + float *key_weight, // [batch_size, num_key, weight_dim] + float *value, // [batch_size, num_key, value_dim] + float *cumulation_value, // [batch_size, num_query, value_dim] + int batch_size, + int num_hash_f, + int num_query, + int num_key, + int value_dim, + int weight_dim +); diff --git a/lib/python3.12/site-packages/transformers/kernels/yoso/fast_lsh_cumulation_torch.cpp b/lib/python3.12/site-packages/transformers/kernels/yoso/fast_lsh_cumulation_torch.cpp new file mode 100644 index 0000000000000000000000000000000000000000..e150a2be604b28f600ab345a8cc9e97819cca416 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/kernels/yoso/fast_lsh_cumulation_torch.cpp @@ -0,0 +1,128 @@ +#include +#include +#include "fast_lsh_cumulation.h" +#include "common_cuda.h" +#include + +std::vector fast_hash( + at::Tensor query_mask, + at::Tensor query_vector, + at::Tensor key_mask, + at::Tensor key_vector, + int num_hash_f, + int hash_code_len, + bool use_cuda, + int version +) { + return fast_hash_ver1_kernel( + query_mask, + query_vector, + key_mask, + key_vector, + num_hash_f, + hash_code_len, + use_cuda + ); +} + +at::Tensor lsh_cumulation( + at::Tensor query_mask, // [batch_size, num_query] + at::Tensor query_hash_code, // [batch_size, num_query, num_hash_f] + at::Tensor key_mask, // [batch_size, num_key] + at::Tensor key_hash_code, // [batch_size, num_key, num_hash_f] + at::Tensor value, // [batch_size, num_key, value_dim] + int hashtable_capacity, + bool use_cuda, + int version +) { + return lsh_cumulation_ver1_kernel( + query_mask, + query_hash_code, + key_mask, + key_hash_code, + value, + hashtable_capacity, + use_cuda + ); +} + +at::Tensor lsh_weighted_cumulation( + at::Tensor query_mask, // [batch_size, num_query] + at::Tensor query_hash_code, // [batch_size, num_query, num_hash_f] + at::Tensor query_weight, // [batch_size, num_query, weight_dim] + at::Tensor key_mask, // [batch_size, num_key] + at::Tensor key_hash_code, // [batch_size, num_key, num_hash_f] + at::Tensor key_weight, // [batch_size, num_key, weight_dim] + at::Tensor value, // [batch_size, num_key, value_dim] + int hashtable_capacity, + bool use_cuda, + int version +) { + if (version == 1) { + return lsh_weighted_cumulation_ver1_kernel( + query_mask, + query_hash_code, + query_weight, + key_mask, + key_hash_code, + key_weight, + value, + hashtable_capacity, + use_cuda + ); + } else if (version == 2) { + return lsh_weighted_cumulation_ver2_kernel( + query_mask, + query_hash_code, + query_weight, + key_mask, + key_hash_code, + key_weight, + value, + hashtable_capacity, + use_cuda + ); + } else if (version == 3) { + return lsh_weighted_cumulation_ver3_kernel( + query_mask, + query_hash_code, + query_weight, + key_mask, + key_hash_code, + key_weight, + value, + hashtable_capacity, + use_cuda + ); + } else if (version == 4) { + return lsh_weighted_cumulation_ver4_kernel( + query_mask, + query_hash_code, + query_weight, + key_mask, + key_hash_code, + key_weight, + value, + hashtable_capacity, + use_cuda + ); + } else { + return lsh_weighted_cumulation_ver3_kernel( + query_mask, + query_hash_code, + query_weight, + key_mask, + key_hash_code, + key_weight, + value, + hashtable_capacity, + use_cuda + ); + } +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("fast_hash", &fast_hash, "Fast Hash (CUDA)"); + m.def("lsh_cumulation", &lsh_cumulation, "LSH Cumulation (CUDA)"); + m.def("lsh_weighted_cumulation", &lsh_weighted_cumulation, "LSH Weighted Cumulation (CUDA)"); +} diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..43e6f3a89b471f062b07cbd6fddd9f1398175df9 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/audio_classification.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/audio_classification.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..01f73abb06bfb50e248d09010d3f03eb495bcf5c Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/audio_classification.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/audio_utils.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/audio_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5c576b9768a33c804173549f0b70b97a9676992b Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/audio_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/automatic_speech_recognition.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/automatic_speech_recognition.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2570ba30ebce32576697c585c4114b8a32d0cc54 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/automatic_speech_recognition.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/base.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7a9fab2c8264f76dff22817a6035177d6fd84914 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/depth_estimation.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/depth_estimation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5e1944bcbf0e94c5cf6cf7e1476851c2effe1891 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/depth_estimation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/document_question_answering.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/document_question_answering.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..21c13c0f58394c0a53f8d3f4b3506613ed40a412 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/document_question_answering.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/feature_extraction.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/feature_extraction.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..33e30c37d8b45707d54ea795dd0ca197298b736c Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/feature_extraction.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/fill_mask.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/fill_mask.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b763301395fdf38e5023964ae2ba1b78ae238c0e Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/fill_mask.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/image_classification.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/image_classification.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a19544e20626e71cbb12bbf134d48474a06d69cb Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/image_classification.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/image_feature_extraction.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/image_feature_extraction.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7f91a0ad62db67ee699282d86cd44f80568b67f1 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/image_feature_extraction.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/image_segmentation.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/image_segmentation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..50b65bdeb8a361051d9fab5892dc76dac67c81ce Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/image_segmentation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/image_text_to_text.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/image_text_to_text.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e7dfb9ccf6e72fa8664d77e85d8e60cc81ca792e Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/image_text_to_text.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/image_to_image.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/image_to_image.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5b70c3aeccfc328228c47cd19ac28b8c8b6e57b3 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/image_to_image.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/image_to_text.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/image_to_text.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b95325d42bb4c0cea9977be4f1f88ccf143e882c Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/image_to_text.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/mask_generation.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/mask_generation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f611e470933c2b13952e1c8c71621d2341dab1db Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/mask_generation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/object_detection.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/object_detection.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..36e2a96f7c1953b6a95e40f6d6128edcd2666341 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/object_detection.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/pt_utils.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/pt_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7974db4e256c4934858523db07fa2a558edaaa9e Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/pt_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/question_answering.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/question_answering.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..73ec6eb938f994b79a8b3ac2e2687872a4f119ef Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/question_answering.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/table_question_answering.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/table_question_answering.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c6563b6fa34962fbad060e843787ff39e6a19545 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/table_question_answering.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/text2text_generation.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/text2text_generation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..35bfcff316c1812afe5311e6aa79f4a4c30c6082 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/text2text_generation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/text_classification.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/text_classification.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cf53886c0e66a470da858bc5f8f413e0547fd47b Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/text_classification.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/text_generation.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/text_generation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..300de31069e65fe932855e9e69f2b4574db730a5 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/text_generation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/text_to_audio.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/text_to_audio.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..97598914613a5742295f950836818ad616fb3bc0 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/text_to_audio.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/token_classification.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/token_classification.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..63b41aff142a36ec368f8b3d866c6c3d808ec0c8 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/token_classification.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/video_classification.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/video_classification.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9abfdbdb658f94cd102b5afcc854956c15047470 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/video_classification.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/visual_question_answering.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/visual_question_answering.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0a7c1aa9983abffc4d2297c9918d2f7e7ff063d6 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/visual_question_answering.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/zero_shot_audio_classification.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/zero_shot_audio_classification.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6ae4e312201c7efc7e57645cda962011a516fd9e Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/zero_shot_audio_classification.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/zero_shot_classification.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/zero_shot_classification.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c514310140c7f0e7f18bfe9cedbd28ceb9951fed Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/zero_shot_classification.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/zero_shot_image_classification.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/zero_shot_image_classification.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1de200cc3f4ad819fdbe796dd73fe7a248dad215 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/zero_shot_image_classification.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/__pycache__/zero_shot_object_detection.cpython-312.pyc b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/zero_shot_object_detection.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f00b30e81baba37ee22b501e705d651e91ae4a7f Binary files /dev/null and b/lib/python3.12/site-packages/transformers/pipelines/__pycache__/zero_shot_object_detection.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/pipelines/automatic_speech_recognition.py b/lib/python3.12/site-packages/transformers/pipelines/automatic_speech_recognition.py new file mode 100644 index 0000000000000000000000000000000000000000..66a9c49ea5f3516053fa9d5835109dcd53e3ff1a --- /dev/null +++ b/lib/python3.12/site-packages/transformers/pipelines/automatic_speech_recognition.py @@ -0,0 +1,766 @@ +# Copyright 2021 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import warnings +from collections import defaultdict +from typing import TYPE_CHECKING, Dict, Optional, Union + +import numpy as np +import requests + +from ..tokenization_utils import PreTrainedTokenizer +from ..utils import is_torch_available, is_torchaudio_available, logging +from .audio_utils import ffmpeg_read +from .base import ChunkPipeline + + +if TYPE_CHECKING: + from pyctcdecode import BeamSearchDecoderCTC + + from ..feature_extraction_sequence_utils import SequenceFeatureExtractor + from ..modeling_utils import PreTrainedModel + +logger = logging.get_logger(__name__) + +if is_torch_available(): + import torch + + from ..models.auto.modeling_auto import MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES + + +def rescale_stride(stride, ratio): + """ + Rescales the stride values from audio space to tokens/logits space. + + (160_000, 16_000, 16_000) -> (2000, 200, 200) for instance. + """ + # Shape is [B, SEQ] for tokens + # [B, SEQ, V] for logits + + new_strides = [] + for input_n, left, right in stride: + token_n = int(round(input_n * ratio)) + left = int(round(left / input_n * token_n)) + right = int(round(right / input_n * token_n)) + new_stride = (token_n, left, right) + new_strides.append(new_stride) + + return new_strides + + +def chunk_iter(inputs, feature_extractor, chunk_len, stride_left, stride_right, dtype=None): + inputs_len = inputs.shape[0] + step = chunk_len - stride_left - stride_right + for chunk_start_idx in range(0, inputs_len, step): + chunk_end_idx = chunk_start_idx + chunk_len + chunk = inputs[chunk_start_idx:chunk_end_idx] + processed = feature_extractor(chunk, sampling_rate=feature_extractor.sampling_rate, return_tensors="pt") + if dtype is not None: + processed = processed.to(dtype=dtype) + _stride_left = 0 if chunk_start_idx == 0 else stride_left + is_last = chunk_end_idx >= inputs_len + _stride_right = 0 if is_last else stride_right + + chunk_len = chunk.shape[0] + stride = (chunk_len, _stride_left, _stride_right) + if chunk.shape[0] > _stride_left: + yield {"is_last": is_last, "stride": stride, **processed} + if is_last: + break + + +def _fast_find_longest_common_sequence(sequence_left, sequence_right): + seq_len_left = len(sequence_left) + seq_len_right = len(sequence_right) + counter = [[0] * (seq_len_right + 1) for _ in range(seq_len_left + 1)] + longest = 0 + for i in range(seq_len_left): + for j in range(seq_len_right): + if sequence_left[i] == sequence_right[j]: + previous_counter = counter[i][j] + 1 + counter[i + 1][j + 1] = previous_counter + if previous_counter > longest: + longest = previous_counter + + counter = np.array(counter) + # we return the idx of the first element of the longest common sequence in the left sequence + index_left = np.argwhere(counter == longest)[-1][0] - longest if longest != 0 else -1 + index_right = np.argwhere(counter == longest)[-1][1] - longest if longest != 0 else -1 + return index_left, index_right, longest + + +def _find_longest_common_sequence(sequences, tokenizer): + # TODO Use a faster algorithm this can probably be done in O(n) + # using suffix array. + # It might be tedious to do because of fault tolerance. + # We actually have a really good property which is that the total sequence + # MUST be those subsequences in order. + # Also the algorithm should be more tolerant to errors. + sequence = [tok_id for tok_id in sequences[0][0].tolist() if tok_id not in tokenizer.all_special_ids] + for new_seq in sequences[1:]: + new_sequence = [tok_id for tok_id in new_seq[0].tolist() if tok_id not in tokenizer.all_special_ids] + + index = 0 + max_ = 0.0 + for i in range(1, len(new_sequence) + 1): + # epsilon to favor long perfect matches + eps = i / 10000.0 + matches = np.sum(np.array(sequence[-i:]) == np.array(new_sequence[:i])) + matching = matches / i + eps + if matches > 1 and matching > max_: + index = i + max_ = matching + sequence.extend(new_sequence[index:]) + return np.array(sequence) + + +class AutomaticSpeechRecognitionPipeline(ChunkPipeline): + """ + Pipeline that aims at extracting spoken text contained within some audio. + + The input can be either a raw waveform or a audio file. In case of the audio file, ffmpeg should be installed for + to support multiple audio formats + + Example: + + ```python + >>> from transformers import pipeline + + >>> transcriber = pipeline(model="openai/whisper-base") + >>> transcriber("https://huggingface.co/datasets/Narsil/asr_dummy/resolve/main/1.flac") + {'text': ' He hoped there would be stew for dinner, turnips and carrots and bruised potatoes and fat mutton pieces to be ladled out in thick, peppered flour-fatten sauce.'} + ``` + + Learn more about the basics of using a pipeline in the [pipeline tutorial](../pipeline_tutorial) + + Arguments: + model ([`PreTrainedModel`] or [`TFPreTrainedModel`]): + The model that will be used by the pipeline to make predictions. This needs to be a model inheriting from + [`PreTrainedModel`] for PyTorch and [`TFPreTrainedModel`] for TensorFlow. + feature_extractor ([`SequenceFeatureExtractor`]): + The feature extractor that will be used by the pipeline to encode waveform for the model. + tokenizer ([`PreTrainedTokenizer`]): + The tokenizer that will be used by the pipeline to encode data for the model. This object inherits from + [`PreTrainedTokenizer`]. + decoder (`pyctcdecode.BeamSearchDecoderCTC`, *optional*): + [PyCTCDecode's + BeamSearchDecoderCTC](https://github.com/kensho-technologies/pyctcdecode/blob/2fd33dc37c4111417e08d89ccd23d28e9b308d19/pyctcdecode/decoder.py#L180) + can be passed for language model boosted decoding. See [`Wav2Vec2ProcessorWithLM`] for more information. + chunk_length_s (`float`, *optional*, defaults to 0): + The input length for in each chunk. If `chunk_length_s = 0` then chunking is disabled (default). + + + + For more information on how to effectively use `chunk_length_s`, please have a look at the [ASR chunking + blog post](https://huggingface.co/blog/asr-chunking). + + + + stride_length_s (`float`, *optional*, defaults to `chunk_length_s / 6`): + The length of stride on the left and right of each chunk. Used only with `chunk_length_s > 0`. This enables + the model to *see* more context and infer letters better than without this context but the pipeline + discards the stride bits at the end to make the final reconstitution as perfect as possible. + + + + For more information on how to effectively use `stride_length_s`, please have a look at the [ASR chunking + blog post](https://huggingface.co/blog/asr-chunking). + + + + framework (`str`, *optional*): + The framework to use, either `"pt"` for PyTorch or `"tf"` for TensorFlow. The specified framework must be + installed. If no framework is specified, will default to the one currently installed. If no framework is + specified and both frameworks are installed, will default to the framework of the `model`, or to PyTorch if + no model is provided. + device (Union[`int`, `torch.device`], *optional*): + Device ordinal for CPU/GPU supports. Setting this to `None` will leverage CPU, a positive will run the + model on the associated CUDA device id. + torch_dtype (Union[`int`, `torch.dtype`], *optional*): + The data-type (dtype) of the computation. Setting this to `None` will use float32 precision. Set to + `torch.float16` or `torch.bfloat16` to use half-precision in the respective dtypes. + + """ + + def __init__( + self, + model: "PreTrainedModel", + feature_extractor: Union["SequenceFeatureExtractor", str] = None, + tokenizer: Optional[PreTrainedTokenizer] = None, + decoder: Optional[Union["BeamSearchDecoderCTC", str]] = None, + device: Union[int, "torch.device"] = None, + torch_dtype: Optional[Union[str, "torch.dtype"]] = None, + **kwargs, + ): + # set the model type so we can check we have the right pre- and post-processing parameters + if model.config.model_type == "whisper": + self.type = "seq2seq_whisper" + elif model.__class__.__name__ in MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES.values(): + self.type = "seq2seq" + elif ( + feature_extractor._processor_class + and feature_extractor._processor_class.endswith("WithLM") + and decoder is not None + ): + self.decoder = decoder + self.type = "ctc_with_lm" + else: + self.type = "ctc" + + super().__init__(model, tokenizer, feature_extractor, device=device, torch_dtype=torch_dtype, **kwargs) + + def __call__( + self, + inputs: Union[np.ndarray, bytes, str], + **kwargs, + ): + """ + Transcribe the audio sequence(s) given as inputs to text. See the [`AutomaticSpeechRecognitionPipeline`] + documentation for more information. + + Args: + inputs (`np.ndarray` or `bytes` or `str` or `dict`): + The inputs is either : + - `str` that is either the filename of a local audio file, or a public URL address to download the + audio file. The file will be read at the correct sampling rate to get the waveform using + *ffmpeg*. This requires *ffmpeg* to be installed on the system. + - `bytes` it is supposed to be the content of an audio file and is interpreted by *ffmpeg* in the + same way. + - (`np.ndarray` of shape (n, ) of type `np.float32` or `np.float64`) + Raw audio at the correct sampling rate (no further check will be done) + - `dict` form can be used to pass raw audio sampled at arbitrary `sampling_rate` and let this + pipeline do the resampling. The dict must be in the format `{"sampling_rate": int, "raw": + np.array}` with optionally a `"stride": (left: int, right: int)` than can ask the pipeline to + treat the first `left` samples and last `right` samples to be ignored in decoding (but used at + inference to provide more context to the model). Only use `stride` with CTC models. + return_timestamps (*optional*, `str` or `bool`): + Only available for pure CTC models (Wav2Vec2, HuBERT, etc) and the Whisper model. Not available for + other sequence-to-sequence models. + + For CTC models, timestamps can take one of two formats: + - `"char"`: the pipeline will return timestamps along the text for every character in the text. For + instance, if you get `[{"text": "h", "timestamp": (0.5, 0.6)}, {"text": "i", "timestamp": (0.7, + 0.9)}]`, then it means the model predicts that the letter "h" was spoken after `0.5` and before + `0.6` seconds. + - `"word"`: the pipeline will return timestamps along the text for every word in the text. For + instance, if you get `[{"text": "hi ", "timestamp": (0.5, 0.9)}, {"text": "there", "timestamp": + (1.0, 1.5)}]`, then it means the model predicts that the word "hi" was spoken after `0.5` and + before `0.9` seconds. + + For the Whisper model, timestamps can take one of two formats: + - `"word"`: same as above for word-level CTC timestamps. Word-level timestamps are predicted + through the *dynamic-time warping (DTW)* algorithm, an approximation to word-level timestamps + by inspecting the cross-attention weights. + - `True`: the pipeline will return timestamps along the text for *segments* of words in the text. + For instance, if you get `[{"text": " Hi there!", "timestamp": (0.5, 1.5)}]`, then it means the + model predicts that the segment "Hi there!" was spoken after `0.5` and before `1.5` seconds. + Note that a segment of text refers to a sequence of one or more words, rather than individual + words as with word-level timestamps. + generate_kwargs (`dict`, *optional*): + The dictionary of ad-hoc parametrization of `generate_config` to be used for the generation call. For a + complete overview of generate, check the [following + guide](https://huggingface.co/docs/transformers/en/main_classes/text_generation). + + Return: + `Dict`: A dictionary with the following keys: + - **text** (`str`): The recognized text. + - **chunks** (*optional(, `List[Dict]`) + When using `return_timestamps`, the `chunks` will become a list containing all the various text + chunks identified by the model, *e.g.* `[{"text": "hi ", "timestamp": (0.5, 0.9)}, {"text": + "there", "timestamp": (1.0, 1.5)}]`. The original full text can roughly be recovered by doing + `"".join(chunk["text"] for chunk in output["chunks"])`. + """ + return super().__call__(inputs, **kwargs) + + def _sanitize_parameters( + self, + chunk_length_s=None, + stride_length_s=None, + ignore_warning=None, + decoder_kwargs=None, + return_timestamps=None, + return_language=None, + generate_kwargs=None, + max_new_tokens=None, + ): + # No parameters on this pipeline right now + preprocess_params = {} + if chunk_length_s is not None: + if self.type == "seq2seq" and not ignore_warning: + logger.warning( + "Using `chunk_length_s` is very experimental with seq2seq models. The results will not necessarily" + " be entirely accurate and will have caveats. More information:" + " https://github.com/huggingface/transformers/pull/20104. Ignore this warning with pipeline(...," + " ignore_warning=True)" + ) + preprocess_params["chunk_length_s"] = chunk_length_s + if stride_length_s is not None: + preprocess_params["stride_length_s"] = stride_length_s + + forward_params = defaultdict(dict) + if max_new_tokens is not None: + warnings.warn( + "`max_new_tokens` is deprecated and will be removed in version 4.49 of Transformers. To remove this warning, pass `max_new_tokens` as a key inside `generate_kwargs` instead.", + FutureWarning, + ) + forward_params["max_new_tokens"] = max_new_tokens + if generate_kwargs is not None: + if max_new_tokens is not None and "max_new_tokens" in generate_kwargs: + raise ValueError( + "`max_new_tokens` is defined both as an argument and inside `generate_kwargs` argument, please use" + " only 1 version" + ) + forward_params.update(generate_kwargs) + + postprocess_params = {} + if decoder_kwargs is not None: + postprocess_params["decoder_kwargs"] = decoder_kwargs + if return_timestamps is not None: + # Check whether we have a valid setting for return_timestamps and throw an error before we perform a forward pass + if self.type == "seq2seq" and return_timestamps: + raise ValueError("We cannot return_timestamps yet on non-CTC models apart from Whisper!") + if self.type == "ctc_with_lm" and return_timestamps != "word": + raise ValueError("CTC with LM can only predict word level timestamps, set `return_timestamps='word'`") + if self.type == "ctc" and return_timestamps not in ["char", "word"]: + raise ValueError( + "CTC can either predict character level timestamps, or word level timestamps. " + "Set `return_timestamps='char'` or `return_timestamps='word'` as required." + ) + if self.type == "seq2seq_whisper" and return_timestamps == "char": + raise ValueError( + "Whisper cannot return `char` timestamps, only word level or segment level timestamps. " + "Use `return_timestamps='word'` or `return_timestamps=True` respectively." + ) + forward_params["return_timestamps"] = return_timestamps + postprocess_params["return_timestamps"] = return_timestamps + if return_language is not None: + if self.type != "seq2seq_whisper": + raise ValueError("Only Whisper can return language for now.") + postprocess_params["return_language"] = return_language + + if self.assistant_model is not None: + forward_params["assistant_model"] = self.assistant_model + if self.assistant_tokenizer is not None: + forward_params["tokenizer"] = self.tokenizer + forward_params["assistant_tokenizer"] = self.assistant_tokenizer + + return preprocess_params, forward_params, postprocess_params + + def preprocess(self, inputs, chunk_length_s=0, stride_length_s=None): + if isinstance(inputs, str): + if inputs.startswith("http://") or inputs.startswith("https://"): + # We need to actually check for a real protocol, otherwise it's impossible to use a local file + # like http_huggingface_co.png + inputs = requests.get(inputs).content + else: + with open(inputs, "rb") as f: + inputs = f.read() + + if isinstance(inputs, bytes): + inputs = ffmpeg_read(inputs, self.feature_extractor.sampling_rate) + + stride = None + extra = {} + if isinstance(inputs, dict): + stride = inputs.pop("stride", None) + # Accepting `"array"` which is the key defined in `datasets` for + # better integration + if not ("sampling_rate" in inputs and ("raw" in inputs or "array" in inputs)): + raise ValueError( + "When passing a dictionary to AutomaticSpeechRecognitionPipeline, the dict needs to contain a " + '"raw" key containing the numpy array representing the audio and a "sampling_rate" key, ' + "containing the sampling_rate associated with that array" + ) + + _inputs = inputs.pop("raw", None) + if _inputs is None: + # Remove path which will not be used from `datasets`. + inputs.pop("path", None) + _inputs = inputs.pop("array", None) + in_sampling_rate = inputs.pop("sampling_rate") + extra = inputs + inputs = _inputs + if in_sampling_rate != self.feature_extractor.sampling_rate: + if is_torchaudio_available(): + from torchaudio import functional as F + else: + raise ImportError( + "torchaudio is required to resample audio samples in AutomaticSpeechRecognitionPipeline. " + "The torchaudio package can be installed through: `pip install torchaudio`." + ) + + inputs = F.resample( + torch.from_numpy(inputs), in_sampling_rate, self.feature_extractor.sampling_rate + ).numpy() + ratio = self.feature_extractor.sampling_rate / in_sampling_rate + else: + ratio = 1 + if stride is not None: + if stride[0] + stride[1] > inputs.shape[0]: + raise ValueError("Stride is too large for input") + + # Stride needs to get the chunk length here, it's going to get + # swallowed by the `feature_extractor` later, and then batching + # can add extra data in the inputs, so we need to keep track + # of the original length in the stride so we can cut properly. + stride = (inputs.shape[0], int(round(stride[0] * ratio)), int(round(stride[1] * ratio))) + if not isinstance(inputs, np.ndarray): + raise TypeError(f"We expect a numpy ndarray as input, got `{type(inputs)}`") + if len(inputs.shape) != 1: + raise ValueError("We expect a single channel audio input for AutomaticSpeechRecognitionPipeline") + + if chunk_length_s: + if stride_length_s is None: + stride_length_s = chunk_length_s / 6 + + if isinstance(stride_length_s, (int, float)): + stride_length_s = [stride_length_s, stride_length_s] + + # XXX: Carefuly, this variable will not exist in `seq2seq` setting. + # Currently chunking is not possible at this level for `seq2seq` so + # it's ok. + align_to = getattr(self.model.config, "inputs_to_logits_ratio", 1) + chunk_len = int(round(chunk_length_s * self.feature_extractor.sampling_rate / align_to) * align_to) + stride_left = int(round(stride_length_s[0] * self.feature_extractor.sampling_rate / align_to) * align_to) + stride_right = int(round(stride_length_s[1] * self.feature_extractor.sampling_rate / align_to) * align_to) + + if chunk_len < stride_left + stride_right: + raise ValueError("Chunk length must be superior to stride length") + + for item in chunk_iter( + inputs, self.feature_extractor, chunk_len, stride_left, stride_right, self.torch_dtype + ): + yield {**item, **extra} + else: + if self.type == "seq2seq_whisper" and inputs.shape[0] > self.feature_extractor.n_samples: + processed = self.feature_extractor( + inputs, + sampling_rate=self.feature_extractor.sampling_rate, + truncation=False, + padding="longest", + return_tensors="pt", + return_attention_mask=True, + ) + else: + if self.type == "seq2seq_whisper" and stride is None: + processed = self.feature_extractor( + inputs, + sampling_rate=self.feature_extractor.sampling_rate, + return_tensors="pt", + return_token_timestamps=True, + return_attention_mask=True, + ) + extra["num_frames"] = processed.pop("num_frames") + else: + processed = self.feature_extractor( + inputs, + sampling_rate=self.feature_extractor.sampling_rate, + return_tensors="pt", + return_attention_mask=True, + ) + if self.torch_dtype is not None: + processed = processed.to(dtype=self.torch_dtype) + if stride is not None: + if self.type == "seq2seq": + raise ValueError("Stride is only usable with CTC models, try removing it !") + + processed["stride"] = stride + yield {"is_last": True, **processed, **extra} + + def _forward(self, model_inputs, return_timestamps=False, **generate_kwargs): + attention_mask = model_inputs.pop("attention_mask", None) + stride = model_inputs.pop("stride", None) + num_frames = model_inputs.pop("num_frames", None) + is_last = model_inputs.pop("is_last") + + if stride is not None and num_frames is not None: + raise ValueError("num_frames must be used only when stride is None") + + if self.type in {"seq2seq", "seq2seq_whisper"}: + # Consume values so we can let extra information flow freely through + # the pipeline (important for `partial` in microphone) + if "input_features" in model_inputs: + inputs = model_inputs.pop("input_features") + elif "input_values" in model_inputs: + inputs = model_inputs.pop("input_values") + else: + raise ValueError( + "Seq2Seq speech recognition model requires either a " + f"`input_features` or `input_values` key, but only has {model_inputs.keys()}" + ) + + # custom processing for Whisper timestamps and word-level timestamps + if return_timestamps and self.type == "seq2seq_whisper": + generate_kwargs["return_timestamps"] = return_timestamps + if return_timestamps == "word": + generate_kwargs["return_token_timestamps"] = True + generate_kwargs["return_segments"] = True + + if stride is not None: + if isinstance(stride, tuple): + generate_kwargs["num_frames"] = stride[0] // self.feature_extractor.hop_length + else: + generate_kwargs["num_frames"] = [s[0] // self.feature_extractor.hop_length for s in stride] + else: + generate_kwargs["num_frames"] = num_frames + + # User-defined `generation_config` passed to the pipeline call take precedence + if "generation_config" not in generate_kwargs: + generate_kwargs["generation_config"] = self.generation_config + + tokens = self.model.generate( + inputs=inputs, + attention_mask=attention_mask, + **generate_kwargs, + ) + # whisper longform generation stores timestamps in "segments" + if return_timestamps == "word" and self.type == "seq2seq_whisper": + if "segments" not in tokens: + out = {"tokens": tokens["sequences"], "token_timestamps": tokens["token_timestamps"]} + else: + token_timestamps = [ + torch.cat([segment["token_timestamps"] for segment in segment_list]) + for segment_list in tokens["segments"] + ] + out = {"tokens": tokens["sequences"], "token_timestamps": token_timestamps} + else: + out = {"tokens": tokens} + if self.type == "seq2seq_whisper": + if stride is not None: + out["stride"] = stride + + else: + inputs = { + self.model.main_input_name: model_inputs.pop(self.model.main_input_name), + "attention_mask": attention_mask, + } + outputs = self.model(**inputs) + logits = outputs.logits + + if self.type == "ctc_with_lm": + out = {"logits": logits} + else: + out = {"tokens": logits.argmax(dim=-1)} + if stride is not None: + # Send stride to `postprocess`. + # it needs to be handled there where + # the pieces are to be concatenated. + ratio = 1 / self.model.config.inputs_to_logits_ratio + if isinstance(stride, tuple): + out["stride"] = rescale_stride([stride], ratio)[0] + else: + out["stride"] = rescale_stride(stride, ratio) + # Leftover + extra = model_inputs + return {"is_last": is_last, **out, **extra} + + def postprocess( + self, model_outputs, decoder_kwargs: Optional[Dict] = None, return_timestamps=None, return_language=None + ): + # Optional return types + optional = {} + + final_items = [] + key = "logits" if self.type == "ctc_with_lm" else "tokens" + stride = None + for outputs in model_outputs: + if self.framework == "pt" and outputs[key].dtype in (torch.bfloat16, torch.float16): + items = outputs[key].to(torch.float32).numpy() + else: + items = outputs[key].numpy() + stride = outputs.get("stride", None) + if stride is not None and self.type in {"ctc", "ctc_with_lm"}: + total_n, left, right = stride + # Total_n might be < logits.shape[1] + # because of padding, that's why + # we need to reconstruct this information + # This won't work with left padding (which doesn't exist right now) + right_n = total_n - right + items = items[:, left:right_n] + final_items.append(items) + + if stride and self.type == "seq2seq": + items = _find_longest_common_sequence(final_items, self.tokenizer) + elif self.type == "seq2seq_whisper": + time_precision = self.feature_extractor.chunk_length / self.model.config.max_source_positions + # Send the chunking back to seconds, it's easier to handle in whisper + sampling_rate = self.feature_extractor.sampling_rate + for output in model_outputs: + if "stride" in output: + chunk_len, stride_left, stride_right = output["stride"] + # Go back in seconds + chunk_len /= sampling_rate + stride_left /= sampling_rate + stride_right /= sampling_rate + output["stride"] = chunk_len, stride_left, stride_right + + text, optional = self.tokenizer._decode_asr( + model_outputs, + return_timestamps=return_timestamps, + return_language=return_language, + time_precision=time_precision, + ) + else: + items = np.concatenate(final_items, axis=1) + items = items.squeeze(0) + + if self.type == "ctc_with_lm": + if decoder_kwargs is None: + decoder_kwargs = {} + beams = self.decoder.decode_beams(items, **decoder_kwargs) + text = beams[0][0] + if return_timestamps: + # Simply cast from pyctcdecode format to wav2vec2 format to leverage + # pre-existing code later + chunk_offset = beams[0][2] + offsets = [] + for word, (start_offset, end_offset) in chunk_offset: + offsets.append({"word": word, "start_offset": start_offset, "end_offset": end_offset}) + elif self.type != "seq2seq_whisper": + skip_special_tokens = self.type != "ctc" + text = self.tokenizer.decode(items, skip_special_tokens=skip_special_tokens) + if return_timestamps: + offsets = self.tokenizer.decode( + items, skip_special_tokens=skip_special_tokens, output_char_offsets=True + )["char_offsets"] + if return_timestamps == "word": + offsets = self.tokenizer._get_word_offsets(offsets, self.tokenizer.replace_word_delimiter_char) + + if return_timestamps and self.type not in {"seq2seq", "seq2seq_whisper"}: + chunks = [] + for item in offsets: + start = item["start_offset"] * self.model.config.inputs_to_logits_ratio + start /= self.feature_extractor.sampling_rate + + stop = item["end_offset"] * self.model.config.inputs_to_logits_ratio + stop /= self.feature_extractor.sampling_rate + + chunks.append({"text": item[return_timestamps], "timestamp": (start, stop)}) + optional["chunks"] = chunks + + extra = defaultdict(list) + for output in model_outputs: + output.pop("tokens", None) + output.pop("logits", None) + output.pop("is_last", None) + output.pop("stride", None) + output.pop("token_timestamps", None) + for k, v in output.items(): + extra[k].append(v) + return {"text": text, **optional, **extra} + + +def _find_timestamp_sequence(sequences, tokenizer, feature_extractor, max_source_positions): + """ + Computes the final sequences by merging the end of the nth sequence with the beginning of the n+1th sequence. Since + `WhisperForConditionalGeneration` produces the timestamps pairwise, we filter the consecutive timestamps and only + iterate over them. We keep track of the `time` which indicates the actual starting time of the chunk that is + processed. We need to make sure to offset the timestamps tokens by the `time` in order for the tokenizer to + properly compute the final `offset`. + """ + # index of the first timestamp token + timestamp_begin = tokenizer.convert_tokens_to_ids("<|notimestamps|>") + 1 + items = [] + # approximation of the token to time ratio : ~0.2seconds + time_precision = feature_extractor.chunk_length / max_source_positions + time = 0 + for seq_idx, item in enumerate(sequences): + sequence, stride = item + if isinstance(sequence, list): + sequence = np.array(sequence) + chunk_len, stride_left, stride_right = stride + sequence = sequence.squeeze(0) + # get rid of the `forced_decoder_idx` that are use to parametrize the generation + begin_idx = np.where(sequence == timestamp_begin)[0][0] if timestamp_begin in sequence else 0 + sequence = sequence[begin_idx:] + + timestamp_tokens = sequence >= timestamp_begin + if seq_idx != 0 and sum(timestamp_tokens) > 0: + consecutive = np.where(timestamp_tokens[:-1] & timestamp_tokens[1:])[0] + 1 + last_timestamp = np.where(timestamp_tokens)[0][-1] + consecutive = np.append(consecutive, last_timestamp) if last_timestamp not in consecutive else consecutive + time -= stride_left + stride_right + offset = int((time / feature_extractor.sampling_rate) / time_precision) + overlap_time = int((stride_left / feature_extractor.sampling_rate) / time_precision) + # relevant timestamps are in the overlapping part + relevant_timestamp = np.where(sequence[consecutive] >= timestamp_begin + overlap_time)[0] + if relevant_timestamp.shape[0] > 0: + relevant_timestamp = ( + consecutive[relevant_timestamp[0] - 1] if relevant_timestamp[0] > 0 else consecutive[0] + ) + # if a big stride is used, we need to check some of the previous items for the best overlap + best_match = 0 + sliced_sequence = [] + for idx, previous_sequence in enumerate(reversed(items)): + previous_tokens = previous_sequence[1:-1] + if previous_sequence[0] < (timestamp_begin + offset - overlap_time) and idx != 0: + break # the previous sequence is too far in the past + if len(previous_tokens) > 0: + # find the longest common sequence between the overlapping parts + index_left, index_right, match_length = _fast_find_longest_common_sequence( + sequence[1:relevant_timestamp], previous_tokens + ) + # don't do anything if only 1 token was matched + if match_length > 1 and match_length > best_match: + best_match = match_length + best_idx = idx + end_of_curr_sequence_idx = ( + np.where(sequence[index_left + 1 :] >= timestamp_begin)[0][0] + 1 + ) + end_of_curr_sequence_idx = end_of_curr_sequence_idx + 1 + index_left + # if all the tokens are matched, suffix + if index_left == 0 and match_length == len(previous_tokens): + sliced_sequence = np.insert( + sequence[index_left + 1 : end_of_curr_sequence_idx], 0, previous_sequence[0] + ) + sliced_sequence[-1] = previous_sequence[-1] + # if part of the previous sequence is not taken + elif index_left >= 0: + sliced_sequence = sequence[index_left + 1 : end_of_curr_sequence_idx] + # let's insert the missing part of the previous sequence + previous_slice = ( + previous_sequence[: index_right + 1] if index_right > 0 else [previous_sequence[0]] + ) + sliced_sequence = np.insert(sliced_sequence, 0, previous_slice) + sliced_sequence[-1] += offset + + if len(sliced_sequence) > 0: + items[len(items) - best_idx - 1] = sliced_sequence + items = items[: len(items) - best_idx] + sequence = sequence[end_of_curr_sequence_idx:] + + # sequence might have changed + timestamp_tokens = sequence >= timestamp_begin + consecutive = np.where(timestamp_tokens[:-1] & timestamp_tokens[1:])[0] + 1 + if sum(timestamp_tokens) > 0: + last_timestamp = np.where(timestamp_tokens)[0][-1] + consecutive = ( + np.append(consecutive, last_timestamp + 1) if last_timestamp not in consecutive else consecutive + ) + + if len(consecutive) > 0: + last_slice = 0 + for current_slice in consecutive: + actual_offset = items[-1][-1] if seq_idx != 0 or last_slice != 0 else sequence[0] + sliced_tokens = sequence[last_slice:current_slice] + duration = sliced_tokens[-1] - sliced_tokens[0] + sliced_tokens[0] = actual_offset + sliced_tokens[-1] = actual_offset + duration + items.append(sliced_tokens) + last_slice = current_slice + + time += chunk_len + result = [] + for i in range(len(items)): + result += items[i].tolist() + return result diff --git a/lib/python3.12/site-packages/transformers/pipelines/fill_mask.py b/lib/python3.12/site-packages/transformers/pipelines/fill_mask.py new file mode 100644 index 0000000000000000000000000000000000000000..c14f54118486b971f64b0985fe2dc688de52f863 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/pipelines/fill_mask.py @@ -0,0 +1,273 @@ +from typing import Dict + +import numpy as np + +from ..utils import add_end_docstrings, is_tf_available, is_torch_available, logging +from .base import GenericTensor, Pipeline, PipelineException, build_pipeline_init_args + + +if is_tf_available(): + import tensorflow as tf + + from ..tf_utils import stable_softmax + + +if is_torch_available(): + import torch + + +logger = logging.get_logger(__name__) + + +@add_end_docstrings( + build_pipeline_init_args(has_tokenizer=True), + r""" + top_k (`int`, *optional*, defaults to 5): + The number of predictions to return. + targets (`str` or `List[str]`, *optional*): + When passed, the model will limit the scores to the passed targets instead of looking up in the whole + vocab. If the provided targets are not in the model vocab, they will be tokenized and the first resulting + token will be used (with a warning, and that might be slower). + tokenizer_kwargs (`dict`, *optional*): + Additional dictionary of keyword arguments passed along to the tokenizer.""", +) +class FillMaskPipeline(Pipeline): + """ + Masked language modeling prediction pipeline using any `ModelWithLMHead`. See the [masked language modeling + examples](../task_summary#masked-language-modeling) for more information. + + Example: + + ```python + >>> from transformers import pipeline + + >>> fill_masker = pipeline(model="google-bert/bert-base-uncased") + >>> fill_masker("This is a simple [MASK].") + [{'score': 0.042, 'token': 3291, 'token_str': 'problem', 'sequence': 'this is a simple problem.'}, {'score': 0.031, 'token': 3160, 'token_str': 'question', 'sequence': 'this is a simple question.'}, {'score': 0.03, 'token': 8522, 'token_str': 'equation', 'sequence': 'this is a simple equation.'}, {'score': 0.027, 'token': 2028, 'token_str': 'one', 'sequence': 'this is a simple one.'}, {'score': 0.024, 'token': 3627, 'token_str': 'rule', 'sequence': 'this is a simple rule.'}] + ``` + + Learn more about the basics of using a pipeline in the [pipeline tutorial](../pipeline_tutorial) + + This mask filling pipeline can currently be loaded from [`pipeline`] using the following task identifier: + `"fill-mask"`. + + The models that this pipeline can use are models that have been trained with a masked language modeling objective, + which includes the bi-directional models in the library. See the up-to-date list of available models on + [huggingface.co/models](https://huggingface.co/models?filter=fill-mask). + + + + This pipeline only works for inputs with exactly one token masked. Experimental: We added support for multiple + masks. The returned values are raw model output, and correspond to disjoint probabilities where one might expect + joint probabilities (See [discussion](https://github.com/huggingface/transformers/pull/10222)). + + + + + + This pipeline now supports tokenizer_kwargs. For example try: + + ```python + >>> from transformers import pipeline + + >>> fill_masker = pipeline(model="google-bert/bert-base-uncased") + >>> tokenizer_kwargs = {"truncation": True} + >>> fill_masker( + ... "This is a simple [MASK]. " + "...with a large amount of repeated text appended. " * 100, + ... tokenizer_kwargs=tokenizer_kwargs, + ... ) + ``` + + + + + + """ + + def get_masked_index(self, input_ids: GenericTensor) -> np.ndarray: + if self.framework == "tf": + masked_index = tf.where(input_ids == self.tokenizer.mask_token_id).numpy() + elif self.framework == "pt": + masked_index = torch.nonzero(input_ids == self.tokenizer.mask_token_id, as_tuple=False) + else: + raise ValueError("Unsupported framework") + return masked_index + + def _ensure_exactly_one_mask_token(self, input_ids: GenericTensor) -> np.ndarray: + masked_index = self.get_masked_index(input_ids) + numel = np.prod(masked_index.shape) + if numel < 1: + raise PipelineException( + "fill-mask", + self.model.base_model_prefix, + f"No mask_token ({self.tokenizer.mask_token}) found on the input", + ) + + def ensure_exactly_one_mask_token(self, model_inputs: GenericTensor): + if isinstance(model_inputs, list): + for model_input in model_inputs: + self._ensure_exactly_one_mask_token(model_input["input_ids"][0]) + else: + for input_ids in model_inputs["input_ids"]: + self._ensure_exactly_one_mask_token(input_ids) + + def preprocess( + self, inputs, return_tensors=None, tokenizer_kwargs=None, **preprocess_parameters + ) -> Dict[str, GenericTensor]: + if return_tensors is None: + return_tensors = self.framework + if tokenizer_kwargs is None: + tokenizer_kwargs = {} + + model_inputs = self.tokenizer(inputs, return_tensors=return_tensors, **tokenizer_kwargs) + self.ensure_exactly_one_mask_token(model_inputs) + return model_inputs + + def _forward(self, model_inputs): + model_outputs = self.model(**model_inputs) + model_outputs["input_ids"] = model_inputs["input_ids"] + return model_outputs + + def postprocess(self, model_outputs, top_k=5, target_ids=None): + # Cap top_k if there are targets + if target_ids is not None and target_ids.shape[0] < top_k: + top_k = target_ids.shape[0] + input_ids = model_outputs["input_ids"][0] + outputs = model_outputs["logits"] + + if self.framework == "tf": + masked_index = tf.where(input_ids == self.tokenizer.mask_token_id).numpy()[:, 0] + + outputs = outputs.numpy() + + logits = outputs[0, masked_index, :] + probs = stable_softmax(logits, axis=-1) + if target_ids is not None: + probs = tf.gather_nd(tf.squeeze(probs, 0), target_ids.reshape(-1, 1)) + probs = tf.expand_dims(probs, 0) + + topk = tf.math.top_k(probs, k=top_k) + values, predictions = topk.values.numpy(), topk.indices.numpy() + else: + masked_index = torch.nonzero(input_ids == self.tokenizer.mask_token_id, as_tuple=False).squeeze(-1) + # Fill mask pipeline supports only one ${mask_token} per sample + + logits = outputs[0, masked_index, :] + probs = logits.softmax(dim=-1) + if target_ids is not None: + probs = probs[..., target_ids] + + values, predictions = probs.topk(top_k) + + result = [] + single_mask = values.shape[0] == 1 + for i, (_values, _predictions) in enumerate(zip(values.tolist(), predictions.tolist())): + row = [] + for v, p in zip(_values, _predictions): + # Copy is important since we're going to modify this array in place + tokens = input_ids.numpy().copy() + if target_ids is not None: + p = target_ids[p].tolist() + + tokens[masked_index[i]] = p + # Filter padding out: + tokens = tokens[np.where(tokens != self.tokenizer.pad_token_id)] + # Originally we skip special tokens to give readable output. + # For multi masks though, the other [MASK] would be removed otherwise + # making the output look odd, so we add them back + sequence = self.tokenizer.decode(tokens, skip_special_tokens=single_mask) + proposition = {"score": v, "token": p, "token_str": self.tokenizer.decode([p]), "sequence": sequence} + row.append(proposition) + result.append(row) + if single_mask: + return result[0] + return result + + def get_target_ids(self, targets, top_k=None): + if isinstance(targets, str): + targets = [targets] + try: + vocab = self.tokenizer.get_vocab() + except Exception: + vocab = {} + target_ids = [] + for target in targets: + id_ = vocab.get(target, None) + if id_ is None: + input_ids = self.tokenizer( + target, + add_special_tokens=False, + return_attention_mask=False, + return_token_type_ids=False, + max_length=1, + truncation=True, + )["input_ids"] + if len(input_ids) == 0: + logger.warning( + f"The specified target token `{target}` does not exist in the model vocabulary. " + "We cannot replace it with anything meaningful, ignoring it" + ) + continue + id_ = input_ids[0] + # XXX: If users encounter this pass + # it becomes pretty slow, so let's make sure + # The warning enables them to fix the input to + # get faster performance. + logger.warning( + f"The specified target token `{target}` does not exist in the model vocabulary. " + f"Replacing with `{self.tokenizer.convert_ids_to_tokens(id_)}`." + ) + target_ids.append(id_) + target_ids = list(set(target_ids)) + if len(target_ids) == 0: + raise ValueError("At least one target must be provided when passed.") + target_ids = np.array(target_ids) + return target_ids + + def _sanitize_parameters(self, top_k=None, targets=None, tokenizer_kwargs=None): + preprocess_params = {} + + if tokenizer_kwargs is not None: + preprocess_params["tokenizer_kwargs"] = tokenizer_kwargs + + postprocess_params = {} + + if targets is not None: + target_ids = self.get_target_ids(targets, top_k) + postprocess_params["target_ids"] = target_ids + + if top_k is not None: + postprocess_params["top_k"] = top_k + + if self.tokenizer.mask_token_id is None: + raise PipelineException( + "fill-mask", self.model.base_model_prefix, "The tokenizer does not define a `mask_token`." + ) + return preprocess_params, {}, postprocess_params + + def __call__(self, inputs, **kwargs): + """ + Fill the masked token in the text(s) given as inputs. + + Args: + inputs (`str` or `List[str]`): + One or several texts (or one list of prompts) with masked tokens. + targets (`str` or `List[str]`, *optional*): + When passed, the model will limit the scores to the passed targets instead of looking up in the whole + vocab. If the provided targets are not in the model vocab, they will be tokenized and the first + resulting token will be used (with a warning, and that might be slower). + top_k (`int`, *optional*): + When passed, overrides the number of predictions to return. + + Return: + A list or a list of list of `dict`: Each result comes as list of dictionaries with the following keys: + + - **sequence** (`str`) -- The corresponding input with the mask token prediction. + - **score** (`float`) -- The corresponding probability. + - **token** (`int`) -- The predicted token id (to replace the masked one). + - **token_str** (`str`) -- The predicted token (to replace the masked one). + """ + outputs = super().__call__(inputs, **kwargs) + if isinstance(inputs, list) and len(inputs) == 1: + return outputs[0] + return outputs diff --git a/lib/python3.12/site-packages/transformers/pipelines/image_classification.py b/lib/python3.12/site-packages/transformers/pipelines/image_classification.py new file mode 100644 index 0000000000000000000000000000000000000000..a9b318df5527e4e641e7e66e4ef36d4a9182809d --- /dev/null +++ b/lib/python3.12/site-packages/transformers/pipelines/image_classification.py @@ -0,0 +1,226 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import List, Union + +import numpy as np + +from ..utils import ( + ExplicitEnum, + add_end_docstrings, + is_tf_available, + is_torch_available, + is_vision_available, + logging, + requires_backends, +) +from .base import Pipeline, build_pipeline_init_args + + +if is_vision_available(): + from PIL import Image + + from ..image_utils import load_image + +if is_tf_available(): + from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES + +if is_torch_available(): + import torch + + from ..models.auto.modeling_auto import MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES + +logger = logging.get_logger(__name__) + + +# Copied from transformers.pipelines.text_classification.sigmoid +def sigmoid(_outputs): + return 1.0 / (1.0 + np.exp(-_outputs)) + + +# Copied from transformers.pipelines.text_classification.softmax +def softmax(_outputs): + maxes = np.max(_outputs, axis=-1, keepdims=True) + shifted_exp = np.exp(_outputs - maxes) + return shifted_exp / shifted_exp.sum(axis=-1, keepdims=True) + + +# Copied from transformers.pipelines.text_classification.ClassificationFunction +class ClassificationFunction(ExplicitEnum): + SIGMOID = "sigmoid" + SOFTMAX = "softmax" + NONE = "none" + + +@add_end_docstrings( + build_pipeline_init_args(has_image_processor=True), + r""" + function_to_apply (`str`, *optional*, defaults to `"default"`): + The function to apply to the model outputs in order to retrieve the scores. Accepts four different values: + + - `"default"`: if the model has a single label, will apply the sigmoid function on the output. If the model + has several labels, will apply the softmax function on the output. + - `"sigmoid"`: Applies the sigmoid function on the output. + - `"softmax"`: Applies the softmax function on the output. + - `"none"`: Does not apply any function on the output.""", +) +class ImageClassificationPipeline(Pipeline): + """ + Image classification pipeline using any `AutoModelForImageClassification`. This pipeline predicts the class of an + image. + + Example: + + ```python + >>> from transformers import pipeline + + >>> classifier = pipeline(model="microsoft/beit-base-patch16-224-pt22k-ft22k") + >>> classifier("https://huggingface.co/datasets/Narsil/image_dummy/raw/main/parrots.png") + [{'score': 0.442, 'label': 'macaw'}, {'score': 0.088, 'label': 'popinjay'}, {'score': 0.075, 'label': 'parrot'}, {'score': 0.073, 'label': 'parodist, lampooner'}, {'score': 0.046, 'label': 'poll, poll_parrot'}] + ``` + + Learn more about the basics of using a pipeline in the [pipeline tutorial](../pipeline_tutorial) + + This image classification pipeline can currently be loaded from [`pipeline`] using the following task identifier: + `"image-classification"`. + + See the list of available models on + [huggingface.co/models](https://huggingface.co/models?filter=image-classification). + """ + + function_to_apply: ClassificationFunction = ClassificationFunction.NONE + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + requires_backends(self, "vision") + self.check_model_type( + TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES + if self.framework == "tf" + else MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES + ) + + def _sanitize_parameters(self, top_k=None, function_to_apply=None, timeout=None): + preprocess_params = {} + if timeout is not None: + preprocess_params["timeout"] = timeout + postprocess_params = {} + if top_k is not None: + postprocess_params["top_k"] = top_k + if isinstance(function_to_apply, str): + function_to_apply = ClassificationFunction(function_to_apply.lower()) + if function_to_apply is not None: + postprocess_params["function_to_apply"] = function_to_apply + return preprocess_params, {}, postprocess_params + + def __call__(self, inputs: Union[str, List[str], "Image.Image", List["Image.Image"]] = None, **kwargs): + """ + Assign labels to the image(s) passed as inputs. + + Args: + inputs (`str`, `List[str]`, `PIL.Image` or `List[PIL.Image]`): + The pipeline handles three types of images: + + - A string containing a http link pointing to an image + - A string containing a local path to an image + - An image loaded in PIL directly + + The pipeline accepts either a single image or a batch of images, which must then be passed as a string. + Images in a batch must all be in the same format: all as http links, all as local paths, or all as PIL + images. + function_to_apply (`str`, *optional*, defaults to `"default"`): + The function to apply to the model outputs in order to retrieve the scores. Accepts four different + values: + + If this argument is not specified, then it will apply the following functions according to the number + of labels: + + - If the model has a single label, will apply the sigmoid function on the output. + - If the model has several labels, will apply the softmax function on the output. + + Possible values are: + + - `"sigmoid"`: Applies the sigmoid function on the output. + - `"softmax"`: Applies the softmax function on the output. + - `"none"`: Does not apply any function on the output. + top_k (`int`, *optional*, defaults to 5): + The number of top labels that will be returned by the pipeline. If the provided number is higher than + the number of labels available in the model configuration, it will default to the number of labels. + timeout (`float`, *optional*, defaults to None): + The maximum time in seconds to wait for fetching images from the web. If None, no timeout is set and + the call may block forever. + + Return: + A dictionary or a list of dictionaries containing result. If the input is a single image, will return a + dictionary, if the input is a list of several images, will return a list of dictionaries corresponding to + the images. + + The dictionaries contain the following keys: + + - **label** (`str`) -- The label identified by the model. + - **score** (`int`) -- The score attributed by the model for that label. + """ + # After deprecation of this is completed, remove the default `None` value for `images` + if "images" in kwargs: + inputs = kwargs.pop("images") + if inputs is None: + raise ValueError("Cannot call the image-classification pipeline without an inputs argument!") + return super().__call__(inputs, **kwargs) + + def preprocess(self, image, timeout=None): + image = load_image(image, timeout=timeout) + model_inputs = self.image_processor(images=image, return_tensors=self.framework) + if self.framework == "pt": + model_inputs = model_inputs.to(self.torch_dtype) + return model_inputs + + def _forward(self, model_inputs): + model_outputs = self.model(**model_inputs) + return model_outputs + + def postprocess(self, model_outputs, function_to_apply=None, top_k=5): + if function_to_apply is None: + if self.model.config.problem_type == "multi_label_classification" or self.model.config.num_labels == 1: + function_to_apply = ClassificationFunction.SIGMOID + elif self.model.config.problem_type == "single_label_classification" or self.model.config.num_labels > 1: + function_to_apply = ClassificationFunction.SOFTMAX + elif hasattr(self.model.config, "function_to_apply") and function_to_apply is None: + function_to_apply = self.model.config.function_to_apply + else: + function_to_apply = ClassificationFunction.NONE + + if top_k > self.model.config.num_labels: + top_k = self.model.config.num_labels + + outputs = model_outputs["logits"][0] + if self.framework == "pt" and outputs.dtype in (torch.bfloat16, torch.float16): + outputs = outputs.to(torch.float32).numpy() + else: + outputs = outputs.numpy() + + if function_to_apply == ClassificationFunction.SIGMOID: + scores = sigmoid(outputs) + elif function_to_apply == ClassificationFunction.SOFTMAX: + scores = softmax(outputs) + elif function_to_apply == ClassificationFunction.NONE: + scores = outputs + else: + raise ValueError(f"Unrecognized `function_to_apply` argument: {function_to_apply}") + + dict_scores = [ + {"label": self.model.config.id2label[i], "score": score.item()} for i, score in enumerate(scores) + ] + dict_scores.sort(key=lambda x: x["score"], reverse=True) + if top_k is not None: + dict_scores = dict_scores[:top_k] + + return dict_scores diff --git a/lib/python3.12/site-packages/transformers/pipelines/image_feature_extraction.py b/lib/python3.12/site-packages/transformers/pipelines/image_feature_extraction.py new file mode 100644 index 0000000000000000000000000000000000000000..391eb2b3aec714dbac61fe46bddc7ee74f10cd2f --- /dev/null +++ b/lib/python3.12/site-packages/transformers/pipelines/image_feature_extraction.py @@ -0,0 +1,112 @@ +from typing import Dict + +from ..utils import add_end_docstrings, is_vision_available +from .base import GenericTensor, Pipeline, build_pipeline_init_args + + +if is_vision_available(): + from ..image_utils import load_image + + +@add_end_docstrings( + build_pipeline_init_args(has_image_processor=True), + """ + image_processor_kwargs (`dict`, *optional*): + Additional dictionary of keyword arguments passed along to the image processor e.g. + {"size": {"height": 100, "width": 100}} + pool (`bool`, *optional*, defaults to `False`): + Whether or not to return the pooled output. If `False`, the model will return the raw hidden states. + """, +) +class ImageFeatureExtractionPipeline(Pipeline): + """ + Image feature extraction pipeline uses no model head. This pipeline extracts the hidden states from the base + transformer, which can be used as features in downstream tasks. + + Example: + + ```python + >>> from transformers import pipeline + + >>> extractor = pipeline(model="google/vit-base-patch16-224", task="image-feature-extraction") + >>> result = extractor("https://huggingface.co/datasets/Narsil/image_dummy/raw/main/parrots.png", return_tensors=True) + >>> result.shape # This is a tensor of shape [1, sequence_lenth, hidden_dimension] representing the input image. + torch.Size([1, 197, 768]) + ``` + + Learn more about the basics of using a pipeline in the [pipeline tutorial](../pipeline_tutorial) + + This image feature extraction pipeline can currently be loaded from [`pipeline`] using the task identifier: + `"image-feature-extraction"`. + + All vision models may be used for this pipeline. See a list of all models, including community-contributed models on + [huggingface.co/models](https://huggingface.co/models). + """ + + def _sanitize_parameters(self, image_processor_kwargs=None, return_tensors=None, pool=None, **kwargs): + preprocess_params = {} if image_processor_kwargs is None else image_processor_kwargs + + postprocess_params = {} + if pool is not None: + postprocess_params["pool"] = pool + if return_tensors is not None: + postprocess_params["return_tensors"] = return_tensors + + if "timeout" in kwargs: + preprocess_params["timeout"] = kwargs["timeout"] + + return preprocess_params, {}, postprocess_params + + def preprocess(self, image, timeout=None, **image_processor_kwargs) -> Dict[str, GenericTensor]: + image = load_image(image, timeout=timeout) + model_inputs = self.image_processor(image, return_tensors=self.framework, **image_processor_kwargs) + if self.framework == "pt": + model_inputs = model_inputs.to(self.torch_dtype) + return model_inputs + + def _forward(self, model_inputs): + model_outputs = self.model(**model_inputs) + return model_outputs + + def postprocess(self, model_outputs, pool=None, return_tensors=False): + pool = pool if pool is not None else False + + if pool: + if "pooler_output" not in model_outputs: + raise ValueError( + "No pooled output was returned. Make sure the model has a `pooler` layer when using the `pool` option." + ) + outputs = model_outputs["pooler_output"] + else: + # [0] is the first available tensor, logits or last_hidden_state. + outputs = model_outputs[0] + + if return_tensors: + return outputs + if self.framework == "pt": + return outputs.tolist() + elif self.framework == "tf": + return outputs.numpy().tolist() + + def __call__(self, *args, **kwargs): + """ + Extract the features of the input(s). + + Args: + images (`str`, `List[str]`, `PIL.Image` or `List[PIL.Image]`): + The pipeline handles three types of images: + + - A string containing a http link pointing to an image + - A string containing a local path to an image + - An image loaded in PIL directly + + The pipeline accepts either a single image or a batch of images, which must then be passed as a string. + Images in a batch must all be in the same format: all as http links, all as local paths, or all as PIL + images. + timeout (`float`, *optional*, defaults to None): + The maximum time in seconds to wait for fetching images from the web. If None, no timeout is used and + the call may block forever. + Return: + A nested list of `float`: The features computed by the model. + """ + return super().__call__(*args, **kwargs) diff --git a/lib/python3.12/site-packages/transformers/pipelines/mask_generation.py b/lib/python3.12/site-packages/transformers/pipelines/mask_generation.py new file mode 100644 index 0000000000000000000000000000000000000000..f87e45b7f8ecb410ba5d0a088188256d59290f0f --- /dev/null +++ b/lib/python3.12/site-packages/transformers/pipelines/mask_generation.py @@ -0,0 +1,287 @@ +from collections import defaultdict +from typing import Optional + +from ..image_utils import load_image +from ..utils import ( + add_end_docstrings, + is_torch_available, + logging, + requires_backends, +) +from .base import ChunkPipeline, build_pipeline_init_args + + +if is_torch_available(): + import torch + + from ..models.auto.modeling_auto import MODEL_FOR_MASK_GENERATION_MAPPING_NAMES + +logger = logging.get_logger(__name__) + + +@add_end_docstrings( + build_pipeline_init_args(has_image_processor=True), + r""" + points_per_batch (*optional*, int, default to 64): + Sets the number of points run simultaneously by the model. Higher numbers may be faster but use more GPU + memory. + output_bboxes_mask (`bool`, *optional*, default to `False`): + Whether or not to output the bounding box predictions. + output_rle_masks (`bool`, *optional*, default to `False`): + Whether or not to output the masks in `RLE` format""", +) +class MaskGenerationPipeline(ChunkPipeline): + """ + Automatic mask generation for images using `SamForMaskGeneration`. This pipeline predicts binary masks for an + image, given an image. It is a `ChunkPipeline` because you can seperate the points in a mini-batch in order to + avoid OOM issues. Use the `points_per_batch` argument to control the number of points that will be processed at the + same time. Default is `64`. + + The pipeline works in 3 steps: + 1. `preprocess`: A grid of 1024 points evenly separated is generated along with bounding boxes and point + labels. + For more details on how the points and bounding boxes are created, check the `_generate_crop_boxes` + function. The image is also preprocessed using the `image_processor`. This function `yields` a minibatch of + `points_per_batch`. + + 2. `forward`: feeds the outputs of `preprocess` to the model. The image embedding is computed only once. + Calls both `self.model.get_image_embeddings` and makes sure that the gradients are not computed, and the + tensors and models are on the same device. + + 3. `postprocess`: The most important part of the automatic mask generation happens here. Three steps + are induced: + - image_processor.postprocess_masks (run on each minibatch loop): takes in the raw output masks, + resizes them according + to the image size, and transforms there to binary masks. + - image_processor.filter_masks (on each minibatch loop): uses both `pred_iou_thresh` and + `stability_scores`. Also + applies a variety of filters based on non maximum suppression to remove bad masks. + - image_processor.postprocess_masks_for_amg applies the NSM on the mask to only keep relevant ones. + + Example: + + ```python + >>> from transformers import pipeline + + >>> generator = pipeline(model="facebook/sam-vit-base", task="mask-generation") + >>> outputs = generator( + ... "http://images.cocodataset.org/val2017/000000039769.jpg", + ... ) + + >>> outputs = generator( + ... "https://huggingface.co/datasets/Narsil/image_dummy/raw/main/parrots.png", points_per_batch=128 + ... ) + ``` + + Learn more about the basics of using a pipeline in the [pipeline tutorial](../pipeline_tutorial) + + This segmentation pipeline can currently be loaded from [`pipeline`] using the following task identifier: + `"mask-generation"`. + + See the list of available models on [huggingface.co/models](https://huggingface.co/models?filter=mask-generation). + """ + + def __init__(self, **kwargs): + super().__init__(**kwargs) + requires_backends(self, "vision") + requires_backends(self, "torch") + + if self.framework != "pt": + raise ValueError(f"The {self.__class__} is only available in PyTorch.") + + self.check_model_type(MODEL_FOR_MASK_GENERATION_MAPPING_NAMES) + + def _sanitize_parameters(self, **kwargs): + preprocess_kwargs = {} + postprocess_kwargs = {} + forward_params = {} + # preprocess args + if "points_per_batch" in kwargs: + preprocess_kwargs["points_per_batch"] = kwargs["points_per_batch"] + if "points_per_crop" in kwargs: + preprocess_kwargs["points_per_crop"] = kwargs["points_per_crop"] + if "crops_n_layers" in kwargs: + preprocess_kwargs["crops_n_layers"] = kwargs["crops_n_layers"] + if "crop_overlap_ratio" in kwargs: + preprocess_kwargs["crop_overlap_ratio"] = kwargs["crop_overlap_ratio"] + if "crop_n_points_downscale_factor" in kwargs: + preprocess_kwargs["crop_n_points_downscale_factor"] = kwargs["crop_n_points_downscale_factor"] + if "timeout" in kwargs: + preprocess_kwargs["timeout"] = kwargs["timeout"] + # postprocess args + if "pred_iou_thresh" in kwargs: + forward_params["pred_iou_thresh"] = kwargs["pred_iou_thresh"] + if "stability_score_offset" in kwargs: + forward_params["stability_score_offset"] = kwargs["stability_score_offset"] + if "mask_threshold" in kwargs: + forward_params["mask_threshold"] = kwargs["mask_threshold"] + if "stability_score_thresh" in kwargs: + forward_params["stability_score_thresh"] = kwargs["stability_score_thresh"] + if "crops_nms_thresh" in kwargs: + postprocess_kwargs["crops_nms_thresh"] = kwargs["crops_nms_thresh"] + if "output_rle_mask" in kwargs: + postprocess_kwargs["output_rle_mask"] = kwargs["output_rle_mask"] + if "output_bboxes_mask" in kwargs: + postprocess_kwargs["output_bboxes_mask"] = kwargs["output_bboxes_mask"] + return preprocess_kwargs, forward_params, postprocess_kwargs + + def __call__(self, image, *args, num_workers=None, batch_size=None, **kwargs): + """ + Generates binary segmentation masks + + Args: + inputs (`np.ndarray` or `bytes` or `str` or `dict`): + Image or list of images. + mask_threshold (`float`, *optional*, defaults to 0.0): + Threshold to use when turning the predicted masks into binary values. + pred_iou_thresh (`float`, *optional*, defaults to 0.88): + A filtering threshold in `[0,1]` applied on the model's predicted mask quality. + stability_score_thresh (`float`, *optional*, defaults to 0.95): + A filtering threshold in `[0,1]`, using the stability of the mask under changes to the cutoff used to + binarize the model's mask predictions. + stability_score_offset (`int`, *optional*, defaults to 1): + The amount to shift the cutoff when calculated the stability score. + crops_nms_thresh (`float`, *optional*, defaults to 0.7): + The box IoU cutoff used by non-maximal suppression to filter duplicate masks. + crops_n_layers (`int`, *optional*, defaults to 0): + If `crops_n_layers>0`, mask prediction will be run again on crops of the image. Sets the number of + layers to run, where each layer has 2**i_layer number of image crops. + crop_overlap_ratio (`float`, *optional*, defaults to `512 / 1500`): + Sets the degree to which crops overlap. In the first crop layer, crops will overlap by this fraction of + the image length. Later layers with more crops scale down this overlap. + crop_n_points_downscale_factor (`int`, *optional*, defaults to `1`): + The number of points-per-side sampled in layer n is scaled down by crop_n_points_downscale_factor**n. + timeout (`float`, *optional*, defaults to None): + The maximum time in seconds to wait for fetching images from the web. If None, no timeout is set and + the call may block forever. + + Return: + `Dict`: A dictionary with the following keys: + - **mask** (`PIL.Image`) -- A binary mask of the detected object as a PIL Image of shape `(width, + height)` of the original image. Returns a mask filled with zeros if no object is found. + - **score** (*optional* `float`) -- Optionally, when the model is capable of estimating a confidence of + the "object" described by the label and the mask. + + """ + return super().__call__(image, *args, num_workers=num_workers, batch_size=batch_size, **kwargs) + + def preprocess( + self, + image, + points_per_batch=64, + crops_n_layers: int = 0, + crop_overlap_ratio: float = 512 / 1500, + points_per_crop: Optional[int] = 32, + crop_n_points_downscale_factor: Optional[int] = 1, + timeout: Optional[float] = None, + ): + image = load_image(image, timeout=timeout) + target_size = self.image_processor.size["longest_edge"] + crop_boxes, grid_points, cropped_images, input_labels = self.image_processor.generate_crop_boxes( + image, target_size, crops_n_layers, crop_overlap_ratio, points_per_crop, crop_n_points_downscale_factor + ) + model_inputs = self.image_processor(images=cropped_images, return_tensors="pt") + if self.framework == "pt": + model_inputs = model_inputs.to(self.torch_dtype) + + with self.device_placement(): + if self.framework == "pt": + inference_context = self.get_inference_context() + with inference_context(): + model_inputs = self._ensure_tensor_on_device(model_inputs, device=self.device) + image_embeddings = self.model.get_image_embeddings(model_inputs.pop("pixel_values")) + model_inputs["image_embeddings"] = image_embeddings + + n_points = grid_points.shape[1] + points_per_batch = points_per_batch if points_per_batch is not None else n_points + + if points_per_batch <= 0: + raise ValueError( + "Cannot have points_per_batch<=0. Must be >=1 to returned batched outputs. " + "To return all points at once, set points_per_batch to None" + ) + + for i in range(0, n_points, points_per_batch): + batched_points = grid_points[:, i : i + points_per_batch, :, :] + labels = input_labels[:, i : i + points_per_batch] + is_last = i == n_points - points_per_batch + yield { + "input_points": batched_points, + "input_labels": labels, + "input_boxes": crop_boxes, + "is_last": is_last, + **model_inputs, + } + + def _forward( + self, + model_inputs, + pred_iou_thresh=0.88, + stability_score_thresh=0.95, + mask_threshold=0, + stability_score_offset=1, + ): + input_boxes = model_inputs.pop("input_boxes") + is_last = model_inputs.pop("is_last") + original_sizes = model_inputs.pop("original_sizes").tolist() + reshaped_input_sizes = model_inputs.pop("reshaped_input_sizes").tolist() + + model_outputs = self.model(**model_inputs) + + # post processing happens here in order to avoid CPU GPU copies of ALL the masks + low_resolution_masks = model_outputs["pred_masks"] + masks = self.image_processor.post_process_masks( + low_resolution_masks, original_sizes, reshaped_input_sizes, mask_threshold, binarize=False + ) + iou_scores = model_outputs["iou_scores"] + masks, iou_scores, boxes = self.image_processor.filter_masks( + masks[0], + iou_scores[0], + original_sizes[0], + input_boxes[0], + pred_iou_thresh, + stability_score_thresh, + mask_threshold, + stability_score_offset, + ) + return { + "masks": masks, + "is_last": is_last, + "boxes": boxes, + "iou_scores": iou_scores, + } + + def postprocess( + self, + model_outputs, + output_rle_mask=False, + output_bboxes_mask=False, + crops_nms_thresh=0.7, + ): + all_scores = [] + all_masks = [] + all_boxes = [] + for model_output in model_outputs: + all_scores.append(model_output.pop("iou_scores")) + all_masks.extend(model_output.pop("masks")) + all_boxes.append(model_output.pop("boxes")) + + all_scores = torch.cat(all_scores) + all_boxes = torch.cat(all_boxes) + output_masks, iou_scores, rle_mask, bounding_boxes = self.image_processor.post_process_for_mask_generation( + all_masks, all_scores, all_boxes, crops_nms_thresh + ) + + extra = defaultdict(list) + for output in model_outputs: + for k, v in output.items(): + extra[k].append(v) + + optional = {} + if output_rle_mask: + optional["rle_mask"] = rle_mask + + if output_bboxes_mask: + optional["bounding_boxes"] = bounding_boxes + + return {"masks": output_masks, "scores": iou_scores, **optional, **extra} diff --git a/lib/python3.12/site-packages/transformers/pipelines/text_to_audio.py b/lib/python3.12/site-packages/transformers/pipelines/text_to_audio.py new file mode 100644 index 0000000000000000000000000000000000000000..d0fbe5f115d88b40ca1cf3edf5859eec0c7b9324 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/pipelines/text_to_audio.py @@ -0,0 +1,221 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License.from typing import List, Union +from typing import List, Union + +from ..utils import is_torch_available +from .base import Pipeline + + +if is_torch_available(): + import torch + + from ..models.auto.modeling_auto import MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING + from ..models.speecht5.modeling_speecht5 import SpeechT5HifiGan + +DEFAULT_VOCODER_ID = "microsoft/speecht5_hifigan" + + +class TextToAudioPipeline(Pipeline): + """ + Text-to-audio generation pipeline using any `AutoModelForTextToWaveform` or `AutoModelForTextToSpectrogram`. This + pipeline generates an audio file from an input text and optional other conditional inputs. + + Example: + + ```python + >>> from transformers import pipeline + + >>> pipe = pipeline(model="suno/bark-small") + >>> output = pipe("Hey it's HuggingFace on the phone!") + + >>> audio = output["audio"] + >>> sampling_rate = output["sampling_rate"] + ``` + + Learn more about the basics of using a pipeline in the [pipeline tutorial](../pipeline_tutorial) + + + + You can specify parameters passed to the model by using [`TextToAudioPipeline.__call__.forward_params`] or + [`TextToAudioPipeline.__call__.generate_kwargs`]. + + Example: + + ```python + >>> from transformers import pipeline + + >>> music_generator = pipeline(task="text-to-audio", model="facebook/musicgen-small", framework="pt") + + >>> # diversify the music generation by adding randomness with a high temperature and set a maximum music length + >>> generate_kwargs = { + ... "do_sample": True, + ... "temperature": 0.7, + ... "max_new_tokens": 35, + ... } + + >>> outputs = music_generator("Techno music with high melodic riffs", generate_kwargs=generate_kwargs) + ``` + + + + This pipeline can currently be loaded from [`pipeline`] using the following task identifiers: `"text-to-speech"` or + `"text-to-audio"`. + + See the list of available models on [huggingface.co/models](https://huggingface.co/models?filter=text-to-speech). + """ + + def __init__(self, *args, vocoder=None, sampling_rate=None, **kwargs): + super().__init__(*args, **kwargs) + + if self.framework == "tf": + raise ValueError("The TextToAudioPipeline is only available in PyTorch.") + + self.vocoder = None + if self.model.__class__ in MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING.values(): + self.vocoder = ( + SpeechT5HifiGan.from_pretrained(DEFAULT_VOCODER_ID).to(self.model.device) + if vocoder is None + else vocoder + ) + + self.sampling_rate = sampling_rate + if self.vocoder is not None: + self.sampling_rate = self.vocoder.config.sampling_rate + + if self.sampling_rate is None: + # get sampling_rate from config and generation config + + config = self.model.config + gen_config = self.model.__dict__.get("generation_config", None) + if gen_config is not None: + config.update(gen_config.to_dict()) + + for sampling_rate_name in ["sample_rate", "sampling_rate"]: + sampling_rate = getattr(config, sampling_rate_name, None) + if sampling_rate is not None: + self.sampling_rate = sampling_rate + + def preprocess(self, text, **kwargs): + if isinstance(text, str): + text = [text] + + if self.model.config.model_type == "bark": + # bark Tokenizer is called with BarkProcessor which uses those kwargs + new_kwargs = { + "max_length": self.generation_config.semantic_config.get("max_input_semantic_length", 256), + "add_special_tokens": False, + "return_attention_mask": True, + "return_token_type_ids": False, + "padding": "max_length", + } + + # priority is given to kwargs + new_kwargs.update(kwargs) + + kwargs = new_kwargs + + output = self.tokenizer(text, **kwargs, return_tensors="pt") + + return output + + def _forward(self, model_inputs, **kwargs): + # we expect some kwargs to be additional tensors which need to be on the right device + kwargs = self._ensure_tensor_on_device(kwargs, device=self.device) + forward_params = kwargs["forward_params"] + generate_kwargs = kwargs["generate_kwargs"] + + if self.model.can_generate(): + # we expect some kwargs to be additional tensors which need to be on the right device + generate_kwargs = self._ensure_tensor_on_device(generate_kwargs, device=self.device) + + # User-defined `generation_config` passed to the pipeline call take precedence + if "generation_config" not in generate_kwargs: + generate_kwargs["generation_config"] = self.generation_config + + # generate_kwargs get priority over forward_params + forward_params.update(generate_kwargs) + + output = self.model.generate(**model_inputs, **forward_params) + else: + if len(generate_kwargs): + raise ValueError( + "You're using the `TextToAudioPipeline` with a forward-only model, but `generate_kwargs` is non " + "empty. For forward-only TTA models, please use `forward_params` instead of `generate_kwargs`. " + f"For reference, the `generate_kwargs` used here are: {generate_kwargs.keys()}" + ) + output = self.model(**model_inputs, **forward_params)[0] + + if self.vocoder is not None: + # in that case, the output is a spectrogram that needs to be converted into a waveform + output = self.vocoder(output) + + return output + + def __call__(self, text_inputs: Union[str, List[str]], **forward_params): + """ + Generates speech/audio from the inputs. See the [`TextToAudioPipeline`] documentation for more information. + + Args: + text_inputs (`str` or `List[str]`): + The text(s) to generate. + forward_params (`dict`, *optional*): + Parameters passed to the model generation/forward method. `forward_params` are always passed to the + underlying model. + generate_kwargs (`dict`, *optional*): + The dictionary of ad-hoc parametrization of `generate_config` to be used for the generation call. For a + complete overview of generate, check the [following + guide](https://huggingface.co/docs/transformers/en/main_classes/text_generation). `generate_kwargs` are + only passed to the underlying model if the latter is a generative model. + + Return: + A `dict` or a list of `dict`: The dictionaries have two keys: + + - **audio** (`np.ndarray` of shape `(nb_channels, audio_length)`) -- The generated audio waveform. + - **sampling_rate** (`int`) -- The sampling rate of the generated audio waveform. + """ + return super().__call__(text_inputs, **forward_params) + + def _sanitize_parameters( + self, + preprocess_params=None, + forward_params=None, + generate_kwargs=None, + ): + if self.assistant_model is not None: + generate_kwargs["assistant_model"] = self.assistant_model + if self.assistant_tokenizer is not None: + generate_kwargs["tokenizer"] = self.tokenizer + generate_kwargs["assistant_tokenizer"] = self.assistant_tokenizer + + params = { + "forward_params": forward_params if forward_params else {}, + "generate_kwargs": generate_kwargs if generate_kwargs else {}, + } + + if preprocess_params is None: + preprocess_params = {} + postprocess_params = {} + + return preprocess_params, params, postprocess_params + + def postprocess(self, waveform): + output_dict = {} + if isinstance(waveform, dict): + waveform = waveform["waveform"] + elif isinstance(waveform, tuple): + waveform = waveform[0] + output_dict["audio"] = waveform.to(device="cpu", dtype=torch.float).numpy() + output_dict["sampling_rate"] = self.sampling_rate + + return output_dict