File size: 1,409 Bytes
9b92c75 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | import json
from PIL import Image
import torch
from objectmodel_v1.data import CocoDetectionDataset, detection_collate
def test_coco_dataset_and_empty_image(tmp_path):
image_dir = tmp_path / "images"
image_dir.mkdir()
Image.new("RGB", (100, 50), "white").save(image_dir / "one.jpg")
Image.new("RGB", (80, 80), "black").save(image_dir / "two.jpg")
annotations = {
"images": [
{"id": 1, "file_name": "one.jpg", "width": 100, "height": 50},
{"id": 2, "file_name": "two.jpg", "width": 80, "height": 80},
],
"annotations": [
{"id": 1, "image_id": 1, "category_id": 7, "bbox": [10, 5, 20, 10], "iscrowd": 0}
],
"categories": [{"id": 7, "name": "object"}],
}
annotation_file = tmp_path / "annotations.json"
annotation_file.write_text(json.dumps(annotations), encoding="utf-8")
dataset = CocoDetectionDataset(image_dir, annotation_file, 64, training=False)
image, target = dataset[0]
empty_image, empty_target = dataset[1]
assert image.shape == (3, 64, 64)
assert target["boxes"].shape == (1, 4)
assert target["labels"].tolist() == [0]
assert empty_target["boxes"].shape == (0, 4)
batch, targets = detection_collate([(image, target), (empty_image, empty_target)])
assert batch.shape == (2, 3, 64, 64)
assert len(targets) == 2
assert torch.isfinite(batch).all()
|