rail_object_detector: show_masks.py Source File

Go to the documentation of this file.
00001 import numpy as np
00002 import matplotlib.pyplot as plt
00003 import random
00004 import cv2
00005 
00006 def show_masks(im, dets, msks, show = True, thresh = 1e-3, scale = 1.0):
00007     plt.cla()
00008     plt.imshow(im)
00009     for det, msk in zip(dets, msks):
00010         color = (random.random(), random.random(), random.random())  # generate a random color
00011         bbox = det[:4] * scale
00012         cod = np.zeros(4).astype(int)
00013         cod[0] = int(bbox[0])
00014         cod[1] = int(bbox[1])
00015         cod[2] = int(bbox[2])
00016         cod[3] = int(bbox[3])
00017         if im[cod[0]:cod[2], cod[1]:cod[3], 0].size > 0:
00018             msk = cv2.resize(msk, im[cod[1]:cod[3], cod[0]:cod[2], 0].T.shape)
00019             bimsk = msk > thresh
00020             bimsk = bimsk.astype(int)
00021             bimsk = np.repeat(bimsk[:, :, np.newaxis], 3, axis=2)
00022             mskd = im[cod[1]:cod[3], cod[0]:cod[2], :] * bimsk
00023             clmsk = np.ones(bimsk.shape) * bimsk
00024             clmsk[:, :, 0] = clmsk[:, :, 0] * color[0] * 256;
00025             clmsk[:, :, 1] = clmsk[:, :, 1] * color[1] * 256;
00026             clmsk[:, :, 2] = clmsk[:, :, 2] * color[2] * 256;
00027             im[cod[1]:cod[3], cod[0]:cod[2], :] = im[cod[1]:cod[3], cod[0]:cod[2], :] + 0.8 * clmsk - 0.8 * mskd
00028     plt.imshow(im)
00029     if(show):
00030         plt.show()
00031     return im
00032