Initial commit

UIED/detect_compo/deprecated/ocr_classify_text.py

@@ -0,0 +1,113 @@
+import pytesseract as pyt
+import cv2
+
+import lib_ip.ip_draw as draw
+from config.CONFIG_UIED import Config
+
+C = Config()
+
+
+def is_text(img, min_word_area, show=False):
+    broad = img.copy()
+    area_word = 0
+    area_total = img.shape[0] * img.shape[1]
+
+    try:
+        # ocr text detection
+        data = pyt.image_to_data(img).split('\n')
+    except:
+        print(img.shape)
+        return -1
+    word = []
+    for d in data[1:]:
+        d = d.split()
+        if d[-1] != '-1':
+            if d[-1] != '-' and d[-1] != '—' and int(d[-3]) < 50 and int(d[-4]) < 100:
+                word.append(d)
+                t_l = (int(d[-6]), int(d[-5]))
+                b_r = (int(d[-6]) + int(d[-4]), int(d[-5]) + int(d[-3]))
+                area_word += int(d[-4]) * int(d[-3])
+                cv2.rectangle(broad, t_l, b_r, (0,0,255), 1)
+
+    if show:
+        for d in word: print(d)
+        print(area_word/area_total)
+        cv2.imshow('a', broad)
+        cv2.waitKey(0)
+        cv2.destroyAllWindows()
+    # no text in this clip or relatively small text area
+    if len(word) == 0 or area_word/area_total < min_word_area:
+        return False
+    return True
+
+
+def text_detection(org, img_clean):
+    try:
+        data = pyt.image_to_data(img_clean).split('\n')
+    except:
+        return org, None
+    corners_word = []
+    for d in data[1:]:
+        d = d.split()
+        if d[-1] != '-1':
+            if d[-1] != '-' and d[-1] != '—' and 5 < int(d[-3]) < 40 and 5 < int(d[-4]) < 100:
+                t_l = (int(d[-6]), int(d[-5]))
+                b_r = (int(d[-6]) + int(d[-4]), int(d[-5]) + int(d[-3]))
+                corners_word.append((t_l, b_r))
+    return corners_word
+
+
+# def text_merge_word_into_line(org, corners_word, max_words_gap=C.THRESHOLD_TEXT_MAX_WORD_GAP):
+#
+#     def is_in_line(word):
+#         for i in range(len(lines)):
+#             line = lines[i]
+#             # at the same row
+#             if abs(line['center'][1] - word['center'][1]) < max_words_gap:
+#                 # small gap between words
+#                 if (abs(line['center'][0] - word['center'][0]) - abs(line['width']/2 + word['width']/2)) < max_words_gap:
+#                     return i
+#         return -1
+#
+#     def merge_line(word, index):
+#         line = lines[index]
+#         # on the left
+#         if word['center'][0] < line['center'][0]:
+#             line['col_min'] = word['col_min']
+#         # on the right
+#         else:
+#             line['col_max'] = word['col_max']
+#         line['row_min'] = min(line['row_min'], word['row_min'])
+#         line['row_max'] = max(line['row_max'], word['row_max'])
+#         line['width'] = line['col_max'] - line['col_min']
+#         line['height'] = line['row_max'] - line['row_min']
+#         line['center'] = ((line['col_max'] + line['col_min'])/2, (line['row_max'] + line['row_min'])/2)
+#
+#     words = []
+#     for corner in corners_word:
+#         word = {}
+#         (top_left, bottom_right) = corner
+#         (col_min, row_min) = top_left
+#         (col_max, row_max) = bottom_right
+#         word['col_min'], word['col_max'], word['row_min'], word['row_max'] = col_min, col_max, row_min, row_max
+#         word['height'] = row_max - row_min
+#         word['width'] = col_max - col_min
+#         word['center'] = ((col_max + col_min)/2, (row_max + row_min)/2)
+#         words.append(word)
+#
+#     lines = []
+#     for word in words:
+#         line_index = is_in_line(word)
+#         # word is in current line
+#         if line_index != -1:
+#             merge_line(word, line_index)
+#         # word is not in current line
+#         else:
+#             # this single word as a new line
+#             lines.append(word)
+#
+#     corners_line = []
+#     for l in lines:
+#         corners_line.append(((l['col_min'], l['row_min']), (l['col_max'], l['row_max'])))
+#     return corners_line
+