Add all necessary files to replicate training run

README.md

@@ -1,5 +1,7 @@
 # Download datasets:  
 * Download and decompress tsv file from here: https://github.com/google-research-datasets/wit/blob/main/DATA.md  
-* Use `prepare_wit.py` to download images from Wikipedia.  
-* Use `discard_incorrect_files` to filter out corrupt files.`TODO: Still some corrupt files are being kept.` `TODO: Make it a CLI`.  
+* Use `prepare_wit.py` to download images from Wikipedia as annotated on each TSV file.  
+* Use `scale_converter.py` to remove corrupt images and resize suitable images to 224x224
+* Use `join_datasets_custom_split.py` to group all JSONs from different subsets of the dataset together
+* Use `discard_incorrect_files.py` to filter out images that we were not able to convert.
 * Finally, use `run-clip.sh` to train.  

discard_incorrect_files.py

@@ -1,23 +1,28 @@
 import json
 import os
+from tqdm import tqdm
 
 import torch
 from torchvision.io import ImageReadMode, read_image
 
-# SUPPORTED_EXTENSIONS = {'PNG', 'JPG', 'png', 'JPEG', 'jpg', 'jpeg'}
+JOINT_JSON_DIRECTORY = f"/home/{os.environ['USER']}/data/wit/all_jsons"
+SCALE_CONVERTED_DIRECTORY = f"/home/{os.environ['USER']}/data/wit_scale_converted"
 
 for split in ["train", "valid", "test"]:
-    with open(f"/home/{os.environ['USER']}/data/wit/prepared_dataset/{split}_dataset.json") as f:
+    print("Reading json")
+    with open(f"{JOINT_JSON_DIRECTORY}/{split}_dataset_all_98_1_1_split.json") as f:
         examples = [json.loads(line) for line in f.readlines()]
-    
+    valid_files = set(os.listdir(SCALE_CONVERTED_DIRECTORY))
+
     supported_examples = []
-    for example in examples:
-        try:
-            image = read_image(example["image_path"], mode=ImageReadMode.RGB)
+    for example in tqdm(examples):
+        directory, filename = os.path.split(example['image_path'])
+        if filename in valid_files:
+            example["image_path"] = os.path.join(SCALE_CONVERTED_DIRECTORY, filename)
             supported_examples.append(json.dumps(example, ensure_ascii=False))
-        except Exception as e:
-            print(f"Excluding file: {example['image_path']} due to error: {e}")
 
     print(f"Total {split} examples: {len(supported_examples)}")
-    with open(f"/home/{os.environ['USER']}/data/wit/prepared_dataset/{split}_dataset_filtered.json", "w") as f:
+    with open(f"{SCALE_CONVERTED_DIRECTORY}/{split}_dataset_scale_converted_98_1_1_split.json", "w") as f:
         f.write("\n".join(supported_examples))
+
+print("DONE!")

join_datasets_custom_split.py

@@ -0,0 +1,48 @@
+import os
+import json
+import random
+
+import pandas as pd
+
+
+DATA_DIR = f"/home/{os.environ['USER']}/data/wit/all_jsons"
+SEED = 0
+PROPORTION_TRAIN = 0.98
+PROPORTION_VALID = 0.01
+
+random.seed(SEED)
+
+all_files = [f"{DATA_DIR}/{file_}" for file_ in os.listdir(DATA_DIR) if ("all" not in file_)]
+
+print(all_files)
+
+examples = []
+for file_ in all_files:
+    print(file_)
+    with open(file_) as f:
+        file_examples = [json.dumps(json.loads(line), ensure_ascii=False) for line in f.readlines()]
+    print(len(file_examples))
+    examples.extend(file_examples)
+
+print(f"Before dedup: {len(examples)}")
+examples = list(set(examples))
+print(f"After dedup: {len(examples)}")
+
+print(examples[0])
+# Shuffle examples
+random.shuffle(examples)
+print(examples[0])
+
+split_dataset = {}
+split_dataset["train"] = examples[:int(len(examples) * PROPORTION_TRAIN)]
+split_dataset["valid"] = examples[int(len(examples) * PROPORTION_TRAIN): int(len(examples) * (PROPORTION_TRAIN + PROPORTION_VALID))]
+split_dataset["test"] = examples[int(len(examples) * (PROPORTION_TRAIN + PROPORTION_VALID)):]
+
+
+for split in ["train", "valid", "test"]:
+    print("-----")
+    print(len(split_dataset[split]))
+    print("-----")
+    with open(f"/home/{os.environ['USER']}/data/wit/all_jsons/{split}_dataset_all_98_1_1_split.json", "w") as f:
+        f.write("\n".join(split_dataset[split]))
+

prepare_wit.py

@@ -33,14 +33,17 @@ def split_and_save_datasets(lines: List[str], output_dir: str, train_proportion:
     with open(f"{output_dir}/test_dataset.json", "w") as f:
         f.write("\n".join(test_lines))
 
-def prepare_wit(tsv: str, language: str, output_dir: str, seed: int, train_proportion: float, valid_proportion: float, backup_period: int, language_col: str="language", caption_col: str="caption_reference_description", url_col: str="image_url", pause=0.1, retries: int=5):
+def prepare_wit(
+    tsv: str, language: str, output_dir: str, seed: int, train_proportion: float, valid_proportion: float, backup_period: int, language_col: str="language", caption_col: str="caption_reference_description", url_col: str="image_url", pause=0.875, retries: int=10):
     os.makedirs(output_dir, exist_ok=True)
     logger.info("Loading dataset")
     df = pd.read_csv(tsv, sep="\t", engine="python")
-    df = df[(df["language"] == language) & (~df["caption_reference_description"].isnull())]
+    existing_files = set(os.listdir(output_dir))
+    not_exists_condition = (~(df[url_col].map(lambda x: x.split("/")[-1][-100:]).isin(existing_files)))
+    df = df[(df["language"] == language) & (~df["caption_reference_description"].isnull()) & not_exists_condition]
     # Shuffle
     df = df.sample(frac=1.0, random_state=seed)
-    logger.info("Download started")
+    logger.info(f"Trying to downloading {df.shape[0]} files")
     lines = []
     count = 0
     try:
@@ -49,7 +52,7 @@ def prepare_wit(tsv: str, language: str, output_dir: str, seed: int, train_propo
                 url = row[url_col]
                 caption = row[caption_col]
                 # Trim image file names so that they are no longer than 100 characters
-                image_filename = url.split('/')[-1][-100:]
+                image_filename = url.split("/")[-1][-100:]
                 image_path = f"{output_dir}/{image_filename}"
                 for retry in range(retries):
                     try:
@@ -59,13 +62,12 @@ def prepare_wit(tsv: str, language: str, output_dir: str, seed: int, train_propo
                         count += 1
                         break
                     except urllib.error.HTTPError as e:
-                        # time.sleep(pause)
-                        pass
+                        time.sleep(pause * 10)
                 if count % backup_period == 0:
                     logger.info(f"Saving dataset backup: Number of lines {len(lines)}")
                     split_and_save_datasets(lines, output_dir, train_proportion, valid_proportion)
-                if retry == retries:
-                    raise ValueError("Rate limit achieved:", e)
+                if retry == retries - 1:
+                    logger.info(f"Skipping {image_filename}")
                 pbar.update(1)
     # Save existing dataset, even upon failure
     finally:

run-clip.sh

@@ -1,12 +1,12 @@
-HUB_TOKEN=`cat $HOME/.huggingface/token`
 python run_hybrid_clip.py \
-    --output_dir "./output_dir" \
+    --output_dir "./output_141230_training_examples" \
     --text_model_name_or_path="dccuchile/bert-base-spanish-wwm-cased" \
     --vision_model_name_or_path="openai/clip-vit-base-patch32" \
     --tokenizer_name="dccuchile/bert-base-spanish-wwm-cased" \
-    --train_file="/home/${USER}/data/wit/prepared_dataset/train_dataset_filtered.json" \
-    --validation_file="/home/${USER}/data/wit/prepared_dataset/valid_dataset_filtered.json" \
-    --do_train --do_eval \
+    --train_file="/home/${USER}/data/wit_scale_converted/train_dataset_scale_converted_98_1_1_split.json" \
+    --validation_file="/home/${USER}/data/wit_scale_converted/valid_dataset_scale_converted_98_1_1_split.json" \
+    --do_train \
+    --do_eval \
     --num_train_epochs="40" \
     --max_seq_length 96 \
     --per_device_train_batch_size="64" \
@@ -14,5 +14,3 @@ python run_hybrid_clip.py \
     --learning_rate="5e-5" --warmup_steps="0" --weight_decay 0.1 \
     --overwrite_output_dir \
     --preprocessing_num_workers 32
-    #--push_to_hub
-

run_hybrid_clip.py

@@ -225,8 +225,9 @@ class ImageTextDataset(VisionDataset):
         self.image_paths = []
 
         for example in examples:
-            self.captions.extend(example["captions"][:captions_per_image])
-            self.image_paths.extend([example["image_path"]] * captions_per_image)
+            captions_subset =  example["captions"][:captions_per_image]
+            self.captions.extend(captions_subset)
+            self.image_paths.extend([example["image_path"]] * len(captions_subset))
 
     def _load_image(self, idx: int):
         path = self.image_paths[idx]

scale_convert.py

@@ -0,0 +1,53 @@
+import glob
+import itertools
+from argparse import ArgumentParser
+from joblib import Parallel, delayed
+import os
+import subprocess
+from collections import Counter
+import shutil
+
+
+parser = ArgumentParser()
+parser.add_argument("in_dir")
+parser.add_argument("out_dir")
+args = parser.parse_args()
+
+os.makedirs(args.out_dir, exist_ok=True)
+
+files = itertools.chain(
+    glob.iglob(f"{args.in_dir}/*/*.jpg"),
+    glob.iglob(f"{args.in_dir}/*/*.JGPG"),
+    glob.iglob(f"{args.in_dir}/*/*.jpeg"),
+    glob.iglob(f"{args.in_dir}/*/*.JPEG"),
+    glob.iglob(f"{args.in_dir}/*/*.png"),
+    glob.iglob(f"{args.in_dir}/*/*.PNG"),
+    glob.iglob(f"{args.in_dir}/*/*.svg"),
+    glob.iglob(f"{args.in_dir}/*/*.SVG"),
+)
+
+def process_file(path):
+    basename = os.path.basename(path)
+    ext = os.path.splitext(basename)[1]
+    name = os.path.splitext(basename)[0]
+
+    dirname = os.path.dirname(path)
+    try:
+        r = subprocess.run(
+            f'convert {path} -resize "224^>" -colorspace RGB -density 1200 {args.out_dir}/{name}.jpg',
+            shell=True,
+            timeout=10
+        )
+        rcode = r.returncode
+    except subprocess.TimeoutExpired:
+        print("conversion timeout expired")
+        rcode = -1
+
+    if rcode == 0:
+        os.remove(path)
+
+    return rcode
+
+codes = Parallel(n_jobs=32, prefer="threads", verbose=1)(delayed(process_file)(f) for f in files)
+print(Counter(codes))
+

test_on_image.py

@@ -1,3 +1,5 @@
+import os
+
 import jax
 import torch
 from torchvision.io import ImageReadMode, read_image
@@ -28,7 +30,7 @@ def run_inference(image_path, text):
     score = jax.nn.sigmoid(logits)
     return score
 
-image_path = "/home/eduardogonzalezponferrada/data/wit/full_dataset/Casa_de_Cultura_%284%29.JPG"
+image_path = f"/home/{os.environ['USER']}/data/wit_scale_converted/Self_Portrait_by_David_Allan.jpg"
 text = "Patio interior de un edificio"
 
-print(run_inference(image_path, text))
+print(run_inference(image_path, text))