produceTreeForValidationProng.py

Go to the documentation of this file.

import sys
import caffe
import matplotlib
import numpy as np
import lmdb
import argparse
import leveldb
import ROOT
import h5py
import eventLabelToLepton as eltl
from collections import defaultdict

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--proto', type=str, required=True)
    parser.add_argument('--model', type=str, required=True)
    parser.add_argument('--leveldb', type=str, required=True)
    parser.add_argument('--hdf5', type=str, required=False)
    parser.add_argument('--output', type=str, required=False)
    args = parser.parse_args()

    count = 0
    correct = 0
    matrix = defaultdict(int) # (real,pred) -> int
    labels_set = set()

    net = caffe.Net(args.proto, args.model, caffe.TEST)
    caffe.set_mode_gpu()
    db = leveldb.LevelDB(args.leveldb)

    if args.hdf5 is not None:
        fileNameH = args.hdf5+".0.h5"
        h_f = h5py.File(fileNameH,  "r")
        a_group_key = h_f.keys()[0]
        h_data = list(h_f[a_group_key])
        h_size=len(h_data)
        h_f.close()
        
    #lmdb_env = lmdb.open(args.lmdb)
    #lmdb_txn = lmdb_env.begin()
    #lmdb_cursor = lmdb_txn.cursor()

    t = ROOT.TTree( 't1', 'tree with histos' )
    if args.output is None:
        fout = ROOT.TFile("output.root", "RECREATE")
    else:
        print("Output name: %s" % args.output)
        fout = ROOT.TFile("%s.root" % args.output, "RECREATE")
    fout.cd()    
    
    arrpidProton=np.zeros(1, dtype=float)
    arrpidPion=np.zeros(1, dtype=float)
    arrpidGamma=np.zeros(1, dtype=float)
    arrpidMuon=np.zeros(1, dtype=float)
    arrpidElectron=np.zeros(1, dtype=float)
    arrpidPiZero=np.zeros(1, dtype=float)
    arrpidNeutron=np.zeros(1, dtype=float)

    truelabel = np.zeros(1, dtype=int)
    truelabelall = np.zeros(1, dtype=int)
    selectedlabel = np.zeros(1, dtype=int)

    nhit = np.zeros(1, dtype=int)

    tVertX = np.zeros(1, dtype=float)
    tVertY = np.zeros(1, dtype=float)
    tVertZ = np.zeros(1, dtype=float)
    
    t.Branch( 'truelabelall', truelabelall, 'truelabelall/I' )
    t.Branch( 'truelabel', truelabel, 'truelabel/I' )
    t.Branch( 'selectedlabel', selectedlabel, 'selectedlabel/I')
    t.Branch( 'nhit', nhit, 'nhit/I' )
    t.Branch( 'tvertx', tVertX, 'tvertx/D' )
    t.Branch( 'tverty', tVertY, 'tverty/D' )
    t.Branch( 'tvertz', tVertZ, 'tvertz/D' )
    t.Branch( 'proton', arrpidProton, 'proton/D' )
    t.Branch( 'pion', arrpidPion, 'pion/D' )
    t.Branch( 'gamma', arrpidGamma, 'gamma/D' )
    t.Branch( 'muon', arrpidMuon, 'muon/D' )
    t.Branch( 'electron', arrpidElectron, 'electron/D' )
    t.Branch( 'pizero', arrpidPiZero, 'pizero/D' )
    t.Branch( 'neutron', arrpidNeutron, 'neutron/D' )
    #how many events do you want to run over?
    nEvents=5000000
    h_count=0
    fileCountH=0
    
    for key, value in db.RangeIter():
        if count > nEvents: 
            break
            
        datum = caffe.proto.caffe_pb2.Datum()
        datum.ParseFromString(value)
        label = int(datum.label)
        truelabelall[0]=label
        truelabel[0]=eltl.labelToLepton(label)
        
        image = caffe.io.datum_to_array(datum)
        image = image.astype(np.uint8)

        nhit[0]=np.count_nonzero(np.asarray([image]))
        
        if args.hdf5 is not None:
            if (count + 1) > h_size:
                fileCountH=fileCountH+1
                fileNameHadd = args.hdf5+"."+str(fileCountH)+".h5"
                hadd_f = h5py.File(fileNameHadd,  "r")
                a_group_key_add = hadd_f.keys()[0]
                h_data = list(hadd_f[a_group_key_add])
                h_count=0
                h_size=h_size+len(h_data)
                hadd_f.close()

            tVertX[0]=h_data[h_count][0]
            tVertY[0]=h_data[h_count][1]
            tVertZ[0]=h_data[h_count][2]
        
        out = net.forward_all(data=np.asarray([image]))
        
        plabel = int(out['prob'][0].argmax(axis=0))
        selectedlabel[0]=plabel 
        
        arrpidProton[0]=out['prob'][0][2]
        arrpidPion[0]=out['prob'][0][4]
        arrpidGamma[0]=out['prob'][0][6]
        arrpidMuon[0]=out['prob'][0][1]
        arrpidElectron[0]=out['prob'][0][0]
        arrpidNeutron[0]=out['prob'][0][3]
        arrpidPiZero[0]=out['prob'][0][5]
 
        t.Fill()

        h_count = h_count + 1
        count = count + 1
        
        iscorrect = label == plabel
        correct = correct + (1 if iscorrect else 0)
        matrix[(label, plabel)] += 1
        labels_set.update([label, plabel])

        # if not iscorrect:
        #     print("\rError: key=%s, expected %i but predicted %i" \
        #                               % (key, label, plabel))

        sys.stdout.write("\rAccuracy: %.1f%% Progress: %.2f%%" % (100.*correct/count,100.*count/nEvents))
        sys.stdout.flush()
                
    print(str(correct) + " out of " + str(count) + " were classified correctly")

    fout.WriteTObject(t,"t")