SDK: Add an adapter layer that presents a CVAT task as a torchvision dataset (#5417)
Roman Donchenko
3 years ago
committed by
GitHub
parent
82adde42aa
commit
487c60ce2b
@ -0,0 +1,359 @@
|
||||
import base64
|
||||
import collections
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import types
|
||||
import zipfile
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
from typing import (
|
||||
Callable,
|
||||
Dict,
|
||||
FrozenSet,
|
||||
List,
|
||||
Mapping,
|
||||
Optional,
|
||||
Sequence,
|
||||
Tuple,
|
||||
Type,
|
||||
TypeVar,
|
||||
)
|
||||
|
||||
import appdirs
|
||||
import attrs
|
||||
import attrs.validators
|
||||
import PIL.Image
|
||||
import torchvision.datasets
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
import cvat_sdk.core
|
||||
import cvat_sdk.core.exceptions
|
||||
from cvat_sdk.api_client.model_utils import to_json
|
||||
from cvat_sdk.core.utils import atomic_writer
|
||||
from cvat_sdk.models import DataMetaRead, LabeledData, LabeledImage, LabeledShape, TaskRead
|
||||
|
||||
_ModelType = TypeVar("_ModelType")
|
||||
|
||||
_CACHE_DIR = Path(appdirs.user_cache_dir("cvat-sdk", "CVAT.ai"))
|
||||
_NUM_DOWNLOAD_THREADS = 4
|
||||
|
||||
|
||||
class UnsupportedDatasetError(cvat_sdk.core.exceptions.CvatSdkException):
|
||||
pass
|
||||
|
||||
|
||||
@attrs.frozen
|
||||
class FrameAnnotations:
|
||||
"""
|
||||
Contains annotations that pertain to a single frame.
|
||||
"""
|
||||
|
||||
tags: List[LabeledImage] = attrs.Factory(list)
|
||||
shapes: List[LabeledShape] = attrs.Factory(list)
|
||||
|
||||
|
||||
@attrs.frozen
|
||||
class Target:
|
||||
"""
|
||||
Non-image data for a dataset sample.
|
||||
"""
|
||||
|
||||
annotations: FrameAnnotations
|
||||
"""Annotations for the frame corresponding to the sample."""
|
||||
|
||||
label_id_to_index: Mapping[int, int]
|
||||
"""
|
||||
A mapping from label_id values in `LabeledImage` and `LabeledShape` objects
|
||||
to an index in the range [0, num_labels), where num_labels is the number of labels
|
||||
defined in the task. This mapping is consistent across all samples for a given task.
|
||||
"""
|
||||
|
||||
|
||||
class TaskVisionDataset(torchvision.datasets.VisionDataset):
|
||||
"""
|
||||
Represents a task on a CVAT server as a PyTorch Dataset.
|
||||
|
||||
This dataset contains one sample for each frame in the task, in the same
|
||||
order as the frames are in the task. Deleted frames are omitted.
|
||||
Before transforms are applied, each sample is a tuple of
|
||||
(image, target), where:
|
||||
|
||||
* image is a `PIL.Image.Image` object for the corresponding frame.
|
||||
* target is a `Target` object containing annotations for the frame.
|
||||
|
||||
This class caches all data and annotations for the task on the local file system
|
||||
during construction. If the task is updated on the server, the cache is updated.
|
||||
|
||||
Limitations:
|
||||
|
||||
* Only tasks with image (not video) data are supported at the moment.
|
||||
* Track annotations are currently not accessible.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: cvat_sdk.core.Client,
|
||||
task_id: int,
|
||||
*,
|
||||
transforms: Optional[Callable] = None,
|
||||
transform: Optional[Callable] = None,
|
||||
target_transform: Optional[Callable] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Creates a dataset corresponding to the task with ID `task_id` on the
|
||||
server that `client` is connected to.
|
||||
|
||||
`transforms`, `transform` and `target_transforms` are optional transformation
|
||||
functions; see the documentation for `torchvision.datasets.VisionDataset` for
|
||||
more information.
|
||||
"""
|
||||
|
||||
self._logger = client.logger
|
||||
|
||||
self._logger.info(f"Fetching task {task_id}...")
|
||||
self._task = client.tasks.retrieve(task_id)
|
||||
|
||||
if not self._task.size or not self._task.data_chunk_size:
|
||||
raise UnsupportedDatasetError("The task has no data")
|
||||
|
||||
if self._task.data_original_chunk_type != "imageset":
|
||||
raise UnsupportedDatasetError(
|
||||
f"{self.__class__.__name__} only supports tasks with image chunks;"
|
||||
f" current chunk type is {self._task.data_original_chunk_type!r}"
|
||||
)
|
||||
|
||||
# Base64-encode the name to avoid FS-unsafe characters (like slashes)
|
||||
server_dir_name = (
|
||||
base64.urlsafe_b64encode(client.api_map.host.encode()).rstrip(b"=").decode()
|
||||
)
|
||||
server_dir = _CACHE_DIR / f"servers/{server_dir_name}"
|
||||
|
||||
self._task_dir = server_dir / f"tasks/{self._task.id}"
|
||||
self._initialize_task_dir()
|
||||
|
||||
super().__init__(
|
||||
os.fspath(self._task_dir),
|
||||
transforms=transforms,
|
||||
transform=transform,
|
||||
target_transform=target_transform,
|
||||
)
|
||||
|
||||
data_meta = self._ensure_model(
|
||||
"data_meta.json", DataMetaRead, self._task.get_meta, "data metadata"
|
||||
)
|
||||
self._active_frame_indexes = sorted(
|
||||
set(range(self._task.size)) - set(data_meta.deleted_frames)
|
||||
)
|
||||
|
||||
self._logger.info("Downloading chunks...")
|
||||
|
||||
self._chunk_dir = self._task_dir / "chunks"
|
||||
self._chunk_dir.mkdir(exist_ok=True, parents=True)
|
||||
|
||||
needed_chunks = {
|
||||
index // self._task.data_chunk_size for index in self._active_frame_indexes
|
||||
}
|
||||
|
||||
with ThreadPoolExecutor(_NUM_DOWNLOAD_THREADS) as pool:
|
||||
for _ in pool.map(self._ensure_chunk, sorted(needed_chunks)):
|
||||
# just need to loop through all results so that any exceptions are propagated
|
||||
pass
|
||||
|
||||
self._logger.info("All chunks downloaded")
|
||||
|
||||
self._label_id_to_index = types.MappingProxyType(
|
||||
{
|
||||
label["id"]: label_index
|
||||
for label_index, label in enumerate(sorted(self._task.labels, key=lambda l: l.id))
|
||||
}
|
||||
)
|
||||
|
||||
annotations = self._ensure_model(
|
||||
"annotations.json", LabeledData, self._task.get_annotations, "annotations"
|
||||
)
|
||||
|
||||
self._frame_annotations: Dict[int, FrameAnnotations] = collections.defaultdict(
|
||||
FrameAnnotations
|
||||
)
|
||||
|
||||
for tag in annotations.tags:
|
||||
self._frame_annotations[tag.frame].tags.append(tag)
|
||||
|
||||
for shape in annotations.shapes:
|
||||
self._frame_annotations[shape.frame].shapes.append(shape)
|
||||
|
||||
# TODO: tracks?
|
||||
|
||||
def _initialize_task_dir(self) -> None:
|
||||
task_json_path = self._task_dir / "task.json"
|
||||
|
||||
try:
|
||||
with open(task_json_path, "rb") as task_json_file:
|
||||
saved_task = TaskRead._new_from_openapi_data(**json.load(task_json_file))
|
||||
except Exception:
|
||||
self._logger.info("Task is not yet cached or the cache is corrupted")
|
||||
|
||||
# If the cache was corrupted, the directory might already be there; clear it.
|
||||
if self._task_dir.exists():
|
||||
shutil.rmtree(self._task_dir)
|
||||
else:
|
||||
if saved_task.updated_date < self._task.updated_date:
|
||||
self._logger.info(
|
||||
"Task has been updated on the server since it was cached; purging the cache"
|
||||
)
|
||||
shutil.rmtree(self._task_dir)
|
||||
|
||||
self._task_dir.mkdir(exist_ok=True, parents=True)
|
||||
|
||||
with atomic_writer(task_json_path, "w", encoding="UTF-8") as task_json_file:
|
||||
json.dump(to_json(self._task._model), task_json_file, indent=4)
|
||||
print(file=task_json_file) # add final newline
|
||||
|
||||
def _ensure_chunk(self, chunk_index: int) -> None:
|
||||
chunk_path = self._chunk_dir / f"{chunk_index}.zip"
|
||||
if chunk_path.exists():
|
||||
return # already downloaded previously
|
||||
|
||||
self._logger.info(f"Downloading chunk #{chunk_index}...")
|
||||
|
||||
with atomic_writer(chunk_path, "wb") as chunk_file:
|
||||
self._task.download_chunk(chunk_index, chunk_file, quality="original")
|
||||
|
||||
def _ensure_model(
|
||||
self,
|
||||
filename: str,
|
||||
model_type: Type[_ModelType],
|
||||
download: Callable[[], _ModelType],
|
||||
model_description: str,
|
||||
) -> _ModelType:
|
||||
path = self._task_dir / filename
|
||||
|
||||
try:
|
||||
with open(path, "rb") as f:
|
||||
model = model_type._new_from_openapi_data(**json.load(f))
|
||||
self._logger.info(f"Loaded {model_description} from cache")
|
||||
return model
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
except Exception:
|
||||
self._logger.warning(f"Failed to load {model_description} from cache", exc_info=True)
|
||||
|
||||
self._logger.info(f"Downloading {model_description}...")
|
||||
model = download()
|
||||
self._logger.info(f"Downloaded {model_description}")
|
||||
|
||||
with atomic_writer(path, "w", encoding="UTF-8") as f:
|
||||
json.dump(to_json(model), f, indent=4)
|
||||
print(file=f) # add final newline
|
||||
|
||||
return model
|
||||
|
||||
def __getitem__(self, sample_index: int):
|
||||
"""
|
||||
Returns the sample with index `sample_index`.
|
||||
|
||||
`sample_index` must satisfy the condition `0 <= sample_index < len(self)`.
|
||||
"""
|
||||
|
||||
frame_index = self._active_frame_indexes[sample_index]
|
||||
chunk_index = frame_index // self._task.data_chunk_size
|
||||
member_index = frame_index % self._task.data_chunk_size
|
||||
|
||||
with zipfile.ZipFile(self._chunk_dir / f"{chunk_index}.zip", "r") as chunk_zip:
|
||||
with chunk_zip.open(chunk_zip.infolist()[member_index]) as chunk_member:
|
||||
sample_image = PIL.Image.open(chunk_member)
|
||||
sample_image.load()
|
||||
|
||||
sample_target = Target(
|
||||
annotations=self._frame_annotations[frame_index],
|
||||
label_id_to_index=self._label_id_to_index,
|
||||
)
|
||||
|
||||
if self.transforms:
|
||||
sample_image, sample_target = self.transforms(sample_image, sample_target)
|
||||
return sample_image, sample_target
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""Returns the number of samples in the dataset."""
|
||||
return len(self._active_frame_indexes)
|
||||
|
||||
|
||||
@attrs.frozen
|
||||
class ExtractSingleLabelIndex:
|
||||
"""
|
||||
A target transform that takes a `Target` object and produces a single label index
|
||||
based on the tag in that object.
|
||||
|
||||
This makes the dataset samples compatible with the image classification networks
|
||||
in torchvision.
|
||||
|
||||
If the annotations contain no tags, or multiple tags, raises a `ValueError`.
|
||||
"""
|
||||
|
||||
def __call__(self, target: Target) -> int:
|
||||
tags = target.annotations.tags
|
||||
if not tags:
|
||||
raise ValueError("sample has no tags")
|
||||
|
||||
if len(tags) > 1:
|
||||
raise ValueError("sample has multiple tags")
|
||||
|
||||
return target.label_id_to_index[tags[0].label_id]
|
||||
|
||||
|
||||
class LabeledBoxes(TypedDict):
|
||||
boxes: Sequence[Tuple[float, float, float, float]]
|
||||
labels: Sequence[int]
|
||||
|
||||
|
||||
_SUPPORTED_SHAPE_TYPES = frozenset(["rectangle", "polygon", "polyline", "points", "ellipse"])
|
||||
|
||||
|
||||
@attrs.frozen
|
||||
class ExtractBoundingBoxes:
|
||||
"""
|
||||
A target transform that takes a `Target` object and returns a dictionary compatible
|
||||
with the object detection networks in torchvision.
|
||||
|
||||
The dictionary contains the following entries:
|
||||
|
||||
"boxes": a sequence of (xmin, ymin, xmax, ymax) tuples, one for each shape
|
||||
in the annotations.
|
||||
"labels": a sequence of corresponding label indices.
|
||||
|
||||
Limitations:
|
||||
|
||||
* Only the following shape types are supported: rectangle, polygon, polyline,
|
||||
points, ellipse.
|
||||
* Rotated shapes are not supported.
|
||||
|
||||
Unsupported shapes will cause a `UnsupportedDatasetError` exception to be
|
||||
raised unless they are filtered out by `include_shape_types`.
|
||||
"""
|
||||
|
||||
include_shape_types: FrozenSet[str] = attrs.field(
|
||||
converter=frozenset,
|
||||
validator=attrs.validators.deep_iterable(attrs.validators.in_(_SUPPORTED_SHAPE_TYPES)),
|
||||
kw_only=True,
|
||||
)
|
||||
"""Shapes whose type is not in this set will be ignored."""
|
||||
|
||||
def __call__(self, target: Target) -> LabeledBoxes:
|
||||
boxes = []
|
||||
labels = []
|
||||
|
||||
for shape in target.annotations.shapes:
|
||||
if shape.type.value not in self.include_shape_types:
|
||||
continue
|
||||
|
||||
if shape.rotation != 0:
|
||||
raise UnsupportedDatasetError("Rotated shapes are not supported")
|
||||
|
||||
x_coords = shape.points[0::2]
|
||||
y_coords = shape.points[1::2]
|
||||
|
||||
boxes.append((min(x_coords), min(y_coords), max(x_coords), max(y_coords)))
|
||||
labels.append(target.label_id_to_index[shape.label_id])
|
||||
|
||||
return LabeledBoxes(boxes=boxes, labels=labels)
|
||||
@ -0,0 +1,207 @@
|
||||
# Copyright (C) 2022 CVAT.ai Corporation
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import io
|
||||
import os
|
||||
from logging import Logger
|
||||
from pathlib import Path
|
||||
from typing import Tuple
|
||||
|
||||
import pytest
|
||||
from cvat_sdk import Client, models
|
||||
from cvat_sdk.core.proxies.tasks import ResourceType
|
||||
|
||||
try:
|
||||
import cvat_sdk.pytorch as cvatpt
|
||||
import PIL.Image
|
||||
import torch
|
||||
import torchvision.transforms
|
||||
import torchvision.transforms.functional as TF
|
||||
from torch.utils.data import DataLoader
|
||||
except ImportError:
|
||||
cvatpt = None
|
||||
|
||||
from shared.utils.helpers import generate_image_files
|
||||
|
||||
|
||||
@pytest.mark.skipif(cvatpt is None, reason="PyTorch dependencies are not installed")
|
||||
class TestTaskVisionDataset:
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Path,
|
||||
fxt_login: Tuple[Client, str],
|
||||
fxt_logger: Tuple[Logger, io.StringIO],
|
||||
fxt_stdout: io.StringIO,
|
||||
):
|
||||
self.tmp_path = tmp_path
|
||||
logger, self.logger_stream = fxt_logger
|
||||
self.stdout = fxt_stdout
|
||||
self.client, self.user = fxt_login
|
||||
self.client.logger = logger
|
||||
|
||||
api_client = self.client.api_client
|
||||
for k in api_client.configuration.logger:
|
||||
api_client.configuration.logger[k] = logger
|
||||
|
||||
monkeypatch.setattr(cvatpt, "_CACHE_DIR", self.tmp_path / "cache")
|
||||
|
||||
self._create_task()
|
||||
|
||||
yield
|
||||
|
||||
def _create_task(self):
|
||||
self.images = generate_image_files(10)
|
||||
|
||||
image_dir = self.tmp_path / "images"
|
||||
image_dir.mkdir()
|
||||
|
||||
image_paths = []
|
||||
for image in self.images:
|
||||
image_path = image_dir / image.name
|
||||
image_path.write_bytes(image.getbuffer())
|
||||
image_paths.append(image_path)
|
||||
|
||||
self.task = self.client.tasks.create_from_data(
|
||||
models.TaskWriteRequest(
|
||||
"PyTorch integration test task",
|
||||
labels=[
|
||||
models.PatchedLabelRequest(name="person"),
|
||||
models.PatchedLabelRequest(name="car"),
|
||||
],
|
||||
),
|
||||
ResourceType.LOCAL,
|
||||
list(map(os.fspath, image_paths)),
|
||||
data_params={"chunk_size": 3},
|
||||
)
|
||||
|
||||
self.label_ids = sorted(l.id for l in self.task.labels)
|
||||
|
||||
self.task.update_annotations(
|
||||
models.PatchedLabeledDataRequest(
|
||||
tags=[
|
||||
models.LabeledImageRequest(frame=5, label_id=self.label_ids[0]),
|
||||
models.LabeledImageRequest(frame=6, label_id=self.label_ids[1]),
|
||||
models.LabeledImageRequest(frame=8, label_id=self.label_ids[0]),
|
||||
models.LabeledImageRequest(frame=8, label_id=self.label_ids[1]),
|
||||
],
|
||||
shapes=[
|
||||
models.LabeledShapeRequest(
|
||||
frame=6,
|
||||
label_id=self.label_ids[1],
|
||||
type=models.ShapeType("rectangle"),
|
||||
points=[1.0, 2.0, 3.0, 4.0],
|
||||
),
|
||||
models.LabeledShapeRequest(
|
||||
frame=7,
|
||||
label_id=self.label_ids[0],
|
||||
type=models.ShapeType("points"),
|
||||
points=[1.1, 2.1, 3.1, 4.1],
|
||||
),
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
def test_basic(self):
|
||||
dataset = cvatpt.TaskVisionDataset(self.client, self.task.id)
|
||||
|
||||
assert len(dataset) == self.task.size
|
||||
|
||||
for index, (sample_image, sample_target) in enumerate(dataset):
|
||||
sample_image_tensor = TF.pil_to_tensor(sample_image)
|
||||
reference_tensor = TF.pil_to_tensor(PIL.Image.open(self.images[index]))
|
||||
assert torch.equal(sample_image_tensor, reference_tensor)
|
||||
|
||||
for index, label_id in enumerate(self.label_ids):
|
||||
assert sample_target.label_id_to_index[label_id] == index
|
||||
|
||||
assert not dataset[0][1].annotations.tags
|
||||
assert not dataset[0][1].annotations.shapes
|
||||
|
||||
assert len(dataset[5][1].annotations.tags) == 1
|
||||
assert dataset[5][1].annotations.tags[0].label_id == self.label_ids[0]
|
||||
assert not dataset[5][1].annotations.shapes
|
||||
|
||||
assert len(dataset[6][1].annotations.tags) == 1
|
||||
assert dataset[6][1].annotations.tags[0].label_id == self.label_ids[1]
|
||||
assert len(dataset[6][1].annotations.shapes) == 1
|
||||
assert dataset[6][1].annotations.shapes[0].type.value == "rectangle"
|
||||
assert dataset[6][1].annotations.shapes[0].points == [1.0, 2.0, 3.0, 4.0]
|
||||
|
||||
assert not dataset[7][1].annotations.tags
|
||||
assert len(dataset[7][1].annotations.shapes) == 1
|
||||
assert dataset[7][1].annotations.shapes[0].type.value == "points"
|
||||
assert dataset[7][1].annotations.shapes[0].points == [1.1, 2.1, 3.1, 4.1]
|
||||
|
||||
def test_deleted_frame(self):
|
||||
self.task.remove_frames_by_ids([1])
|
||||
|
||||
dataset = cvatpt.TaskVisionDataset(self.client, self.task.id)
|
||||
|
||||
assert len(dataset) == self.task.size - 1
|
||||
|
||||
# sample #0 is still frame #0
|
||||
assert torch.equal(
|
||||
TF.pil_to_tensor(dataset[0][0]), TF.pil_to_tensor(PIL.Image.open(self.images[0]))
|
||||
)
|
||||
|
||||
# sample #1 is now frame #2
|
||||
assert torch.equal(
|
||||
TF.pil_to_tensor(dataset[1][0]), TF.pil_to_tensor(PIL.Image.open(self.images[2]))
|
||||
)
|
||||
|
||||
# sample #4 is now frame #5
|
||||
assert len(dataset[4][1].annotations.tags) == 1
|
||||
assert dataset[4][1].annotations.tags[0].label_id == self.label_ids[0]
|
||||
assert not dataset[4][1].annotations.shapes
|
||||
|
||||
def test_extract_single_label_index(self):
|
||||
dataset = cvatpt.TaskVisionDataset(
|
||||
self.client,
|
||||
self.task.id,
|
||||
transform=torchvision.transforms.PILToTensor(),
|
||||
target_transform=cvatpt.ExtractSingleLabelIndex(),
|
||||
)
|
||||
|
||||
assert dataset[5][1] == 0
|
||||
assert dataset[6][1] == 1
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
# no tags
|
||||
_ = dataset[7]
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
# multiple tags
|
||||
_ = dataset[8]
|
||||
|
||||
# make sure the samples can be batched with the default collater
|
||||
loader = DataLoader(dataset, batch_size=2, sampler=[5, 6])
|
||||
|
||||
batch = next(iter(loader))
|
||||
assert torch.equal(batch[0][0], TF.pil_to_tensor(PIL.Image.open(self.images[5])))
|
||||
assert torch.equal(batch[0][1], TF.pil_to_tensor(PIL.Image.open(self.images[6])))
|
||||
assert torch.equal(batch[1], torch.tensor([0, 1]))
|
||||
|
||||
def test_extract_bounding_boxes(self):
|
||||
dataset = cvatpt.TaskVisionDataset(
|
||||
self.client,
|
||||
self.task.id,
|
||||
transform=torchvision.transforms.PILToTensor(),
|
||||
target_transform=cvatpt.ExtractBoundingBoxes(include_shape_types={"rectangle"}),
|
||||
)
|
||||
|
||||
assert dataset[0][1] == {"boxes": [], "labels": []}
|
||||
assert dataset[6][1] == {"boxes": [(1.0, 2.0, 3.0, 4.0)], "labels": [1]}
|
||||
assert dataset[7][1] == {"boxes": [], "labels": []} # points are filtered out
|
||||
|
||||
def test_transforms(self):
|
||||
dataset = cvatpt.TaskVisionDataset(
|
||||
self.client,
|
||||
self.task.id,
|
||||
transforms=lambda x, y: (y, x),
|
||||
)
|
||||
|
||||
assert isinstance(dataset[0][0], cvatpt.Target)
|
||||
assert isinstance(dataset[0][1], PIL.Image.Image)
|
||||
Loading…
Reference in New Issue