Python 读取 SOT 数据集图片

数据集的读取

import numpy as np
import cv2
import os
from PIL import Image


def read_image(filename, color_fmt='RGB'):
    assert color_fmt in Image.MODES
    img = Image.open(filename)
    if not img.mode == color_fmt:
        img = img.convert(color_fmt)
    return np.asarray(img)


def save_image(filename, img, color_fmt='RGB'):
    assert color_fmt in ['RGB', 'BGR']
    if color_fmt == 'BGR' and img.ndim == 3:
        img = img[..., ::-1]
    img = Image.fromarray(img)
    return img.save(filename)
    

def show_frames(sequence_dir, frame_format, sequence_len, predict_bboxes):
    """输出最后的显示结果
    sequence_dir: 序列的文件夹部分
    frame_format: 帧的格式 传入格式化字符串
      '{:08d}.jpg' 表示 GOT-10k 以及 LaSOT 的 frame 格式
      '{}.jpg' 表示 TrackingNet 的格式
    sequence_len: 序列的长度
      设置的原因是：LaSOT 里面有一个 groundtruth.txt，混合在里面
    predict_bboxes: 预测的边界框
    """
    for i in range(1, sequence_len):
        # (1) 获取 image
        img_file = os.path.join(sequence_dir, frame_format.format(i))
        # print(img_file)

        # (2) 获取 predict_bbox
        predict_bbox = predict_bboxes[i]
        # # (3) 展示添加预测框的图片
        # flags[0, 1] 分别表示灰度、彩色图像
        img = cv2.imread(img_file, flags=1)
        img = cv2.rectangle(
            img,
            (int(predict_bbox[0]), int(predict_bbox[1])),
            (int(predict_bbox[0] + predict_bbox[2]),
             int(predict_bbox[1] + predict_bbox[3])),
            (0, 0, 255),
            3,
        )

        if max(img.shape[:2]) > 640:
            scale = 640 / max(img.shape[:2])
            out_size = (int(img.shape[1] * scale), int(img.shape[0] * scale))
            img = cv2.resize(img, out_size)

        # @SAVE_IMAGE
        # n 来自接收的参数
        # img = img[..., ::-1]
        # img = Image.fromarray(img)
        # image_path = f'visualize/tracking_image/{n}'
        # if not os.path.exists(image_path):
        #     os.makedirs(image_path)
        # img.save(f'visualize/tracking_image/{n}/{i}.jpg')
        # END

        cv2.imshow("TEST", img)
        cv2.waitKey(10)
    cv2.destroyAllWindows()


if __name__ == '__main__':
    for n in range(1, 181):
        sequence_dir = '/data/GOT-10k/test/GOT-10k_Test_{:06d}/'.format(n)
        predict_bboxes = []
        with open(
                "/home/guest/XieBailian/proj/SwinTrack/visualize/got-10k-test/Tiny/GOT-10k_Test_{:06d}/GOT-10k_Test_{:06d}_001.txt"
                .format(n, n)) as fp:
            lines = fp.readlines()
            # print(lines)
            for line in lines:
                predict_bbox = np.fromstring(line, dtype=float, sep=",")
                predict_bboxes.append(predict_bbox)

        show_frames(sequence_dir, '{:08d}.jpg',
                    len(os.listdir(sequence_dir)) - 1, predict_bboxes)