class CapacityCounter(PipelineProcessor): 

 

 

    def __init__(self, area_mask, save_image=False, image_dir='./'): 

         super(CapacityCounter, self).__init__() 

      

        self.area_mask = area_mask 

         self.all = np.count_nonzero(area_mask) 

         self.image_dir = image_dir 

        self.save_image = save_image 

 

 

    def calculate_capacity(self, frame, frame_number): 

         base_frame = frame 

         # CLAHE (Contrast Limited Adaptive Histogram Equalization) 

         # this used for noise reduction at night time 

         frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY) 

         clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8)) 

         cl1 = clahe.apply(frame) 

 

 

         # getting edges with Canny filter 

         edges = cv2.Canny(frame,50,70) 

         # invert them to get white background 

         edges = ~edges 

 

 

         # blur with additional use of bilateralFilter to remove color noise 

         blur = cv2.bilateralFilter(cv2.blur(edges,(21,21), 100),9,200,200) 

 

 

         # threshold with ROI overlapping 

         _, threshold = cv2.threshold(blur,230, 255,cv2.THRESH_BINARY) 

        t = cv2.bitwise_and(threshold,threshold,mask = self.area_mask) 

 

 

         # counting capacity area 

         free = np.count_nonzero(t) 

         capacity = 1 - float(free)/self.all 

 

 

         # creating plot for debugging and visualization 

         if self.save_image: 

             img = np.zeros(base_frame.shape, base_frame.dtype) 

             img[:, :] = EXIT_COLOR 

             mask = cv2.bitwise_and(img, img, mask=self.area_mask) 

             cv2.addWeighted(mask, 1, base_frame, 1, 0, base_frame) 

 

 

             fig = plt.figure() 

             fig.suptitle("Capacity: {}%".format(capacity*100), fontsize=16) 

            plt.subplot(221),plt.imshow(base_frame),plt.title('Original') 

            plt.xticks([]), plt.yticks([]) 

            plt.subplot(222),plt.imshow(edges),plt.title('Cany edges') 

             plt.xticks([]), plt.yticks([]) 

             plt.subplot(223),plt.imshow(blur),plt.title('Blur') 

             plt.xticks([]), plt.yticks([]) 

             plt.subplot(224),plt.imshow(t),plt.title('Threshold with ROI mask') 

             plt.xticks([]), plt.yticks([]) 

 

 

             fig.savefig(self.image_dir + ("/processed_%s.png" % frame_number), dpi=500) 

 

 

         return capacity 

      

     def __call__(self, context): 

         frame = context['frame'].copy() 

     frame_number = context['frame_number'] 

                   capacity = self.calculate_capacity(frame, frame_number) 

        

       self.log.debug("Capacity: {}%".format(capacity*100)) 

       context['capacity'] = capacity 

         

        return context