728x90
반응형
import os
import numpy as np
from IPython.display import Image
def get_files(path):
files = os.listdir(path)
img_path =[]
gt_path = []
for file in files:
if file.split(".")[-1] in ["jpg"]:
img_path.append(path+"/"+file)
gt_path.append(path+"/"+file.split(".")[0]+".txt")
return img_path , gt_path
def is_on_same_line(box_a, box_b, min_y_overlap_ratio=0.8):
"""Check if two boxes are on the same line by their y-axis coordinates.
Two boxes are on the same line if they overlap vertically, and the length
of the overlapping line segment is greater than min_y_overlap_ratio * the
height of either of the boxes.
Args:
box_a (list), box_b (list): Two bounding boxes to be checked
min_y_overlap_ratio (float): The minimum vertical overlapping ratio
allowed for boxes in the same line
Returns:
The bool flag indicating if they are on the same line
"""
a_y_min = np.min(box_a[1::2])
b_y_min = np.min(box_b[1::2])
a_y_max = np.max(box_a[1::2])
b_y_max = np.max(box_b[1::2])
# Make sure that box a is always the box above another
if a_y_min > b_y_min:
a_y_min, b_y_min = b_y_min, a_y_min
a_y_max, b_y_max = b_y_max, a_y_max
if b_y_min <= a_y_max:
if min_y_overlap_ratio is not None:
sorted_y = sorted([b_y_min, b_y_max, a_y_max])
overlap = sorted_y[1] - sorted_y[0]
min_a_overlap = (a_y_max - a_y_min) * min_y_overlap_ratio
min_b_overlap = (b_y_max - b_y_min) * min_y_overlap_ratio
return overlap >= min_a_overlap or \
overlap >= min_b_overlap
else:
return True
return False
def main():
img_path , gt_path = get_files("./sort_sample")
#for gt in gt_path:
file = open(gt_path[2], 'r', encoding="utf-8")
lines = file.readlines()
sort_box_mid = []
for line in lines:
polygon = line.split("\t")[:8]
label = line.split("\t")[8]
ner = line.split("\t")[-1].rstrip("\n")
x1, y1, x3, y3 = int(polygon[0]), int(polygon[1]), int(polygon[4]), int(polygon[5])
xm, ym = ((x3+x1)/2), ((y3+y1)/2)
sort_box_mid.append([x1, y1, x3, y3, xm, ym, label, ner])
x_sorted_boxes = sorted(sort_box_mid, key=lambda x: np.min(x[4]))
x_sorted_boxes
dict_val = {}
for index, val in enumerate(x_sorted_boxes):
dict_val[index] = val
print(dict_val)
skip_idxs = set()
min_y_overlap_ratio = 0.8
lines = {}
lines_label = []
i = 0
# locate lines of boxes starting from the leftmost one
# x_sorted_boxes : x좌표 (좌 -> 우)
for i in range(len(x_sorted_boxes)):#한 단어 씩 검토
if i in skip_idxs:
continue
# the rightmost box in the current line
rightmost_box_idx = i
line = [rightmost_box_idx]
line_key = rightmost_box_idx
label = [x_sorted_boxes[rightmost_box_idx][6]]
for j in range(i + 1, len(x_sorted_boxes)):
if j in skip_idxs:
continue
if is_on_same_line(x_sorted_boxes[rightmost_box_idx][0:4], x_sorted_boxes[j][0:4], min_y_overlap_ratio):
line.append(j)
skip_idxs.add(j)
rightmost_box_idx = j
label.append(x_sorted_boxes[j][6])
lines_label.append(label)
# ########TEST
lines[line_key]=line
# 라인단위 첫 글자들만 뽑아서 y위치에 맞춰서 sorting 진행
first = []
for key_idx in lines.keys():
first.append((key_idx, x_sorted_boxes[key_idx]))# 각 라인의 제일 왼쪽 글자의 my 값
y_sorted_boxes = sorted(first, key=lambda first: first[1][5])
y_sorted_boxes = [y[0] for y in y_sorted_boxes]
new = []
for y in y_sorted_boxes:
vlist = lines.get(y)
new.append(vlist)
sorted_result = []
for line in new:
for val in line:
sorted_result.append(dict_val[val])
print(' '.join(x[6] for x in sorted_result))
728x90
반응형
'머신러닝 > Computer Vision' 카테고리의 다른 글
[python/cv2] RGB colour table/ rgb 컬러 triplet (0) | 2022.10.31 |
---|---|
[python] matplotlib 에서 한글 사용하기(파이썬 그래프 그리기) / python graph visualization / networkx 사용법 (0) | 2022.07.11 |
[OpenCV] 이미지 내 특정 컬러 변경(글자 색 바꾸기) (0) | 2022.04.12 |
[OpenCV] 이미지 내 글자 탐지, 글자 삭제(배경 컬러로 합성) (0) | 2022.04.12 |
[cv2/PIL] png 이미지 jpg로 변경 (0) | 2022.03.30 |