python: 3.10.3

依赖:

torchreid==0.2.5
gdown==5.2.1
tensorboard==2.20.0
torch==2.6.0+cu126
torchaudio==2.6.0+cu126
torchvision==0.21.0+cu126
opencv-python==4.11.0.86

class PersonReid(object):
    '''行人重识别 + 实时流跟踪'''

    def __init__(self, yolo_weights_file, reid_model_path, device):
        self.device = device
        self.model = YOLO(yolo_weights_file)
        self.model.to(self.device)
        self.font = ImageFont.truetype(SimHeiPath, 10)
        self._init_reid_model(reid_model_path, device)

        # 底库：存储目标人员的ID和特征向量
        self.target_person = {
            'person_id': '',
            'feature': None
        }

    def _init_reid_model(self, reid_model_path, device):
        name: str = os.path.splitext(os.path.basename(reid_model_path))[0]

        # 自动识别 num_classes
        if 'market' in reid_model_path:
            num_classes = 751
        elif 'dukemtmcreid' in reid_model_path:
            num_classes = 702
        elif 'msmt' in reid_model_path:
            num_classes = 1041
        else:
            num_classes = 1000  # imagenet 预训练不用管

        use_gpu = 'cpu' not in device

        self.reid_model = torchreid.models.build_model(
            name=name,
            num_classes=num_classes,
            loss='softmax',
            pretrained=False,
            use_gpu=use_gpu
        )

        torchreid.utils.load_pretrained_weights(self.reid_model, reid_model_path)
        self.reid_model = self.reid_model.to(device).eval()

        self.transform = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Resize((256, 128)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
        ])

    # def extract_single_feature(self, cropped_image):
    #     '''提取特征向量（已归一化，可直接比对）'''
    #     img_rgb = cv2.cvtColor(cropped_image, cv2.COLOR_BGR2RGB)
    #     img = self.transform(img_rgb).unsqueeze(0).to(self.device)
    #
    #     with torch.no_grad():
    #         feat = self.reid_model(img)
    #         feat = torch.nn.functional.normalize(feat, p=2, dim=1)
    #
    #     return feat.cpu().numpy().squeeze()

    def extract_single_feature(self, cropped_images):
        '''
        批量提取特征向量（已归一化，支持N张图一起推理）
        :param cropped_images: list -> [img1, img2, img3...]
        :return: numpy数组 -> shape = (N, 512)
        '''
        # 保存所有图像的 tensor
        batch_tensors = []

        for img in cropped_images:
            img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            tensor = self.transform(img_rgb)  # 只转tensor，不升维
            batch_tensors.append(tensor)

        # 拼成一个 batch [N, 3, 256, 128]
        batch_tensor = torch.stack(batch_tensors).to(self.device)

        # 一次性推理N张图（GPU 并行，速度极快）
        with torch.no_grad():
            feats = self.reid_model(batch_tensor)
            feats = torch.nn.functional.normalize(feats, p=2, dim=1) # L2 归一化

        return feats.cpu().numpy() # 转回 numpy 返回

    def is_same_person(self, vec1, vec2, threshold=0.65):
        '''比对是否为同一个人'''
        similarity = np.dot(vec1, vec2)
        return similarity > threshold, similarity

    def person_predict(self, image):
        '''YOLO 检测行人'''
        with torch.no_grad():
            boxes = self.model.predict(
                source=image,
                imgsz=640,
                conf=0.5,
                iou=0.5,
                classes=[0],
                verbose=False,
            )[0].boxes

        output = []
        for i in boxes:
            x1, y1, x2, y2 = [int(i) for i in i.xyxy[0].tolist()]  # 绝对位置  检测框的绝对位置
            output.append([x1, y1, x2, y2])
        return output

    def crop_image(self, image, _box):
        '''安全裁剪，防止越界报错'''
        h, w = image.shape[:2]
        x1, y1, x2, y2 = _box
        x1 = max(0, int(x1))
        y1 = max(0, int(y1))
        x2 = min(w, int(x2))
        y2 = min(h, int(y2))
        return image[y1:y2, x1:x2]

    def set_target_person(self, image, person_id):
        '''
        设置目标人员（底库入库）
        :param image: 单张目标人物图片
        :param person_id: 人员ID
        '''
        boxes = self.person_predict(image)
        if len(boxes) != 1:
            logger.error(f"未检测到人员或人员过多!{boxes}")
            return False

        # 取第一个人
        box = boxes[0]
        crop = self.crop_image(image, box)
        feat = self.extract_single_feature([crop, ])[0] # 固定取第一个向量

        self.target_person['person_id'] = person_id
        self.target_person['feature'] = feat
        logger.success(f"目标人物 {person_id} 已入库！")
        return True

    def draw_image(self, image, red_box, green_box):
        '''画图 - 只画四个角，不画整框'''
        rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        pil_image = Image.fromarray(rgb_image)
        draw = ImageDraw.Draw(pil_image)

        # 拐角长度（可以自己调大小，15~20最合适）
        corner_len = 8

        for _box in green_box:
            x1, y1, x2, y2, score = _box
            draw.rectangle((x1, y1, x2, y2), outline=(0, 255, 0), width=1)  # 区域 outline颜色 width粗细
            draw.text((x1, y1 - 11), f"{self.target_person['person_id']}[{score}]", font=self.font,
                      fill=(0, 255, 0))  # 位置. 文字. 字体. 颜色

        for _box in red_box:
            x1, y1, x2, y2 = _box

            # 左上
            draw.line([x1, y1, x1 + corner_len, y1], fill=(255, 0, 0), width=1)
            draw.line([x1, y1, x1, y1 + corner_len], fill=(255, 0, 0), width=1)
            # 右上
            draw.line([x2 - corner_len, y1, x2, y1], fill=(255, 0, 0), width=1)
            draw.line([x2, y1, x2, y1 + corner_len], fill=(255, 0, 0), width=1)
            # 左下
            draw.line([x1, y2 - corner_len, x1, y2], fill=(255, 0, 0), width=1)
            draw.line([x1, y2, x1 + corner_len, y2], fill=(255, 0, 0), width=1)
            # 右下
            draw.line([x2 - corner_len, y2, x2, y2], fill=(255, 0, 0), width=1)
            draw.line([x2, y2 - corner_len, x2, y2], fill=(255, 0, 0), width=1)

        frame = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
        return frame

    def resize_frame(self, frame, target_width):
        '''等比缩放'''
        h, w = frame.shape[:2]
        scale = target_width / w  # 计算缩放比例
        target_height = int(h * scale)  # 高度等比自动计算
        resized_frame = cv2.resize(frame, (target_width, target_height), interpolation=cv2.INTER_LINEAR)
        return resized_frame

    def start_realtime_track(self, frame, threshold, width, height):
        '''
        启动实时流追踪
        @params  frame      -->  帧;
        @params  threshold  -->  相似度阈值;
        @params  width      -->  原始宽;
        @params  height     -->  原始高;
        '''
        frame = self.resize_frame(frame, 640) # 等比缩放
        boxes = self.person_predict(frame)  # yolo检测
        cropped_images = [] # 裁剪后的数组
        red_box, green_box = [], []  # 目标和非目标

        for box in boxes: # 裁剪全部人存入数组
            x1, y1, x2, y2 = box
            crop = self.crop_image(frame, [x1, y1, x2, y2])  # 2. 裁剪 + 提取特征
            cropped_images.append(crop)

        if cropped_images:
            feat_list = self.extract_single_feature(cropped_images)  # 批量提取向量
            for index, feat in enumerate(feat_list):
                is_same, score = self.is_same_person(self.target_person['feature'], feat, threshold)  # 和底库比对

                _box: list = boxes[index] # 当前处理的人员
                if is_same:
                    _box.append(f'{score:.2f}')
                    green_box.append(_box)
                else:
                    red_box.append(_box)

            frame = self.draw_image(frame, red_box, green_box)
        frame = cv2.resize(frame, (width, height)) # 缩放回原始宽高
        return frame

实时流如下

import asyncio
import json
import cv2
from yolo import utils
from loguru import logger




async def stop_event(cap):
    '''监听前端退出接收流'''
    try:
        await asyncio.sleep(0.01)
    except asyncio.CancelledError:
        cap.release()
        logger.error(f'浏览器已被关闭,退出拉流停止')
        raise


async def reid_event_generator(rtsp_url: str, person_id, image, threshold):
    '''响应层'''
    status = utils.person_reid.set_target_person(image, str(person_id))
    max_reconnect = 3 # 最大从连次数
    cap = None

    if not status:
        logger.error(f'底图人物数量有误')
        return


    for con_num in range(0, max_reconnect):
        await stop_event(cap)
        if cap:
            cap.release() # 如果已经打开，先释放

        cap = cv2.VideoCapture(rtsp_url)
        cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)  # 缓冲区只留1帧
        cap.set(cv2.CAP_PROP_FPS, 13)  # 限制读取FPS，减轻压力

        if not cap.isOpened():
            cap.release()
            logger.error(f'第({con_num})次获取流失败!')
            continue

        success, frame = cap.read()
        height, width = frame.shape[:2]  # 取第一帧摄像头原始宽高

        while True:
            await stop_event(cap)
            for _ in range(5):
                cap.grab() # 丢3帧

            success, frame = cap.retrieve()
            if not success:
                cap.release()
                logger.error(f'获取流失败,断开, 准备第{con_num}次重连...')
                break

            frame = utils.person_reid.start_realtime_track(frame, threshold, width, height)
            ret, buffer = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 70])
            jpg_bytes = buffer.tobytes()
            yield (b'--frame\r\n' 
                   b'Content-Type: image/jpeg\r\n\r\n' + jpg_bytes + b'\r\n')

    if cap:
        cap.release()  # 如果已经打开，先释放

接口

async def person_reid_views(request: Request):
    '''
    人体重识别
    GET
    @params  url            -->  摄像头地址;
    @params  person_id      -->  人员id;
    @params  threshold      -->  阈值;

    POST
    @params  image_base64   -->  图像base64;
    '''
    if request.method == 'GET':
        params = request.query_params
        url, person_id, threshold = [params.get(key, None) for key in ('url', 'person_id',  'threshold')]
        with open('1.jpg', 'r') as f:
            image_base64 = f.read()
        return StreamingResponse(
            reid_event_generator(url, person_id, decodeb64_to_np(image_base64), float(threshold)),
            media_type='multipart/x-mixed-replace; boundary=frame')  # "text/event-stream")

    else:
        params = await request.json()  # 请求参数
        image_base64 = params.get('image_base64')
        with open('1.jpg', 'w+', encoding='utf-8') as f:
            f.write(image_base64)
        return {'code': 200, 'msg': 'success'}

@app.api_route('/person_reid', methods=['GET', 'POST']) # 行人重识别
async def person_reid(request: Request) -> dict: return await person_reid_views(request)

模型地址

通过网盘分享的文件：行人追踪-torchreid模型
链接: https://pan.baidu.com/s/10PGpNJZnhxOuURBqNtEU wg 提取码: 1111

页面

<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <!-- 🔥 修复中文乱码 关键 -->
    <meta charset="UTF-8">
</head>
<body>

    RTSP：<input type="text" id="url" value="rtsp://"><br><br>
    人员ID：<input type="text" id="person_id" value="1001"><br><br>
    阈值：<input type="text" id="threshold" value="0.7"><br><br>
    图片：<input type="file" id="file"><br><br>

    <button onclick="upload()">上传图片</button>
    <button onclick="startPlay()">播放</button>
    <button onclick="stopPlay()">停止</button>

    <br><br>
    <img id="video" width="800">

    <script>
        const baseUrl = "http://192.168.8.119:49351";
        const video = document.getElementById('video');

        // 上传图片 POST
        async function upload() {
            const file = document.getElementById('file').files[0];
            if (!file) return alert("请选图片");

            const reader = new FileReader();
            reader.onload = async (e) => {
                const image_base64 = e.target.result.split(',')[1];
                await fetch(baseUrl + "/person_reid", {
                    method: "POST",
                    headers: { "Content-Type": "application/json" },
                    body: JSON.stringify({ image_base64 })
                });
                alert("上传成功");
            };
            reader.readAsDataURL(file);
        }

        // 播放
        function startPlay() {
            const u = encodeURIComponent(document.getElementById('url').value);
            const pid = document.getElementById('person_id').value;
            const t = document.getElementById('threshold').value;
            video.src = `${baseUrl}/person_reid?url=${u}&person_id=${pid}&threshold=${t}`;
        }

        // 停止
        function stopPlay() {
            video.src = "";
        }
    </script>
</body>
</html>