import os from PIL import Image import numpy as np import pandas as pd import cv2 import string import random CV2_FONTS = [ #cv2.FONT_HERSHEY_COMPLEX, cv2.FONT_HERSHEY_COMPLEX_SMALL, cv2.FONT_HERSHEY_DUPLEX, cv2.FONT_HERSHEY_PLAIN, cv2.FONT_HERSHEY_SIMPLEX, cv2.FONT_HERSHEY_TRIPLEX, cv2.FONT_ITALIC, cv2.QT_FONT_BLACK, cv2.QT_FONT_NORMAL ] # рандомный float между x и y def random_float(x, y): return random.random()*(y-x)+x # вычисляет размер текста в пикселях для cv2.putText def get_text_size(text, font, font_scale, thickness): (w, h), baseline = cv2.getTextSize(text, font, font_scale, thickness) return w, h+baseline # вычисляет какой нужен font_scale для определенного размера текста (по высоте) def get_font_scale(needed_height, text, font, thickness): w, h = get_text_size(text, font, 1, thickness) return needed_height/h # добавляет текст на изображение def place_text(image, text, color=(255,255,255), alpha=1, position=(0, 0), angle=0, font=cv2.FONT_HERSHEY_SIMPLEX, font_scale=1.0, thickness=3): image = np.array(image) overlay = np.zeros_like(image) output = image.copy() cv2.putText(overlay, text, position, font, font_scale, color, thickness) if angle != 0: text_w, text_h = get_text_size(text, font, font_scale, thickness) rotate_M = cv2.getRotationMatrix2D((position[0]+text_w//2, position[1]-text_h//2), angle, 1) overlay = cv2.warpAffine(overlay, rotate_M, (overlay.shape[1], overlay.shape[0])) overlay[overlay==0] = image[overlay==0] cv2.addWeighted(overlay, alpha, output, 1-alpha, 0, output) return Image.fromarray(output) def get_random_font_params(text, text_height, fonts, font_thickness_range): font = random.choice(fonts) font_thickness_range_scaled = [int(font_thickness_range[0]*(text_height/35)), int(font_thickness_range[1]*(text_height/85))] try: font_thickness = min(random.randint(*font_thickness_range_scaled), 2) except ValueError: font_thickness = 2 font_scale = get_font_scale(text_height, text, font, font_thickness) return font, font_scale, font_thickness # устанавливает вотермарку в центре изображения с рандомными параметрами def place_random_centered_watermark( pil_image, text, center_point_range_shift=(-0.025, 0.025), random_angle=(0,0), text_height_in_percent_range=(0.15, 0.18), text_alpha_range=(0.23, 0.5), fonts=CV2_FONTS, font_thickness_range=(2, 7), colors=[(255,255,255)] ): w, h = pil_image.size position_shift_x = random_float(*center_point_range_shift) offset_x = int(w*position_shift_x) position_shift_y = random_float(*center_point_range_shift) offset_y = int(w*position_shift_y) text_height = int(h*random_float(*text_height_in_percent_range)) font, font_scale, font_thickness = get_random_font_params(text, text_height, fonts, font_thickness_range) text_width, _ = get_text_size(text, font, font_scale, font_thickness) position_x = int((w/2)-text_width/2+offset_x) position_y = int((h/2)+text_height/2+offset_y) return place_text( pil_image, text, color=random.choice(colors), alpha=random_float(*text_alpha_range), position=(position_x, position_y), angle=random.randint(*random_angle), thickness=font_thickness, font=font, font_scale=font_scale ) def place_random_watermark( pil_image, text, random_angle=(0,0), text_height_in_percent_range=(0.10, 0.18), text_alpha_range=(0.18, 0.4), fonts=CV2_FONTS, font_thickness_range=(2, 6), colors=[(255,255,255)] ): w, h = pil_image.size text_height = int(h*random_float(*text_height_in_percent_range)) font, font_scale, font_thickness = get_random_font_params(text, text_height, fonts, font_thickness_range) text_width, _ = get_text_size(text, font, font_scale, font_thickness) position_x = random.randint(0, max(w-text_width, 10)) position_y = random.randint(text_height, h) return place_text( pil_image, text, color=random.choice(colors), alpha=random_float(*text_alpha_range), position=(position_x, position_y), angle=random.randint(*random_angle), thickness=font_thickness, font=font, font_scale=font_scale ) def center_crop(image, w, h): center = image.shape x = center[1]/2 - w/2 y = center[0]/2 - h/2 return image[int(y):int(y+h), int(x):int(x+w)] # добавляет текст в шахматном порядке на изображение def place_text_checkerboard(image, text, color=(255,255,255), alpha=1, step_x=0.1, step_y=0.1, angle=0, font=cv2.FONT_HERSHEY_SIMPLEX, font_scale=1.0, thickness=3): image_size = image.size image = np.array(image.convert('RGB')) if angle != 0: border_scale = 0.4 overlay_size = [int(i*(1+border_scale)) for i in list(image_size)] else: overlay_size = image_size w, h = overlay_size overlay = np.zeros((overlay_size[1], overlay_size[0], 3)) # change dimensions output = image.copy() text_w, text_h = get_text_size(text, font, font_scale, thickness) c = 0 for rel_pos_x in np.arange(0, 1, step_x): c += 1 for rel_pos_y in np.arange(text_h/h+(c%2)*step_y/2, 1, step_y): position = (int(w*rel_pos_x), int(h*rel_pos_y)) cv2.putText(overlay, text, position, font, font_scale, color, thickness) if angle != 0: rotate_M = cv2.getRotationMatrix2D((w//2, h//2), angle, 1) overlay = cv2.warpAffine(overlay, rotate_M, (overlay.shape[1], overlay.shape[0])) overlay = center_crop(overlay, image_size[0], image_size[1]) overlay[overlay==0] = image[overlay==0] overlay = overlay.astype(np.uint8) cv2.addWeighted(overlay, alpha, output, 1-alpha, 0, output) return Image.fromarray(output) def place_random_diagonal_watermark( pil_image, text, random_step_x=(0.25, 0.4), random_step_y=(0.25, 0.4), random_angle=(-60,60), text_height_in_percent_range=(0.10, 0.18), text_alpha_range=(0.18, 0.4), fonts=CV2_FONTS, font_thickness_range=(2, 6), colors=[(255,255,255)] ): w, h = pil_image.size text_height = int(h*random_float(*text_height_in_percent_range)) font, font_scale, font_thickness = get_random_font_params(text, text_height, fonts, font_thickness_range) text_width, _ = get_text_size(text, font, font_scale, font_thickness) return place_text_checkerboard( pil_image, text, color=random.choice(colors), alpha=random_float(*text_alpha_range), step_x=random_float(*random_step_x), step_y=random_float(*random_step_y), angle=random.randint(*random_angle), thickness=font_thickness, font=font, font_scale=font_scale )