【图像-关键点json标签】的数据集生成器

最新推荐文章于 2026-01-06 11:13:57 发布

原创

最新推荐文章于 2026-01-06 11:13:57 发布 · 1.5k 阅读

9 ·

CC 4.0 BY-SA版权

文章标签：

#计算机视觉 #json #opencv

该博客介绍了如何利用OpenCV处理库和PyTorch的数据加载器来构建一个人体关键点数据集。数据集包括训练和验证集，标签以JSON格式存储，包含图像路径和人体关键点坐标。代码中定义了一个KeyPointsPose类，用于读取和预处理图像及其对应的坐标信息，并实现了多种图像变换，如随机裁剪、翻转、旋转和色彩扭曲。此外，还提供了用于显示图像和关键点坐标的辅助函数。

利用opencv图像处理库制作【图像-关键点json标签】的数据集生成器

背景：关于读取json文件的方法

1. 数据

数据目录：
– /keypoints_pose/

–/train_img/
–/train_label_pose/
–/val_img/
–/val_label_pose/

如图所示：
在这里插入图片描述

2. 数据集结果展示：（图像，坐标）

在这里插入图片描述

3. 代码实现

`dataset_json_by_cv2.py`

# !/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2022/5/24 15:32
# @Author  : Hyan Tan 
# @File    : dataset_json_by_cv2.py.py
import os
from abc import abstractmethod, ABCMeta
import numpy as np
import pandas as pd
import json

import cv2
from torch.utils.data import Dataset, DataLoader

from torchvision import transforms
from transformation import RandomAdd, ImageResize, RandomCrop, RandomFlip, RandomRotate, Distort


# ---人像-人体关键点数据集生成器---
class KeyPointsPose(Dataset):
    def __init__(self, root_dir=r'E:/Datasets/keypoints_pose/', image_set='train', transforms=None):
        """
        初始化数据集
        :param root_dir: 数据目录(.csv和images的根目录)
        :param image_set: train训练,val验证,test测试
        :param transforms（callable,optional）:图像变换-可选
        标签数据文件格式为json_file: 标签json文件(内容：图像相对地址-category类型-标签coordination坐标)
        """
        super(KeyPointsPose, self).__init__()
        self._imgset = image_set
        self._image_paths = []
        self._labels = []
        self._root_dir = root_dir
        self._transform = transforms

        self.__getFileList()  # 图地址列表和标签列表

    def __getFileList(self):
        img_path = os.path.join(self._root_dir, self._imgset + '_img')  # 图像目录
        label_path = os.path.join(self._root_dir, self._imgset + '_label_pose')  # 标签目录
        self._image_paths = os.listdir(img_path)
        for img_ in self._image_paths:
            pose_name = img_.replace('.png', '_keypoints.json').replace('.jpg', '_keypoints.json')  # 标签名和图像对应，后缀不同。
            with open(os.path.join(label_path, pose_name), 'r') as f:
                pose_label = json.load(f)
                try:
                    pose_data = pose_label['people'][0]['pose_keypoints']
                except IndexError:
                    pose_data = [0 for i in range(54)]  # 数据读取失败时，创建18组0,0,0
                pose_data = np.array(pose_data)
                pose_data = pose_data.reshape((-1, 3))
                self._labels.append(pose_data)

    def __getitem__(self, idx):
        label