44 lines
1.5 KiB
Python
44 lines
1.5 KiB
Python
from __future__ import print_function, absolute_import, division
|
|
|
|
import os
|
|
import os.path as osp
|
|
import lmdb
|
|
|
|
|
|
class BaseImgDataset(object):
|
|
def __init__(self):
|
|
self.train_lmdb_path = None
|
|
self.query_lmdb_path = None
|
|
self.gallery_lmdb_path = None
|
|
|
|
def generate_lmdb(self):
|
|
assert isinstance(self.train, list)
|
|
assert isinstance(self.query, list)
|
|
assert isinstance(self.gallery, list)
|
|
|
|
print("Reminder: this function is under development, some datasets might not be applicable yet")
|
|
|
|
self.train_lmdb_path = osp.join(self.dataset_dir, 'train_lmdb')
|
|
self.query_lmdb_path = osp.join(self.dataset_dir, 'query_lmdb')
|
|
self.gallery_lmdb_path = osp.join(self.dataset_dir, 'gallery_lmdb')
|
|
|
|
def _write_lmdb(write_path, data_list):
|
|
if osp.exists(write_path):
|
|
return
|
|
|
|
print("Generating lmdb files to '{}'".format(write_path))
|
|
|
|
num_data = len(data_list)
|
|
max_map_size = int(num_data * 500**2 * 3) # be careful with this
|
|
env = lmdb.open(write_path, map_size=max_map_size)
|
|
|
|
for img_path, pid, camid in data_list:
|
|
with env.begin(write=True) as txn:
|
|
with open(img_path, 'rb') as imgf:
|
|
imgb = imgf.read()
|
|
txn.put(img_path, imgb)
|
|
|
|
_write_lmdb(self.train_lmdb_path, self.train)
|
|
_write_lmdb(self.query_lmdb_path, self.query)
|
|
_write_lmdb(self.gallery_lmdb_path, self.gallery)
|