added data/download_images.py

llava/data/download_images.py

@@ -0,0 +1,49 @@
+import os
+import json
+import shutil
+from tqdm import tqdm
+import tarfile
+import argparse
+from urllib.error import HTTPError
+import urllib.request
+
+
+def main(args):
+  input_data = []
+  with open(args.input_path) as f:
+    for line in f:
+      input_data.append(json.loads(line))
+
+  # Download all PMC articles
+  print('Downloading PMC articles')
+  for idx, sample in enumerate(tqdm(input_data)):
+    try:
+      urllib.request.urlretrieve(sample['pmc_tar_url'], os.path.join(args.pmc_output_path, os.path.basename(sample['pmc_tar_url'])))
+    except HTTPError as e:
+      print('Error downloading PMC article: {}'.format(sample['pmc_tar_url']))
+      continue
+
+
+  # Untar all PMC articles
+  print('Untarring PMC articles')
+  for sample in tqdm(input_data):
+    fname = os.path.join(args.pmc_output_path, os.path.basename(os.path.join(sample['pmc_tar_url'])))
+    tar = tarfile.open(fname, "r:gz")
+    tar.extractall(args.pmc_output_path)
+    tar.close()
+
+  # Copy to images directory
+  print('Copying images')
+  for sample in tqdm(input_data):
+    src = os.path.join(args.pmc_output_path, sample['image_file_path'])
+    dst = os.path.join(args.images_output_path, sample['pair_id']+'.jpg')
+    shutil.copyfile(src, dst)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input_path', type=str, default='data/llava_med_image_urls.jsonl')
+    parser.add_argument('--pmc_output_path', type=str, default='data/pmc_articles/')
+    parser.add_argument('--images_output_path', type=str, default='data/images/')
+    args = parser.parse_args()
+    main(args)