手动切分图片和标注以及转换 COCO 格式到 DOTA 格式

image_patch_split.py

import cv2
from pycocotools.coco import COCO
from tqdm import tqdm
import numpy as np
import os
import os.path as osp
import shapely
from copy import deepcopy
import json
import shutil

PATCH_SIZE = 1280
SLIDING = PATCH_SIZE // 2
image_folder = "/media/ps/extra/Dataset/sh1k/big_img"
output_folder = "/media/ps/extra/Dataset/sh1k/chunibyo_img"
coco = COCO("/media/ps/extra/Dataset/sh1k/trainval_sh_big.json")

rgb_color_list = [
    [0, 0, 255],
    [0, 255, 0],
    [255, 0, 0],
    [0, 255, 255],
]


def find_sector(poly: np.ndarray):
    area, index = 0, 0
    for i in range(len(poly)):
        a, b, c = poly[i - 1], poly[i], poly[(i + 1) % len(poly)]
        area_ = np.linalg.norm(np.cross(b - a, c - b))
        if area_ > area:
            area, index = area_, i

    A, O, B = poly[index - 1], poly[index], poly[(index + 1) % len(poly)]
    return A, O, B


def process(_type, ids):
    new_images = []
    new_annotations = []

    dota_ann_folder = osp.join(output_folder, "..", f"chunibyo_dota_{_type}")
    if osp.exists(dota_ann_folder):
        shutil.rmtree(dota_ann_folder)
    os.makedirs(dota_ann_folder)

    for imgId in tqdm(ids, desc=f"Processing to {_type} data patches"):
        imgInfo = coco.loadImgs(imgId)[0]
        height, width = imgInfo["height"], imgInfo["width"]
        file_name = imgInfo["file_name"]
        img = cv2.imread(osp.join(image_folder, file_name))

        # split images to patches
        top, left, i, j = 0, 0, 0, 0
        while top + PATCH_SIZE < height:
            while left + PATCH_SIZE < width:
                cropped_image = deepcopy(img[top : top + PATCH_SIZE, left : left + PATCH_SIZE])
                roi = shapely.Polygon(
                    [
                        (left, top),
                        (left + PATCH_SIZE, top),
                        (left + PATCH_SIZE, top + PATCH_SIZE),
                        (left, top + PATCH_SIZE),
                    ]
                )
                anns_group_by_image = []
                for ann in coco.imgToAnns[imgId]:
                    if ann["category_id"] >= 4:
                        continue

                    new_ann = deepcopy(ann)
                    tmp = shapely.Polygon(np.array(ann["segmentation"][0]).reshape(-1, 2)).convex_hull
                    intersection = roi.intersection(tmp)
                    if intersection.area < 10:
                        continue
                    if ann["category_id"] == 3 and intersection.area != tmp.area:
                        continue

                    # 转真值为 4 个点
                    pts = np.array(intersection.exterior.coords.xy).T[:-1]
                    if ann["category_id"] == 3:
                        A, O, B = find_sector(pts)
                        r = (np.linalg.norm(B - O) + np.linalg.norm(A - O)) / 2
                        OA, OB = (A - O) / np.linalg.norm(A - O), (B - O) / np.linalg.norm(B - O)
                        dir = (OA + OB) / np.linalg.norm(OA + OB)
                        pts = np.array([O, O + OA * r, O + dir * r, O + OB * r])
                    else:
                        rect = intersection.minimum_rotated_rectangle
                        if intersection.area / rect.area < 0.8:
                            continue
                        pts = np.array(rect.exterior.coords.xy).T[:-1]
                    # 转真值为 4 个点

                    new_ann["segmentation"][0] = (pts - np.array([[left, top]])).flatten().tolist()
                    bbox = intersection.bounds
                    # cx, cy, w, h
                    new_ann["bbox"] = [
                        bbox[0] + SLIDING - left,
                        bbox[1] + SLIDING - top,
                        bbox[2] - bbox[0],
                        bbox[3] - bbox[1],
                    ]
                    new_ann["area"] = intersection.area
                    new_ann["id"] = len(new_annotations) + 1
                    new_ann["image_id"] = len(new_images) + 1
                    anns_group_by_image.append(new_ann)
                    new_annotations.append(new_ann)

                debug_image = deepcopy(cropped_image)
                for new_ann in anns_group_by_image:
                    cv2.polylines(
                        debug_image,
                        [np.array(new_ann["segmentation"][0]).reshape(-1, 2).astype(np.int32)],
                        True,
                        (0, 0, 255),
                        2,
                    )
                    cv2.polylines(
                        debug_image,
                        [np.array(new_ann["segmentation"][0], np.int32).reshape((-1, 1, 2))],
                        True,
                        rgb_color_list[new_ann["category_id"]],
                        thickness=10,
                        lineType=cv2.LINE_AA,
                    )

                    for pts in np.array(new_ann["segmentation"][0], np.int32).reshape(-1, 2):
                        cv2.circle(
                            debug_image,
                            pts,
                            10,
                            [255, 0, 255],
                            -1,
                        )

                basename = f"{file_name.split('.')[0]}_{PATCH_SIZE}_{i}_{j}"
                cv2.imwrite(
                    osp.join(output_folder, f"{basename}.png"),
                    cropped_image,
                )
                cv2.imwrite(
                    osp.join(output_folder, f"{basename}_debug.png"),
                    debug_image,
                )
                new_image = {
                    "height": PATCH_SIZE,
                    "width": PATCH_SIZE,
                    "id": len(new_images) + 1,
                    "file_name": f"{basename}.png",
                }
                new_images.append(new_image)

                with open(osp.join(dota_ann_folder, f"{basename}.txt"), "w") as f:
                    for ann in anns_group_by_image:
                        f.write(
                            f"{' '.join([str(i) for i in np.array(ann['segmentation'][0]).astype(np.int32)])} {coco.loadCats(ann['category_id'])[0]['name']} 0\n"
                        )

                j += 1
                left = SLIDING * j

            i += 1
            top = SLIDING * i

    result = {
        "categories": [coco.cats[i] for i in coco.cats],
        "images": new_images,
        "annotations": new_annotations,
    }
    with open(osp.join(image_folder, "..", f"chunibyo_coco_{_type}.json"), "w") as f:
        f.write(json.dumps(result))


def main():
    ids = sorted(list(coco.imgToAnns.keys()))

    ratio = 0.8

    process("train", ids[: int(len(ids) * ratio)])
    process("val", ids[int(len(ids) * ratio) :])


if __name__ == "__main__":
    main()