计划使用MOT16数据集做基于ConvLSTM的目标检测,其标注数据在gt.txt中,第8列为分类,有以下12类

1. Pedestrian(行人)——主要目标,要跟踪
2. Person on vehicle(车上的人)
3. Car(轿车)
4. Bicycle(自行车)
5. Motorbike(摩托车)
6. Non motorized vehicle(非机动车)
7. Static person(静止的人)
8. Distractor(干扰物)
9. Occluder(遮挡物)
10. Occluder on the ground(地面遮挡物)
11. Occluder full(完全遮挡)
12. Reflection(倒影)

如果做目标检测的话,我们只关注前面7类,后面的就不关注了,预览代码如下:

import os
import cv2
import numpy as np
from collections import defaultdict

mot16_root = "MOT16/train"
seqs = [d for d in os.listdir(mot16_root) if os.path.isdir(os.path.join(mot16_root, d))]

colors = {
    1: (0,255,0),
    2: (255,0,0),
    3: (0,0,255),
    4: (255,255,0),
    5: (255,0,255),
    6: (0,255,255),
    7: (128,0,128)
}

paused = False

for seq in seqs:
    gt_path = os.path.join(mot16_root, seq, "gt/gt.txt")
    img_dir = os.path.join(mot16_root, seq, "img1")
    
    if not os.path.exists(gt_path) or not os.path.exists(img_dir):
        continue
    
    gts = np.loadtxt(gt_path, delimiter=',')
    frame_gts = defaultdict(list)
    for gt in gts:
        frame_gts[int(gt[0])].append(gt)
    
    for frame_id in sorted(frame_gts.keys()):
        img = cv2.imread(os.path.join(img_dir, f"{frame_id:06d}.jpg"))
        
        for gt in frame_gts[frame_id]:
            cls = int(gt[7])
            
            if cls >= 8:
                continue
            
            x, y, w, h = gt[2:6]
            x1, y1 = int(x), int(y)
            x2, y2 = int(x + w), int(y + h)
            color = colors.get(cls, (200,200,200))
            label = f"{cls}"
            cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
            cv2.putText(img, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
        
        cv2.imshow(seq, img)
        
        while True:
            key = cv2.waitKey(30) & 0xFF
            if key == ord(' '):
                paused = not paused
                break
            if key == ord('q'):
                break
            if not paused:
                break
        if key == ord('q'):
            break
    if key == ord('q'):
        break

    cv2.destroyAllWindows()

Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐