半自动标注(使用自己的分割或检测模型推理完得到的矩阵再生成json文件)

前言
都知道标注很麻烦、很累,不然先训练一批,然后推理得到它的掩码图,先生成自动标注,再人工手动修改也许会快很多

半自动标注代码

这是我自己写的,是labelme的格式,大家想要修改成自己的json格式可以修改json_dict_init与dict_init函数

默认大家已经得到自己图片经过模型推理之后的掩码图了,掩码图上面生成的像素应该是0,1,2,3. 分别对应自己训练的时候的标签

import cv2
import sys
import base64
import cv2
import json
import time
import os
from tqdm import tqdm
from PIL import Image
import numpy as np

def dict_init(label_name):
    temp_dict = {
        "label": label_name,
        "line_color": None,
        "fill_color": None,
        "points":[],
        "shape_type": "polygon",
        "flags": {}
    }

    return  temp_dict

def json_dict_init(num,height,width,img):
    temp_dict = {
        "version": "3.16.2",
        "flags": {},
        "shapes": [],
        "lineColor": [
            0,
            255,
            0,
            128
        ],
        "fillColor": [
            255,
            0,
            0,
            128
        ],
        "imagePath": "..\\img\\{}.jpg".format(num),
        #原图像数据通过b64编码生成的字符串数据,可以再次解码成图片
        "imageData":img,
        "imageHeight": height,
        "imageWidth": width
    }

    return  temp_dict

def image_to_base64(image_path):
    # 读取二进制图片,获得原始字节码
    with open(image_path, 'rb') as jpg_file:
        byte_content = jpg_file.read()

    # 把原始字节码编码成base64字节码
    base64_bytes = base64.b64encode(byte_content)

    # 把base64字节码解码成utf-8格式的字符串
    base64_string = base64_bytes.decode('utf-8')

    return base64_string

# 定义一个函数,用于判断两个点是否相邻
def is_adjacent(pt1, pt2):
    x1, y1 = pt1
    x2, y2 = pt2
    return abs(x1 - x2) <= 8 and abs(y1 - y2) <= 8

def main():
    img_path = "D:/data_val/new/temp/img"
    label_path = "D:/data_val/new/temp/label"
    json_save_path = 'D:/data_val/new/temp/json'
    os.makedirs(json_save_path,exist_ok=True)

    # 放入你的标签名 例如 你的掩码图 像素是1 1对应的就是phone 像素是2 2就是line
    label_name_list = ['phone','line']

    img_list = os.listdir(img_path)
    pbar = tqdm(total=len(img_list))

    for filename in img_list:
        name = filename.split(".")[0]
        # print("*"*10)
        # print(f'{filename}')

        #label图
        label = Image.open(f'{label_path}/{name}.png')

        label = np.array(label)
        #输出自己的label有多少种像素
        # unique_values = np.unique(label)
        # print('')
        # print(unique_values)
        # print('')

        height  = label.shape[0]
        width   = label.shape[1]
        base64_string = image_to_base64(f'{img_path}/{filename}')
        json_dict = json_dict_init(name,height,width,base64_string)

        contours, hierarchy = cv2.findContours(label, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

        background_flag =False
        if (len(contours) == 0):
            print("如果未识别到 则变成背景")
            shapes_dict = dict_init("_background_")
            left_top, left_botton = [10, 10], [10, height - 10]
            right_botton, right_top = [width - 10, height - 10], [width - 10, 10]
            shapes_dict['points'] = [left_top, left_botton, right_botton, right_top]
            json_dict["shapes"].append(shapes_dict)

        else:
           for i,clasee_name in enumerate(label_name_list):
               temp = label.copy()
               temp[temp!=(i+1)] = 0
               contours, hierarchy = cv2.findContours(temp, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
               if (len(contours) == 0):
                   continue
               else:
                    count = 0
                    for contour in contours:
                        shapes_dict = dict_init(clasee_name)
                        count = count + 1
                        #10 20
                        
                        #语义分割
                        approxCourve = cv2.approxPolyDP(contour, 3, True)
                        for Courve in approxCourve:
                            point = Courve[0]
                            x, y = point
                            x = int(x)
                            y = int(y)
                            temp_list = []
                            temp_list.append(x)
                            temp_list.append(y)
                            shapes_dict["points"].append(temp_list)
                            
                        #目标检测 这里只写到生成检测框的四个点 根据自己需求填进去
                        # rect = cv2.minAreaRect(contour)
                        # box = cv2.boxPoints(rect)
                        # # print(box)
                        # # 轮廓必须是整数, 不能是小数, 所以转化为整数
                        # box = np.round(box).astype('int64')
                        # left_point_x = np.min(box[:, 0])
                        # right_point_x = np.max(box[:, 0])
                        # top_point_y = np.min(box[:, 1])
                        # bottom_point_y = np.max(box[:, 1])
                        
                        
                        json_dict["shapes"].append(shapes_dict)

        with open(f'{json_save_path}/{name}.json', "w", encoding='utf-8') as f:
            f.write(json.dumps(json_dict, ensure_ascii=False))
        pbar.update(1)


if __name__ == "__main__":
    main()

欢迎大家点赞或收藏哦~
你们的点赞或收藏可以鼓励作者加快更新哟~

  • 2
    点赞
  • 2
    收藏
    觉得还不错? 一键收藏
  • 7
    评论
评论 7
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值