目标检测

from ultralytics import YOLO

# 加载预训练模型
model = YOLO('yolov8n.pt')

# 训练
model.train(
    data='robocon.yaml',  # 数据集配置
    epochs=100,
    imgsz=640,
    batch=16,
    device='cuda'
)

# 导出为 TensorRT
model.export(format='engine', device='cuda')

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit

class YOLODetector:
    def __init__(self, engine_path):
        self.logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, 'rb') as f:
            runtime = trt.Runtime(self.logger)
            self.engine = runtime.deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()
        
    def detect(self, image):
        # 预处理
        input_image = self.preprocess(image)
        
        # 推理
        output = self.infer(input_image)
        
        # 后处理
        results = self.postprocess(output)
        return results

import cv2
import numpy as np

def color_detect(frame, lower_hsv, upper_hsv):
    """基于颜色的目标检测"""
    # 转换到HSV
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    
    # 阈值分割
    mask = cv2.inRange(hsv, lower_hsv, upper_hsv)
    
    # 形态学操作
    kernel = np.ones((5, 5), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
    
    # 找轮廓
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # 筛选最大轮廓
    if contours:
        largest = max(contours, key=cv2.contourArea)
        x, y, w, h = cv2.boundingRect(largest)
        return (x, y, w, h)
    return None

def shape_detect(contour, template_contours):
    """形状匹配"""
    scores = []
    for template in template_contours:
        score = cv2.matchShapes(contour, template, cv2.CONTOURS_MATCH_I1, 0)
        scores.append(score)
    return np.argmin(scores)