trf_learn: recognize_3d_density

Go to the documentation of this file.
00001 #!/usr/bin/python
00002 import roslib; roslib.load_manifest('trf_learn')
00003 import trf_learn.recognize_3d as r3d
00004 import hrl_lib.util as ut
00005 import pylab as pb
00006 from PIL import Image
00007 import os.path as pt
00008 import pdb
00009 import numpy as np
00010 import os
00011 
00012 def minmax(mat):
00013     return (np.min(mat[0,:]), np.max(mat[0,:]), 
00014             np.min(mat[1,:]), np.max(mat[1,:]))
00015 
00016 def num_bins(points, bin_size):
00017     minx, maxx, miny, maxy = minmax(points)
00018 
00019     rangex = maxx-minx
00020     xbins = np.ceil(rangex/bin_size)
00021 
00022     rangey = maxy-miny
00023     ybins = np.ceil(rangey/bin_size)
00024 
00025     print 'XBINS', xbins, 'YBINS', ybins
00026     return xbins, ybins
00027 
00028 ##
00029 #   ['neg_pred', 'image', 'pos_pred', 'tried', 'center']
00030 def density_plot(pickle_file_name):
00031     print 'density_plot: processing', pickle_file_name
00032     BIN_SIZE = 20
00033     #PICKLE_FOLDER = 'pickle_files'
00034 
00035     data_dict = ut.load_pickle(pickle_file_name)
00036     orig_pickle_folder, _ = pt.split(pickle_file_name)
00037     folder_name, img_name = pt.split(data_dict['image'])
00038     nimg_path = pt.join(orig_pickle_folder, img_name)
00039     img_obj = Image.open(nimg_path)
00040     w, h = img_obj.size
00041     pb.imshow(img_obj, origin='lower')
00042 
00043     data_dict['neg_pred'][1,:] = h - data_dict['neg_pred'][1,:]
00044     data_dict['pos_pred'][1,:] = h - data_dict['pos_pred'][1,:]
00045 
00046     all_pts = np.column_stack((data_dict['neg_pred'], data_dict['pos_pred']))
00047     Hall, xedges, yedges = np.histogram2d(all_pts[0,:].A1, all_pts[1,:].A1, 
00048                                           bins=num_bins(all_pts, BIN_SIZE))
00049     Hneg, xedges, yedges = np.histogram2d(data_dict['neg_pred'][0,:].A1, data_dict['neg_pred'][1,:].A1, 
00050                                           bins=[xedges, yedges])
00051 
00052     extent = [xedges[0], xedges[-1], yedges[-1], yedges[0]]
00053     Himage = (Hall-Hneg).T
00054     max_val, min_val = np.max(Himage), np.min(Himage)
00055     Hrgba = np.zeros((Himage.shape[0], Himage.shape[1], 4), dtype='uint8')
00056     Hrgba[:,:,0] = 0
00057     Hrgba[:,:,1] = 255 #Himage*80
00058     Hrgba[:,:,2] = 0
00059     Hrgba[:,:,3] = 255
00060     r,c = np.where(Himage == 0)
00061     Hrgba[r,c,3] = 0
00062 
00063     print 'max', max_val, 'min', min_val
00064     pb.imshow(Hrgba, extent=extent, interpolation='spline36', origin='upper', alpha = .7)
00065     if data_dict['tried'][1] == r3d.NEGATIVE:
00066         pb.plot(data_dict['tried'][0][0,0], h-data_dict['tried'][0][1,0], 'bx')
00067     else:
00068         pb.plot(data_dict['tried'][0][0,0], h-data_dict['tried'][0][1,0], 'bx')
00069 
00070     #pdb.set_trace()
00071     #pb.plot(data_dict['neg_pred'][0,:].A1, data_dict['neg_pred'][1,:].A1, 'rx')
00072     #pb.plot(data_dict['pos_pred'][0,:].A1, data_dict['pos_pred'][1,:].A1, 'x')
00073     min_x, max_x, min_y, max_y = minmax(all_pts)
00074     pb.axis([max(min_x-100,0), min(max_x+100,w), max(min_y-100, 0), min(max_y+100, h)])
00075     #pb.axis([0, w, 0, h])
00076     #name, extension = pt.splitext(img_name)
00077     name = pt.splitext(pt.split(pickle_file_name)[1])[0]
00078 
00079     if data_dict['tried'][1] == r3d.NEGATIVE:
00080         figname = pt.join(orig_pickle_folder, name + '_plot_FAIL.png')
00081     else:
00082         figname = pt.join(orig_pickle_folder, name + '_plot_SUCC.png')
00083 
00084     pb.savefig(figname)
00085     
00086     #pb.show()
00087 
00088 if __name__ == '__main__':
00089     import sys
00090     import optparse
00091     if len(sys.argv) > 1:
00092         for i in range(1, len(sys.argv)):
00093             density_plot(sys.argv[i])
00094 
00095