test

data/lmdbify.py

@@ -1,30 +1,31 @@
 import os
 import pickle
-from io import BytesIO
 import argparse
-
+from PIL import Image
 import lmdb
-from data.dataset import CARS, ImprovedImageFolder
+from data.dataset import ImprovedImageFolder
 from tqdm import tqdm
 
 
 def content_loader(path):
-    with open(path, "rb") as f:
-        return f.read()
+    im = Image.open(path)
+    im = im.resize((256, 256))
+    if im.mode != "RGB":
+        im = im.convert("RGB")
+    return im
 
 
 def dataset_to_lmdb(dataset, lmdb_path):
     env = lmdb.open(lmdb_path, map_size=1099511627776*2, subdir=os.path.isdir(lmdb_path))
     with env.begin(write=True) as txn:
         for i in tqdm(range(len(dataset)), ncols=50):
-            txn.put("{}".format(i).encode(), bytearray(dataset[i]))
+            txn.put("{}".format(i).encode(), pickle.dumps(dataset[i]))
         txn.put(b"classes_list", pickle.dumps(dataset.classes_list))
         txn.put(b"__len__", pickle.dumps(len(dataset)))
 
 
 def transform(save_path, dataset_path):
     print(save_path, dataset_path)
-    # origin_dataset = CARS("/data/few-shot/STANFORD-CARS/", loader=content_loader)
     origin_dataset = ImprovedImageFolder(dataset_path, loader=content_loader)
     dataset_to_lmdb(origin_dataset, save_path)