基于labelme创建自己的仿cityscapes的训练数据集（视频转图片训练集）

最新推荐文章于 2024-06-21 23:55:20 发布

理想国の糕

最新推荐文章于 2024-06-21 23:55:20 发布

阅读量1.4k

点赞数 1

本文链接：https://blog.csdn.net/qq_45751990/article/details/127233401

版权

课程设计记录专栏收录该内容

4 篇文章 1 订阅

订阅专栏

课设要求基于原有的模型做一个增量训练，要求我们使用自己创建的校园场景数据集……比较有趣的一个作业~理直气壮地出去玩了半天😄

第一步：去拍校园的道路场景_{（不提倡单手骑车拍视频哈）}
第二步：处理视频数据，分割图片，筛选出比较好的_{（没拍糊的）}
第三步：人为标注数据集（这里使用的是labelme）
第四步：对生成的json数据处理，生成ids类型的灰度图片
第五步：增量训练

1、安装labelme

激活自己的虚拟空间后，在命令行安装即可，推荐python版本3.6、3.7：

pip install labelme

2、labelme的使用

安装完成后在命令行输入labelme即可打开labelme的图形化界面
接着主要是标注数据
在这里插入图片描述

3、json数据的转化

labelme标注的数据是json格式，我们需要将这个数据转化成labelid的类别图进行训练。
转化脚本如下：

import os
import os.path as osp
import io
import math
import base64
import json
import uuid

import numpy as np
import imgviz
import PIL.Image
import cv2


def img_data_to_pil(img_data):
    f = io.BytesIO()
    f.write(img_data)
    img_pil = PIL.Image.open(f)
    return img_pil


def img_data_to_arr(img_data):
    img_pil = img_data_to_pil(img_data)
    img_arr = np.array(img_pil)
    return img_arr


def img_b64_to_arr(img_b64):
    img_data = base64.b64decode(img_b64)
    img_arr = img_data_to_arr(img_data)
    return img_arr


def shape_to_mask(img_shape, points, shape_type=None, line_width=10, point_size=5):
    mask = np.zeros(img_shape[:2], dtype=np.uint8)
    mask = PIL.Image.fromarray(mask)
    draw = PIL.ImageDraw.Draw(mask)
    xy = [tuple(point) for point in points]
    if shape_type == "circle":
        assert len(xy) == 2, "Shape of shape_type=circle must have 2 points"
        (cx, cy), (px, py) = xy
        d = math.sqrt((cx - px) ** 2 + (cy - py) ** 2)
        draw.ellipse([cx - d, cy - d, cx + d, cy + d], outline=1, fill=1)
    elif shape_type == "rectangle":
        assert len(xy) == 2, "Shape of shape_type=rectangle must have 2 points"
        draw.rectangle(xy, outline=1, fill=1)
    elif shape_type == "line":
        assert len(xy) == 2, "Shape of shape_type=line must have 2 points"
        draw.line(xy=xy, fill=1, width=line_width)
    elif shape_type == "linestrip":
        draw.line(xy=xy, fill=1, width=line_width)
    elif shape_type == "point":
        assert len(xy) == 1, "Shape of shape_type=point must have 1 points"
        cx, cy = xy[0]
        r = point_size
        draw.ellipse([cx - r, cy - r, cx + r, cy + r], outline=1, fill=1)
    else:
        assert len(xy) > 2, "Polygon must have points more than 2"
        draw.polygon(xy=xy, outline=1, fill=1)
    mask = np.array(mask, dtype=bool)
    return mask


def shapes_to_label(img_shape, shapes, label_name_to_id):
    cls = np.zeros(img_shape[:2], dtype=np.int32)
    ins = np.zeros_like(cls)
    instances = []
    for shape in shapes:
        points = shape["points"]
        label = shape["label"]
        group_id = shape.get("group_id")
        if group_id is None:
            group_id = uuid.uuid1()
        shape_type = shape.get("shape_type", None)

        cls_name = label
        instance = (cls_name, group_id)

        if instance not in instances:
            instances.append(instance)
        ins_id = instances.index(instance) + 1
        cls_id = label_name_to_id[cls_name]

        mask = shape_to_mask(img_shape[:2], points, shape_type)
        cls[mask] = cls_id
        ins[mask] = ins_id

    return cls, ins


def lblsave(filename, lbl, palette=True):
    if osp.splitext(filename)[1] != ".png":
        filename += ".png"
    # Assume label ranses [-1, 254] for int32, and [0, 255] for uint8 as VOC.
    if lbl.min() >= -1 and lbl.max() < 255:
        lbl_pil = PIL.Image.fromarray(lbl.astype(np.uint8), mode="P")
        if palette:
            colormap = imgviz.label_colormap()
            lbl_pil.putpalette(colormap.flatten())
        lbl_pil.save(filename)
    else:
        raise ValueError("[%s] Cannot save the pixel-wise class label as PNG. Please consider using the .npy format." % filename)


def json_to_dataset(json_file, out_dir=None, demo=False):
    state = True
    data = json.load(open(json_file))
    imageData = data.get("imageData")
    if not imageData:
        imagePath = os.path.join(os.path.dirname(json_file), data["imagePath"])
        with open(imagePath, "rb") as f:
            imageData = f.read()
            imageData = base64.b64encode(imageData).decode("utf-8")
    img = img_b64_to_arr(imageData)
    for shape in sorted(data["shapes"], key=lambda x: x["label"]):
        label_name = shape["label"]
        if label_name in label_name_to_id:
            label_value = label_name_to_id[label_name]
        else:
            print("标签名错误！{}".format(label_name))
            state = False
            label_value = len(label_name_to_id)
            label_name_to_id[label_name] = label_value
    try:
        lbl, _ = shapes_to_label(img.shape, data["shapes"], label_name_to_id)
    except Exception as e:
        print(e)
        return False, None, None
    if demo:
        if out_dir is None:
            out_dir = osp.basename(json_file).replace(".", "_")
            out_dir = osp.join(osp.dirname(json_file), out_dir)
        if not osp.exists(out_dir):
            os.mkdir(out_dir)
        print("Saved to: {}".format(out_dir))
        PIL.Image.fromarray(img).save(osp.join(out_dir, "img.png"))
        lblsave(osp.join(out_dir, "label.png"), lbl, palette=False)
        lblsave(osp.join(out_dir, "label_color.png"), lbl)
        label_names = [None] * (max(label_name_to_id.values()) + 1)
        for name, value in label_name_to_id.items():
            label_names[value] = name
        print(label_names)    
        lbl_viz = imgviz.label2rgb(lbl, imgviz.asgray(img), label_names=label_names, loc="rb")
        PIL.Image.fromarray(lbl_viz).save(osp.join(out_dir, "label_viz.png"))
        with open(osp.join(out_dir, "label_names.txt"), "w") as f:
            for lbl_name in label_names:
                f.write(lbl_name + "\n")
        return state, None, None
    else:
        # lbl_pil = PIL.Image.fromarray(lbl.astype(np.uint8), mode="P")
        # img_pil = PIL.Image.fromarray(img)
        img_cv = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        lbl_cv = lbl.astype(np.uint8)
        return state, img_cv, lbl_cv


label_id_to_name = {
    0: '_background_',
    1: "ego vehicle",
    2:"rectification border",
    3:"out of roi",
    4: "static",
    5: "dynamic",
    6: "ground",
    7: "road",
    8: "sidewalk",
    9:"parking",

    10:"rail track",
    11:"building",
    12:"wall",
    13:"fence",
    14:"guard rail",
    15:"bridge",
    16:"tunnel",
    17:"pole",
    18:"polegroup",
    19:"traffic light",

    20:"traffic sign",
    21:"vegetation",
    22:"terrain",
    23:"sky",
    24:"person",
    25:"rider",
    26:"car",
    27:"truck",
    28:"bus",
    29:"caravan",
    
    30:"trailer",
    31:"train",
    32:"motorcycle",
    33:"bicycle",
    
}
label_name_to_id = dict([val,key] for key,val in label_id_to_name.items())


if __name__ == "__main__":
    # 输出详细的label信息到指定目录
    #json_to_dataset("./images/ghy1/00000001.json", "./images/ghy1/out", demo=True)
    # 返回图片和对应的标签map，转化是否成功
    #state, img_cv, lbl_cv = json_to_dataset("./images/ghy1/00000001.json")
    #print(state)
    #cv2.imwrite("img.png",img_cv)
    #cv2.imwrite("lbl.png",lbl_cv)
    #img_cv = cv2.imread("img.png")
    #lbl_cv = cv2.imread("lbl.png", 0)
    #批量操作
    json_dir=".\images\ghy1"# 自己放json的目录
    out_path="./images/ghy1/out"#无所谓
    list_path = os.listdir(json_dir)
    print('freedom =', json_dir)
    for i in range(0, len(list_path)):
        if "json" not in list_path[i]:
            continue
        path = os.path.join(json_dir, list_path[i])
        (filename,_)=os.path.splitext(list_path[i])
        # print(filename)
        json_to_dataset(path, out_path, demo=True)
        state, img_cv, lbl_cv = json_to_dataset(path)
        print("state=",state)
        cv2.imwrite(filename+"img.png",img_cv)# 生成的原始图片png格式
        cv2.imwrite(filename+"lbl.png",lbl_cv)# 生成的labelid图片路径
        # img_cv = cv2.imread("img.png")
        # lbl_cv = cv2.imread("lbl.png", 0)
    pass