forked from umautobots/gta-postprocessing
/
Db_Manager.py
142 lines (110 loc) · 6.27 KB
/
Db_Manager.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
from configparser import ConfigParser
import postgresql as pg
from itertools import groupby
from pathlib import Path
import numpy as np
from Snapshot import Snapshot
from Detection import Detection
class db_manager:
def __init__(self):
self.ini_filename = "gta-postprocessing.ini.actual"
self.db_uri = None
def open_connection(self):
self.parse_ini()
self.conn = pg.open(self.db_uri)
def close_connection(self):
self.conn.close()
def setFalseProcessed(self):
self.open_connection()
query = "update snapshots set processed=FALSE;"
self.conn.query(query)
self.close_connection()
def parse_ini(self):
CONFIG = ConfigParser()
CONFIG.read(self.ini_filename)
self.db_uri = CONFIG["Database"]["URI"]
def getAllRuns(self, sess_id):
query = "SELECT run_id, runguid FROM runs WHERE session_id="+str(sess_id)+"order by run_id asc"
self.open_connection()
prep_stmt = self.conn.query(query)
runIds = []
runguIds = []
for p in prep_stmt:
runIds.append(p['run_id'])
runguIds.append(p['runguid'])
self.close_connection()
return runIds, runguIds
# fetches data from database that belongs to particular run of particular session and form them in data structure.
def getSnapShotsFromSessionAndRun(self, sess_id, r_id):
# sess_id : session id to be processed.
# r_id : run id that belongs to sess_id
# forms data in to datastructures.
def form_snapshots_and_detections(db_objs):
snapshot_tmp = Snapshot()
snapshot_tmp.snapshot_id = db_objs[0]['snapshot_id']
snapshot_tmp.runguid = db_objs[0]['runguid']
snapshot_tmp.imagepath = db_objs[0]['imagepath']
snapshot_tmp.view_matrix = np.array(db_objs[0]['view_matrix'], dtype=np.float64)
snapshot_tmp.proj_matrix = np.array(db_objs[0]['proj_matrix'], dtype=np.float64)
snapshot_tmp.width = db_objs[0]['width']
snapshot_tmp.height = db_objs[0]['height']
detectionList = []
for obj in db_objs:
detection_tmp = Detection()
detection_tmp.detection_id = obj['detection_id']
detection_tmp.type = obj['type']
detection_tmp.pos = obj['pos'] # need to process this position
detection_tmp.bbox = obj['bbox'] # geometry.Box
detection_tmp.vehicle_class = obj['class']
detection_tmp.handle = obj['handle']
detection_tmp.best_bbox = obj['best_bbox']
detection_tmp.fullbox = obj['fullbox']
detection_tmp.bbox3d_min = obj['bbox3d_min']
detection_tmp.bbox3d_max = obj['bbox3d_max']
detection_tmp.rotation = obj['rot']
detection_tmp.coverage = obj['coverage']
detectionList.append(detection_tmp)
snapshot_tmp.detections = detectionList
return snapshot_tmp
query = "SELECT snapshot_id, detection_id, coverage, type, class, best_bbox, runguid::text, imagepath, view_matrix," \
"width, height, proj_matrix, handle, pos::bytea, rot::bytea, bbox, ngv_box3dpolygon(bbox3d)::bytea as fullbox," \
"ST_MakePoint(ST_XMin(bbox3d), ST_YMin(bbox3d), ST_ZMin(bbox3d))::bytea as bbox3d_min," \
"ST_MakePoint(ST_XMax(bbox3d), ST_YMax(bbox3d), ST_ZMax(bbox3d))::bytea as bbox3d_max FROM detections JOIN snapshots USING (snapshot_id) JOIN runs USING (run_id) JOIN sessions USING(session_id) WHERE session_id=" + str(
sess_id) + " and run_id="+str(r_id)+" and processed=false and camera_pos <-> pos < 200 order by snapshot_id desc"
self.open_connection()
prep_stmt = self.conn.query(query)
snapshotList = []
for snapshot_id, db_objs in groupby(prep_stmt, key=lambda x: x['snapshot_id']):
db_objs = list(db_objs)
snapshot_tmp = form_snapshots_and_detections(db_objs)
snapshotList.append(snapshot_tmp)
self.close_connection()
return snapshotList
def getSnapShotsFromSession(self, sess_id):
query = "SELECT snapshot_id, detection_id, coverage, type, class, best_bbox, runguid::text, imagepath, view_matrix," \
"width, height, proj_matrix, handle, pos::bytea, rot::bytea, bbox, ngv_box3dpolygon(bbox3d)::bytea as fullbox," \
"ST_MakePoint(ST_XMin(bbox3d), ST_YMin(bbox3d), ST_ZMin(bbox3d))::bytea as bbox3d_min," \
"ST_MakePoint(ST_XMax(bbox3d), ST_YMax(bbox3d), ST_ZMax(bbox3d))::bytea as bbox3d_max FROM detections JOIN snapshots USING (snapshot_id) JOIN runs USING (run_id) JOIN sessions USING(session_id) WHERE session_id="+str(sess_id)+" and processed=false and camera_pos <-> pos < 200 order by snapshot_id desc"
self.open_connection()
prep_stmt = self.conn.query(query)
snapshotList = []
for snapshot_id, db_objs in groupby(prep_stmt, key=lambda x: x['snapshot_id']):
db_objs = list(db_objs)
snapshot_tmp = self.form_snapshots_and_detections(db_objs)
snapshotList.append(snapshot_tmp)
self.close_connection()
return snapshotList
# def update_bestBoxes_processed(self, snapshot:Snapshot, pixel_path: Path):
def update_bestBoxes_processed(self, snapshot:Snapshot):
self.open_connection()
update_query = self.conn.prepare("UPDATE detections SET best_bbox=$1, coverage=$2 WHERE detection_id = $3")
done_query = self.conn.prepare("UPDATE snapshots set processed=$1 where snapshot_id=$2")
done_query(snapshot.processed, snapshot.snapshot_id)
if snapshot.detections is not None:
for detectionObj in snapshot.detections:
assert isinstance(detectionObj, Detection)
update_query(detectionObj.best_bbox, detectionObj.coverage, detectionObj.detection_id) # just updating bestbounding boxes found earlier
# if snapshot.refined_stencil_coded is not None:
# np.savez_compressed(str(pixel_path / (str(snapshot.snapshot_id) + ".npz")),
# snapshot.refined_stencil_coded) # saves filtered stencil map in compressed format
self.close_connection()