【比赛】红绿灯

适用于有规则的红绿灯

思路

1. yolo检测红绿灯，返回对应的框图像
2. 根据像素坐标筛选我们要的那一个（去除别的红绿灯）
3. 转灰度图
4. orb特征点检测
5. 计算特征点在红绿灯的位置（上，中，下）
6. 根据位置和规则得出红绿灯状态
7. 众数滤波

优缺点

优点：针对比赛而言，效果牛逼
缺点：通用性不强

效果如图

下面是代码，懂得都懂

三部分
yolo
orb
filter

# -*- coding: UTF-8 -*-
import os
import sys
sys.path.append('/home/qm/light_ws/src/yolo_class/scripts/')
# sys.path.remove('/opt/ros/kinetic/lib/python2.7/dist-packages')

import cv2
import numpy as np
import random
import time
from pathlib import Path
import time
import threading

import torch
import torch.backends.cudnn as cudnn

from models.experimental import attempt_load
from utils.datasets import LoadStreams, LoadImages, letterbox
from utils.general import (
    check_img_size, non_max_suppression, apply_classifier, scale_coords,
    xyxy2xywh, plot_one_box, strip_optimizer, set_logging)
from utils.torch_utils import select_device, load_classifier, time_synchronized
from matplotlib import pyplot as plt


def create_random_color():
    # 功能：产生随机RGB颜色
    # 输出：color <class 'tuple'> 颜色

    r = random.randint(0, 255)
    g = random.randint(0, 255)
    b = random.randint(0, 255)

    color = (r, g, b)

    return color


def draw_one_box(img, classname, score, xyxy, color=[0, 255, 0], line_thickness=3):
    # 功能：绘制box
    # 输入：img <class 'numpy.ndarray'> (frame_height, frame_width, 3)
    #      classname <class 'str'>
    #      score <class 'float'>
    #      xyxy <class 'numpy.ndarray'> (4,)
    # 输出：img <class 'numpy.ndarray'> (frame_height, frame_width, 3)

    c1, c2 = (int(xyxy[0]), int(xyxy[1])), (int(xyxy[2]), int(xyxy[3]))
    cv2.rectangle(img, c1, c2, color, thickness=line_thickness, lineType=cv2.LINE_AA)

    tf = max(line_thickness - 1, 1)  # font thickness
    label = f'{classname} {score:.2f}'
    t_size = cv2.getTextSize(label, 0, fontScale=line_thickness / 3, thickness=tf)[0]
    c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
    cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filled
    cv2.putText(img, label, (c1[0], c1[1] - 2), 0, line_thickness / 3, [225, 255, 255], thickness=tf,
                lineType=cv2.LINE_AA)


class Yolov5Detector():
    def __init__(self):
        # Initialize
        self.device = select_device()
        self.half = self.device.type != 'cpu'  # half precision only supported on CUDA

        # Load model
        weights_path = './yolov5s.pt'
        self.model = attempt_load(weights_path, map_location=self.device)  # load FP32 model

        self.stride = int(self.model.stride.max())  # model stride
        self.img_size = check_img_size(img_size=640, s=self.stride)  # check img_size

        if self.half:
            self.model.half()  # to FP16

        # Get names
        self.names = self.model.module.names if hasattr(self.model, 'module') else self.model.names

        # Run inference
        img = torch.zeros((1, 3, self.img_size, self.img_size), device=self.device)  # init img
        _ = self.model(img.half() if self.half else img) if self.device.type != 'cpu' else None  # run once

    def run(self, img0, items=None, conf_thres=0.25):
        img_src = img0.copy()
        all_light = []
        # Padded resize
        img = letterbox(img0, new_shape=self.img_size)[0]
        # Convert
        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
        img = np.ascontiguousarray(img)

        with torch.no_grad():
            img = torch.from_numpy(img).to(self.device)
            img = img.half() if self.half else img.float()  # uint8 to fp16/32
            img /= 255.0  # 0 - 255 to 0.0 - 1.0
            if img.ndimension() == 3:
                img = img.unsqueeze(0)

            # Inference
            pred = self.model(img, augment=False)[0]

            # Apply NMS
            pred = non_max_suppression(pred, conf_thres=conf_thres, iou_thres=0.45, classes=items)

            det = pred[0]  # det <class 'torch.Tensor'> (N, 6) N(>=0)为目标数量

            all_light = []
            if det is not None:
                classes, scores, boxes = [], [], []

                # Rescale boxes from img_size to img0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round()

                for *xyxy, conf, cls in reversed(det):  # 按置信度从高到低
                    if 1:  # Add bbox to image
                        # label = '%s %.2f' % (names[int(cls)], conf)
                        # plot_one_box(xyxy, img_src, label=label, color=colors[int(cls)], line_thickness=3)
                        # print(xyxy[0], xyxy[1], xyxy[2], xyxy[3])
                        h1 = int(xyxy[1].cpu().numpy())
                        h2 = int(xyxy[3].cpu().numpy())
                        w1 = int(xyxy[0].cpu().numpy())
                        w2 = int(xyxy[2].cpu().numpy())
                        center_x = (w1 + w2) / 2
                        one_light = [h1, h2, w1, w2, center_x]
                        all_light.append(one_light)
                ######
                all_light = np.array(all_light)
                all_light = all_light[np.lexsort(all_light.T)]
                ######
                return all_light
            else:
                return None


class LightDetector:
    def __init__(self):
        self.detector=Yolov5Detector()
    def detect_light(self,input_img):
        img_height=input_img.shape[0]
        img=input_img
        src_img=img.copy()
        boxes = self.detector.run(img, items=[9], conf_thres=0.5)
        if boxes is None:
            return "No light",30,30
        if boxes is not None:
            if(len(boxes)<3):
                return "light<3",30,30
            if(len(boxes)>=3):
                ####y高度相似滤波
                l0_h = int(boxes[0][0])
                l1_h = int(boxes[1][0])
                l2_h = int(boxes[2][0])
                delta_y=np.abs(l0_h-l1_h)+np.abs(l1_h-l2_h)
                ####y高度相似滤波
                if (delta_y > int(img_height / 10)):
                    return "wrong list", 30, 30
                if(delta_y<int(img_height/10)):
                    h1=int(boxes[1][0])
                    h2 = int(boxes[1][1])
                    w1 = int(boxes[1][2])
                    w2 = int(boxes[1][3])
                    light_two=src_img[h1:h2,w1:w2]
                    light_two_x=int((w1+w2)/2)
                    light_two_y = int((h1 + h2) / 2)
                    if light_two is not None:
                        h=light_two.shape[0]
                        h1=int(4*h/10)
                        h2=int(h/2)
                        h3=int(4*h/5)
                        # 4.4之前cv2.xfeatures2d.SIFT_create()
                        sift = cv2.SIFT_create()
                        light_two_gray= cv2.cvtColor(light_two, cv2.COLOR_BGR2GRAY)
                        kp, des = sift.detectAndCompute(light_two_gray, None)
                        y_list=[]
                        for i in range(0,len(kp)):
                            y_list.append(kp[i].pt[1])
                        target_y=np.mean(y_list)
                        d1=np.abs(target_y-h1)
                        d2=np.abs(target_y-h2)
                        d3=np.abs(target_y-h3)
                        if(d1<=d2 and d1<=d3):
                            return "red",light_two_x,light_two_y
                        elif(d2<=d1 and d2<=d3):
                            return "yellow",light_two_x,light_two_y
                        elif(d3<=d1 and d3<=d2):
                            return "green",light_two_x,light_two_y
                        else:
                            return "light black",30,30
                    if light_two is None:
                        return "not light",30,30
class QmDetector:
    def __init__(self):
        self.light_detector=LightDetector()
        self.statue_array=np.zeros(20).tolist()
        self.result="green"
        self.x=20
        self.y=20
    def detect_one(self,frame):
        tmp_result,self.x,self.y=self.light_detector.detect_light(frame)
        self.statue_array=self.statue_array[1:]
        self.statue_array.append(tmp_result)
    def caculate_statue(self):
        cnt_g=0
        cnt_y=0
        cnt_r=0

        for i in range(0,len(self.statue_array)):
            if(self.statue_array[i]=="red"):
                cnt_r=cnt_r+1
            if(self.statue_array[i]=="green"):
                cnt_g=cnt_g+1
            if(self.statue_array[i]=="yellow"):
                cnt_y=cnt_y+1
        if(cnt_y>0 or cnt_r>0 or cnt_g>0):
            if(cnt_r>=cnt_g and cnt_r>=cnt_y):
                self.result= "red"
            elif (cnt_y >= cnt_g and cnt_y >= cnt_r):
                self.result= "yellow"
            elif (cnt_g >= cnt_r and cnt_g >= cnt_y):
                self.result= "green"
            else:
                self.result = "green"
        else:
            self.result ="green"


    def detect_light(self,frame):
        #计算一帧，更新状态数组
        self.detect_one(frame)
        #更新结果
        self.caculate_statue()
        return self.result,self.x,self.y

def main():
    detector=QmDetector()

    cap = cv2.VideoCapture('./target_video/light3.mp4')

    ###保存结果
    fps = 30
    size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    videoWriter = cv2.VideoWriter('./target_video/result.mp4', cv2.VideoWriter_fourcc('M', 'P', 'E', 'G'), fps, size)

    while (1):
        ret, frame = cap.read()
        t1=time.time()
        result,x,y=detector.detect_light(frame)
        print("time used: ", time.time() - t1)

        if result is not None:
            cv2.putText(frame,result,(20,20),cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255),2)
        cv2.imshow("capture", frame)
        videoWriter.write(frame)
        if cv2.waitKey(33) & 0xFF == ord('q'):
            break
    cap.release()
    videoWriter.release()
if __name__ == '__main__':
    main()