Working download script for COCO dataset

src/config.py

@@ -1,10 +1,16 @@
+import pathlib
+
 import pydantic
 
+MAX_DOWNLOAD_TIME = 0.2
+
+IMAGE_DOWNLOAD_PATH = pathlib.Path("/tmp/images")
+
 
 class DataConfig(pydantic.BaseModel):
     buffer_size: int = 1000
-    data_len: int = 100000
-    train_len: int = 90000
+    data_len: int = 100
+    train_len: int = 90
     small_dataset: str = "laion/220k-gpt4vision-captions-from-livis"
     large_dataset: str = "laion/laion400m"
     dataset: str = small_dataset
@@ -16,7 +22,7 @@ class ModelConfig(pydantic.BaseModel):
     projection_layers: int = 3
     embed_dim: int = 256
     transformer_embed_dim: int = 768
-    max_len: int = 77  # maximum length of text in CLIP
+    max_len: int = 128  # 77
     cls_type: bool = True
     freeze_vision_base: bool = False
     freeze_text_base: bool = False
@@ -36,5 +42,5 @@ class TrainerConfig(pydantic.BaseModel):
 
     run_openai_clip: bool = False
 
-    model_config: ModelConfig = ModelConfig()
-    data_config: DataConfig = DataConfig()
+    _model_config: ModelConfig = ModelConfig()
+    _data_config: DataConfig = DataConfig()

src/download.py

@@ -0,0 +1,54 @@
+from io import BytesIO
+import pathlib
+from functools import partial
+from typing import Any
+
+import datasets
+from PIL import Image
+from loguru import logger
+import requests
+from tqdm.auto import tqdm
+
+from src import config
+
+
+def _save_resized_image(example: dict[str, Any], size: tuple[int, int], path: pathlib.Path):
+    # Download the image
+    image_url = example["url"]
+    image_path = path / image_url.rsplit("/", 1)[-1]
+    if image_path.exists():
+        return
+
+    response = requests.get(image_url)
+    image = Image.open(BytesIO(response.content))
+    # Resize the image
+    image_resized = image.resize(size)
+    image_resized.save(image_path)
+
+
+def _get_images(dataset: datasets.Dataset, path: pathlib.Path):
+    save_resized_image = partial(_save_resized_image, path=path, size=(256, 256))
+    dataset.map(save_resized_image, num_proc=128)
+
+
+def _check_corrupt_images(image_file: pathlib.Path):
+    try:
+        with Image.open(image_file) as img:
+            img.verify()  # Verify the integrity of the image
+    except (IOError, SyntaxError) as e:
+        logger.error(f"Corrupt image: {image_file}")
+
+
+if __name__ == "__main__":
+    hyper_parameters = config.TrainerConfig()
+
+    dataset = datasets.load_dataset(
+        hyper_parameters._data_config.dataset,
+        split="train",
+    )
+
+    config.IMAGE_DOWNLOAD_PATH.mkdir(parents=True, exist_ok=True)
+    _get_images(dataset, config.IMAGE_DOWNLOAD_PATH)  # type: ignore
+
+    for image in tqdm(config.IMAGE_DOWNLOAD_PATH.iterdir()):
+        _check_corrupt_images(image)