import os, sys
import pickle
from easydict import EasyDict as edict
from collections import OrderedDict as odict
import cv2
import numpy as np
from basic.common import env, Open, add_path
from basic.util import load_yaml
from numpy_db import npy_table, npy_db, dtype_summary, reorder_dtype
add_path('../anno_db_v2')
from db_type import viewpoint, proj_info, image_info, object_anno # , pose_hypo
from util_v2 import rescale_anno
conf = load_yaml('config.yml') # odict
Pascal3D_root = os.path.expanduser(conf['Pascal3D_release_root'])
# Pascal3D_root = env.Home+'/working/cvpr17pose/dataset/PASCAL3D+_release1.1' # Image_sets/%s_imagenet_%s.txt
PascalVOC_root = Pascal3D_root + '/PASCAL'
assert os.path.exists(
Pascal3D_root), "Please replace with your path/to/PASCAL3D+_release1.1"
assert os.path.exists(PascalVOC_root), "Cannot find %s" % PascalVOC_root
protocol = conf[
'pkl_protocol'] # pickle dump protocol. Change -1 to 2 for python2.x compatibility.
if True:
print('\n\n--------- viewpoint --------')
dtype_summary(np.dtype(viewpoint))
print('\n\n--------- proj_info --------')
dtype_summary(np.dtype(proj_info))
def __init__(self,
db_path,
mode='r',
map_N=30000,
max_readers=256,
always_load_color=True,
silent=False,
**kwargs): # map_N=250000
""" [kwargs examples]:
remap=dict(png_dtype='uint16',npy_dtype='float32', min=[0,0], max=[1,1] ) # for png handler of opt-flow (2-channel float32 img)
"""
self.db_path = os.path.abspath(
db_path.rstrip("/")) # remove last '/' if there's any.
# assert os.path.exists(self.db_path)
self.mode = mode
if self.mode == 'w':
os.system(
'rm -rf %s' % self.db_path
) # overwrite if overwrite and os.path.exists(self.db_path):
elif self.mode in ['r', 'a+']:
assert os.path.exists(
self.db_path), "[Path not exists] %s" % self.db_path
else:
raise NotImplementedError
# self.lmdb_env = lmdb.open(self.db_path, map_size=self.map_size)
if self.mode == 'r':
self.map_size = (map_N * 256 * 256 * 3 * 4)
self.lmdb_env = lmdb.open(self.db_path,
map_size=self.map_size,
max_readers=max_readers,
readahead=True,
readonly=True,
lock=False) #
else:
self.map_size = (map_N * 256 * 256 * 3 * 4) * 10
self.lmdb_env = lmdb.open(self.db_path,
map_size=self.map_size,
max_readers=max_readers) # lock=True
if self.db_path.endswith('.Rawjpg.lmdb'):
if not silent: print("[Using] Handle_Rawjpg")
self.handle = Handle_Rawjpg(always_load_color=always_load_color
) # bytes data handler (pack/unpack)
elif self.db_path.endswith('.Rawpng.lmdb'):
if not silent: print("[Using] Handle_Rawpng")
yamlfile = os.path.join(db_path, 'remap.yml')
if self.mode in ['r', 'a']:
''' e.g. {dtype: 32FC2, min : 0.0, max : 1.0} '''
remap = load_yaml(yamlfile) if os.path.exists(
yamlfile) else None
print('---------> remap yaml: ', remap)
else: # write mode
remap = kwargs.get('remap', None)
print("Write png with remap: %s" % remap)
dump_yaml(remap, yamlfile)
self.handle = Handle_Rawpng(
always_load_color=always_load_color,
remap=remap) # bytes data handler (pack/unpack)
elif self.db_path.endswith('.Npyarr.lmdb'):
if not silent: print("[Using] Handle_Npyarr")
self.handle = Handle_Npyarr()
elif self.db_path.endswith('.Rawebp.lmdb'):
if not silent: print("[Using] Handle_Rawebp")
if self.mode in ['w', 'a']:
QUALITY = kwargs.get('QUALITY', 95)
print("Compress QUALITY: ", QUALITY)
self.handle = Handle_Rawebp(QUALITY=QUALITY)
else:
self.handle = Handle_Rawebp()
else:
print('Unrecognized imagedata_lmdb extension:\n[db_path] %s' %
self.db_path)
raise NotImplementedError
if not silent: print(self)
# print(self.len)
""" --- patch for rename keys ---
In lmdb, "keys are always lexicographically sorted".
This prevent us to shuffle the storage order of images, which is necessary when the training dataset size is really large, e.g ImageNet (45G~120G).
Pre-shuffling image order favor image data loader in training code, as it can read sequentially along the physical storage.
To do this, we re-name all image keys to '0000XXXXX' ('%09d' % image_id) format as it would be sorted by lmdb (same trick in caffe).
So when imgId2dbId.pkl, we need to map the actual image_id to db_id for retrieve a data.
"""
if os.path.exists(os.path.join(self.db_path, 'imgId2dbId.pkl')):
assert self.mode not in ['w', 'a+'], 'Not implement renamed key '
self.key_renamed = True
self.imgId2dbId = pickle.load(
open(os.path.join(self.db_path, 'imgId2dbId.pkl'),
'rb')) # OrderedDict
else:
self.key_renamed = False