main.py 2.38 KB
# -*- coding: utf-8 -*-
# @Author        : Lyu Kui
# @Email         : 9428.al@gmail.com
# @Create Date   : 2022-04-23 13:15:18
# @Last Modified : 2022-04-23 13:59:50
# @Description   : 

import os
import csv
import json

import os
import cv2
import json
import numpy as np

def save(save_file, shapes, imagePath, imageHeight, imageWidth, imageData=None, lineColor=None, fillColor=None, flags=None):

    if flags is None:
        flags = {}
    data = dict(
        version="3.16.2",
        flags=flags,
        shapes=shapes,
        lineColor=lineColor,
        fillColor=fillColor,
        imagePath=imagePath,
        imageData=imageData,
        imageHeight=imageHeight,
        imageWidth=imageWidth,
    )

    json.dump(data, save_file, ensure_ascii=False, indent=2)

    return

def writeJson(imagePath, im_size, lines, save_file):
    def format_shape(s):
        return dict(
            label='1',
            line_color=None,
            fill_color=None,
            points=[(float(p[0]), float(p[1])) for p in s],
            shape_type="polygon",
            flags={}
        )
    shapes = [format_shape(shape) for shape in lines]

    save(save_file, shapes=shapes, imagePath=imagePath, imageData=None, imageHeight=im_size[0], imageWidth=im_size[1], lineColor=[0,255,0,128], fillColor=[255,0,0,128], flags=None)
    
    return

f = open('/home/lk/MyProject/文字检测相关/inter_ocr/data/Xeon1OCR_round1_train_20210524.csv', 'r')
lines = csv.reader(f)
next(lines)

for line in lines:
    # index, label = line.split(',', 1)
    i, j, k = line

    tfspath = eval(j)['tfspath']
    fn = os.path.basename(tfspath)

    m, n = eval(k)

    shapes = []
    for text in m:
        label = eval(text['text'])
        lable = label['text']
        points = np.array(text['coord'], dtype=np.float32).reshape((-1, 2)).tolist()

        shape = {
            "label": lable,
            "points": points,
            "group_id": None,
            "shape_type": "polygon",
            "flags": {}
        }
        shapes.append(shape)

    labelme = {
      "version": "4.4.0",
      "flags": {},
      "shapes": shapes,
      "imagePath": f"../image/{fn}",
      "imageData": None,
      "imageHeight": None,
      "imageWidth": None
    }

    save_file = f"./json/{fn.replace('.jpg', '.json')}"
    labelme = json.dumps(labelme, ensure_ascii=False, indent=2)

    with open(save_file, 'w') as f:
        f.write(labelme)