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