telekinesis-iris 0.0.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (57) hide show
  1. telekinesis_iris-0.0.3/PKG-INFO +179 -0
  2. telekinesis_iris-0.0.3/README.md +133 -0
  3. telekinesis_iris-0.0.3/pyproject.toml +84 -0
  4. telekinesis_iris-0.0.3/setup.cfg +4 -0
  5. telekinesis_iris-0.0.3/src/telekinesis/iris/__init__.py +1 -0
  6. telekinesis_iris-0.0.3/src/telekinesis/iris/dataset.py +199 -0
  7. telekinesis_iris-0.0.3/src/telekinesis/iris/deploy.py +479 -0
  8. telekinesis_iris-0.0.3/src/telekinesis/iris/export.py +151 -0
  9. telekinesis_iris-0.0.3/src/telekinesis/iris/generate.py +323 -0
  10. telekinesis_iris-0.0.3/src/telekinesis/iris/logger.py +100 -0
  11. telekinesis_iris-0.0.3/src/telekinesis/iris/models/__init__.py +15 -0
  12. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/__init__.py +111 -0
  13. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/_assignment.py +332 -0
  14. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/_defaults.py +115 -0
  15. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/_namespace.py +155 -0
  16. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/_types.py +106 -0
  17. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/assets/__init__.py +17 -0
  18. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/assets/model_weights.py +168 -0
  19. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/backbone/__init__.py +129 -0
  20. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/backbone/backbone.py +275 -0
  21. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/backbone/base.py +22 -0
  22. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/backbone/dinov2.py +298 -0
  23. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/backbone/dinov2_with_windowed_attn.py +1444 -0
  24. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/backbone/projector.py +342 -0
  25. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/config.py +1925 -0
  26. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/criterion.py +1483 -0
  27. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/heads/__init__.py +20 -0
  28. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/heads/keypoints.py +654 -0
  29. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/heads/segmentation.py +529 -0
  30. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/lwdetr.py +1080 -0
  31. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/matcher.py +1233 -0
  32. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/math.py +98 -0
  33. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/ops/__init__.py +7 -0
  34. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/ops/functions/__init__.py +22 -0
  35. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/ops/functions/ms_deform_attn_func.py +97 -0
  36. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/ops/modules/__init__.py +19 -0
  37. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/ops/modules/ms_deform_attn.py +282 -0
  38. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/position_encoding.py +279 -0
  39. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/postprocess.py +583 -0
  40. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/transformer.py +1682 -0
  41. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/utilities/__init__.py +59 -0
  42. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/utilities/box_ops.py +387 -0
  43. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/utilities/compiler.py +32 -0
  44. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/utilities/distributed.py +202 -0
  45. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/utilities/files.py +65 -0
  46. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/utilities/io.py +110 -0
  47. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/utilities/keypoints.py +234 -0
  48. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/utilities/logger.py +145 -0
  49. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/utilities/state_dict.py +430 -0
  50. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/utilities/tensors.py +812 -0
  51. telekinesis_iris-0.0.3/src/telekinesis/iris/models/rfdetr/weights.py +732 -0
  52. telekinesis_iris-0.0.3/src/telekinesis/iris/trainer.py +952 -0
  53. telekinesis_iris-0.0.3/src/telekinesis_iris.egg-info/PKG-INFO +179 -0
  54. telekinesis_iris-0.0.3/src/telekinesis_iris.egg-info/SOURCES.txt +55 -0
  55. telekinesis_iris-0.0.3/src/telekinesis_iris.egg-info/dependency_links.txt +1 -0
  56. telekinesis_iris-0.0.3/src/telekinesis_iris.egg-info/requires.txt +29 -0
  57. telekinesis_iris-0.0.3/src/telekinesis_iris.egg-info/top_level.txt +1 -0
@@ -0,0 +1,179 @@
1
+ Metadata-Version: 2.4
2
+ Name: telekinesis-iris
3
+ Version: 0.0.3
4
+ Summary: Model training and inference repository
5
+ Author-email: Telekinesis <support@telekinesis.ai>
6
+ License-Expression: LicenseRef-Proprietary
7
+ Project-URL: Homepage, https://telekinesis.ai
8
+ Project-URL: Documentation, https://docs.telekinesis.ai
9
+ Project-URL: Telekinesis Examples Repository, https://github.com/telekinesis-ai/telekinesis-examples
10
+ Classifier: Programming Language :: Python :: 3
11
+ Classifier: Programming Language :: Python :: 3.10
12
+ Classifier: Programming Language :: Python :: 3.11
13
+ Classifier: Programming Language :: Python :: 3.12
14
+ Classifier: Operating System :: OS Independent
15
+ Classifier: Intended Audience :: Developers
16
+ Classifier: Topic :: Software Development :: Libraries
17
+ Requires-Python: >=3.11
18
+ Description-Content-Type: text/markdown
19
+ Requires-Dist: loguru>=0.7.0
20
+ Requires-Dist: opencv-python>=4.7.0.72
21
+ Requires-Dist: pycocotools>=2.0.6
22
+ Requires-Dist: numpy>=1.24.0
23
+ Requires-Dist: pydantic>=2.0
24
+ Requires-Dist: requests>=2.31.0
25
+ Requires-Dist: rerun-sdk>=0.34.0
26
+ Requires-Dist: scipy>=1.10.0
27
+ Requires-Dist: torch>=2.2.0
28
+ Requires-Dist: torch-hungarian==0.1.0rc0
29
+ Requires-Dist: tensorboard>=2.16.0
30
+ Requires-Dist: tqdm>=4.66.0
31
+ Requires-Dist: transformers<6,>=5.0.0
32
+ Requires-Dist: torchvision>=0.17.0
33
+ Requires-Dist: onnx>=1.16.0
34
+ Requires-Dist: onnxscript>=0.1.0
35
+ Requires-Dist: onnxruntime>=1.18.0
36
+ Requires-Dist: torchmetrics[detection]<1.9.0,>=1.8.2
37
+ Provides-Extra: docs
38
+ Requires-Dist: mkdocs<2,>=1.6; extra == "docs"
39
+ Requires-Dist: mkdocstrings[python]<1,>=0.25; extra == "docs"
40
+ Requires-Dist: mkdocs-awesome-pages-plugin<3,>=2; extra == "docs"
41
+ Provides-Extra: dev
42
+ Requires-Dist: ruff==0.15.8; extra == "dev"
43
+ Requires-Dist: pylint==4.0.5; extra == "dev"
44
+ Requires-Dist: pytest>=8.0.0; extra == "dev"
45
+ Requires-Dist: pytest-cov>=5.0.0; extra == "dev"
46
+
47
+ <div align="center">
48
+ <p>
49
+ <a href="https://telekinesis.ai" target="_blank">
50
+ <img
51
+ width="100%"
52
+ src="https://assets.telekinesis.ai/logo/telekinesis-banner.png"
53
+ alt="Telekinesis"
54
+ >
55
+ </a>
56
+ </p>
57
+
58
+ <br>
59
+
60
+ [![PyPI version](https://img.shields.io/pypi/v/telekinesis-iris)](https://pypi.org/project/telekinesis-iris/)
61
+ [![License](https://img.shields.io/pypi/l/telekinesis-iris)](https://pypi.org/project/telekinesis-iris/)
62
+ [![Python versions](https://img.shields.io/pypi/pyversions/telekinesis-iris)](https://pypi.org/project/telekinesis-iris/)
63
+ </div>
64
+
65
+ <p align="center">
66
+ <a href="https://github.com/telekinesis-ai">GitHub</a>
67
+ &nbsp;•&nbsp;
68
+ <a href="https://www.linkedin.com/company/telekinesis-ai/">LinkedIn</a>
69
+ &nbsp;•&nbsp;
70
+ <a href="https://x.com/telekinesis_ai">X</a>
71
+ &nbsp;•&nbsp;
72
+ <a href="https://discord.gg/7NnQ3bQHqm">Discord</a>
73
+ </p>
74
+
75
+ # Telekinesis Iris
76
+
77
+ Telekinesis Iris is a computer vision library for creating COCO datasets and
78
+ training, exporting, and deploying object detection and instance segmentation
79
+ models.
80
+
81
+ It includes:
82
+
83
+ - COCO dataset generation and loading with bounding boxes and segmentation masks
84
+ - Local RF-DETR detection and segmentation model implementations
85
+ - PyTorch training, validation, checkpointing, and TensorBoard logging
86
+ - ONNX export and ONNX Runtime inference
87
+ - Rerun visualization for datasets and predictions
88
+
89
+ ## Release Model
90
+
91
+ Telekinesis Iris is currently in active development (pre-1.0). APIs may evolve
92
+ between minor releases. Install the latest package version for the newest
93
+ features and fixes.
94
+
95
+ ## Installation
96
+
97
+ 1. Install [Miniconda](https://docs.conda.io/projects/miniconda/en/latest/).
98
+
99
+ 2. Create a Python 3.11 environment:
100
+
101
+ ```bash
102
+ conda create -n telekinesis-iris python=3.11
103
+ ```
104
+
105
+ 3. Activate the environment:
106
+
107
+ ```bash
108
+ conda activate telekinesis-iris
109
+ ```
110
+
111
+ 4. Install PyTorch and TorchVision for CUDA 12.8, then install the package:
112
+
113
+ ```bash
114
+ pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128
115
+ ```
116
+
117
+ ```bash
118
+ pip install telekinesis-iris
119
+ ```
120
+
121
+ The Python package is imported as `telekinesis.iris`, while the package published on PyPI is named `telekinesis-iris`.
122
+
123
+ ## Example
124
+
125
+ Train an RF-DETR segmentation model on a COCO dataset:
126
+
127
+ ```python
128
+ from telekinesis.iris.dataset import COCODataset, ResizeSample
129
+ from telekinesis.iris.models import (
130
+ RFDETRSegNanoConfig,
131
+ SegmentationTrainConfig,
132
+ build_criterion_from_config,
133
+ build_model_from_config,
134
+ load_pretrain_weights,
135
+ )
136
+ from telekinesis.iris.trainer import Trainer
137
+
138
+ model_config = RFDETRSegNanoConfig(num_classes=3)
139
+ train_config = SegmentationTrainConfig(
140
+ dataset_dir="dataset/train",
141
+ output_dir="results/seg-nano",
142
+ epochs=10,
143
+ batch_size=4,
144
+ )
145
+
146
+ model = build_model_from_config(model_config, train_config)
147
+ load_pretrain_weights(model, model_config)
148
+ criterion, _ = build_criterion_from_config(model_config, train_config)
149
+ dataset = COCODataset(
150
+ "dataset/train",
151
+ transforms=ResizeSample(model_config.resolution),
152
+ include_masks=True,
153
+ )
154
+
155
+ Trainer(
156
+ model=model,
157
+ criterion=criterion,
158
+ dataset=dataset,
159
+ output_dir="results/seg-nano",
160
+ epochs=10,
161
+ batch_size=4,
162
+ evaluate_masks=True,
163
+ ).train()
164
+ ```
165
+
166
+ ## Resources
167
+
168
+ - [Telekinesis Examples](https://github.com/telekinesis-ai/telekinesis-examples)
169
+ - [Telekinesis Documentation](https://docs.telekinesis.ai)
170
+ - [Development guide](DEVELOPMENT.md)
171
+ - [Changelog](CHANGELOG.md)
172
+
173
+ ## Support
174
+
175
+ For issues and questions:
176
+
177
+ - Create an [issue](https://github.com/telekinesis-ai/telekinesis-examples/issues).
178
+ - Contact `support@telekinesis.ai`.
179
+ - Join the [Telekinesis Discord](https://discord.com/invite/7NnQ3bQHqm).
@@ -0,0 +1,133 @@
1
+ <div align="center">
2
+ <p>
3
+ <a href="https://telekinesis.ai" target="_blank">
4
+ <img
5
+ width="100%"
6
+ src="https://assets.telekinesis.ai/logo/telekinesis-banner.png"
7
+ alt="Telekinesis"
8
+ >
9
+ </a>
10
+ </p>
11
+
12
+ <br>
13
+
14
+ [![PyPI version](https://img.shields.io/pypi/v/telekinesis-iris)](https://pypi.org/project/telekinesis-iris/)
15
+ [![License](https://img.shields.io/pypi/l/telekinesis-iris)](https://pypi.org/project/telekinesis-iris/)
16
+ [![Python versions](https://img.shields.io/pypi/pyversions/telekinesis-iris)](https://pypi.org/project/telekinesis-iris/)
17
+ </div>
18
+
19
+ <p align="center">
20
+ <a href="https://github.com/telekinesis-ai">GitHub</a>
21
+ &nbsp;•&nbsp;
22
+ <a href="https://www.linkedin.com/company/telekinesis-ai/">LinkedIn</a>
23
+ &nbsp;•&nbsp;
24
+ <a href="https://x.com/telekinesis_ai">X</a>
25
+ &nbsp;•&nbsp;
26
+ <a href="https://discord.gg/7NnQ3bQHqm">Discord</a>
27
+ </p>
28
+
29
+ # Telekinesis Iris
30
+
31
+ Telekinesis Iris is a computer vision library for creating COCO datasets and
32
+ training, exporting, and deploying object detection and instance segmentation
33
+ models.
34
+
35
+ It includes:
36
+
37
+ - COCO dataset generation and loading with bounding boxes and segmentation masks
38
+ - Local RF-DETR detection and segmentation model implementations
39
+ - PyTorch training, validation, checkpointing, and TensorBoard logging
40
+ - ONNX export and ONNX Runtime inference
41
+ - Rerun visualization for datasets and predictions
42
+
43
+ ## Release Model
44
+
45
+ Telekinesis Iris is currently in active development (pre-1.0). APIs may evolve
46
+ between minor releases. Install the latest package version for the newest
47
+ features and fixes.
48
+
49
+ ## Installation
50
+
51
+ 1. Install [Miniconda](https://docs.conda.io/projects/miniconda/en/latest/).
52
+
53
+ 2. Create a Python 3.11 environment:
54
+
55
+ ```bash
56
+ conda create -n telekinesis-iris python=3.11
57
+ ```
58
+
59
+ 3. Activate the environment:
60
+
61
+ ```bash
62
+ conda activate telekinesis-iris
63
+ ```
64
+
65
+ 4. Install PyTorch and TorchVision for CUDA 12.8, then install the package:
66
+
67
+ ```bash
68
+ pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128
69
+ ```
70
+
71
+ ```bash
72
+ pip install telekinesis-iris
73
+ ```
74
+
75
+ The Python package is imported as `telekinesis.iris`, while the package published on PyPI is named `telekinesis-iris`.
76
+
77
+ ## Example
78
+
79
+ Train an RF-DETR segmentation model on a COCO dataset:
80
+
81
+ ```python
82
+ from telekinesis.iris.dataset import COCODataset, ResizeSample
83
+ from telekinesis.iris.models import (
84
+ RFDETRSegNanoConfig,
85
+ SegmentationTrainConfig,
86
+ build_criterion_from_config,
87
+ build_model_from_config,
88
+ load_pretrain_weights,
89
+ )
90
+ from telekinesis.iris.trainer import Trainer
91
+
92
+ model_config = RFDETRSegNanoConfig(num_classes=3)
93
+ train_config = SegmentationTrainConfig(
94
+ dataset_dir="dataset/train",
95
+ output_dir="results/seg-nano",
96
+ epochs=10,
97
+ batch_size=4,
98
+ )
99
+
100
+ model = build_model_from_config(model_config, train_config)
101
+ load_pretrain_weights(model, model_config)
102
+ criterion, _ = build_criterion_from_config(model_config, train_config)
103
+ dataset = COCODataset(
104
+ "dataset/train",
105
+ transforms=ResizeSample(model_config.resolution),
106
+ include_masks=True,
107
+ )
108
+
109
+ Trainer(
110
+ model=model,
111
+ criterion=criterion,
112
+ dataset=dataset,
113
+ output_dir="results/seg-nano",
114
+ epochs=10,
115
+ batch_size=4,
116
+ evaluate_masks=True,
117
+ ).train()
118
+ ```
119
+
120
+ ## Resources
121
+
122
+ - [Telekinesis Examples](https://github.com/telekinesis-ai/telekinesis-examples)
123
+ - [Telekinesis Documentation](https://docs.telekinesis.ai)
124
+ - [Development guide](DEVELOPMENT.md)
125
+ - [Changelog](CHANGELOG.md)
126
+
127
+ ## Support
128
+
129
+ For issues and questions:
130
+
131
+ - Create an [issue](https://github.com/telekinesis-ai/telekinesis-examples/issues).
132
+ - Contact `support@telekinesis.ai`.
133
+ - Join the [Telekinesis Discord](https://discord.com/invite/7NnQ3bQHqm).
@@ -0,0 +1,84 @@
1
+ [build-system]
2
+ requires = [
3
+ "setuptools",
4
+ ]
5
+ build-backend = "setuptools.build_meta"
6
+
7
+ [project]
8
+ name = "telekinesis-iris"
9
+ version = "0.0.3"
10
+ description = "Model training and inference repository"
11
+ readme = "README.md"
12
+ requires-python = ">=3.11"
13
+ license = "LicenseRef-Proprietary"
14
+ license-files = [
15
+ "LICEN[CS]E*",
16
+ ]
17
+ authors = [
18
+ { name = "Telekinesis", email = "support@telekinesis.ai" },
19
+ ]
20
+ classifiers = [
21
+ "Programming Language :: Python :: 3",
22
+ "Programming Language :: Python :: 3.10",
23
+ "Programming Language :: Python :: 3.11",
24
+ "Programming Language :: Python :: 3.12",
25
+ "Operating System :: OS Independent",
26
+ "Intended Audience :: Developers",
27
+ "Topic :: Software Development :: Libraries",
28
+ ]
29
+ dependencies = [
30
+ "loguru>=0.7.0",
31
+ "opencv-python>=4.7.0.72",
32
+ "pycocotools>=2.0.6",
33
+ "numpy>=1.24.0",
34
+ "pydantic>=2.0",
35
+ "requests>=2.31.0",
36
+ "rerun-sdk>=0.34.0",
37
+ "scipy>=1.10.0",
38
+ "torch>=2.2.0",
39
+ "torch-hungarian==0.1.0rc0",
40
+ "tensorboard>=2.16.0",
41
+ "tqdm>=4.66.0",
42
+ "transformers>=5.0.0,<6",
43
+ "torchvision>=0.17.0",
44
+ "onnx>=1.16.0",
45
+ "onnxscript>=0.1.0",
46
+ "onnxruntime>=1.18.0",
47
+ "torchmetrics[detection]>=1.8.2,<1.9.0",
48
+ ]
49
+
50
+ [project.optional-dependencies]
51
+ docs = [
52
+ "mkdocs>=1.6,<2",
53
+ "mkdocstrings[python]>=0.25,<1",
54
+ "mkdocs-awesome-pages-plugin>=2,<3",
55
+ ]
56
+ dev = [
57
+ "ruff==0.15.8",
58
+ "pylint==4.0.5",
59
+ "pytest>=8.0.0",
60
+ "pytest-cov>=5.0.0",
61
+ ]
62
+
63
+ [project.urls]
64
+ Homepage = "https://telekinesis.ai"
65
+ Documentation = "https://docs.telekinesis.ai"
66
+ "Telekinesis Examples Repository" = "https://github.com/telekinesis-ai/telekinesis-examples"
67
+
68
+ [tool.ruff]
69
+ line-length = 100
70
+
71
+ [tool.setuptools.packages.find]
72
+ where = [
73
+ "src",
74
+ ]
75
+ namespaces = true
76
+
77
+ [tool.pytest.ini_options]
78
+ testpaths = [
79
+ "tests",
80
+ ]
81
+ python_files = [
82
+ "*_test.py",
83
+ ]
84
+ addopts = "-ra"
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1 @@
1
+ """Computer vision dataset, training, export, and inference tools."""
@@ -0,0 +1,199 @@
1
+ """COCO dataset loading and sample transformation utilities."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from pathlib import Path
6
+ from typing import Any
7
+
8
+ import numpy as np
9
+ import torch
10
+ from torch import Tensor
11
+ from torch.utils.data import Dataset
12
+ from PIL import Image
13
+ from pycocotools.coco import COCO
14
+ from torchvision.transforms import InterpolationMode
15
+ from torchvision.transforms import functional as transform_functional
16
+
17
+
18
+ class ResizeSample:
19
+ """Resize an image and its spatial annotations to a square resolution."""
20
+
21
+ def __init__(self, resolution: int) -> None:
22
+ """Initializes a square resize transform.
23
+
24
+ Args:
25
+ resolution: Output image height and width in pixels.
26
+
27
+ """
28
+ self.resolution = resolution
29
+
30
+ def __call__(self, image: Tensor, target: dict[str, Any]) -> tuple[Tensor, dict[str, Any]]:
31
+ """Resizes an image and its spatial annotations.
32
+
33
+ Args:
34
+ image: Image tensor in ``CHW`` layout.
35
+ target: Detection target containing boxes and optional masks.
36
+
37
+ Returns:
38
+ The resized image and updated target.
39
+
40
+ """
41
+ source_height, source_width = image.shape[-2:]
42
+ scale_x = self.resolution / source_width
43
+ scale_y = self.resolution / source_height
44
+ target = target.copy()
45
+
46
+ image = transform_functional.resize(
47
+ image,
48
+ [self.resolution, self.resolution],
49
+ interpolation=InterpolationMode.BILINEAR,
50
+ antialias=True,
51
+ )
52
+
53
+ boxes = target.get("boxes")
54
+ if isinstance(boxes, Tensor):
55
+ boxes = boxes.clone()
56
+ boxes[:, 0::2] *= scale_x
57
+ boxes[:, 1::2] *= scale_y
58
+ target["boxes"] = boxes
59
+
60
+ keypoints = target.get("keypoints")
61
+ if isinstance(keypoints, Tensor):
62
+ keypoints = keypoints.clone()
63
+ keypoints[..., 0] *= scale_x
64
+ keypoints[..., 1] *= scale_y
65
+ target["keypoints"] = keypoints
66
+
67
+ masks = target.get("masks")
68
+ if isinstance(masks, Tensor):
69
+ if masks.shape[0] == 0:
70
+ target["masks"] = masks.new_empty((0, self.resolution, self.resolution))
71
+ else:
72
+ target["masks"] = transform_functional.resize(
73
+ masks,
74
+ [self.resolution, self.resolution],
75
+ interpolation=InterpolationMode.NEAREST,
76
+ )
77
+
78
+ area = target.get("area")
79
+ if isinstance(area, Tensor):
80
+ target["area"] = area * scale_x * scale_y
81
+ target["size"] = torch.tensor([self.resolution, self.resolution])
82
+ return image, target
83
+
84
+
85
+ class COCODataset(Dataset):
86
+ """Load a COCO detection dataset for model training."""
87
+
88
+ def __init__(
89
+ self,
90
+ root: str | Path,
91
+ transforms=None,
92
+ *,
93
+ include_masks: bool = False,
94
+ ) -> None:
95
+ """Initializes a COCO detection dataset.
96
+
97
+ Args:
98
+ root: Root directory containing images and COCO annotations.
99
+ transforms: Optional callable applied to each image and target.
100
+ include_masks: Decode segmentation masks for segmentation models.
101
+
102
+ Raises:
103
+ FileNotFoundError: If the annotations file or image directory is missing.
104
+
105
+ """
106
+ self.root = Path(root)
107
+ rfdetr_annotations = self.root / "_annotations.coco.json"
108
+ if rfdetr_annotations.is_file():
109
+ self.images_dir = self.root
110
+ self.annotations_path = rfdetr_annotations
111
+ else:
112
+ self.images_dir = self.root / "images"
113
+ self.annotations_path = self.root / "annotations.json"
114
+
115
+ if not self.annotations_path.is_file():
116
+ raise FileNotFoundError(f"COCO annotations not found: {self.annotations_path}")
117
+ if not self.images_dir.is_dir():
118
+ raise FileNotFoundError(f"Image directory not found: {self.images_dir}")
119
+
120
+ self.coco = COCO(str(self.annotations_path))
121
+
122
+ self.image_ids = sorted(self.coco.getImgIds())
123
+
124
+ self.transforms = transforms
125
+ self.include_masks = include_masks
126
+
127
+ def __len__(self) -> int:
128
+ """Returns the number of images in the dataset."""
129
+ return len(self.image_ids)
130
+
131
+ def __getitem__(self, index: int) -> tuple[torch.Tensor, dict]:
132
+ """Loads an image and its detection target.
133
+
134
+ Args:
135
+ index: Zero-based dataset index.
136
+
137
+ Returns:
138
+ An image tensor and target dictionary.
139
+
140
+ Raises:
141
+ ValueError: If masks are requested but an annotation has no segmentation.
142
+
143
+ """
144
+ image_id = self.image_ids[index]
145
+
146
+ image_info = self.coco.loadImgs([image_id])[0]
147
+ image_path = self.images_dir / image_info["file_name"]
148
+ image = Image.open(image_path).convert("RGB")
149
+ image = transform_functional.pil_to_tensor(image).float() / 255.0
150
+ annotation_ids = self.coco.getAnnIds(imgIds=[image_id])
151
+ annotations = self.coco.loadAnns(annotation_ids)
152
+
153
+ boxes = []
154
+ labels = []
155
+ areas = []
156
+ iscrowd = []
157
+ masks = []
158
+
159
+ for annotation in annotations:
160
+ x, y, width, height = annotation["bbox"]
161
+
162
+ boxes.append([x, y, x + width, y + height])
163
+
164
+ labels.append(annotation["category_id"])
165
+
166
+ areas.append(annotation["area"])
167
+
168
+ iscrowd.append(annotation.get("iscrowd", 0))
169
+
170
+ if self.include_masks:
171
+ if not annotation.get("segmentation"):
172
+ raise ValueError(f"Annotation {annotation['id']} has no segmentation mask")
173
+ masks.append(self.coco.annToMask(annotation))
174
+
175
+ boxes = torch.tensor(boxes, dtype=torch.float32).reshape(-1, 4)
176
+
177
+ labels = torch.tensor(labels, dtype=torch.int64)
178
+ areas = torch.tensor(areas, dtype=torch.float32)
179
+ iscrowd = torch.tensor(iscrowd, dtype=torch.int64)
180
+
181
+ target = {
182
+ "boxes": boxes,
183
+ "labels": labels,
184
+ "image_id": torch.tensor(image_id, dtype=torch.int64),
185
+ "area": areas,
186
+ "iscrowd": iscrowd,
187
+ }
188
+
189
+ if self.include_masks:
190
+ target["masks"] = (
191
+ torch.tensor(np.stack(masks), dtype=torch.uint8)
192
+ if masks
193
+ else torch.empty((0, image.shape[-2], image.shape[-1]), dtype=torch.uint8)
194
+ )
195
+
196
+ if self.transforms is not None:
197
+ image, target = self.transforms(image, target)
198
+
199
+ return image, target