htr_demo / helper /examples /create_examples.py
Gabriel's picture
test docker
080d81a
raw
history blame
2.7 kB
import os
import tarfile
import datasets
import pandas as pd
_CITATION = """\
@InProceedings{huggingface:dataset,
title = {Small htr examples images},
author={Gabriel Borg},
year={2023}
}
"""
_DESCRIPTION = """\
Demo dataset for the htr demo.
"""
_HOMEPAGE = "https://huggingface.co/datasets/Riksarkivet/test_images_demo"
_LICENSE = ""
_REPO = "https://huggingface.co/datasets/Riksarkivet/test_images_demo"
class ExampleImages(datasets.GeneratorBasedBuilder):
"""Small sample of image-text pairs"""
def _info(self):
return datasets.DatasetInfo(
description=_DESCRIPTION,
features=datasets.Features(
{
"text": datasets.Value("string"),
"image": datasets.Image(),
}
),
supervised_keys=None,
homepage=_HOMEPAGE,
citation=_CITATION,
)
def _split_generators(self, dl_manager):
images_archive = dl_manager.download(f"{_REPO}/resolve/main/images.tar.gz")
metadata_path = dl_manager.download(f"{_REPO}/resolve/main/images.txt")
image_iters = dl_manager.iter_archive(images_archive)
return [
datasets.SplitGenerator(
name=datasets.Split.TRAIN, gen_kwargs={"images": image_iters, "metadata_path": metadata_path}
),
]
def _generate_examples(self, images, metadata_path):
"""Generate images and text."""
with open(metadata_path, encoding="utf-8") as f:
metadata_list = f.read().split("\n")
for idx, (img_obj, meta_txt) in enumerate(zip(images, metadata_list)):
filepath, image = img_obj
text_value = meta_txt.split("= ")[-1].strip()
yield idx, {
"image": {"path": filepath, "bytes": image.read()},
"text": text_value,
}
def txt_to_csv(file_name):
text_file_path = f"{file_name}.txt"
df = pd.read_csv(text_file_path, delimiter="=", header=None, names=["Key", "Label"], encoding="utf-8")
df["Key"] = df["Key"].str.strip()
df["Label"] = df["Label"].str.strip()
print(df)
df.to_csv(f"{file_name}.csv", index=False)
def sort_and_compress_images(images_folder, tar_file):
sorted_images = sorted(os.listdir(images_folder))
with tarfile.open(tar_file, "w:gz") as tar:
for image_name in sorted_images:
image_path = os.path.join(images_folder, image_name)
tar.add(image_path, arcname=image_name)
print("Images sorted and compressed into tar.gz archive.")
if __name__ == "__main__":
txt_to_csv("info")
sort_and_compress_images("images", "sorted_images.tar.gz")