60 lines
2.0 KiB
Python
60 lines
2.0 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from torch.utils.data import Dataset
|
|
|
|
from data.schema import OBJECT_TO_ID, PREDICATE_TO_ID, SUBJECT_TO_ID, Fact
|
|
|
|
|
|
def fact_to_ids(value: dict[str, str]) -> tuple[int, int, int]:
|
|
fact = Fact.from_dict(value)
|
|
return (
|
|
SUBJECT_TO_ID[fact.subject],
|
|
PREDICATE_TO_ID[fact.predicate],
|
|
OBJECT_TO_ID[fact.object],
|
|
)
|
|
|
|
|
|
class EvidenceDataset(Dataset):
|
|
def __init__(self, manifest: Path, image_size: int = 224):
|
|
self.root = manifest.parent
|
|
self.image_size = image_size
|
|
self.records = [json.loads(line) for line in manifest.read_text(encoding="utf-8").splitlines()]
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.records)
|
|
|
|
def __getitem__(self, index: int) -> dict[str, Any]:
|
|
record = self.records[index]
|
|
image = Image.open(self.root / record["image"]).convert("RGB")
|
|
image = image.resize((self.image_size, self.image_size))
|
|
image_tensor = torch.from_numpy(np.asarray(image).copy()).permute(2, 0, 1).float() / 255.0
|
|
return {
|
|
"image": image_tensor,
|
|
"rules_text": record["rules_text"],
|
|
"visual_facts": torch.tensor(
|
|
[fact_to_ids(value) for value in record["visual_facts"]], dtype=torch.long
|
|
),
|
|
"text_facts": torch.tensor(
|
|
[fact_to_ids(value) for value in record["text_facts"]], dtype=torch.long
|
|
),
|
|
"scene_id": record["scene_id"],
|
|
}
|
|
|
|
|
|
def collate_evidence(batch: list[dict[str, Any]]) -> dict[str, Any]:
|
|
return {
|
|
"images": torch.stack([item["image"] for item in batch]),
|
|
"rules_text": [item["rules_text"] for item in batch],
|
|
"visual_facts": [item["visual_facts"] for item in batch],
|
|
"text_facts": [item["text_facts"] for item in batch],
|
|
"scene_id": [item["scene_id"] for item in batch],
|
|
}
|
|
|