CC12M downloader script added

data/CC12M_downloader.py

@@ -0,0 +1,91 @@
+# Luke Melas-Kyriazi's code. (https://twitter.com/lukemelas)
+
+#%% 
+import sys
+import os
+from datetime import datetime
+import pandas as pd
+import contexttimer
+from urllib.request import urlopen
+import requests
+from PIL import Image
+import torch
+from torchvision.transforms import functional as TF
+from multiprocessing import Pool
+from tqdm import tqdm
+import logging
+
+# Setup
+logging.basicConfig(filename='download.log', filemode='w', level=logging.INFO)
+requests.packages.urllib3.disable_warnings(requests.packages.urllib3.exceptions.InsecureRequestWarning)
+
+
+# # For downloading SVG images (I can't get this to work)
+# from io import BytesIO
+# import cairosvg
+
+#%% 
+# Load data
+print(f'Starting to load at {datetime.now().isoformat(timespec="minutes")}')
+with contexttimer.Timer(prefix="Loading from tsv"):
+    df = pd.read_csv('./cc12m.tsv', delimiter='\t', header=None)
+
+url_to_idx_map = {url: index for index, url, caption in df.itertuples()}
+print(f'Loaded {len(url_to_idx_map)} urls')
+
+#%% 
+df.head()
+
+#%% 
+
+# Note: it seems that there are no SVG images
+df.sample(10000)[1].str.contains('.svg').sum()
+
+#%% 
+# Resize function
+def resize(img):
+    max_size_of_short_side = 512
+    if min(img.size) > max_size_of_short_side:
+        img = TF.resize(img, size=max_size_of_short_side, interpolation=Image.LANCZOS) 
+    return img
+
+base_dir = os.path.join(os.getcwd(), 'images')
+
+def process(item):
+    url, image_id = item
+    try:
+        base_url = os.path.basename(url)  # extract base url
+        stem, ext = os.path.splitext(base_url)  # split into stem and extension
+        filename = f'{image_id:08d}---{stem}.jpg'  # create filename
+        filepath = os.path.join(base_dir, filename)  # concat to get filepath
+        if not os.path.isfile(filepath):
+            # if filepath.endswith('.svg'):
+            #     raise NotImplementedError()
+            #     image_bytes = BytesIO()  # create a bytestream
+            #     cairosvg.svg2png(url=url, write_to=image_bytes)  # convert svg into image
+            # else:
+            req = requests.get(url, stream=True, timeout=1, verify=False).raw
+            image = Image.open(req).convert('RGB')
+            if min(image.size) > 512:
+                image = TF.resize(image, size=512, interpolation=Image.LANCZOS)
+            # image = resize(image)  # resize PIL image
+            image.save(filepath)  # save PIL image
+    except Exception as e:
+        logging.info(" ".join(repr(e).splitlines()))
+        logging.error(url)
+
+#%% 
+#for i, item in enumerate(tqdm(url_to_idx_map.items(), total=len(url_to_idx_map))):
+#    process(item)
+#    if i > 100:
+#        break
+
+# Use multiprocessing for speed
+list_of_items = list(url_to_idx_map.items())
+print(len(list_of_items))
+list_of_items = list_of_items[10_000_000:]
+print(len(list_of_items))
+with Pool(128) as p:
+    r = list(tqdm(p.imap(process, list_of_items), total=len(list_of_items)))
+    print('DONE')
+