File size: 4,082 Bytes
002bd9b |
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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
import json
import os
import datasets
logger = datasets.logging.get_logger(__name__)
_BASE_IMAGE_METADATA_FEATURES = {
"image_id": datasets.Value("int32"),
"width": datasets.Value("int32"),
"height": datasets.Value("int32"),
"file_name": datasets.Value("string"),
"task_type": datasets.Value("string"),
}
_BASE_REGION_FEATURES = {
"region_id": datasets.Value("int32"),
"image_id": datasets.Value("int32"),
"phrases": [datasets.Value("string")],
"x": datasets.Value("int32"),
"y": datasets.Value("int32"),
"width": datasets.Value("int32"),
"height": datasets.Value("int32"),
}
_ANNOTATION_FEATURES = {
"regions": [_BASE_REGION_FEATURES],
}
class CustomDatasetConfig(datasets.BuilderConfig):
def __init__(
self,
name,
splits,
data_dir: str = None,
with_image: bool = True,
task_type: str = "recognition",
**kwargs,
):
super().__init__(name, **kwargs)
self.splits = splits
self.with_image = with_image
self.task_type = task_type
self.data_dir = data_dir
@property
def features(self):
return datasets.Features(
{
**({"image": datasets.Image()} if self.with_image else {}),
**_BASE_IMAGE_METADATA_FEATURES,
**_ANNOTATION_FEATURES,
}
)
class CustomDataset(datasets.GeneratorBasedBuilder):
VERSION = datasets.Version("0.0.1")
BUILDER_CONFIG_CLASS = CustomDatasetConfig
BUILDER_CONFIGS = [
CustomDatasetConfig(name="custom", splits=["train", "validation"]),
]
DEFAULT_CONFIG_NAME = "custom"
config: CustomDatasetConfig
def _info(self):
return datasets.DatasetInfo(features=self.config.features)
def _split_generators(self, dl_manager):
"""Returns SplitGenerators."""
data_dir = self.config.data_dir
if data_dir is None:
raise ValueError(
"This script is supposed to work with a local dataset. The argument `data_dir` in `load_dataset()` is required."
)
splits = []
for split in self.config.splits:
if split == "train":
dataset = datasets.SplitGenerator(
name=datasets.Split.TRAIN,
gen_kwargs={
"json_path": os.path.join(data_dir, "train.json"),
},
)
elif split in ["val", "valid", "validation", "dev"]:
dataset = datasets.SplitGenerator(
name=datasets.Split.VALIDATION,
gen_kwargs={
"json_path": os.path.join(data_dir, "test.json"), # Using test.json for validation
},
)
else:
continue
splits.append(dataset)
return splits
def _generate_examples(self, json_path):
with open(json_path, 'r') as f:
data = json.load(f)
for idx, image in enumerate(data["images"]):
image_id = image["image_id"]
image_metadata = {
"file_name": image["file_name"],
"height": image["height"],
"width": image["width"],
"image_id": image["image_id"],
}
annotations = [
ann for ann in data["annotations"] if ann["image_id"] == image_id
]
regions = [{
"region_id": ann["region_id"],
"image_id": ann["image_id"],
"phrases": ann["phrases"],
"x": ann["x"],
"y": ann["y"],
"width": ann["width"],
"height": ann["height"],
} for ann in annotations]
image_dict = {"image": os.path.join(self.config.data_dir, image["file_name"])} if self.config.with_image else {}
yield idx, {**image_dict, **image_metadata, "regions": regions, "task_type": self.config.task_type} |