db.py

Go to the documentation of this file.

import roslib
roslib.load_manifest('semantic_model_web_interface')

import mongo_ros as mr
import semanticmodel.msg as smm
import sensor_msgs.msg as sm
import cv
import cv_bridge
import os
from os import path
from . import site_root, lock, DB_PORT, DB_NAME, DB_HOST
import logging
from logging import config
import pymongo as pm
import semanticmodel.db as sdb

config.fileConfig(path.join(site_root, 'logging.conf'))
log = logging.getLogger('semantic_model_web.db')

cluster_collections = {}
image_collections = {}
bridge = cv_bridge.CvBridge()

POSITIONS = {
  'green_room': (20, 20, 8),
  'pool_room': (17, 35, 7),
  'white_lab': (34, 43, 10),
  'cafe': (30, 15, 7),
  'research_hallway': (5, 22, 7)
}

def filter_query(q):
    #q['$or'] = [{'num_vertical': {'$gt': 5}}, {'num_vertical': {'$exists':
    #    False}}]
    return q

def get_cluster_collection(run):
    with lock:
        return get_cached_collection(cluster_collections, run, 'blobs',
                                     smm.BlobMessage)

def get_image_collection(run):
    with lock:
        return get_cached_collection(image_collections, 
                                     run, 'images', sm.Image)

def get_cached_collection(colls, run, c, t):
    if run not in colls:
        colls[run] =\
            mr.MessageCollection(DB_NAME, sdb.collection_name(run, c), t)
    return colls[run]

def list_runs():
    conn = pm.Connection(host=DB_HOST, port=DB_PORT)
    db = conn[DB_NAME]
    runs = db['runs']
    return [('run{0}'.format(run['id']),
            '{0}.bag'.format(run['bag_file'])) for run in runs.find({})]

def run_info(run):
    conn = pm.Connection(host=DB_HOST, port=DB_PORT)
    db = conn[DB_NAME]
    runs = db['runs']
    id = int(run[3:])
    cursor = runs.find({'id': id})
    return next(cursor, None)

def image_file_loc(i, run):
    """
    @param i: Id of cluster
    @param run: Run id

    Ensure that the file for the image of the cluster is available (download
    if necessary), and return its site-relative URL
    """
    with lock:
        coll = get_image_collection(run)
        meta = coll.find_one({'cluster_id': i}, True, sort_by='creation_time',
                             ascending=False)
        relative_path = path.join('static/images', 
                                  run, 
                                  '{0}.jpg'.format(meta['blob_id']))
        filename = path.join(site_root, relative_path)
        log.debug("Looking for image file {0}".format(relative_path))
        if not path.exists(filename):
            log.info("{0} not found: getting from db".format(relative_path))
            msg, _ = coll.find_one(meta)
            img = bridge.imgmsg_to_cv(msg)
            d = path.dirname(filename)
            if not path.exists(d):
                os.makedirs(d)
            cv.SaveImage(filename, img)
        else:
            log.debug(" File found; not downloading")
        return '/'+relative_path    

def next_id(coll, id):
    cursor = coll.find(filter_query({'id': {'$gt': id}})).sort('id', pm.ASCENDING)
    item = next(cursor, None)
    if item:
        return item['id']

def prev_id(coll, id):
    cursor = coll.find(filter_query({'id': {'$lt': id}})).sort('id', pm.DESCENDING)
    item = next(cursor, None)
    if item:
        return item['id']

def semantic_query(args):
    q = {}
    c = args.get('color')
    if c:
        q['color'] = c
    s = args.get('size')
    if s == 'small':
        q['diameter'] = {'$lt': 0.2}
    elif s == 'medium':
        q['diameter'] = {'$gte': 0.2, '$lt': 0.8}
    elif s == 'large':
        q['diameter'] = {'$gte': 0.8}
    shape = args.get('shape')
    if shape == 'planar':
        q['c01'] = {'$gt': 0.5}
    elif shape == 'curved':
        q['c01'] = {'$lte': 0.5}
    s = args.get('loc')
    if s:
        pos = POSITIONS.get(s)
        if pos:
            radius = pos[2]
            q['x'] = {'$lt': pos[0]+radius, '$gt': pos[0]-radius}
            q['y'] = {'$lt': pos[1]+radius, '$gt': pos[1]-radius}

    return q