Merge pull request #425 from sh8/master

Fix to cache image sizes data

lib/roi_data_layer/roidb.py

@@ -3,6 +3,9 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import os
+import pickle
+
 import datasets
 import numpy as np
 from model.utils.config import cfg
@@ -20,8 +23,18 @@ def prepare_roidb(imdb):
 
   roidb = imdb.roidb
   if not (imdb.name.startswith('coco')):
-    sizes = [PIL.Image.open(imdb.image_path_at(i)).size
-         for i in range(imdb.num_images)]
+    cache_file = os.path.join(imdb.cache_path, imdb.name + '_sizes.pkl')
+    if os.path.exists(cache_file):
+      print('Image sizes loaded from %s' % cache_file)
+      with open(cache_file, 'rb') as f:
+        sizes = pickle.load(f)
+    else:
+      print('Extracting image sizes... (It may take long time)')
+      sizes = [PIL.Image.open(imdb.image_path_at(i)).size
+                for i in range(imdb.num_images)]
+      with open(cache_file, 'wb') as f:
+        pickle.dump(sizes, f)
+      print('Done!!')
          
   for i in range(len(imdb.image_index)):
     roidb[i]['img_id'] = imdb.image_id_at(i)