grasping_learner.py

Go to the documentation of this file.

#!/usr/bin/env python
import roslib; roslib.load_manifest('iri_grasping_learner')
import rospy

import numpy as np
from scipy.spatial.distance import cdist
import sys
from os.path import exists
import cPickle
import cv
from cv_bridge import CvBridge, CvBridgeError

from iri_grasping_learner.srv import *
from iri_learning_msgs.msg import GraspInfo

import ipdb

def cv_select_point(event, x, y, flags, param):
    # param is list where point will be added
    if event and cv.CV_EVENT_RBUTTONDOWN:
        param.append((x,y))
#-----

def user_select_point(cvimage):
    """Select point based on user input."""
    cv.NamedWindow('select_point')
    cv.ShowImage('select_point', cvimage)
    selected_point=[]
    cv.SetMouseCallback('select_point', cv_select_point, selected_point)
    while len(selected_point)==0:
        cv.WaitKey(100)
    cv.DestroyWindow('select_point')

    return selected_point
#------

class GraspsMemory:
    def __init__(self):
        self.descriptors = None
        self.grasps = []
        self.threshold = 0.3

    def add_grasp(self, desc, ginfo):
        self.grasps.append(ginfo)
        if self.descriptors == None:
            self.descriptors = desc
        else:
            self.descriptors = np.vstack((self.descriptors,desc))
        return True

    def compute_threshold(self, desc_set):
        pass
#------

class LearnedGrasps:
    def __init__(self):
        self.grasps_memory = None
        self.path_memory = None

    def select_learned_grasp_point_cb(self, req):
        # analyze finddd vectors in data OR ask user to select a point
        # FOR now only ask user
        
        bridge = CvBridge()
        try:
            cv_image = bridge.imgmsg_to_cv(req.ima, "bgr8")
        except CvBridgeError, e:
            print e

        # if no point found, ask user to select a point
        # get image from pointcloud
        selected_point = np.array(user_select_point(cv_image)[0])

        ##### search memory and chose grasp_info

        #preprocess descriptors
        descs = np.array([d.descriptor for d in req.descs.descriptors])
        uvs   = np.array([(d.u, d.v) for d in req.descs.descriptors])
        #find linear index and take descriptor
        pos_image = np.argmin(cdist(np.matrix(selected_point), uvs)[0])
        selected_desc = descs[pos_image,:]

        if len(self.grasps_memory.grasps) != 0:
            # compare with memory, and do something
            dists_to_memory = cdist(np.matrix(selected_desc), np.matrix(self.grasps_memory.descriptors))
            # select closest one
            pos_memory = np.argmin(dists_to_memory)
            score_memory = np.min(dists_to_memory)
            #filter by distance (?)
            if score_memory < self.grasps_memory.threshold:
                knows = True
                grasp_info = self.grasps_memory.grasps[pos_memory]
            else:
                knows = False
                grasp_info = GraspInfo()
        else:
            knows = False
            grasp_info = GraspInfo()

        # fill response and return
        res = FindGoodPointResponse()
        res.knows = knows
        res.grasp_info = grasp_info
        res.desc_index = pos_image
        return res

    def learn_new_grasp_cb(self, req):
        the_desc = np.array(req.desc.descriptors[req.desc_index].descriptor)
        update_ok = self.grasps_memory.add_grasp(the_desc, req.grasp_info)
        if update_ok:
            pf = open(self.path_memory,'w')
            cPickle.dump(self.grasps_memory,pf)
            pf.close()
        res = LearnNewGraspResponse()
        res.update_ok = update_ok
        return res

    def load_memory(self, path_memory=None):
        self.path_memory = path_memory
        if exists(self.path_memory):
            pf=open(self.path_memory)
            self.grasps_memory = cPickle.load(pf)
            pf.close()
        else:
            self.grasps_memory = GraspsMemory()

    def listener(self):
        rospy.init_node('grasping_learner')
        serv = rospy.Service('get_grasping_point', FindGoodPoint, self.select_learned_grasp_point_cb)
        serv2 = rospy.Service('learn_new_grasping', LearnNewGrasp, self.learn_new_grasp_cb)
        # check params
        if (not rospy.has_param("~path_memory")):
            rospy.logerr("Node has no param path_memory. Please set it up on launch file")
            sys.exit()
        path_memory = rospy.get_param("~path_memory")
        print "Launching grasp learner. Loading/saving from/to %s."%path_memory
        self.load_memory(path_memory)
        rospy.spin()



if __name__=="__main__":
    LfD = LearnedGrasps()
    LfD.listener()