Spaces:
Sleeping
Sleeping
# -*- coding: utf-8 -*- | |
import os | |
import numpy as np | |
import cv2 | |
import matplotlib.pyplot as plt | |
import crop_images | |
import imgproc | |
import matplotlib.pyplot as plt | |
# borrowed from https://github.com/lengstrom/fast-style-transfer/blob/master/src/utils.py | |
def get_files(img_dir): | |
imgs, masks, xmls = list_files(img_dir) | |
return imgs, masks, xmls | |
def list_files(in_path): | |
img_files = [] | |
mask_files = [] | |
gt_files = [] | |
for (dirpath, dirnames, filenames) in os.walk(in_path): | |
for file in filenames: | |
filename, ext = os.path.splitext(file) | |
ext = str.lower(ext) | |
if ext == '.jpg' or ext == '.jpeg' or ext == '.gif' or ext == '.png' or ext == '.pgm': | |
img_files.append(os.path.join(dirpath, file)) | |
elif ext == '.bmp': | |
mask_files.append(os.path.join(dirpath, file)) | |
elif ext == '.xml' or ext == '.gt' or ext == '.txt': | |
gt_files.append(os.path.join(dirpath, file)) | |
elif ext == '.zip': | |
continue | |
# img_files.sort() | |
# mask_files.sort() | |
# gt_files.sort() | |
return img_files, mask_files, gt_files | |
def saveResult(img_file, img, boxes, dirname='Results', verticals=None, texts=None): | |
""" save text detection result one by one | |
Args: | |
img_file (str): image file name | |
img (array): raw image context | |
boxes (array): array of result file | |
Shape: [num_detections, 4] for BB output / [num_detections, 4] for QUAD output | |
Return: | |
None | |
""" | |
img = np.array(img) | |
# make result file list: tên ảnh và đuôi ảnh | |
filename, file_ext = os.path.splitext(os.path.basename(img_file)) | |
# result directory | |
res_file = dirname + "res_" + filename + '.txt' | |
res_img_file = dirname + "res_" + filename + '.jpg' | |
if not os.path.isdir(dirname): | |
os.mkdir(dirname) | |
with open(res_file, 'w') as f: | |
for i, box in enumerate(boxes): | |
poly = np.array(box).astype(np.int32).reshape((-1)) | |
strResult = ','.join([str(p) for p in poly]) + '\r\n' | |
f.write(strResult) | |
# bỏ comment đoạn này để có thể bounding box lại b | |
poly = poly.reshape(-1, 2) | |
# cv2.polylines(img, [poly.reshape((-1, 1, 2))], True, color=(0, 0, 255), thickness=2) | |
# ptColor = (0, 255, 255) | |
xmin = min(poly[:, 0]) | |
xmax = max(poly[:, 0]) | |
ymin = min(poly[:, 1]) | |
ymax = max(poly[:, 1]) | |
width = xmax - xmin | |
height = ymax - ymin | |
# các điểm này từ file txt | |
pts = np.array([[xmin, ymax], [xmax, ymax], [xmax, ymin], [xmin, ymin]]) | |
word = crop_images.crop(pts, img) | |
folder = '/'.join(filename.split('/')[:-1]) | |
# đầu tiên đây là folder cropWord | |
dir = 'cropWord/' | |
if os.path.isdir(os.path.join(dir + folder)) == False: | |
os.makedirs(os.path.join(dir + folder)) | |
try: | |
file_name = os.path.join(dir + filename) | |
cv2.imwrite(file_name + str(i) + '.jpg', word) | |
except: | |
continue | |
if verticals is not None: | |
if verticals[i]: | |
ptColor = (255, 0, 0) | |
if texts is not None: | |
font = cv2.FONT_HERSHEY_SIMPLEX | |
font_scale = 0.5 | |
cv2.putText(img, "{}".format(texts[i]), (poly[0][0] + 1, poly[0][1] + 1), font, font_scale, (0, 0, 0), | |
thickness=1) | |
cv2.putText(img, "{:.2f}".format(texts[i]), tuple(poly[0]), font, font_scale, (0, 255, 255), | |
thickness=1) | |
# Save result image | |
# cv2.imwrite(res_img_file, img) | |