FORMAT_CONVERSION_MAXCOMPUTE/isr/pre_process.py

import colorsys
import time
import numpy as np
import cv2


def count_red_pixel(image_np, cnt=1000):
    # 红色像素计数
    start_time = time.time()
    image_hsv = cv2.cvtColor(image_np, cv2.COLOR_BGR2HSV)
    # minus_1 = image_np[:, :, 2].astype('int32') - image_np[:, :, 0].astype('int32')
    # minus_2 = image_np[:, :, 2].astype('int32') - image_np[:, :, 1].astype('int32')
    # red_mask = (image_np[:, :, 2] >= 180) & (minus_1 >= 60) & (minus_2 >= 60)
    red_mask = ((image_hsv[:, :, 0] >= 0) & (image_hsv[:, :, 0] <= 10) | (image_hsv[:, :, 0] <= 180) & (image_hsv[:, :, 0] >= 156)) \
               & (image_hsv[:, :, 1] <= 255) & (image_hsv[:, :, 1] >= 43) \
               & (image_hsv[:, :, 2] <= 255) & (image_hsv[:, :, 2] >= 100)
    red_cnt = np.sum(red_mask != 0)
    print("red_cnt", red_cnt, time.time()-start_time)
    if red_cnt >= cnt:
        return True
    else:
        return False


def get_classes(classes_path):
    """loads the classes"""
    with open(classes_path) as f:
        class_names = f.readlines()
    class_names = [c.strip() for c in class_names]
    return class_names


def get_anchors(anchors_path):
    """loads the anchors from a file"""
    with open(anchors_path) as f:
        anchors = f.readline()
    anchors = [float(x) for x in anchors.split(',')]
    return np.array(anchors).reshape(-1, 2)


def get_colors(number, bright=True):
    """
    Generate random colors for drawing bounding boxes.
    To get visually distinct colors, generate them in HSV space then
    convert to RGB.
    """
    if number <= 0:
        return []

    brightness = 1.0 if bright else 0.7
    hsv_tuples = [(x / number, 1., brightness)
                  for x in range(number)]
    colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
    colors = list(
        map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
            colors))
    np.random.seed(10101)  # Fixed seed for consistent colors across runs.
    np.random.shuffle(colors)  # Shuffle colors to decorrelate adjacent classes.
    np.random.seed(None)  # Reset seed to default.
    return colors