segment plot

from ROOT import larcv
from matplotlib import colors
import numpy as np
import matplotlib
import matplotlib.pyplot as plt

matplotlib.rcParams['font.size'] = 20
matplotlib.rcParams['font.family'] = 'serif'

larcv.load_pyutil

cmap = colors.ListedColormap(['blue','red','green','yellow','pink','orange'])
bounds=[0,1,2,3,4,5,6]
norm = colors.BoundaryNorm(bounds, cmap.N)

iom = larcv.IOManager()
iom.add_in_file("/data/vgentyvalid_hires_100.root")
iom.add_in_file("/data/vgenty/UBFCN/retraining/alex/fcn16/weights/ana/alex_fcn16_weights_21696_writeana.root.root")
iom.initialize()
iom.read_entry(397)

imgs=iom.get_data(larcv.kProductImage2D,"tpc_hires_crop").Image2DArray()
gts=iom.get_data(larcv.kProductImage2D,"segment_hires_crop").Image2DArray()
segs=iom.get_data(larcv.kProductImage2D,"fcn").Image2DArray()

# create 6 channel score tensor
ss = np.zeros([6,576,576])
for s in xrange(6):
    seg = larcv.as_ndarray(segs[s])
    ss[s,:,:] = seg

# argmax to get the scoremap
scoremap=ss.argmax(axis=0)

#Get plane 2 image and ground truth labels
p2img = larcv.as_ndarray(imgs[2])
p2seg = larcv.as_ndarray(gts[2])

# Threshold low pixel values
picut = p2img < 10.0

# Make them background in scoremap
scoremap[picut] = 0.0

# Order classes 1 to 5 (0 is background)
for ix,c in enumerate([3.0,4.0,6.0,8.0,9.0]):
    p2seg[p2seg == c] = ix+1

fig,_=plt.subplots(figsize=(30,20))

#Ground Truth
ax = plt.subplot(2,1,1)
img = plt.imshow(p2seg,cmap=cmap,norm=norm,interpolation='none')
ax.set_title("Ground Truth",fontweight='bold')

cbar = plt.colorbar(img, cmap=cmap, norm=norm, boundaries=bounds, ticks=[0,1,2,3,4,5,6],fraction=0.046, pad=0.04)
plt.axis("off")

#Segmentation
ax = plt.subplot(2,1,2)
img = plt.imshow(scoremap,cmap=cmap,norm=norm,interpolation='none')
ax.set_title("Inference",fontweight='bold')
cbar = plt.colorbar(img, cmap=cmap, norm=norm, boundaries=bounds, ticks=[0,1,2,3,4,5,6],fraction=0.046, pad=0.04)

plt.axis("off")
plt.show()