152 lines
4.9 KiB
Python
152 lines
4.9 KiB
Python
import os
|
|
import datetime
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw
|
|
|
|
|
|
def save_images(img_list, folder):
|
|
if not os.path.exists(folder):
|
|
os.makedirs(folder)
|
|
now = datetime.datetime.now()
|
|
date_str = now.strftime("%Y-%m-%d")
|
|
folder_path = os.path.join(folder, date_str)
|
|
if not os.path.exists(folder_path):
|
|
os.makedirs(folder_path)
|
|
time_str = now.strftime("%H_%M_%S")
|
|
for idx, img in enumerate(img_list):
|
|
image_number = idx + 1
|
|
filename = f"{time_str}_{image_number}.jpg"
|
|
save_path = os.path.join(folder_path, filename)
|
|
cv2.imwrite(save_path, img[..., ::-1])
|
|
|
|
|
|
def check_channels(image):
|
|
channels = image.shape[2] if len(image.shape) == 3 else 1
|
|
if channels == 1:
|
|
image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
|
|
elif channels > 3:
|
|
image = image[:, :, :3]
|
|
return image
|
|
|
|
|
|
def resize_image(img, max_length=768):
|
|
height, width = img.shape[:2]
|
|
max_dimension = max(height, width)
|
|
|
|
if max_dimension > max_length:
|
|
scale_factor = max_length / max_dimension
|
|
new_width = int(round(width * scale_factor))
|
|
new_height = int(round(height * scale_factor))
|
|
new_size = (new_width, new_height)
|
|
img = cv2.resize(img, new_size)
|
|
height, width = img.shape[:2]
|
|
img = cv2.resize(img, (width - (width % 64), height - (height % 64)))
|
|
return img
|
|
|
|
|
|
def insert_spaces(string, nSpace):
|
|
if nSpace == 0:
|
|
return string
|
|
new_string = ""
|
|
for char in string:
|
|
new_string += char + " " * nSpace
|
|
return new_string[:-nSpace]
|
|
|
|
|
|
def draw_glyph(font, text):
|
|
g_size = 50
|
|
W, H = (512, 80)
|
|
new_font = font.font_variant(size=g_size)
|
|
img = Image.new(mode="1", size=(W, H), color=0)
|
|
draw = ImageDraw.Draw(img)
|
|
left, top, right, bottom = new_font.getbbox(text)
|
|
text_width = max(right - left, 5)
|
|
text_height = max(bottom - top, 5)
|
|
ratio = min(W * 0.9 / text_width, H * 0.9 / text_height)
|
|
new_font = font.font_variant(size=int(g_size * ratio))
|
|
|
|
text_width, text_height = new_font.getsize(text)
|
|
offset_x, offset_y = new_font.getoffset(text)
|
|
x = (img.width - text_width) // 2
|
|
y = (img.height - text_height) // 2 - offset_y // 2
|
|
draw.text((x, y), text, font=new_font, fill="white")
|
|
img = np.expand_dims(np.array(img), axis=2).astype(np.float64)
|
|
return img
|
|
|
|
|
|
def draw_glyph2(
|
|
font, text, polygon, vertAng=10, scale=1, width=512, height=512, add_space=True
|
|
):
|
|
enlarge_polygon = polygon * scale
|
|
rect = cv2.minAreaRect(enlarge_polygon)
|
|
box = cv2.boxPoints(rect)
|
|
box = np.int0(box)
|
|
w, h = rect[1]
|
|
angle = rect[2]
|
|
if angle < -45:
|
|
angle += 90
|
|
angle = -angle
|
|
if w < h:
|
|
angle += 90
|
|
|
|
vert = False
|
|
if abs(angle) % 90 < vertAng or abs(90 - abs(angle) % 90) % 90 < vertAng:
|
|
_w = max(box[:, 0]) - min(box[:, 0])
|
|
_h = max(box[:, 1]) - min(box[:, 1])
|
|
if _h >= _w:
|
|
vert = True
|
|
angle = 0
|
|
|
|
img = np.zeros((height * scale, width * scale, 3), np.uint8)
|
|
img = Image.fromarray(img)
|
|
|
|
# infer font size
|
|
image4ratio = Image.new("RGB", img.size, "white")
|
|
draw = ImageDraw.Draw(image4ratio)
|
|
_, _, _tw, _th = draw.textbbox(xy=(0, 0), text=text, font=font)
|
|
text_w = min(w, h) * (_tw / _th)
|
|
if text_w <= max(w, h):
|
|
# add space
|
|
if len(text) > 1 and not vert and add_space:
|
|
for i in range(1, 100):
|
|
text_space = insert_spaces(text, i)
|
|
_, _, _tw2, _th2 = draw.textbbox(xy=(0, 0), text=text_space, font=font)
|
|
if min(w, h) * (_tw2 / _th2) > max(w, h):
|
|
break
|
|
text = insert_spaces(text, i - 1)
|
|
font_size = min(w, h) * 0.80
|
|
else:
|
|
shrink = 0.75 if vert else 0.85
|
|
font_size = min(w, h) / (text_w / max(w, h)) * shrink
|
|
new_font = font.font_variant(size=int(font_size))
|
|
|
|
left, top, right, bottom = new_font.getbbox(text)
|
|
text_width = right - left
|
|
text_height = bottom - top
|
|
|
|
layer = Image.new("RGBA", img.size, (0, 0, 0, 0))
|
|
draw = ImageDraw.Draw(layer)
|
|
if not vert:
|
|
draw.text(
|
|
(rect[0][0] - text_width // 2, rect[0][1] - text_height // 2 - top),
|
|
text,
|
|
font=new_font,
|
|
fill=(255, 255, 255, 255),
|
|
)
|
|
else:
|
|
x_s = min(box[:, 0]) + _w // 2 - text_height // 2
|
|
y_s = min(box[:, 1])
|
|
for c in text:
|
|
draw.text((x_s, y_s), c, font=new_font, fill=(255, 255, 255, 255))
|
|
_, _t, _, _b = new_font.getbbox(c)
|
|
y_s += _b
|
|
|
|
rotated_layer = layer.rotate(angle, expand=1, center=(rect[0][0], rect[0][1]))
|
|
|
|
x_offset = int((img.width - rotated_layer.width) / 2)
|
|
y_offset = int((img.height - rotated_layer.height) / 2)
|
|
img.paste(rotated_layer, (x_offset, y_offset), rotated_layer)
|
|
img = np.expand_dims(np.array(img.convert("1")), axis=2).astype(np.float64)
|
|
return img
|