From 4d47ad72c19b70ebc2f238660280ced6a4416f31 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EB=B0=95=EC=84=A0=EA=B7=9C=5FT4082?= Date: Thu, 15 Dec 2022 08:55:06 +0000 Subject: [PATCH] :sparkles:Add polygon to quadrangle converter --- src/polygon_to_quadrangle.ipynb | 569 ++++++++++++++++++++++++++++++++ 1 file changed, 569 insertions(+) create mode 100644 src/polygon_to_quadrangle.ipynb diff --git a/src/polygon_to_quadrangle.ipynb b/src/polygon_to_quadrangle.ipynb new file mode 100644 index 0000000..7b2011e --- /dev/null +++ b/src/polygon_to_quadrangle.ipynb @@ -0,0 +1,569 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import os.path as osp\n", + "import json\n", + "import math\n", + "from tqdm import tqdm\n", + "from glob import glob\n", + "from pprint import pprint\n", + "import copy\n", + "\n", + "import torch\n", + "import matplotlib.pyplot as plt\n", + "from matplotlib.animation import FuncAnimation\n", + "import numpy as np\n", + "import cv2\n", + "import albumentations as A\n", + "import lanms\n", + "from albumentations.pytorch import ToTensorV2\n", + "from imageio import imread\n", + "\n", + "from model import EAST\n", + "from detect import detect" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### polygon -> rectangle(외접하는 직사각형, 회전X)" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "100%|██████████| 1510/1510 [00:00<00:00, 52540.14it/s]\n" + ] + } + ], + "source": [ + "add_data_dir = os.environ.get(\"SM_CHANNEL_TRAIN\", \"../input/data/camper\")\n", + "\n", + "with open(osp.join(add_data_dir, \"ufo/annotation.json\"), \"r\") as f:\n", + " anno = json.load(f)\n", + "\n", + "anno = anno[\"images\"]\n", + "anno_temp = copy.deepcopy(anno)\n", + "\n", + "for img_name, img_info in tqdm(anno.items()):\n", + " # annotation이 없는 이미지 삭제\n", + " if img_info[\"words\"] == {}:\n", + " del anno_temp[img_name]\n", + " continue\n", + " for obj, obj_info in img_info[\"words\"].items():\n", + " if len(img_info[\"words\"][obj][\"points\"]) == 4:\n", + " continue\n", + " # 점이 홀수개인 텍스트박스 삭제\n", + " elif len(img_info[\"words\"][obj][\"points\"]) % 2 == 1:\n", + " del anno_temp[img_name][\"words\"][obj]\n", + "\n", + " if anno_temp[img_name][\"words\"] == {}:\n", + " del anno_temp[img_name]\n", + " continue\n", + "\n", + " # polygon을 x, y의 min, max로 외접하는 직사각형으로 변환\n", + " else:\n", + " x_points = [\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][i][0]\n", + " for i in range(len(anno_temp[img_name][\"words\"][obj][\"points\"]))\n", + " ]\n", + " y_points = [\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][i][1]\n", + " for i in range(len(anno_temp[img_name][\"words\"][obj][\"points\"]))\n", + " ]\n", + " x_min, x_max = min(x_points), max(x_points)\n", + " y_min, y_max = min(y_points), max(y_points)\n", + " anno_temp[img_name][\"words\"][obj][\"points\"] = [\n", + " [x_min, y_min],\n", + " [x_max, y_min],\n", + " [x_min, y_max],\n", + " [x_max, y_max],\n", + " ]\n", + "\n", + " if anno_temp[img_name][\"words\"] == {}:\n", + " del anno_temp[img_name]\n", + " continue\n", + "\n", + "anno = {\"images\": anno_temp}\n", + "\n", + "ufo_dir = osp.join(\"../input/data/camper\", \"ufo\")\n", + "with open(osp.join(ufo_dir, \"train.json\"), \"w\") as f:\n", + " json.dump(anno, f, indent=4)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### polygon 삭제" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "100%|██████████| 1510/1510 [00:00<00:00, 79271.53it/s]\n" + ] + } + ], + "source": [ + "add_data_dir = os.environ.get(\"SM_CHANNEL_TRAIN\", \"../input/data/camper\")\n", + "\n", + "with open(osp.join(add_data_dir, \"ufo/annotation.json\"), \"r\") as f:\n", + " anno = json.load(f)\n", + "\n", + "anno = anno[\"images\"]\n", + "anno_temp = copy.deepcopy(anno)\n", + "\n", + "for img_name, img_info in tqdm(anno.items()):\n", + " # annotation이 없는 이미지 삭제\n", + " if img_info[\"words\"] == {}:\n", + " del anno_temp[img_name]\n", + " continue\n", + " for obj, obj_info in img_info[\"words\"].items():\n", + " if len(img_info[\"words\"][obj][\"points\"]) == 4:\n", + " continue\n", + " # 점이 홀수개인 텍스트박스 삭제\n", + " elif len(img_info[\"words\"][obj][\"points\"]) % 2 == 1:\n", + " del anno_temp[img_name][\"words\"][obj]\n", + "\n", + " if anno_temp[img_name][\"words\"] == {}:\n", + " del anno_temp[img_name]\n", + " continue\n", + "\n", + " # polygon을 삭제\n", + " else:\n", + " del anno_temp[img_name][\"words\"][obj]\n", + " if anno_temp[img_name][\"words\"] == {}:\n", + " del anno_temp[img_name]\n", + " continue\n", + "\n", + "anno = {\"images\": anno_temp}\n", + "\n", + "ufo_dir = osp.join(\"../input/data/camper\", \"ufo\")\n", + "with open(osp.join(ufo_dir, \"train.json\"), \"w\") as f:\n", + " json.dump(anno, f, indent=4)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 양 끝 4개의 점으로 사각형 만들기\n", + "* point들이 annotation 가이드에 맞게 찍혀있다고 가정\n", + "* 모두 가로쓰기로 가정하고 사각형 위, 아래에 동일한 개수의 점이 있다고 가정" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "100%|██████████| 1510/1510 [00:00<00:00, 82433.93it/s]\n" + ] + } + ], + "source": [ + "add_data_dir = os.environ.get(\"SM_CHANNEL_TRAIN\", \"../input/data/camper\")\n", + "\n", + "with open(osp.join(add_data_dir, \"ufo/annotation.json\"), \"r\") as f:\n", + " anno = json.load(f)\n", + "\n", + "anno = anno[\"images\"]\n", + "anno_temp = copy.deepcopy(anno)\n", + "\n", + "for img_name, img_info in tqdm(anno.items()):\n", + " # annotation이 없는 이미지 삭제\n", + " if img_info[\"words\"] == {}:\n", + " del anno_temp[img_name]\n", + " continue\n", + " for obj, obj_info in img_info[\"words\"].items():\n", + " if len(img_info[\"words\"][obj][\"points\"]) == 4:\n", + " continue\n", + " # 점이 홀수개인 텍스트박스 삭제\n", + " elif len(img_info[\"words\"][obj][\"points\"]) % 2 == 1:\n", + " del anno_temp[img_name][\"words\"][obj]\n", + "\n", + " if anno_temp[img_name][\"words\"] == {}:\n", + " del anno_temp[img_name]\n", + " continue\n", + "\n", + " # polygon의 양 끝 4점을 추출하여 사각형 생성\n", + " else:\n", + " anno_temp[img_name][\"words\"][obj][\"points\"] = [\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][0],\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][\n", + " (len(anno_temp[img_name][\"words\"][obj][\"points\"]) // 2 - 1)\n", + " ],\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][\n", + " (len(anno_temp[img_name][\"words\"][obj][\"points\"]) // 2)\n", + " ],\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][-1],\n", + " ]\n", + "\n", + " if anno_temp[img_name][\"words\"] == {}:\n", + " del anno_temp[img_name]\n", + " continue\n", + "\n", + "anno = {\"images\": anno_temp}\n", + "\n", + "ufo_dir = osp.join(\"../input/data/camper\", \"ufo\")\n", + "with open(osp.join(ufo_dir, \"train.json\"), \"w\") as f:\n", + " json.dump(anno, f, indent=4)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 양 끝 4개의 점으로 사각형 만들기, 세로쓰기 구분" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "add_data_dir = os.environ.get(\"SM_CHANNEL_TRAIN\", \"../input/data/camper\")\n", + "\n", + "with open(osp.join(add_data_dir, \"ufo/annotation.json\"), \"r\") as f:\n", + " anno = json.load(f)\n", + "\n", + "anno = anno[\"images\"]\n", + "anno_temp = copy.deepcopy(anno)\n", + "\n", + "for img_name, img_info in tqdm(anno.items()):\n", + " # annotation이 없는 이미지 삭제\n", + " if img_info[\"words\"] == {}:\n", + " del anno_temp[img_name]\n", + " continue\n", + " for obj, obj_info in img_info[\"words\"].items():\n", + " if len(img_info[\"words\"][obj][\"points\"]) == 4:\n", + " continue\n", + " # 점이 홀수개인 텍스트박스 삭제\n", + " elif len(img_info[\"words\"][obj][\"points\"]) % 2 == 1:\n", + " del anno_temp[img_name][\"words\"][obj]\n", + "\n", + " if anno_temp[img_name][\"words\"] == {}:\n", + " del anno_temp[img_name]\n", + " continue\n", + "\n", + " # polygon의 양 끝 4점을 추출하여 사각형 생성\n", + " else:\n", + " if anno_temp[img_name][\"words\"][obj][\"orientation\"] == \"Vertical\":\n", + " anno_temp[img_name][\"words\"][obj][\"points\"] = [\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][0],\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][1],\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][\n", + " (len(anno_temp[img_name][\"words\"][obj][\"points\"]) // 2)\n", + " ],\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][\n", + " (len(anno_temp[img_name][\"words\"][obj][\"points\"]) // 2 + 1)\n", + " ],\n", + " ]\n", + " else:\n", + " anno_temp[img_name][\"words\"][obj][\"points\"] = [\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][0],\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][\n", + " (len(anno_temp[img_name][\"words\"][obj][\"points\"]) // 2 - 1)\n", + " ],\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][\n", + " (len(anno_temp[img_name][\"words\"][obj][\"points\"]) // 2)\n", + " ],\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][-1],\n", + " ]\n", + "\n", + " if anno_temp[img_name][\"words\"] == {}:\n", + " del anno_temp[img_name]\n", + " continue\n", + "\n", + "anno = {\"images\": anno_temp}\n", + "\n", + "ufo_dir = osp.join(\"../input/data/camper\", \"ufo\")\n", + "with open(osp.join(ufo_dir, \"train.json\"), \"w\") as f:\n", + " json.dump(anno, f, indent=4)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 양 끝 4개의 점으로 사각형 -> 사각형 외부에 점이 존재하는 경우 사각형 확대\n", + "* 각 점과 사각형의 거리를 계산하여 사각형에서 가장 많이 벗어난 점을 찾아 변을 평행이동하여 사각형 확대" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "100%|██████████| 1510/1510 [00:00<00:00, 19938.61it/s]\n" + ] + } + ], + "source": [ + "def cal_dist(x1, y1, x2, y2, x0, y0):\n", + " \"\"\"각 point가 새로 생성된 사각형 외부에 위치하는지 확인하기 위한 거리 계산\"\"\"\n", + " area = abs((x1 - x0) * (y2 - y0) - (y1 - y0) * (x2 - x0))\n", + " AB = ((x1 - x2) ** 2 + (y1 - y2) ** 2) ** 0.5\n", + " distance = area / AB\n", + " return distance\n", + "\n", + "\n", + "def get_crosspt(x0, y0, x11, y11, x12, y12, x21, y21, x22, y22):\n", + " \"\"\"사각형의 변을 외부에 위치한 점을 지나게 평행이동하여 사각형의 꼭짓점을 반환\"\"\"\n", + " if x12 == x11 or x22 == x21:\n", + " if x12 == x11:\n", + " cx = x0\n", + " m2 = (y22 - y21) / (x22 - x21)\n", + " cy = m2 * (cx - x21) + y21\n", + " return [cx, cy]\n", + " if x22 == x21:\n", + " cx = x22\n", + " m1 = (y12 - y11) / (x12 - x11)\n", + " cy = m1 * (cx - x0) + y0\n", + " return [cx, cy]\n", + " m1 = (y12 - y11) / (x12 - x11)\n", + " m2 = (y22 - y21) / (x22 - x21)\n", + " if m1 == m2:\n", + " return None\n", + " cx = (x0 * m1 - y0 - x21 * m2 + y21) / (m1 - m2)\n", + " cy = m1 * (cx - x0) + y0\n", + " return [cx, cy]\n", + "\n", + "\n", + "add_data_dir = os.environ.get(\"SM_CHANNEL_TRAIN\", \"../input/data/camper\")\n", + "\n", + "with open(osp.join(add_data_dir, \"ufo/annotation.json\"), \"r\") as f:\n", + " anno = json.load(f)\n", + "\n", + "anno = anno[\"images\"]\n", + "anno_temp = copy.deepcopy(anno)\n", + "\n", + "for img_name, img_info in tqdm(anno.items()):\n", + " # annotation이 없는 이미지 삭제\n", + " if img_info[\"words\"] == {}:\n", + " del anno_temp[img_name]\n", + " continue\n", + " for obj, obj_info in img_info[\"words\"].items():\n", + " if len(img_info[\"words\"][obj][\"points\"]) == 4:\n", + " continue\n", + " # 점이 홀수개인 텍스트박스 삭제\n", + " elif len(img_info[\"words\"][obj][\"points\"]) % 2 == 1:\n", + " del anno_temp[img_name][\"words\"][obj]\n", + "\n", + " if anno_temp[img_name][\"words\"] == {}:\n", + " del anno_temp[img_name]\n", + " continue\n", + "\n", + " # polygon을 텍스트 시작과 끝에서 각각 2개의 점 추출\n", + " else:\n", + " if anno_temp[img_name][\"words\"][obj][\"orientation\"] == \"Vertical\":\n", + " anno_temp[img_name][\"words\"][obj][\"points\"] = [\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][0],\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][1],\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][\n", + " (len(anno_temp[img_name][\"words\"][obj][\"points\"]) // 2)\n", + " ],\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][\n", + " (len(anno_temp[img_name][\"words\"][obj][\"points\"]) // 2 + 1)\n", + " ],\n", + " ]\n", + " else:\n", + " [topleft, topright, bottomright, bottomleft] = [\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][0],\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][\n", + " (len(anno_temp[img_name][\"words\"][obj][\"points\"]) // 2 - 1)\n", + " ],\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][\n", + " (len(anno_temp[img_name][\"words\"][obj][\"points\"]) // 2)\n", + " ],\n", + " anno_temp[img_name][\"words\"][obj][\"points\"][-1],\n", + " ]\n", + "\n", + " left_side = (\n", + " (topleft[0] - bottomleft[0]) ** 2\n", + " + (topleft[1] - bottomleft[1]) ** 2\n", + " ) ** 0.5\n", + " right_side = (\n", + " (topright[0] - bottomright[0]) ** 2\n", + " + (topright[1] - bottomright[1]) ** 2\n", + " ) ** 0.5\n", + "\n", + " idx = None\n", + " max_dist = 0\n", + " out_of_topside, out_of_bottomside = None, None\n", + "\n", + " # 윗변에서 사각형 바깥에 존재하는 점의 존재 여부 확인\n", + " for i in range(len(anno_temp[img_name][\"words\"][obj][\"points\"]) // 2):\n", + " point = anno_temp[img_name][\"words\"][obj][\"points\"][i]\n", + " dist = cal_dist(\n", + " bottomright[0],\n", + " bottomright[1],\n", + " bottomleft[0],\n", + " bottomleft[1],\n", + " point[0],\n", + " point[1],\n", + " )\n", + " if dist > max(left_side, right_side) and dist > max_dist:\n", + " idx = i\n", + " max_dist = dist\n", + "\n", + " # 외부에 점이 존재하면 사각형 확장\n", + " if idx:\n", + " out_of_topside = anno_temp[img_name][\"words\"][obj][\"points\"][idx]\n", + " new_topleft = get_crosspt(\n", + " out_of_topside[0],\n", + " out_of_topside[1],\n", + " topleft[0],\n", + " topleft[1],\n", + " topright[0],\n", + " topright[1],\n", + " topleft[0],\n", + " topleft[1],\n", + " bottomleft[0],\n", + " bottomleft[1],\n", + " )\n", + " new_topright = get_crosspt(\n", + " out_of_topside[0],\n", + " out_of_topside[1],\n", + " topleft[0],\n", + " topleft[1],\n", + " topright[0],\n", + " topright[1],\n", + " topright[0],\n", + " topright[1],\n", + " bottomright[0],\n", + " bottomright[1],\n", + " )\n", + "\n", + " # 아랫변에서 사각형 바깥에 존재하는 점의 존재 여부 확인\n", + " idx = None\n", + " max_dist = 0\n", + " for i in range(\n", + " len(anno_temp[img_name][\"words\"][obj][\"points\"]) // 2,\n", + " len(anno_temp[img_name][\"words\"][obj][\"points\"]),\n", + " ):\n", + " point = anno_temp[img_name][\"words\"][obj][\"points\"][i]\n", + " dist = cal_dist(\n", + " topleft[0],\n", + " topleft[1],\n", + " topright[0],\n", + " topright[1],\n", + " point[0],\n", + " point[1],\n", + " )\n", + " if dist > max(left_side, right_side) and dist > max_dist:\n", + " idx = i\n", + " max_dist = dist\n", + "\n", + " # 외부에 점이 존재하면 사각형 확장\n", + " if idx:\n", + " out_of_bottomside = anno_temp[img_name][\"words\"][obj][\"points\"][idx]\n", + " new_bottomright = get_crosspt(\n", + " out_of_bottomside[0],\n", + " out_of_bottomside[1],\n", + " bottomright[0],\n", + " bottomright[1],\n", + " bottomleft[0],\n", + " bottomleft[1],\n", + " topright[0],\n", + " topright[1],\n", + " bottomright[0],\n", + " bottomright[1],\n", + " )\n", + " new_bottomleft = get_crosspt(\n", + " out_of_bottomside[0],\n", + " out_of_bottomside[1],\n", + " bottomright[0],\n", + " bottomright[1],\n", + " bottomleft[0],\n", + " bottomleft[1],\n", + " topleft[0],\n", + " topleft[1],\n", + " bottomleft[0],\n", + " bottomleft[1],\n", + " )\n", + "\n", + " if out_of_topside:\n", + " topleft = new_topleft\n", + " topright = new_topright\n", + " if out_of_bottomside:\n", + " bottomright = new_bottomright\n", + " bottomleft = new_bottomleft\n", + "\n", + " anno_temp[img_name][\"words\"][obj][\"points\"] = [\n", + " topleft,\n", + " topright,\n", + " bottomright,\n", + " bottomleft,\n", + " ]\n", + "\n", + "anno = {\"images\": anno_temp}\n", + "\n", + "ufo_dir = osp.join(\"../input/data/camper\", \"ufo\")\n", + "with open(osp.join(ufo_dir, \"train.json\"), \"w\") as f:\n", + " json.dump(anno, f, indent=4)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "base", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.8.5 (default, Sep 4 2020, 07:30:14) \n[GCC 7.3.0]" + }, + "orig_nbformat": 4, + "vscode": { + "interpreter": { + "hash": "d4d1e4263499bec80672ea0156c357c1ee493ec2b1c70f0acce89fc37c4a6abe" + } + } + }, + "nbformat": 4, + "nbformat_minor": 2 +}