LSUN object dataset download & multiple category option

README.md

@@ -1,6 +1,7 @@
 # LSUN
 
 Please check [LSUN webpage](http://www.yf.io/p/lsun) for more information about the dataset.
+This forked repo includes the code to download LSUN object dataset.
 
 ## Data Release
 
@@ -26,13 +27,21 @@ If you find LSUN dataset useful in your research, please consider citing:
 Please make sure you have cURL installed
 ```bash
 # Download the whole latest data set
-python3 download.py
+python3 download_scenes.py
 # Download the whole latest data set to <data_dir>
-python3 download.py -o <data_dir>
+python3 download_scenes.py -o <data_dir>
 # Download data for bedroom
-python3 download.py -c bedroom
+python3 download_scenes.py -c bedroom
+# It supports the option for multiple categories
+python3 download_scenes.py -c bedroom church_outdoor
 # Download testing set
-python3 download.py -c test
+python3 download_scenes.py -c test
+
+# The options work in the download script of the object dataset.
+python3 download_objects.py
+python3 download_objects.py -o <data_dir>
+python3 download_objects.py -c cat
+python3 download_objects.py -c cat dog
 ```
 
 ## Demo code

data.py

@@ -39,14 +39,14 @@ def export_images(db_path, out_dir, flat=False, limit=-1):
         cursor = txn.cursor()
         for key, val in cursor:
             if not flat:
-                image_out_dir = join(out_dir, '/'.join(key[:6]))
+                image_out_dir = join(out_dir, '/'.join(str(key[:6])))
             else:
                 image_out_dir = out_dir
             if not exists(image_out_dir):
                 os.makedirs(image_out_dir)
-            image_out_path = join(image_out_dir, key + '.webp')
+            image_out_path = join(image_out_dir, str(key) + '.webp')
             with open(image_out_path, 'w') as fp:
-                fp.write(val)
+                fp.write(str(val))
             count += 1
             if count == limit:
                 break

download_objects.py

@@ -0,0 +1,46 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import print_function, division
+import argparse
+from os.path import join
+from os import makedirs
+
+import subprocess
+from urllib.request import Request, urlopen
+
+__author__ = 'Fisher Yu'
+__email__ = 'fy@cs.princeton.edu'
+__license__ = 'MIT'
+
+
+def download(out_dir, category):
+    url = 'http://dl.yf.io/lsun/objects/{category}.zip'.format(**locals())
+    print(url)
+    out_name = '{category}.zip'.format(**locals())
+    out_path = join(out_dir, out_name)
+    cmd = ['curl', url, '-o', out_path]
+    print('Downloading', category, 'set')
+    subprocess.call(cmd)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-c', '--category', nargs='+', type=str, default=None)
+    parser.add_argument('-o', '--out_dir', default='')
+    args = parser.parse_args()
+    makedirs(args.out_dir, exist_ok=True)
+    categories = ['airplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'dining_table', 'dog', 'horse', 'motorbike', 'person', 'potted_plant', 'sheep', 'sofa', 'train' ,'tv-monitor']
+
+    if args.category is not None:
+        for category in args.category:
+            if category not in categories:
+                raise AttributeError('Error:', args.category, "doesn't exist in", 'LSUN release')
+        categories = args.category
+
+    print('Downloading', len(categories), 'categories')
+    for category in categories:
+        download(args.out_dir, category)
+
+if __name__ == '__main__':
+    main()
+

download.py → download_scenes.py

@@ -3,6 +3,7 @@
 from __future__ import print_function, division
 import argparse
 from os.path import join
+from os import makedirs
 
 import subprocess
 from urllib.request import Request, urlopen
@@ -35,8 +36,9 @@ def download(out_dir, category, set_name):
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('-o', '--out_dir', default='')
-    parser.add_argument('-c', '--category', default=None)
+    parser.add_argument('-c', '--category', nargs='+', type=str, default=None)
     args = parser.parse_args()
+    makedirs(args.out_dir, exist_ok=True)
 
     categories = list_categories()
     if args.category is None:
@@ -45,15 +47,27 @@ def main():
             download(args.out_dir, category, 'train')
             download(args.out_dir, category, 'val')
         download(args.out_dir, '', 'test')
-    else:
+
+    elif len(args.category) == 1:
+        args.category = args.category[0]
+        if args.category not in categories:
+            print('Error:', args.category, "doesn't exist in", 'LSUN release')
         if args.category == 'test':
             download(args.out_dir, '', 'test')
-        elif args.category not in categories:
-            print('Error:', args.category, "doesn't exist in", 'LSUN release')
         else:
             download(args.out_dir, args.category, 'train')
             download(args.out_dir, args.category, 'val')
 
+    else:
+        print('Downloading', len(args.category), 'categories')
+        for category in args.category:
+            if category not in categories:
+                raise AttributeError('Error:', category, "doesn't exist in", 'LSUN release')
+        for category in args.category:
+            download(args.out_dir, category, 'train')
+            download(args.out_dir, category, 'val')
+
 
 if __name__ == '__main__':
     main()
+