| 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() |
|
|
