00001 import sqlite3
00002 import numpy as np
00003 import pickle as pk
00004 import hrl_lib.util as ut
00005
00006 def create_raw_table(cur):
00007 cur.execute("CREATE TABLE raw_data (id text PRIMARY KEY," \
00008 + "surface_id integer, "\
00009 + "surface_height real, " \
00010 + "unrotated_points matrix, " \
00011 + "normal matrix, " \
00012 + "labels matrix, " \
00013 + "image text, " \
00014 + "intensities matrix, " \
00015 + "points matrix)")
00016
00017
00018 def matrix_adapter(some_mat):
00019 return pk.dumps(some_mat)
00020
00021 def matrix_converter(matrix_string):
00022 try:
00023 if len(matrix_string) == 0:
00024 print "matrix_converter: WARNING -- encountered string of length 0"
00025 return np.matrix([])
00026 return pk.loads(matrix_string)
00027 except EOFError, e:
00028 import pdb
00029 pdb.set_trace()
00030 print 'error'
00031 except ValueError, e:
00032 import pdb
00033 pdb.set_trace()
00034 print 'error'
00035 return np.matrix([])
00036
00037
00038 def store_raw_data_from_individual_files(cur, basedir='clutter_db'):
00039 import clutter_util as cu
00040 scans = cu.list_files('pkl', basedir)
00041 print '>> starting'
00042 for filename in scans:
00043 store_raw_data_from_file(filename)
00044 print '>> done.'
00045
00046 def store_raw_data_from_file(filename):
00047 print '>> loading and storing', filename
00048 r = ut.load_pickle(filename)
00049 cur.execute('INSERT INTO raw_data (id, surface_id, surface_height, unrotated_points, normal, labels, image, intensities, points) '\
00050 + 'values (?, ?, ?, ?, ?, ?, ?, ?, ?)',\
00051 (r['id'], r['surface_id'], r['surface_height'], r['unrotated_points'],
00052 r['normal'], np.matrix(r['labels']), r['image_name'], r['intensities'], r['points']))
00053
00054 def db_iterator(name=None, con=None, cur=None):
00055 started = False
00056 if name != None:
00057 con, cur = start(name)
00058 started = True
00059
00060 data = cur.execute("SELECT * from raw_data")
00061 for i, r in enumerate(data):
00062
00063 d = convert_record_to_dict(r)
00064 yield d
00065
00066 if started:
00067 cur.close()
00068 con.close()
00069
00070 def convert_record_to_dict(record):
00071 keys = ["id", "surface_id", "surface_height", "unrotated_points", "normal", "labels", "image", "intensities", "points"]
00072 a = dict()
00073 for i, k in enumerate(keys):
00074 a[k] = record[i]
00075 return a
00076
00077 def start(name):
00078 sqlite3.register_adapter(np.matrix, matrix_adapter)
00079 sqlite3.register_converter('matrix', matrix_converter)
00080 con = sqlite3.connect(name, detect_types=sqlite3.PARSE_DECLTYPES)
00081 cur = con.cursor()
00082 return con, cur
00083
00084 def create_db(name):
00085 con, cur = start(name)
00086 create_raw_table(cur)
00087 store_raw_data_from_individual_files(cur)
00088 con.commit()
00089 cur.close()
00090 con.close()
00091 print '>> db closed.'
00092
00093 def select(name):
00094 con, cur = start(name)
00095 selected = cur.execute('SELECT id, surface_id FROM raw_data')
00096 id, surface_id = zip(*selected)
00097 print "print len(set(id))"
00098 print len(set(id))
00099 print id
00100 print "print len(set(surface_id))"
00101 print len(set(surface_id))
00102 print surface_id
00103
00104 def recreate_table(name):
00105 con, cur = start(name)
00106 cur.execute('DROP TABLE bkup')
00107 create_bkup(cur)
00108 cur.execute('INSERT INTO bkup SELECT id, surface_id, surface_height, unrotated_points, normal, labels, image, points FROM raw_data')
00109 cur.execute('DROP TABLE raw_data')
00110 create_raw_table(cur)
00111 cur.execute('INSERT INTO raw_data SELECT id, surface_id, surface_height, unrotated_points, normal, labels, image, points FROM bkup')
00112 cur.execute('DROP TABLE bkup')
00113 con.commit()
00114 cur.close()
00115 con.close()
00116
00117 def test_converter_adapter():
00118 con, cur = start(':memory:')
00119 cur.execute("create table test(m matrix)")
00120 cur.execute("insert into test(m) values (?)", (np.matrix([1,2,3.])))
00121 cur.execute("select m from test")
00122 print "with declared types:", cur.fetchone()
00123 cur.close()
00124 con.close()
00125
00126
00127 if __name__ == '__main__':
00128 name = 'clutter_database.sqlite'
00129
00130
00131 con, cur = start(name)
00132
00133
00134
00135 result = cur.execute('SELECT * FROM raw_data')
00136 for i, r in enumerate(result):
00137 scan = convert_record_to_dict(r)
00138 print i, scan['id'], scan['points'].shape, np.matrix(scan['labels']).shape, scan['intensities'].shape, scan['unrotated_points'].shape,
00139 print scan['normal'].shape,
00140 pts = np.concatenate((scan['points'], np.matrix(scan['labels']), scan['intensities']), 0)
00141 print pts.shape
00142
00143
00144 cur.close()
00145 con.close()
00146
00147
00148
00149
00150
00151
00152
00153
00154
00155
00156
00157
00158
00159
00160
00161
00162
00163
00164
00165
00166
00167
00168
00169
00170
00171
00172
00173
00174
00175
00176
00177
00178
00179
00180