PaddleX/static/examples/meter_reader/reader_infer.py

# coding: utf8
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import os.path as osp
import numpy as np
import math
import cv2
import argparse

from paddlex.seg import transforms
import paddlex as pdx

# The size of inputting images (METER_SHAPE x METER_SHAPE) of the segmenter,
# also the size of circular meters.
METER_SHAPE = 512
# Center of a circular meter
CIRCLE_CENTER = [256, 256]
# Radius of a circular meter
CIRCLE_RADIUS = 250
PI = 3.1415926536
# During the postprocess phase, annulus formed by the radius from
# 130 to 250 of a circular meter will be converted to a rectangle.
# So the height of the rectangle is 120.
LINE_HEIGHT = 120
# The width of the rectangle is 1570, that is to say the perimeter of a circular meter
LINE_WIDTH = 1570
# The type of a meter is estimated by a threshold. If the number of scales in a meter is
# greater than or equal to the threshold, the meter is belong to the former type.
# Otherwize, the latter.
TYPE_THRESHOLD = 40
# The configuration information of a meter, composed of scale value, range, unit.
METER_CONFIG = [{
    'scale_value': 25.0 / 50.0,
    'range': 25.0,
    'unit': "(MPa)"
}, {
    'scale_value': 1.6 / 32.0,
    'range': 1.6,
    'unit': "(MPa)"
}]


def parse_args():
    parser = argparse.ArgumentParser(description='Meter Reader Infering')
    parser.add_argument(
        '--detector_dir',
        dest='detector_dir',
        help='The directory of models to do detection',
        type=str)
    parser.add_argument(
        '--segmenter_dir',
        dest='segmenter_dir',
        help='The directory of models to do segmentation',
        type=str)
    parser.add_argument(
        '--image_dir',
        dest='image_dir',
        help='The directory of images to be infered',
        type=str,
        default=None)
    parser.add_argument(
        '--image',
        dest='image',
        help='The image to be infered',
        type=str,
        default=None)
    parser.add_argument(
        '--use_camera',
        dest='use_camera',
        help='Whether use camera or not',
        action='store_true')
    parser.add_argument(
        '--camera_id',
        dest='camera_id',
        type=int,
        help='The camera id',
        default=0)
    parser.add_argument(
        '--use_erode',
        dest='use_erode',
        help='Whether erode the predicted lable map',
        action='store_true')
    parser.add_argument(
        '--erode_kernel',
        dest='erode_kernel',
        help='Erode kernel size',
        type=int,
        default=4)
    parser.add_argument(
        '--save_dir',
        dest='save_dir',
        help='The directory for saving the inference results',
        type=str,
        default='./output/result')
    parser.add_argument(
        '--score_threshold',
        dest='score_threshold',
        help="Detected bbox whose score is lower than this threshlod is filtered",
        type=float,
        default=0.5)
    parser.add_argument(
        '--seg_batch_size',
        dest='seg_batch_size',
        help="Segmentation batch size",
        type=int,
        default=2)

    return parser.parse_args()


def is_pic(img_name):
    valid_suffix = ['JPEG', 'jpeg', 'JPG', 'jpg', 'BMP', 'bmp', 'PNG', 'png']
    suffix = img_name.split('.')[-1]
    if suffix not in valid_suffix:
        return False
    return True


class MeterReader:
    """Find the meters in images and provide a digital readout of each meter.

    Args:
        detector_dir(str): directory of the detector.
        segmenter_dir(str): directory of the segmenter.

    """

    def __init__(self, detector_dir, segmenter_dir):
        if not osp.exists(detector_dir):
            raise Exception("Model path {} does not exist".format(
                detector_dir))
        if not osp.exists(segmenter_dir):
            raise Exception("Model path {} does not exist".format(
                segmenter_dir))
        self.detector = pdx.load_model(detector_dir)
        self.segmenter = pdx.load_model(segmenter_dir)
        # Because we will resize images with (METER_SHAPE, METER_SHAPE) before fed into the segmenter,
        # here the transform is composed of normalization only.
        self.seg_transforms = transforms.Compose([transforms.Normalize()])

    def predict(self,
                im_file,
                save_dir='./',
                use_erode=True,
                erode_kernel=4,
                score_threshold=0.5,
                seg_batch_size=2):
        """Detect meters in a image, segment scales and points in these meters, the postprocess are
        done to provide a digital readout according to scale and point location.

        Args:
            im_file (str):  the path of a image to be predicted.
            save_dir (str): the directory to save the visual prediction. Default: './'.
            use_erode (bool, optional): whether to do image erosion by using a specific structuring element for
                the label map output from the segmenter. Default: True.
            erode_kernel (int, optional): structuring element used for erosion. Default: 4.
            score_threshold (float, optional): detected meters whose scores are not lower than `score_threshold`
                will be fed into the following segmenter. Default: 0.5.
            seg_batch_size (int, optional): batch size of meters when do segmentation. Default: 2.

        """
        if isinstance(im_file, str):
            im = cv2.imread(im_file).astype('float32')
        else:
            im = im_file.copy()
        # Get detection results
        det_results = self.detector.predict(im)
        # Filter bbox whose score is lower than score_threshold
        filtered_results = list()
        for res in det_results:
            if res['score'] > score_threshold:
                filtered_results.append(res)

        resized_meters = list()
        for res in filtered_results:
            # Crop the bbox area
            xmin, ymin, w, h = res['bbox']
            xmin = max(0, int(xmin))
            ymin = max(0, int(ymin))
            xmax = min(im.shape[1], int(xmin + w - 1))
            ymax = min(im.shape[0], int(ymin + h - 1))
            sub_image = im[ymin:(ymax + 1), xmin:(xmax + 1), :]

            # Resize the image with shape (METER_SHAPE, METER_SHAPE)
            meter_shape = sub_image.shape
            scale_x = float(METER_SHAPE) / float(meter_shape[1])
            scale_y = float(METER_SHAPE) / float(meter_shape[0])
            meter_meter = cv2.resize(
                sub_image,
                None,
                None,
                fx=scale_x,
                fy=scale_y,
                interpolation=cv2.INTER_LINEAR)
            meter_meter = meter_meter.astype('float32')
            resized_meters.append(meter_meter)

        meter_num = len(resized_meters)
        seg_results = list()
        for i in range(0, meter_num, seg_batch_size):
            im_size = min(meter_num, i + seg_batch_size)
            meter_images = list()
            for j in range(i, im_size):
                meter_images.append(resized_meters[j])
            # Segment scales and point in each meter area
            result = self.segmenter.batch_predict(
                transforms=self.seg_transforms, img_file_list=meter_images)
            # Do image erosion for the predicted label map of each meter
            if use_erode:
                kernel = np.ones((erode_kernel, erode_kernel), np.uint8)
                for i in range(len(result)):
                    result[i]['label_map'] = cv2.erode(result[i]['label_map'],
                                                       kernel)
            seg_results.extend(result)

        results = list()
        # The postprocess are done to get the point location relative to the scales
        for i, seg_result in enumerate(seg_results):
            result = self.read_process(seg_result['label_map'])
            results.append(result)

        # Provide a digital readout according to point location relative to the scales
        meter_values = list()
        for i, result in enumerate(results):
            if result['scale_num'] > TYPE_THRESHOLD:
                value = result['scales'] * METER_CONFIG[0]['scale_value']
            else:
                value = result['scales'] * METER_CONFIG[1]['scale_value']
            meter_values.append(value)
            print("-- Meter {} -- result: {} --\n".format(i, value))

        # Visualize the results
        visual_results = list()
        for i, res in enumerate(filtered_results):
            # Use `score` to represent the meter value
            res['score'] = meter_values[i]
            visual_results.append(res)
        pdx.det.visualize(im_file, visual_results, -1, save_dir=save_dir)

    def read_process(self, label_maps):
        """Get the pointer location relative to the scales.

        Args:
            label_maps (np.array): the label map output from a segmeter for a meter.

        """
        # Convert the circular meter into a rectangular meter
        line_images = self.creat_line_image(label_maps)
        # Get two one-dimension data where 0 represents background and >0 represents
        # a scale or a pointer
        scale_data, pointer_data = self.convert_1d_data(line_images)
        # Fliter scale data whose value is lower than the mean value
        self.scale_mean_filtration(scale_data)
        # Get the number of scales，the pointer location relative to the scales, the ratio between
        # the distance from the pointer to the starting scale and distance from the ending scale to the
        # starting scale.
        result = self.get_meter_reader(scale_data, pointer_data)
        return result

    def creat_line_image(self, meter_image):
        """Convert the circular meter into a rectangular meter.

        The minimum scale value is at the bottom left, the maximum scale value
        is at the bottom right, so the vertical down axis is the starting axis and
        rotates around the meter ceneter counterclockwise.

        Args:
            meter_image (np.array): the label map output from a segmeter for a meter.

        Returns:
            line_image (np.array): a rectangular meter.
        """

        line_image = np.zeros((LINE_HEIGHT, LINE_WIDTH), dtype=np.uint8)
        for row in range(LINE_HEIGHT):
            for col in range(LINE_WIDTH):
                theta = PI * 2 / LINE_WIDTH * (col + 1)
                rho = CIRCLE_RADIUS - row - 1
                y = int(CIRCLE_CENTER[0] + rho * math.cos(theta) + 0.5)
                x = int(CIRCLE_CENTER[1] - rho * math.sin(theta) + 0.5)
                line_image[row, col] = meter_image[y, x]
        return line_image

    def convert_1d_data(self, meter_image):
        """Get two one-dimension data where 0 represents background and >0 represents
           a scale or a pointer from the rectangular meter.

        Args:
            meter_image (np.array): the two-dimension rectangular meter output
                from function creat_line_image().

        Returns:
            scale_data (np.array): a one-dimension data where 0 represents background and
                >0 represents scales.
            pointer_data (np.array): a one-dimension data where 0 represents background and
                >0 represents a pointer.
        """

        scale_data = np.zeros((LINE_WIDTH), dtype=np.uint8)
        pointer_data = np.zeros((LINE_WIDTH), dtype=np.uint8)
        # Accumulte the number of positions whose label is 1 along the height axis.
        # Accumulte the number of positions whose label is 2 along the height axis.
        for col in range(LINE_WIDTH):
            for row in range(LINE_HEIGHT):
                if meter_image[row, col] == 1:
                    pointer_data[col] += 1
                elif meter_image[row, col] == 2:
                    scale_data[col] += 1
        return scale_data, pointer_data

    def scale_mean_filtration(self, scale_data):
        """Set the element in the scale data which is lower than its mean value to 0.

        Args:
            scale_data (np.array): the scale data output from function convert_1d_data().
        """
        mean_data = np.mean(scale_data)
        for col in range(LINE_WIDTH):
            if scale_data[col] < mean_data:
                scale_data[col] = 0

    def get_meter_reader(self, scale_data, pointer_data):
        """Calculate the number of scales，the pointer location relative to the scales, the ratio between
        the distance from the pointer to the starting scale and distance from the ending scale to the
        starting scale.

        Args:
            scale_data (np.array): a scale data output from function scale_mean_filtration().
            pointer_data (np.array): a pointer data output from function convert_1d_data().

        Returns:
            Dict (keys: 'scale_num', 'scales', 'ratio'):
                The value of key 'scale_num' (int): the number of scales;
                The value of 'scales' (float): the pointer location relative to the scales;
                the value of 'ratio' (float): the ratio between from the pointer to the starting scale and
                distance from the ending scale to the starting scale.

        """
        scale_flag = False
        pointer_flag = False
        one_scale_start = 0
        one_scale_end = 0
        one_pointer_start = 0
        one_pointer_end = 0
        scale_location = list()
        pointer_location = 0
        for i in range(LINE_WIDTH - 1):
            if scale_data[i] > 0 and scale_data[i + 1] > 0:
                if scale_flag == False:
                    one_scale_start = i
                    scale_flag = True
            if scale_flag:
                if scale_data[i] == 0 and scale_data[i + 1] == 0:
                    one_scale_end = i - 1
                    one_scale_location = (one_scale_start + one_scale_end) / 2
                    scale_location.append(one_scale_location)
                    one_scale_start = 0
                    one_scale_end = 0
                    scale_flag = False
            if pointer_data[i] > 0 and pointer_data[i + 1] > 0:
                if pointer_flag == False:
                    one_pointer_start = i
                    pointer_flag = True
            if pointer_flag:
                if pointer_data[i] == 0 and pointer_data[i + 1] == 0:
                    one_pointer_end = i - 1
                    pointer_location = (
                        one_pointer_start + one_pointer_end) / 2
                    one_pointer_start = 0
                    one_pointer_end = 0
                    pointer_flag = False

        scale_num = len(scale_location)
        scales = -1
        ratio = -1
        if scale_num > 0:
            for i in range(scale_num - 1):
                if scale_location[
                        i] <= pointer_location and pointer_location < scale_location[
                            i + 1]:
                    scales = i + (pointer_location - scale_location[i]) / (
                        scale_location[i + 1] - scale_location[i] + 1e-05) + 1
            ratio = (pointer_location - scale_location[0]) / (
                scale_location[scale_num - 1] - scale_location[0] + 1e-05)
        result = {'scale_num': scale_num, 'scales': scales, 'ratio': ratio}
        return result


def infer(args):
    image_lists = list()
    if args.image is not None:
        if not osp.exists(args.image):
            raise Exception("Image {} does not exist.".format(args.image))
        if not is_pic(args.image):
            raise Exception("{} is not a picture.".format(args.image))
        image_lists.append(args.image)
    elif args.image_dir is not None:
        if not osp.exists(args.image_dir):
            raise Exception("Directory {} does not exist.".format(
                args.image_dir))
        for im_file in os.listdir(args.image_dir):
            if not is_pic(im_file):
                continue
            im_file = osp.join(args.image_dir, im_file)
            image_lists.append(im_file)

    meter_reader = MeterReader(args.detector_dir, args.segmenter_dir)
    if len(image_lists) > 0:
        for im_file in image_lists:
            meter_reader.predict(im_file, args.save_dir, args.use_erode,
                                 args.erode_kernel, args.score_threshold,
                                 args.seg_batch_size)
    elif args.use_camera:
        cap_video = cv2.VideoCapture(args.camera_id)
        if not cap_video.isOpened():
            raise Exception(
                "Error opening video stream, please make sure the camera is working"
            )

        while cap_video.isOpened():
            ret, frame = cap_video.read()
            if ret:
                meter_reader.predict(frame, args.save_dir, args.use_erode,
                                     args.erode_kernel, args.score_threshold,
                                     args.seg_batch_size)
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break
            else:
                break
        cap_video.release()


if __name__ == '__main__':
    args = parse_args()
    infer(args)