1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788278927902791279227932794279527962797279827992800280128022803280428052806280728082809281028112812281328142815281628172818281928202821282228232824282528262827282828292830283128322833283428352836283728382839284028412842284328442845284628472848284928502851285228532854285528562857285828592860286128622863286428652866286728682869287028712872287328742875287628772878287928802881288228832884288528862887288828892890289128922893289428952896289728982899290029012902290329042905290629072908290929102911291229132914291529162917291829192920292129222923292429252926292729282929293029312932293329342935293629372938293929402941294229432944294529462947294829492950295129522953295429552956295729582959296029612962296329642965296629672968296929702971297229732974297529762977297829792980298129822983298429852986298729882989299029912992299329942995299629972998299930003001300230033004300530063007300830093010301130123013301430153016301730183019302030213022302330243025302630273028302930303031303230333034 |
- import copy
- import gc
- import io
- import json
- import math
- import os
- import random
- import re
- import sys
- import threading
- import time
- import traceback
- from glob import glob
- from itertools import combinations, product
- import matplotlib
- # 不显示图
- # matplotlib.use('agg')
- # print(matplotlib.get_backend())
- import matplotlib.pyplot as plt
- import matplotlib.patheffects as path_effects
- import chardet
- import cv2
- import jieba
- import numpy as np
- from PIL import ImageFont, ImageDraw, Image, ImageEnhance, ImageFilter
- from captcha.image import ImageCaptcha
- from keras_preprocessing.sequence import pad_sequences
- from matplotlib import _pylab_helpers
- from matplotlib.colors import rgb_to_hsv, hsv_to_rgb
- from click_captcha.model import u_net_denoise
- sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../")
- sys.path.append(os.path.dirname(os.path.abspath(__file__)))
- from click_captcha.utils import np2pil, pil2np, pil_resize, pil_rotate, pil2np_a, np2pil_a, pil_resize_a, pil_rotate_a
- fig = plt.figure(figsize=(1, 1), facecolor='none')
- def gen_siamese(paths, batch_size=32, shape=(40, 40), cls_num=1):
- num = len(paths)
- data_path = os.path.dirname(os.path.abspath(__file__)) + "/../data/click/"
- i = 0
- while True:
- if i >= num:
- i = 0
- random.shuffle(paths)
- height, width = shape[:2]
- X1 = np.zeros((batch_size, height, width, 1))
- X2 = np.zeros((batch_size, height, width, 1))
- Y = np.zeros((batch_size, 2))
- for j in range(batch_size):
- # 生成标注数据
- img1, img2, label = paths[i][:-1].split("\t")
- # print(img1, img2, label)
- img1 = cv2.imread(data_path + img1)
- img1 = pil_resize(img1, shape[0], shape[1])
- img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
- img1 = np.expand_dims(img1, axis=-1)
- img2 = cv2.imread(data_path + img2)
- img2 = pil_resize(img2, shape[0], shape[1])
- img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
- img2 = np.expand_dims(img2, axis=-1)
- if label == "1":
- label = np.array([0, 1])
- else:
- label = np.array([1, 0])
- X1[j] = img1
- X2[j] = img2
- Y[j] = label
- yield {"input_1": X1, "input_2": X2}, {"output": Y}
- def gen_char(paths, batch_size=32, shape=(40, 40), cls_num=6270, data_path="click"):
- num = len(paths)
- data_path = os.path.dirname(os.path.abspath(__file__)) + "/../data/" + data_path + "/"
- i = 0
- random.shuffle(paths)
- while True:
- if i >= num:
- i = 0
- random.shuffle(paths)
- height, width = shape[:2]
- if len(shape) > 2 and shape[2] == 3:
- channel = 3
- else:
- channel = 1
- X = np.zeros((batch_size, height, width, channel))
- Y = np.zeros((batch_size, cls_num))
- j = 0
- random.shuffle(paths)
- while j < batch_size:
- if i >= num:
- random.shuffle(paths)
- i = 0
- path = paths[i].split(os.sep)[-1]
- char_index = int(path.split("_")[0])
- label = np.zeros(cls_num)
- # print("char_index", char_index)
- label[char_index] = 1
- # print("label", np.argmax(label), char_index)
- img1 = cv2.imread(data_path + path)
- img1 = pil_resize(img1, shape[0], shape[1])
- # img2 = copy.deepcopy(img1)
- # 数据增强
- if random.choice([0, 1]):
- img1_pil = np2pil(img1)
- aug = random.choice([0, 1, 2, 4, 5])
- # aug = 6
- if aug == 0:
- img1_pil = image_enhance_color(img1_pil)
- elif aug == 1:
- img1_pil = image_enhance_brightness(img1_pil)
- elif aug == 3:
- img1_pil = image_enhance_contrast(img1_pil)
- elif aug == 4:
- img1_pil = image_enhance_sharpness(img1_pil)
- elif aug == 5:
- img1_pil = image_enhance_blur(img1_pil)
- img1 = pil2np(img1_pil)
- if aug == 6:
- img1 = image_enhance_distort(img1)
- # print(aug)
- # cv2.imshow("origin", img2)
- # cv2.imshow("gen_char", img1)
- # cv2.waitKey(0)
- img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
- img1 = np.expand_dims(img1, axis=-1)
- img1 = img1 / 255.
- X[j] = img1
- Y[j] = label
- i += 1
- j += 1
- # print("error_num", error_num)
- yield X, Y
- def gen_yolo_char(paths, batch_size, input_shape, anchors, num_classes, box_num=6):
- """data generator for fit_generator"""
- n = len(paths)
- data_path = os.path.dirname(os.path.abspath(__file__)) + "/../data/detect/"
- i = 0
- while True:
- image_data = []
- box_data = []
- batch_cnt = 0
- while batch_cnt < batch_size:
- try:
- if i == 0:
- np.random.shuffle(paths)
- ss = paths[i][:-1].split(" ")
- image_path = ss[0]
- image = cv2.imread(data_path+image_path)
- origin_h, origin_w = image.shape[:2]
- image = pil_resize(image, input_shape[0], input_shape[1])
- # 数据增强
- if random.choice([0, 0, 1]):
- img1_pil = np2pil(image)
- aug = random.choice([0, 1, 2, 4, 5])
- if aug == 0:
- img1_pil = image_enhance_color(img1_pil)
- elif aug == 1:
- img1_pil = image_enhance_brightness(img1_pil)
- elif aug == 3:
- img1_pil = image_enhance_contrast(img1_pil)
- elif aug == 4:
- img1_pil = image_enhance_sharpness(img1_pil)
- elif aug == 5:
- img1_pil = image_enhance_blur(img1_pil)
- image = pil2np(img1_pil)
- image_show = copy.deepcopy(image)
- image = image / 255.
- box = np.array([np.array(list(map(int, box.split(',')))) for box in ss[1:]])
- # box数不同,复制
- # print("box.shape", box.shape)
- if box.shape[0] < box_num:
- # box = np.concatenate([box]+[np.zeros(5)]*(box_num-box.shape[0]), axis=0)
- if box_num-box.shape[0] >= box.shape[0]:
- box = np.concatenate([box]+[box[0, :]]*(box_num-box.shape[0]), axis=0)
- else:
- box = np.concatenate([box, box[:box_num-box.shape[0], :]], axis=0)
- # show
- # box_show = box.tolist()
- # for b in box_show:
- # print("box", b)
- # cv2.rectangle(image_show, (b[0], b[1]), (b[2], b[3]), (255, 0, 0), 2)
- # cv2.imshow("image_show", image_show)
- # cv2.waitKey(0)
- # print("box.shape", box.shape)
- image_data.append(image)
- box_data.append(box)
- i = (i+1) % n
- batch_cnt += 1
- except:
- i = (i+1) % n
- continue
- # print
- # print(image.shape)
- # image_show = (image*255).astype(np.uint8)
- # print("annotation_lines[i]", annotation_lines[i])
- # for _b in box:
- # print(_b)
- # cv2.rectangle(image_show, (int(_b[0]), int(_b[1])), (int(_b[2]), int(_b[3])), (0, 255, 0), 1)
- # cv2.imshow("image", image_show)
- # cv2.waitKey(0)
- image_data = np.array(image_data)
- box_data = np.array(box_data)
- # print(image_data.shape, box_data.shape)
- y_true = preprocess_true_boxes(box_data, input_shape, anchors, num_classes)
- yield [image_data, *y_true], np.zeros(batch_size)
- def gen_yolo_puzzle(paths, batch_size, input_shape, anchors, num_classes, box_num=1):
- """data generator for fit_generator"""
- n = len(paths)
- data_path = os.path.dirname(os.path.abspath(__file__)) + "/../data/detect2/"
- i = 0
- while True:
- image_data = []
- box_data = []
- batch_cnt = 0
- while batch_cnt < batch_size:
- try:
- if i == 0:
- np.random.shuffle(paths)
- ss = paths[i][:-1].split(" ")
- image_path = ss[0]
- image = cv2.imread(data_path+image_path)
- image = pil_resize(image, input_shape[0], input_shape[1])
- image_show = copy.deepcopy(image)
- image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
- image = 255. - image
- image = np.uint8(image)
- # cv2.imshow("image", image)
- # cv2.waitKey(0)
- image = np.expand_dims(image, -1)
- image = image / 255.
- box = np.array([np.array(list(map(int, box.split(',')))) for box in ss[1:]])
- # box数不同,复制
- if box.shape[0] < box_num:
- box = np.concatenate([box, box[:2, :]], axis=0)
- # show
- # box_show = box.tolist()
- # for b in box_show:
- # print("box", b)
- # cv2.rectangle(image_show, (b[0], b[1]), (b[2], b[3]), (0, 0, 255), 2)
- # cv2.imshow("image_show", image_show)
- # cv2.waitKey(0)
- image_data.append(image)
- box_data.append(box)
- i = (i+1) % n
- batch_cnt += 1
- except:
- i = (i+1) % n
- continue
- # print
- # print(image.shape)
- # image_show = (image*255).astype(np.uint8)
- # for _b in box:
- # print(_b)
- # cv2.rectangle(image_show, (int(_b[0]), int(_b[1])), (int(_b[2]), int(_b[3])), (0, 255, 0), 1)
- # cv2.imshow("image", image_show)
- # cv2.waitKey(0)
- image_data = np.array(image_data)
- box_data = np.array(box_data)
- # print(image_data.shape, box_data.shape)
- y_true = preprocess_true_boxes(box_data, input_shape, anchors, num_classes)
- yield [image_data, *y_true], np.zeros(batch_size)
- def gen_drag(paths, batch_size=32, shape=(128, 256), cls_num=2):
- num = len(paths)
- data_path = os.path.dirname(os.path.abspath(__file__)) + "/../data/drag/"
- map_path = data_path+"map.txt"
- with open(map_path, "r") as f:
- _list = f.readlines()
- map_dict = {}
- for s in _list:
- ss = s[:-1].split(" ")
- map_dict[ss[0]] = ss[1]
- i = 0
- random.shuffle(paths)
- while True:
- if i >= num:
- i = 0
- random.shuffle(paths)
- height, width = shape[:2]
- if len(shape) > 2:
- channel = 3
- else:
- channel = 1
- X = np.zeros((batch_size, height, width, channel))
- Y = np.zeros((batch_size, height, width, 1))
- for j in range(batch_size):
- if i >= num:
- random.shuffle(paths)
- i = 0
- path = paths[i].split(os.sep)[-1]
- w_index = int(map_dict.get(path))
- # label = np.zeros(cls_num)
- # print("char_index", char_index)
- # label[w_index] = 1
- # print("label", np.argmax(label), char_index)
- img1 = cv2.imread(data_path + path)
- img1 = pil_resize(img1, shape[0], shape[1])
- # cv2.imshow("image", img1)
- label = np.full((shape[0], shape[1], 1), 0, dtype='uint8')
- label[:, w_index, 0] = 1
- # label[:, w_index, 1] = 1
- # cv2.imshow("label", np.expand_dims(label[..., 0], -1))
- # cv2.waitKey(0)
- img1 = img1 / 255.
- # img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
- # img1 = np.expand_dims(img1, axis=-1)
- i += 1
- X[j] = img1
- Y[j] = label
- yield X, Y
- def gen_phrase(map_list, batch_size=32, shape=(5707, 3)):
- voc_dim, timesteps = shape[:2]
- data_list = []
- for line in map_list:
- data_list.append([eval(line[:-3]), int(line[-2:-1])])
- num = len(data_list)
- i = 0
- random.shuffle(data_list)
- while True:
- X = np.zeros((batch_size, timesteps))
- Y = np.zeros((batch_size, 1))
- for j in range(batch_size):
- if i >= num:
- random.shuffle(data_list)
- i = 0
- data = data_list[i]
- d_list = [x for x in data[0]]
- d_list = d_list + [voc_dim]*(timesteps-len(d_list))
- X[j] = np.array(d_list)
- Y[j] = data[1]
- i += 1
- yield X, Y
- def gen_equation(paths, batch_size=32, shape=(40, 40), input_len=21, label_len=8, cls_num=6270, data_path='equation'):
- num = len(paths)
- data_path = os.path.dirname(os.path.abspath(__file__)) + "/../data/" + data_path + "/"
- map_path = os.path.dirname(os.path.abspath(__file__)) + "/../data/equation.txt"
- with open(map_path, "r") as f:
- map_list = f.readlines()
- map_str = "".join(map_list)
- map_str = re.sub("\n", "", map_str)
- char_map_dict = {
- "星": '*',
- "斜": "/",
- "问": "?",
- 'x': '×',
- '?': '?'
- }
- i = 0
- random.shuffle(paths)
- while True:
- if i >= num:
- i = 0
- random.shuffle(paths)
- height, width = shape[:2]
- if len(shape) > 2 and shape[2] == 3:
- channel = 3
- else:
- channel = 1
- X = np.zeros((batch_size, height, width, channel))
- Y = np.zeros((batch_size, label_len))
- input_length = np.ones(batch_size) * input_len
- label_length = np.ones(batch_size) * label_len
- j = 0
- while j < batch_size:
- if i >= num:
- random.shuffle(paths)
- i = 0
- path = paths[i].split(os.sep)[-1]
- char_index_list = []
- for c in path.split(".")[0].split('_')[1:]:
- if c in char_map_dict.keys():
- c = char_map_dict.get(c)
- if not c:
- continue
- # print("c", c)
- char_index_list.append(map_str.index(c)+1)
- char_index_list.extend([0] * (label_len-len(char_index_list)))
- label = np.array(char_index_list)
- img1 = cv2.imread(data_path + path)
- img1 = pil_resize(img1, shape[0], shape[1])
- # img2 = copy.deepcopy(img1)
- # cv2.imshow("origin", img2)
- # cv2.imshow("gen_char", img1)
- # cv2.waitKey(0)
- img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
- img1 = np.expand_dims(img1, axis=-1)
- img1 = img1 / 255.
- X[j] = img1
- Y[j] = label
- i += 1
- j += 1
- yield [X, Y, input_length, label_length], np.ones(batch_size)
- def gen_equation2(paths, batch_size=32, shape=(40, 40), input_len=21, label_len=8, cls_num=6270, data_path='equation'):
- os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
- image_shape = (32, 192, 1)
- weights_path = "./models/e153-loss53.97-denoise.h5"
- model = u_net_denoise(input_shape=image_shape, class_num=image_shape[2])
- model.load_weights(weights_path)
- os.environ["CUDA_VISIBLE_DEVICES"] = "0"
- data_path = os.path.dirname(os.path.abspath(__file__)) + "/../data/" + data_path + "/"
- map_path = os.path.dirname(os.path.abspath(__file__)) + "/../data/equation.txt"
- with open(map_path, "r") as f:
- map_list = f.readlines()
- map_str = "".join(map_list)
- map_str = re.sub("\n", "", map_str)
- blank_index = len(map_str)+1
- char_map_dict = {
- "星": '*',
- "斜": "/",
- "问": "?",
- 'x': '×',
- '?': '?'
- }
- i = 0
- while True:
- height, width = shape[:2]
- if len(shape) > 2 and shape[2] == 3:
- channel = 3
- else:
- channel = 1
- X = np.zeros((batch_size, height, width, channel))
- Y = np.zeros((batch_size, label_len))
- input_length = np.ones(batch_size) * input_len
- label_length = np.ones(batch_size) * label_len
- noise = random.choice([False, True, True])
- result_list = generate_data_equation2(batch_size, noise=noise)
- # 模型降噪
- img_list = [np.expand_dims(cv2.cvtColor(pil_resize(x[0], image_shape[0], image_shape[1]), cv2.COLOR_BGR2GRAY), axis=-1) for x in result_list]
- img_data = np.array(img_list) / 255.
- pred = model.predict(img_data)
- img_list = np.uint8(pred*255.)
- # for j in range(len(result_list)):
- # img1 = img_list[j]
- # img1 = add_contrast(img1)
- # ratio = get_image_legal(img1)
- # if ratio >= 0.02:
- for j in range(len(result_list)):
- if random.choice([0, 1, 1, 1]):
- img1 = img_list[j]
- img1 = add_contrast(img1)
- else:
- img1 = result_list[j][0]
- img1 = pil_resize(img1, shape[0], shape[1])
- img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
- img1 = np.expand_dims(img1, axis=-1)
- # img1 = cv2.adaptiveThreshold(img1, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,
- # 11, 20)
- # img1 = np.expand_dims(img1, axis=-1)
- # cv2.imshow("origin", img1)
- # if noise:
- # if random.choice([0, 1]):
- # img1 = 255 - img1
- # _, img1 = cv2.threshold(img1, 110, 255, cv2.THRESH_BINARY)
- char_list = result_list[j][1]
- if char_list[0] is None:
- # char_index_list = [0] * len(char_list)
- # print("len(map_str)+1", len(map_str)+1)
- char_index_list = [blank_index] * len(char_list)
- else:
- char_index_list = []
- for c in char_list:
- if c in char_map_dict.keys():
- c = char_map_dict.get(c)
- if not c:
- continue
- # print("c", c)
- # char_index_list.append(map_str.index(c)+1)
- char_index_list.append(map_str.index(c))
- char_index_list.extend([blank_index] * (label_len-len(char_index_list)))
- label = np.array(char_index_list)
- ratio = get_image_legal(img1)
- # img1 = pil_resize(img, shape[0], shape[1])
- # img2 = copy.deepcopy(img1)
- # print(char_list)
- # print(char_index_list)
- # print("ratio", ratio)
- # cv2.imshow("gen_char", img1)
- # cv2.waitKey(0)
- # img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
- # img1 = np.expand_dims(img1, axis=-1)
- img1 = img1 / 255.
- X[j] = img1
- Y[j] = label
- i += 1
- yield [X, Y, input_length, label_length], np.ones(batch_size)
- def gen_equation_denoise(paths, batch_size=32, shape=(40, 40)):
- while True:
- height, width = shape[:2]
- if len(shape) > 2 and shape[2] == 3:
- channel = 3
- else:
- channel = 1
- X = np.zeros((batch_size, height, width, channel))
- Y = np.zeros((batch_size, height, width, channel))
- result_list = generate_data_denoise(batch_size)
- j = 0
- for img, noise in result_list:
- img1 = pil_resize(img, shape[0], shape[1])
- img2 = pil_resize(noise, shape[0], shape[1])
- # img1 = add_contrast(img1)
- # img2 = add_contrast(img2)
- # cv2.imshow("origin", img2)
- img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
- img1 = np.expand_dims(img1, axis=-1)
- img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
- img2 = np.expand_dims(img2, axis=-1)
- # cv2.imshow("bg", img1)
- # cv2.imshow("noise", img2)
- # cv2.waitKey(0)
- img1 = img1 / 255.
- img2 = img2 / 255.
- X[j] = img2
- Y[j] = img1
- j += 1
- yield X, Y
- def generate_data_siamese(char_num=6, char_shape=(40, 40)):
- bg_paths = glob("../data/base/*")
- char_path = "../data/chinese_2500.txt"
- with open(char_path, "r") as f:
- char_str = f.read()
- data_dir = "../data/click/"
- for i in range(1000):
- bg = cv2.imread(random.sample(bg_paths, 1)[0])
- char_list = [char_str[x] for x in random.sample(range(len(char_str)), char_num)]
- image_np, position_list, tips_image_list = char_on_image(bg, char_list, char_shape)
- # for j in range(len(char_list)):
- # char = char_list[j]
- # p = position_list[j]
- # print(char)
- # cv2.rectangle(image_np, [p[1], p[0]], [p[1]+char_shape[1], p[0]+char_shape[0]], (255, 0, 0), 2)
- # cv2.imshow("generate_data", image_np)
- # cv2.waitKey(0)
- # 保存tips图片
- tips_path_list = []
- for k in range(len(tips_image_list)):
- tips_image = tips_image_list[k]
- tips_path = str(i) + "_" + str(k) + ".jpg"
- tips_path_list.append(tips_path)
- cv2.imwrite(data_dir+tips_path, tips_image)
- # 保存文字区域图片
- char_path_list = []
- for j in range(len(char_list)):
- p = position_list[j]
- char_path = str(i) + "_" + str(len(tips_image_list)+j) + ".jpg"
- char_path_list.append(char_path)
- cv2.imwrite(data_dir+char_path, image_np[p[0]:p[0]+char_shape[0], p[1]:p[1]+char_shape[1], :])
- # 生成映射数据
- with open("../data/click/map.txt", "a") as f:
- for j in range(len(tips_path_list)):
- tips_path = tips_path_list[j]
- for k in range(len(char_path_list)):
- char_path = char_path_list[k]
- if j == k:
- f.write(tips_path + "\t" + char_path + "\t" + str(1) + "\n")
- else:
- f.write(tips_path + "\t" + char_path + "\t" + str(0) + "\n")
- def generate_data_char(char_num=6, char_shape=(40, 40), image_shape=(160, 260)):
- # (40,40) (160, 260)
- # (80,80) (360, 590)
- bg_paths = glob("../data/base/*") + glob("../data/base1/*")*100
- char_path = "../data/chinese_simple_5649.txt"
- with open(char_path, "r") as f:
- char_list = f.readlines()
- # random.shuffle(char_list)
- char_str = "".join(char_list)
- char_str = re.sub("\n", "", char_str)
- data_dir = "../data/click_simple/"
- # 每个字生成多张图片
- start_time = time.time()
- # 0- 1520 4000- 5760
- for i in range(0, 6000):
- if i % 20 == 0:
- print("Loop", i, time.time()-start_time)
- start_time = time.time()
- char = char_str[i]
- # 生成带背景图数
- image_cnt = 1
- char_list = [char] * image_cnt
- tips_cnt = 0
- image_cnt = 265
- tips_list = []
- for l in range(10):
- try:
- # 背景图
- bg = cv2.imread(random.sample(bg_paths, 1)[0])
- if random.choice([0, 0, 1]):
- bg = image_enhance_distort(bg)
- if random.choice([0, 0, 1]):
- bg = image_enhance_flip(bg)
- # 生成4张tips图,6张带背景的旋转图
- image_np, p_list, t_list = char_on_image(bg, char_list, char_shape, image_shape,
- char_stretch=False)
- # cv2.imshow("char_on_image", image_np)
- # cv2.waitKey(0)
- tips_list += t_list
- for p in p_list:
- char_path = str(i) + "_" + str(image_cnt) + "_2" + ".jpg"
- # 放大缩小box
- if random.choice([1, 1, 1, 0]):
- threshold = random.randint(12, 25)
- y1 = max(0, p[0]-random.choice([0, threshold]))
- y2 = min(image_np.shape[0], p[0]+char_shape[0]+random.choice([0, threshold]))
- x1 = max(0, p[1]-random.choice([0, threshold]))
- x2 = min(image_np.shape[1], p[1]+char_shape[1]+random.choice([0, threshold]))
- # print(threshold)
- else:
- threshold = random.randint(4, 8)
- y1 = p[0]+random.choice([0, threshold])
- y2 = p[0]+char_shape[0]-random.choice([0, threshold])
- x1 = p[1]+random.choice([0, threshold])
- x2 = p[1]+char_shape[1]-random.choice([0, threshold])
- # 平移box
- if random.choice([0, 0, 1]):
- threshold = random.randint(8, 13)
- h_flag, w_flag = 0, 0
- if random.choice([0, 1]):
- threshold = -threshold
- if random.choice([0, 1, 1]):
- h_flag = 1
- if random.choice([0, 1, 1]):
- w_flag = 1
- y1 = p[0]+threshold*h_flag
- y2 = p[0]+char_shape[0]+threshold*h_flag
- x1 = p[1]+threshold*w_flag
- x2 = p[1]+char_shape[1]+threshold*w_flag
- sub_image_np = image_np[y1:y2,
- x1:x2,
- :]
- if sub_image_np.shape[0] == 0 or sub_image_np.shape[1] == 0:
- continue
- sub_image_np = pil_resize(sub_image_np, char_shape[0], char_shape[1])
- # cv2.imshow("sub_image_np", sub_image_np)
- # cv2.waitKey(0)
- cv2.imwrite(data_dir+char_path, sub_image_np)
- image_cnt += 1
- except:
- print(1)
- continue
- for tips_image in tips_list[:1]:
- tips_path = str(i) + "_" + str(tips_cnt) + "_1" + ".jpg"
- # cv2.imshow("tips_image", tips_image)
- # cv2.waitKey(0)
- cv2.imwrite(data_dir+tips_path, tips_image)
- tips_cnt += 1
- def generate_data_char_from_yolo():
- _dir = "D:/Project/captcha/data/detect/"
- map_path = _dir + "map.txt"
- with open(map_path, "r") as f:
- map_list = f.readlines()
- char_image_list = []
- for line in map_list[:10]:
- image_path = _dir + line.split(" ")[0]
- boxes = line[:-1].split(" ")[1:]
- boxes = [list(map(int, box.split(','))) for box in boxes]
- print(image_path, boxes)
- image = cv2.imread(image_path)
- for box in boxes:
- char_image = image[box[1]:box[3], box[0]:box[2], :]
- char_image_list.append(char_image)
- return
- def generate_data_yolo_char(image_shape=(160, 256)):
- bg_paths = glob("../data/base1/*")
- char_path = "../data/chinese_6270.txt"
- with open(char_path, "r") as f:
- char_map_list = f.readlines()
- random.shuffle(char_map_list)
- data_dir = "../data/detect/"
- # with open(data_dir+"map.txt", "w") as f:
- # f.write("")
- j = 0
- start_time = time.time()
- for i in range(46121, 46500):
- if i % 50 == 0:
- print("Loop", i, time.time()-start_time)
- start_time = time.time()
- if j >= len(char_map_list)-1:
- j = 0
- random.shuffle(char_map_list)
- # 随机底图大小
- r = random.randint(160, 256)
- # random_image_shape = (r, random.choice([r*3, int(r*2.5), int(r*3.5), r*4]))
- random_image_shape = (r, r*3)
- # 随机shape和num
- r = random.randint(30, 40)
- char_shape = (r, r)
- char_num = random.randint(4, 6)
- tips_char_num = min(random.randint(3, char_num), 4)
- # 背景
- try:
- bg = cv2.imread(random.sample(bg_paths, 1)[0])
- # bg = pil_resize(bg, random_image_shape[0], random_image_shape[1])
- if random.choice([0, 0, 0]):
- bg = image_enhance_distort(bg)
- if random.choice([0, 0, 0]):
- bg = image_enhance_flip(bg)
- except:
- print(1)
- continue
- if bg is None:
- print(2)
- continue
- # 随机取char
- char_list = [x[:-1] for x in char_map_list[j:j+char_num]]
- # print("char_list", char_list)
- j = j+char_num
- # 生成图
- image_np, position_list, tips_image_list = char_on_image(bg, char_list, char_shape, random_image_shape, tips_char_num)
- if i < 5000:
- tips_image_np, tips_position_list = get_tips_image(tips_image_list, char_shape, random_image_shape)
- image_np_path = str(i) + ".jpg"
- # print(image_np.shape)
- image_np = pil_resize(image_np, image_shape[0], image_shape[1])
- cv2.imwrite(data_dir+image_np_path, image_np)
- if i < 5000:
- tips_image_np_path = str(i) + "_0.jpg"
- tips_image_np = pil_resize(tips_image_np, image_shape[0], image_shape[1])
- cv2.imwrite(data_dir+tips_image_np_path, tips_image_np)
- # 生成映射数据
- with open(data_dir+"map.txt", "a") as f:
- box_str = ""
- for p in position_list:
- x1 = int(p[1] * image_shape[1] / random_image_shape[1])
- y1 = int(p[0] * image_shape[0] / random_image_shape[0])
- x2 = int((p[1]+char_shape[1]) * image_shape[1] / random_image_shape[1])
- y2 = int((p[0]+char_shape[0]) * image_shape[0] / random_image_shape[0])
- threshold = random.randint(2, 4)
- if random.choice([0, 1]):
- y1 = max(0, y1 - threshold)
- y2 = min(image_np.shape[0], y2 + threshold)
- x1 = max(0, x1 - threshold)
- x2 = min(image_np.shape[1], x2 + threshold)
- # print(threshold)
- box_str += str(x1) + "," + str(y1) + "," + \
- str(x2) + "," + str(y2) + "," + \
- str(0) + " "
- # cv2.rectangle(image_np, (x1, y1), (x2, y2), (255, 0, 0), 2)
- # cv2.imshow("image_np", image_np)
- # cv2.waitKey(0)
- box_str = box_str[:-1]
- f.write(image_np_path + " " + box_str + "\n")
- if i < 5000:
- box_str = ""
- x1 = int(p[1] * image_shape[1] / random_image_shape[1])
- y1 = int(p[0] * image_shape[0] / random_image_shape[0])
- x2 = int((p[1]+char_shape[1]) * image_shape[1] / random_image_shape[1])
- y2 = int((p[0]+char_shape[0]) * image_shape[0] / random_image_shape[0])
- for p in tips_position_list:
- box_str += str(x1) + "," + str(y1) + "," + \
- str(x2) + "," + str(y2) + "," + \
- str(0) + " "
- cv2.rectangle(tips_image_np, (x1, y1), (x2, y2), (255, 0, 0), 2)
- # cv2.imshow("tips_image_np", tips_image_np)
- box_str = box_str[:-1]
- f.write(tips_image_np_path + " " + box_str + "\n")
- def generate_data_yolo_char_1(image_shape=(160, 256)):
- bg_paths = glob("../data/base/*")
- char_path = "../data/chinese_6270.txt"
- with open(char_path, "r") as f:
- char_map_list = f.readlines()
- data_dir = "../data/detect/"
- # with open(data_dir+"map.txt", "w") as f:
- # f.write("")
- j = 0
- for i in range(20000, 25000):
- if i % 1000 == 0:
- print("Loop", i)
- if j >= len(char_map_list)-1:
- j = 0
- random.shuffle(char_map_list)
- # 随机shape和num
- r = random.randint(31, 51)
- char_shape = (r, r)
- char_num = random.randint(3, 6)
- tips_char_num = 0
- # 背景
- bg = cv2.imread(random.sample(bg_paths, 1)[0])
- if random.choice([0, 0, 1]):
- bg = image_enhance_distort(bg)
- if random.choice([0, 0, 1]):
- bg = image_enhance_flip(bg)
- # 随机取char
- char_list = [x[:-1] for x in char_map_list[j:j+char_num]]
- # print("char_list", char_list)
- j = j+char_num
- # 生成图
- image_np, position_list, tips_image_list = char_on_image(bg, char_list, char_shape,
- image_shape, tips_char_num,
- only_color=(0, 0, 0))
- image_np_path = str(i) + ".jpg"
- cv2.imwrite(data_dir+image_np_path, image_np)
- # 生成映射数据
- with open(data_dir+"map.txt", "a") as f:
- box_str = ""
- for p in position_list:
- box_str += str(p[1]) + "," \
- + str(p[0]) + "," \
- + str(p[1]) + "," \
- + str(p[0]) + "," \
- + str(0) + " "
- # cv2.rectangle(image_np, (p[1], p[0]), (p[1]+char_shape[1], p[0]+char_shape[0]), (255, 0, 0), 2)
- # cv2.imshow("image_np", image_np)
- # cv2.waitKey(0)
- box_str = box_str[:-1]
- f.write(image_np_path + " " + box_str + "\n")
- def generate_data_yolo_puzzle(image_shape=(160, 256)):
- bg_paths = glob("../data/base/*.jpeg")
- data_dir = "../data/detect2/"
- with open(data_dir+"map.txt", "w") as f:
- f.write("")
- for i in range(0, 10000):
- if i % 1000 == 0:
- print("Loop", i)
- bg = cv2.imread(random.sample(bg_paths, 1)[0])
- if random.choice([0, 0, 1]):
- bg = image_enhance_distort(bg)
- if random.choice([0, 0, 1]):
- bg = image_enhance_flip(bg)
- r = random.randint(35, 60)
- puzzle_shape = (r, r)
- image_np, position_list = puzzle_on_image(bg, puzzle_shape, image_shape)
- image_np_path = str(i) + ".jpg"
- cv2.imwrite(data_dir+image_np_path, image_np)
- # 生成映射数据
- with open(data_dir+"map.txt", "a") as f:
- box_str = ""
- for p in position_list:
- box_str += str(p[1]) + "," + str(p[0]) + "," + \
- str(p[1]+puzzle_shape[1]) + "," + str(p[0]+puzzle_shape[0]) + \
- "," + str(0) + " "
- # cv2.rectangle(image_np, (p[1], p[0]), (p[1]+puzzle_shape[1], p[0]+puzzle_shape[0]), (255, 0, 0), 2)
- # cv2.imshow("image_np", image_np)
- # cv2.waitKey(0)
- box_str = box_str[:-1]
- f.write(image_np_path + " " + box_str + "\n")
- def generate_data_drag_image(image_shape=(160, 260)):
- bg_paths = glob("../data/base/*")
- data_dir = "../data/drag/"
- with open(data_dir+"map.txt", "w") as f:
- f.write("")
- for i in range(10000):
- if i % 1000 == 0:
- print("Loop", i)
- bg = cv2.imread(random.sample(bg_paths, 1)[0])
- bg = pil_resize(bg, image_shape[0], image_shape[1])
- if random.choice([0, 0, 1]):
- bg = image_enhance_distort(bg)
- if random.choice([0, 0, 1]):
- bg = image_enhance_flip(bg)
- image_np, clip_line = get_drag_image(bg)
- image_np_path = str(i) + ".jpg"
- cv2.imwrite(data_dir+image_np_path, image_np)
- # 生成映射数据
- with open(data_dir+"map.txt", "a") as f:
- f.write(image_np_path + " " + str(clip_line[0][0]) + "\n")
- def generate_data_phrase():
- data_path = os.path.dirname(os.path.abspath(__file__)) + "/../data/phrase/"
- char_path = data_path+"char.txt"
- with open(char_path, "r") as f:
- char_list = f.readlines()
- char_dict = {}
- for i in range(len(char_list)):
- char_dict[char_list[i]] = i
- phrase_list = []
- phrase_path = data_path+"phrase3.txt"
- with open(phrase_path, "r") as f:
- phrase_list += f.readlines()
- phrase_path = data_path+"phrase4.txt"
- with open(phrase_path, "r") as f:
- phrase_list += f.readlines()
- phrase_path = data_path+"phrase5.txt"
- with open(phrase_path, "r") as f:
- phrase_list += f.readlines()
- phrase_set = set(phrase_list)
- map_path = data_path+"map3.txt"
- with open(map_path, "w") as f:
- f.write("")
- data_list = []
- start_time = time.time()
- i = 0
- negative_way_flag = False
- for phrase in phrase_list:
- if i % 500000 == 0:
- with open(map_path, "a") as f:
- f.writelines(data_list)
- data_list = []
- print("Loop", i, len(phrase_list), time.time()-start_time)
- start_time = time.time()
- i += 1
- # 正样本
- index_list = []
- for char in phrase[:-1]:
- index_list.append(char_dict.get(char+"\n"))
- data_list.append(str(index_list) + " 1\n")
- # 负样本
- if negative_way_flag:
- index1 = random.randint(0, len(index_list)-1)
- find_flag = False
- while not find_flag:
- index2 = random.randint(0, len(index_list)-1)
- if index1 != index2:
- find_flag = True
- temp = index_list[index1]
- index_list[index1] = index_list[index2]
- index_list[index2] = temp
- if "".join([char_list[x][:-1] for x in index_list]) + "\n" not in phrase_set:
- data_list.append(str(index_list) + " 0\n")
- else:
- products = list(product(index_list, repeat=len(index_list)))
- random.shuffle(products)
- negative_cnt = 0
- for p in products:
- if negative_cnt >= 2:
- break
- p = list(p)
- if len(set(p)) != len(p):
- continue
- if p != index_list and "".join([char_list[x][:-1] for x in p]) + "\n" not in phrase_set:
- data_list.append(str(p) + " 0\n")
- negative_cnt += 1
- with open(map_path, "a") as f:
- f.writelines(data_list)
- def generate_data_phrase_raw(word_len=5):
- paths = glob("D:/Chinese_corpus/answer/*/*.txt")
- phrase_path = "../data/phrase/phrase" + str(word_len) + "_new.txt"
- triple_list = []
- reg = "[^\u4e00-\u9fa5]"
- start_time = time.time()
- for i in range(len(paths)):
- if i % 1000 == 0:
- with open(phrase_path, "w") as f:
- f.writelines(triple_list)
- print("Loop", i, len(paths), time.time()-start_time)
- start_time = time.time()
- triple_list = []
- with open(paths[i], "rb") as f:
- _b = f.read()
- try:
- text = _b.decode("gbk")
- except:
- try:
- text = _b.decode("gb2312")
- except:
- try:
- text = _b.decode("gb18030")
- except:
- print(chardet.detect(_b), "is None")
- filter_word = ["的"]
- for word in filter_word:
- text = re.sub(word, "#"*len(word), text)
- word_list = jieba.lcut(text, cut_all=False, HMM=True)
- for j in range(1, len(word_list)):
- current = word_list[j]
- current_re = re.search(reg, current)
- last = word_list[j-1]
- last_re = re.search(reg, last)
- if current_re:
- continue
- if len(current) == word_len:
- triple_list.append(current + "\n")
- elif len(current) + len(last) == word_len and not last_re:
- triple_list.append(last+current + "\n")
- triple_list = list(set(triple_list))
- print("len(triple_list)", len(triple_list))
- with open(phrase_path, "w") as f:
- f.writelines(triple_list)
- def generate_data_equation():
- char_dict = {
- 1: ['1', '一'],
- 2: ['2', '二'],
- 3: ['3', '三'],
- 4: ['4', '四'],
- 5: ['5', '五'],
- 6: ['6', '六'],
- 7: ['7', '七'],
- 8: ['8', '八'],
- 9: ['9', '九'],
- 0: ['0', '零'],
- '+': ['+', '加', '加上'],
- '-': ['-', '减', '减去'],
- '*': ["*", "×", 'x', '乘', '乘以'],
- '/': ["/", '除', '÷'],
- '?': ['?', '?'],
- '去': ['去'],
- '上': ['上'],
- '以': ['以'],
- }
- only_op_dict = {
- '+': ['+', '加'],
- '-': ['-', '减'],
- '*': ["*", "×", 'x', '乘'],
- '/': ["/", '除', '÷'],
- }
- file_name_dict = {
- '*': "星",
- "/": "斜",
- "?": "问",
- }
- data_dir = "../data/equation/"
- # 每个字生成多张图片
- start_time = time.time()
- # 0- 1520 4000- 5760
- for i in range(0, 100):
- if i % 20 == 0:
- print("Loop", i, time.time()-start_time)
- start_time = time.time()
- if random.choice([1, 1]):
- # 随机生成算式
- char_list = []
- result = -1
- op = random.choice(['+', '-', '*', '/'])
- while result < 0:
- n1 = random.choice([random.randint(0, 9), random.randint(10, 99)])
- n2 = random.choice([random.randint(0, 9), random.randint(10, 99)])
- if op == '-':
- result = n1 - n2
- else:
- result = 1
- if len(str(n1)) > 1:
- n1 = [x for x in str(n1)]
- else:
- n1 = [random.choice(char_dict[n1])]
- if len(str(n2)) > 1:
- n2 = [x for x in str(n2)]
- else:
- n2 = [random.choice(char_dict[n2])]
- op = random.choice(char_dict[op])
- if len(op) > 1:
- op = [x for x in op]
- else:
- op = [op]
- char_list.extend(n1)
- char_list.extend(op)
- char_list.extend(n2)
- if random.choice([0, 0, 0, 1]):
- char_list.append('=')
- char_list.append(random.choice([''] + char_dict['?']))
- else:
- # 随机生成非算式
- char_list = random.sample(list(char_dict.keys())+list(only_op_dict.keys())*2, random.randint(3, 6))
- char_list = [random.choice(char_dict[x]) for x in char_list]
- if random.choice([0, 0]):
- # 生成背景图
- big_flag = 0
- if random.choice([0, 1]):
- h = random.randint(50, 60)
- w = random.randint(200, 280)
- else:
- h = random.randint(120, 150)
- w = random.randint(400, 680)
- big_flag = 1
- bg = np.full((h, w, 3), random.randint(180, 255), dtype=np.uint8)
- current_w = random.randint(0, int(w/3))
- for char in char_list:
- # 获取单字图片
- if big_flag:
- r = random.randint(80, 120)
- else:
- r = random.randint(35, 45)
- char_image_pil = get_char_image2(char, (r, r), rotate=False)
- # 字体添加到背景图上
- position_h = random.randint(0, bg.shape[0]-r)
- try:
- position_w = random.randint(current_w-15, min(current_w+random.randint(0, 5), bg.shape[1]-r))
- except:
- position_w = current_w
- current_w = position_w+r
- bg = get_image_paste(bg, char_image_pil, position_h, position_w)
- # 增加噪声
- bg = create_noise(bg)
- else:
- fig_size = random.choice([(70, 25), (100, 26), (100, 40)])
- line = (1, 4)
- wavy = (1, 3)
- font_size = random.choice([(18, 25), (35, 40)])
- # 定制
- fig_size = (90, 34)
- line = (1, 4)
- wavy = (0, 0)
- font_size = (38, 50)
- offset_w = (-10, -2)
- bg = generate_data_equation_lsm("".join(char_list),
- fig_size=fig_size,
- font_color=random.choice([(20, 230, 20, 230, 20, 230),
- (70, 100),
- (10, 230, 10, 230, 10, 230),
- ]),
- font_size=font_size,
- rotate=random.choice([0, ]),
- shortline=(1, 5),
- line=line,
- line_width=(1, 3),
- wavy=wavy,
- offset_w=offset_w,
- offset_h=5,
- same_color=0,
- )
- bg = pil2np(bg)
- bg = create_noise_short_lines(bg, random.randint(3, 8))
- # 数据增强
- if random.choice([0, 1]):
- bg = np2pil(bg)
- aug = random.choice([0, 1, 2, 3, 5])
- if aug == 0:
- bg = image_enhance_color(bg)
- elif aug == 1:
- bg = image_enhance_brightness(bg)
- elif aug == 2:
- bg = image_enhance_contrast(bg)
- elif aug == 3:
- bg = image_enhance_sharpness(bg)
- elif aug == 4:
- bg = image_enhance_blur(bg)
- bg = pil2np(bg)
- if aug == 5:
- bg = image_enhance_distort(bg)
- # print("aug", aug)
- # show
- cv2.imshow("bg", bg)
- cv2.waitKey(0)
- # write
- for j in range(len(char_list)):
- if char_list[j] in file_name_dict.keys():
- char_list[j] = file_name_dict.get(char_list[j])
- image_path = str(i) + "_" + "_".join(char_list) + ".jpg"
- cv2.imwrite(data_dir+image_path, bg)
- def generate_data_equation2(batch_size, noise=True):
- char_dict = {
- 1: ['1', '一'],
- 2: ['2', '二'],
- 3: ['3', '三'],
- 4: ['4', '四'],
- 5: ['5', '五'],
- 6: ['6', '六'],
- 7: ['7', '七'],
- 8: ['8', '八'],
- 9: ['9', '九'],
- 0: ['0', '零'],
- '+': ['+', '加', '加上'],
- '-': ['-', '减', '减去'],
- '*': ["*", "×", 'x', '乘', '乘以'],
- '/': ['除', '÷'],
- '?': ['?', '?'],
- '去': ['去'],
- '上': ['上'],
- '以': ['以'],
- }
- only_op_dict = {
- '+': ['+', '加', '加上'],
- '-': ['-', '减', '减去'],
- '*': ["*", "×", 'x', '乘', '乘以'],
- '/': ['除', '÷'],
- }
- # 每个字生成多张图片
- result_list = []
- i = 0
- while i < batch_size:
- # for i in range(0, batch_size):
- simple_flag = random.choice([0, 1])
- # empty_flag = random.choice([0, 0, 0, 0, 0, 1, 0, 0, 0, 0])
- empty_flag = 0
- if random.choice([1, 1]):
- # 随机生成算式
- char_list = []
- # result = -1
- op = random.choice(['+', '-', '*'])
- # while result < 0:
- if not simple_flag:
- n1 = random.choice([random.randint(0, 9), random.randint(10, 99)])
- n2 = random.choice([random.randint(0, 9), random.randint(10, 99)])
- # if op == '-':
- # result = n1 - n2
- # else:
- # result = 1
- if len(str(n1)) > 1:
- n1 = [x for x in str(n1)]
- else:
- n1 = [random.choice(char_dict[n1])]
- if len(str(n2)) > 1:
- n2 = [x for x in str(n2)]
- else:
- n2 = [random.choice(char_dict[n2])]
- else:
- n1 = random.randint(0, 9)
- n2 = random.randint(0, 9)
- n1 = [random.choice(char_dict[n1])]
- n2 = [random.choice(char_dict[n2])]
- op = random.choice(char_dict[op])
- # if simple_flag:
- # op = op[0]
- if len(op) > 1:
- op = [x for x in op]
- else:
- op = [op]
- char_list.extend(n1)
- char_list.extend(op)
- char_list.extend(n2)
- if random.choice([0, 0, 0, 1]):
- char_list.append('=')
- if random.choice([0, 1]) and not simple_flag:
- char_list.append(random.choice(char_dict['?']))
- else:
- # 随机生成非算式
- char_list = random.sample(list(char_dict.keys())+list(only_op_dict.keys())*2, random.randint(3, 6))
- char_list = [random.choice(char_dict[x]) for x in char_list]
- new_char_list = []
- for c in char_list:
- if len(c) > 1:
- new_char_list.extend([x for x in c])
- else:
- new_char_list.append(c)
- char_list = new_char_list
- char_list = char_list[:8]
- if random.choice([0, 0]):
- # 生成背景图
- big_flag = 0
- if random.choice([0, 1]):
- h = random.randint(50, 60)
- w = random.randint(200, 280)
- else:
- h = random.randint(120, 150)
- w = random.randint(400, 680)
- big_flag = 1
- bg = np.full((h, w, 3), random.randint(180, 255), dtype=np.uint8)
- current_w = random.randint(-10, 1)
- for char in char_list:
- # 获取单字图片
- if big_flag:
- r = random.randint(80, 120)
- else:
- r = random.randint(35, 45)
- char_image_pil = get_char_image2(char, (r, r), rotate=False)
- # 字体添加到背景图上
- position_h = random.randint(0, bg.shape[0]-r)
- try:
- position_w = random.randint(current_w-15, min(current_w+random.randint(0, 5), bg.shape[1]-r))
- except:
- position_w = current_w
- current_w = position_w+r
- bg = get_image_paste(bg, char_image_pil, position_h, position_w)
- # 增加噪声
- bg = create_noise(bg)
- else:
- if random.choice([0, 1]):
- fig_size = random.choice([(70, 25), (100, 26), (100, 40)])
- line = (0, 0)
- wavy = (0, 0)
- font_size = random.choice([(18, 25), (35, 40)])
- # offset_w = (-20, 6)
- else:
- # 定制
- fig_size = (90, 34)
- line = (0, 0)
- wavy = (0, 0)
- font_size = (38, 50)
- offset_w = (0, 1)
- if simple_flag:
- fig_size = (random.randint(120, 240), 32)
- font_size = (18, 28)
- if empty_flag:
- bg = np.full((fig_size[1], fig_size[0], 3), fill_value=random.randint(125, 255), dtype=np.uint8)
- char_list = [None] * random.randint(3, 8)
- else:
- bg = generate_data_equation_lsm("".join(char_list),
- fig_size=fig_size,
- font_color=random.choice([(0, 100)]),
- font_size=font_size,
- rotate=random.choice([0, ]),
- shortline=(0, 0),
- line=line,
- line_width=(1, 3),
- wavy=wavy,
- offset_w=offset_w,
- offset_h=5,
- same_color=0,
- simple_flag=simple_flag
- )
- bg = pil2np(bg)
- if get_image_legal(bg) < 0.02:
- continue
- if noise:
- bg = create_noise_point(bg, random.randint(15, 50))
- bg = create_noise_short_lines(bg, random.randint(1, 5))
- # 数据增强
- if random.choice([0, 1]):
- bg = np2pil(bg)
- aug = random.choice([0, 1, 2, 3, 5])
- if aug == 0:
- bg = image_enhance_color(bg)
- elif aug == 1:
- bg = image_enhance_brightness(bg)
- elif aug == 2:
- bg = image_enhance_contrast(bg)
- elif aug == 3:
- bg = image_enhance_sharpness(bg)
- elif aug == 4:
- bg = image_enhance_blur(bg)
- bg = pil2np(bg)
- if aug == 5:
- bg = image_enhance_distort(bg)
- # print("aug", aug)
- # show
- # cv2.imshow("bg", bg)
- # gray = cv2.cvtColor(bg, cv2.COLOR_BGR2GRAY)
- # gray = eight_neighbour(gray, 4)
- # cv2.imshow("gray", gray)
- # cv2.waitKey(0)
- result_list.append([bg, char_list])
- i += 1
- return result_list
- def generate_data_equation_lsm(text, fig_size=(200, 70), fonts=glob("../font/*"), font_color=(10, 100),
- same_color=1, font_size=(25, 35), rotate=0, font_noise=0,
- offset_w=(0, 0), offset_h=0, line=(0, 0), shortline=(0, 0),
- line_width=(0, 1), line_color=(200, 250), point=(0, 500),
- point_color=(150, 250), frame_color=None, wavy=(0, 0),
- bg=(200, 255), simple_flag=0):
- """
- text:验证码文本
- size:验证码图片宽高
- fonts:字体列表,随机选择一个
- font_noise: 字体散点干扰,0不加干扰,1加干扰
- offset_hor: 左右偏移值
- offset_var: 上下偏移值
- fill:字体颜色范围
- rotate:字体旋转角度
- line:干扰线条数范围
- point:干扰点数范围
- wavy:波浪线数范围
- color:干扰线、点 颜色
- bg:背景色范围
- """
- def random_xy(width, height):
- """
- 随机位置函数,返回指定范围随机位置坐标
- 参数:width:图片宽,height:图片高
- """
- x = random.randint(0, width)
- y = random.randint(0, height)
- return x, y
- def random_color(color_tuple):
- """
- 随机颜色函数,返回指定范围随机颜色值
- 参数:start:颜色最低值,end:颜色最高值
- """
- if len(color_tuple)==2:
- rs, re = color_tuple
- gs = bs = rs
- ge = be = re
- else:
- rs, re, gs, ge, bs, be = color_tuple
- red = random.randint(rs, re)
- green = random.randint(gs, ge)
- blue = random.randint(bs, be)
- return (red, green, blue)
- def Asin(x, A=8, w=0.05, b=6, k=40):
- """
- y=Asin(ωx+φ)+k在直角坐标系上的图象
- A——振幅,当物体作轨迹符合正弦曲线的直线往复运动时,其值为行程的1/2。
- (ωx+φ)——相位,反映变量y所处的状态。
- φ——初相,x=0时的相位;反映在坐标系上则为图像的左右移动。
- k——偏距,反映在坐标系上则为图像的上移或下移。
- ω——角速度, 控制正弦周期(单位弧度内震动的次数)。
- """
- return A*math.sin(w*x+b)+k
- def get_wavy_line(w=(0, 100), h=(30, 50)):
- """产生波浪线坐标"""
- import random
- n = 50
- x = 0
- y = random.randint(h[0],h[1])
- flag = random.randint(0,2)
- xy = [(x, y)]
- while x < w[1]:
- temp_y = random.randint(1, 3)
- temp_x = random.randint(5, 10)
- if flag == 0:
- if y + temp_y > h[1]:
- y -= temp_y
- flag = 1
- else:
- y += temp_y
- else:
- if y - temp_y < h[0]:
- y += temp_y
- flag = 0
- else:
- y -= temp_y
- x = x+temp_x if x+temp_x < w[1] else w[1]
- xy.append((x, y))
- return xy
- def get_char_img(char, font, font_color, rotate, bg, font_noise=0):
- """
- 生成单个字符图片,随机颜色加随机旋转
- """
- w, h = draw.textsize(char, font=font)
- im = Image.new('RGBA', (w, h), color=bg)
- ImageDraw.Draw(im).text((0,0), char, font=font, fill=font_color)
- if rotate and char not in ['+', '-', '×']:
- im = im.rotate(random.randint(-rotate, rotate), Image.BILINEAR, expand=1)
- im = im.crop(im.getbbox())
- if font_noise:
- im_draw = ImageDraw.Draw(im)
- # for i in range(random.randint(1,20)):
- for i in range(random.randint(int(w*h*0.01),min(int(w*h*0.05), 5))):
- im_draw.point(xy=(random.randint(0, w), random.randint(0, h)),fill=bg)
- table = []
- for i in range(256):
- table.append(i * 97) # 5.97
- mask = im.convert('L').point(table)
- return (im, mask)
- bg = random_color(bg)
- img = Image.new(mode='RGB', size=fig_size, color=bg)
- draw = ImageDraw.Draw(im=img, mode='RGB')
- font_path = random.choice(fonts)
- # print("font_path", font_path)
- font_size1 = random.randint(font_size[0], font_size[1])
- font = ImageFont.truetype(font_path, size=font_size1) # font=None, size=10, index=0, encoding=""
- rotate = random.randint(0, rotate)
- char_color = random_color(font_color)
- re_s = re.search('(\d+|\?)(\+|-|\*|×)(\d+|\?)(=)(-?\d+|\?)?', text)
- if re_s:
- # print(re_s.group(0))
- char_imgs = []
- char_list = []
- if same_color:
- for i in range(1, 6):
- if re_s.group(i) is not None:
- char_list.append(re_s.group(i))
- char_imgs.append(get_char_img(re_s.group(i), font, font_color=char_color, rotate=rotate, bg=bg, font_noise=font_noise))
- else:
- for i in range(1, 6):
- if re_s.group(i) is not None:
- char_list.append(re_s.group(i))
- char_imgs.append(get_char_img(re_s.group(i), font, font_color=random_color(font_color), rotate=rotate, bg=bg, font_noise=font_noise))
- else:
- if same_color:
- char_imgs = [get_char_img(char, font, font_color=char_color, rotate=rotate, bg=bg, font_noise=font_noise) for char in text]
- else:
- char_imgs = [get_char_img(char, font, font_color=random_color(font_color), rotate=rotate, bg=bg, font_noise=font_noise) for char in text]
- ws = [img[0].size[0] for img in char_imgs]
- hs = [img[0].size[1] for img in char_imgs]
- w = max(sum(ws), fig_size[0])
- h = max(max(hs), fig_size[1])
- if w>fig_size[0] or h>fig_size[1]:
- img = Image.new('RGB', (w+6, h+6), color=bg)
- draw = ImageDraw.Draw(im=img, mode='RGB') # im, mode=None
- w, h = img.size
- fig_size = img.size
- # 短线
- for i in range(random.randint(shortline[0], shortline[1])):
- x0, y0 = random_xy(w, h)
- x1 = x0 + random.randint(2, 5)
- y1 = y0 + random.randint(2, 5)
- draw.line(xy=((x0,y0),(x1,y1)),
- fill=random_color(line_color),
- width=random.randint(line_width[0], line_width[1])) # xy, fill=None, width=0
- if rotate:
- temp_x = random.randint(0, min(570, int((fig_size[0]-sum(ws))/2+1))) # int((fig_size[0]-sum(ws))/5)
- # temp_y = random.randint(int((fig_size[1]-hs[0])/8), int((fig_size[1]-hs[0])/2+1))
- temp_y = random.randint(0, int(fig_size[1]/4))
- # print('len(char_imgs):',len(char_imgs))
- for i in range(len(char_imgs)):
- # tmp_offset = random.randint(offset_w[0], offset_w[1]) if sum(ws)+(len(ws)-1)*offset_w[1]<fig_size[0] else 0
- tmp_offset = random.randint(offset_w[0], offset_w[1])
- if i > 0:
- temp_x = new_x
- temp_y = new_y
- # temp_x = new_x if new_x+ws[i] < fig_size[0] else temp_x+ws[i-1]
- # temp_y = new_y if 0 < new_y and new_y+hs[i] < fig_size[1] else random.randint(0, h-hs[i]+1)
- table = []
- for _i in range(256):
- table.append(_i * 97) # 5.97
- mask = img.crop((temp_x, temp_y, temp_x+char_imgs[i][0].size[0], temp_y+char_imgs[i][0].size[1])).convert('L').point(table)
- img.paste(char_imgs[i][0], box=(temp_x, temp_y), mask=mask)
- new_x = temp_x+ws[i]+tmp_offset
- new_y = temp_y+random.randint(-offset_h, offset_h)
- # 直线
- for i in range(random.randint(line[0], line[1])):
- x0, y0 = random_xy(w, h)
- x1, y1 = random_xy(w, h)
- draw.line(xy=((x0, y0), (x1, y1)),
- fill=random_color(line_color),
- width=random.randint(line_width[0], line_width[1]))
- # 散点
- # for i in range(random.randint(point[0], point[1])):
- # if font_color == point_color:
- # draw.point(xy=(random.randint(0, fig_size[0]), random.randint(0, fig_size[1])),
- # fill=char_color)
- # else:
- # draw.point(xy=(random.randint(0, fig_size[0]), random.randint(0, fig_size[1])),
- # fill=random_color(point_color))
- if random.random() >= 0.5:
- A_ = random.uniform(hs[1]*0.1, hs[1]*0.2)
- w_ = math.pi*4/w # random.uniform(0.04, 0.06)
- b_ = random.random()*math.pi
- k_ = random.uniform(h*0.5, h*0.7)
- # 波浪线
- for _ in range(random.randint(wavy[0],wavy[1])):
- draw.line(xy=[(x, Asin(x, A_, w_, b_, k_)) for x in range(int(w))],
- fill=char_color, width=random.randint(line_width[0], line_width[1]))
- else:
- # 波浪线
- for _ in range(random.randint(wavy[0], wavy[1])):
- draw.line(xy=get_wavy_line(w=(0, w), h=(min(hs)-5, max(hs)+5)),
- fill=char_color, width=random.randint(line_width[0], line_width[1]))
- # 边框
- if frame_color is not None:
- draw.line(xy=[(0, 0), (0, h), (0, 0), (w, 0), (w-1, 0), (w-1, h), (0, h-1), (w-1, h-1)], fill=random_color(frame_color))
- if simple_flag:
- offset_w = (int(1/15*fig_size[0]), int(1/12*fig_size[0]))
- else:
- offset_w = (int(-1/5*font_size1), int(-1/10*font_size1))
- if not rotate:
- # temp_x = random.randint(0, min(70, int((fig_size[0]-sum(ws))/2+1))) #int((fig_size[0]-sum(ws))/5)
- # temp_y = random.randint(int((fig_size[1]-hs[0])/8), int((fig_size[1]-hs[0])/2+1))
- if simple_flag:
- temp_x = random.randint(0, int(fig_size[0]/8))
- else:
- temp_x = random.randint(0, int(fig_size[0]/4))
- temp_y = 0
- for i in range(len(char_imgs)):
- # tmp_offset = random.randint(offset_w[0], offset_w[1]) if sum(ws)+(len(ws)-1)*offset_w[1]<fig_size[0] else 0
- tmp_offset = random.randint(offset_w[0], offset_w[1])
- if i > 0:
- temp_x = new_x
- temp_y = new_y
- # temp_x = new_x if new_x+ws[i] < fig_size[0] else temp_x+ws[i-1]
- # temp_y = new_y if 0 < new_y and new_y+hs[i] < fig_size[1] else temp_y
- if same_color:
- if re_s:
- draw.text((temp_x, temp_y), char_list[i], font=font, fill=char_color)
- else:
- draw.text((temp_x, temp_y), text[i], font=font, fill=char_color)
- else:
- if re_s:
- draw.text((temp_x, temp_y), char_list[i], font=font, fill=random_color(font_color))
- else:
- draw.text((temp_x, temp_y), text[i], font=font, fill=random_color(font_color)) #
- new_x = temp_x+ws[i]+tmp_offset
- # print("new_x", new_x, temp_x, ws[i], tmp_offset)
- new_y = temp_y+random.randint(-offset_h, offset_h)
- # filename = font_name+"_"+str(uuid.uuid1())+"_"+text
- # img.save('/data/python/lsm/gen_captcha/{}.jpg'.format(filename))
- # return img
- return img
- def generate_data_denoise(batch_size):
- char_dict = {
- 1: ['1', '一'],
- 2: ['2', '二'],
- 3: ['3', '三'],
- 4: ['4', '四'],
- 5: ['5', '五'],
- 6: ['6', '六'],
- 7: ['7', '七'],
- 8: ['8', '八'],
- 9: ['9', '九'],
- 0: ['0', '零'],
- '+': ['+', '加', '加上'],
- '-': ['-', '减', '减去'],
- '*': ["*", "×", 'x', '乘', '乘以'],
- '/': ["/", '除', '÷'],
- '?': ['?', '?'],
- '去': ['去'],
- '上': ['上'],
- '以': ['以'],
- }
- only_op_dict = {
- '+': ['+', '加', '加上'],
- '-': ['-', '减', '减去'],
- '*': ["*", "×", 'x', '乘', '乘以'],
- '/': ["/", '除', '÷'],
- }
- # 每个字生成多张图片
- result_list = []
- for i in range(0, batch_size):
- if random.choice([1, 1, 0]):
- # 随机生成算式
- char_list = []
- result = -1
- op = random.choice(['+', '-', '*', '/'])
- while result < 0:
- n1 = random.choice([random.randint(0, 9), random.randint(10, 99)])
- n2 = random.choice([random.randint(0, 9), random.randint(10, 99)])
- if op == '-':
- result = n1 - n2
- else:
- result = 1
- if len(str(n1)) > 1:
- n1 = [x for x in str(n1)]
- else:
- n1 = [random.choice(char_dict[n1])]
- if len(str(n2)) > 1:
- n2 = [x for x in str(n2)]
- else:
- n2 = [random.choice(char_dict[n2])]
- op = random.choice(char_dict[op])
- if len(op) > 1:
- op = [x for x in op]
- else:
- op = [op]
- char_list.extend(n1)
- char_list.extend(op)
- char_list.extend(n2)
- if random.choice([0, 0, 0, 1]):
- char_list.append('=')
- char_list.append(random.choice([''] + char_dict['?']))
- else:
- # 随机生成非算式
- char_list = random.sample(list(char_dict.keys())+list(only_op_dict.keys())*2, random.randint(3, 6))
- char_list = [random.choice(char_dict[x]) for x in char_list]
- new_char_list = []
- for c in char_list:
- if len(c) > 1:
- new_char_list.extend([x for x in c])
- else:
- new_char_list.append(c)
- char_list = new_char_list
- char_list = char_list[:8]
- if random.choice([0, 1]):
- fig_size = random.choice([(70, 25), (100, 26), (100, 40)])
- line = (0, 0)
- wavy = (0, 0)
- font_size = random.choice([(18, 25), (35, 40)])
- # offset_w = (-20, 6)
- else:
- # 定制
- fig_size = (90, 34)
- line = (0, 0)
- wavy = (0, 0)
- font_size = (38, 50)
- offset_w = (-8, -2)
- bg = generate_data_equation_lsm("".join(char_list),
- fig_size=fig_size,
- font_color=random.choice([(20, 230, 20, 230, 20, 230),
- (70, 100),
- (10, 230, 10, 230, 10, 230),
- ]),
- font_size=font_size,
- rotate=random.choice([0, ]),
- shortline=(0, 0),
- line=line,
- line_width=(0, 0),
- wavy=wavy,
- offset_w=offset_w,
- offset_h=5,
- same_color=0,
- )
- bg = pil2np(bg)
- # print(1)
- noise = create_noise_point(bg, random.randint(15, 50))
- # print(2)
- noise = create_noise_short_lines(noise, random.randint(1, 5))
- # 数据增强
- if random.choice([0, 1]):
- bg = np2pil(bg)
- noise = np2pil(noise)
- aug = random.choice([0, 1, 2, 3])
- if aug == 0:
- r = random.uniform(.5, 6.)
- bg = image_enhance_color(bg, _range=(r, r))
- noise = image_enhance_color(noise, _range=(r, r))
- elif aug == 1:
- r = random.uniform(.3, 2.)
- bg = image_enhance_brightness(bg, (r, r))
- noise = image_enhance_brightness(noise, (r, r))
- elif aug == 2:
- r = random.uniform(.3, 6.)
- bg = image_enhance_contrast(bg, (r, r))
- noise = image_enhance_contrast(noise, (r, r))
- elif aug == 3:
- r = random.uniform(4., 8.)
- bg = image_enhance_sharpness(bg, (r, r))
- noise = image_enhance_sharpness(noise, (r, r))
- elif aug == 4:
- bg = image_enhance_blur(bg)
- noise = image_enhance_blur(noise)
- bg = pil2np(bg)
- noise = pil2np(noise)
- if aug == 5:
- bg = image_enhance_distort(bg)
- # print("aug", aug)
- # show
- # cv2.imshow("bg", bg)
- # cv2.imshow("noise", noise)
- # cv2.waitKey(0)
- result_list.append([bg, noise])
- return result_list
- def char_on_image(image_np, char_list, char_shape, image_shape, tip_char_num=1, only_color=None, char_stretch=False):
- position_list = []
- for char in char_list:
- # 获取单字图片
- char_image_pil = get_char_image2(char, char_shape, char_color=only_color)
- image_np = pil_resize(image_np, image_shape[0], image_shape[1])
- # h, w
- fg_w, fg_h = char_image_pil.size[:2]
- bg_h, bg_w = image_np.shape[:2]
- # 字体放置的位置,且位置不重叠
- find_flag = 0
- while not find_flag:
- position_h = random.randint(0, bg_h-fg_h)
- position_w = random.randint(0, bg_w-fg_w)
- if len(position_list) < 1:
- find_flag = 1
- break
- for p in position_list:
- if get_iou(position_w, position_h, position_w+fg_w, position_h+fg_h,
- p[1], p[0], p[1]+fg_w, p[0]+fg_h) > 0:
- find_flag = 0
- break
- else:
- find_flag = 1
- position_list.append([position_h, position_w])
- # 字体添加到背景图上
- # image_np = get_image_roi(image_np, char_image_np, position_h, position_w)
- image_np = get_image_paste(image_np, char_image_pil, position_h, position_w, stretch=char_stretch)
- # 生成提示图片
- image_list = []
- for char in char_list[:tip_char_num]:
- char_image_pil = get_char_image2(char, char_shape, rotate=True, is_tips=True)
- char_image_np = pil2np_a(char_image_pil)
- # char_image_np = pil_resize(char_image_np, char_shape[0], char_shape[1])
- image_list.append(char_image_np)
- if image_list:
- tips_image_np = np.concatenate(image_list, axis=1)
- # 加干扰
- tips_image_np = create_noise(tips_image_np)
- # 切割
- image_list = []
- for i in range(tip_char_num):
- image_list.append(tips_image_np[:, i*char_shape[1]:(i+1)*char_shape[1], :])
- return image_np, position_list, image_list
- def get_char_image(char, char_shape, rotate=True, bg_color=(0, 0, 0, 0)):
- # 创建空图
- image_pil = Image.new('RGBA', (80, 80), bg_color)
- # 空图上写字
- # font_size = 35 # (40, 40)
- font_size = 75 # (80, 80)
- font_type_list = glob("../font/*")
- font_type = random.sample(font_type_list, 1)[0]
- font_config = ImageFont.truetype(font_type, int(font_size))
- dr = ImageDraw.Draw(image_pil)
- fill_color = random_color()
- fill_color = (fill_color[0], fill_color[1], fill_color[2])
- dr.text((3, -6), char, font=font_config, fill=fill_color)
- if rotate:
- if random.choice([0, 1]):
- angle = random.randint(0, 80)
- else:
- angle = random.randint(280, 360)
- image_pil = image_pil.rotate(angle, expand=False, fillcolor=bg_color)
- # image_pil.show("1")
- image_pil = image_pil.resize(char_shape)
- # cv2.imshow("get_char_image", pil2np(image_pil))
- # cv2.waitKey(0)
- return image_pil
- def get_char_image2(char, char_shape, rotate=True, char_color=None, is_tips=False):
- # plt.ion()
- # plt.ioff()
- font_files = ['simhei', 'simsun', "Microsoft YaHei"]
- # 修改matplotlib中font_manager.py中defaultFamily
- plt.rcParams["font.sans-serif"] = ["Microsoft YaHei"]
- font_colors = [
- 'aqua',
- 'aquamarine',
- 'bisque',
- 'black',
- 'blue',
- 'blueviolet',
- 'brown',
- 'burlywood',
- 'cadetblue',
- 'chartreuse',
- 'chocolate',
- 'coral',
- 'cornflowerblue',
- 'cornsilk',
- 'crimson',
- 'cyan',
- 'darkblue',
- 'darkcyan',
- 'darkgoldenrod',
- 'darkgray',
- 'darkgreen',
- 'darkmagenta',
- 'darkolivegreen',
- 'darkorange',
- 'darkorchid',
- 'darkred',
- 'darksalmon',
- 'darkseagreen',
- 'darkslategray',
- 'darkturquoise',
- 'darkviolet',
- 'deeppink',
- 'deepskyblue',
- 'dimgray',
- 'dodgerblue',
- 'firebrick',
- 'forestgreen',
- 'fuchsia',
- 'gold',
- 'goldenrod',
- 'gray',
- 'green',
- 'greenyellow',
- 'hotpink',
- 'indianred',
- 'indigo',
- 'lawngreen',
- 'lightseagreen',
- 'lightskyblue',
- 'lightslategray',
- 'lightsteelblue',
- 'lime',
- 'limegreen',
- 'magenta',
- 'maroon',
- 'mediumaquamarine',
- 'mediumblue',
- 'mediumorchid',
- 'mediumpurple',
- 'mediumseagreen',
- 'mediumslateblue',
- 'mediumspringgreen',
- 'mediumturquoise',
- 'mediumvioletred',
- 'midnightblue',
- 'navy',
- 'olive',
- 'olivedrab',
- 'orange',
- 'orangered',
- 'orchid',
- 'palegoldenrod',
- 'palegreen',
- 'paleturquoise',
- 'palevioletred',
- 'peachpuff',
- 'peru',
- 'pink',
- 'plum',
- 'purple',
- 'red',
- 'rosybrown',
- 'royalblue',
- 'saddlebrown',
- 'salmon',
- 'sandybrown',
- 'seagreen',
- 'sienna',
- 'silver',
- 'skyblue',
- 'slateblue',
- 'slategray',
- 'springgreen',
- 'steelblue',
- 'tan',
- 'teal',
- 'tomato',
- 'turquoise',
- 'violet',
- 'yellow',
- 'yellowgreen'
- ]
- line_width = random.choice([random.randint(1, 6), 0, 0])
- # print("char_color", char_color)
- if char_color == (0, 0, 0):
- font_color = 'black'
- line_color = 'black'
- font_weight = random.randint(1, 300)
- else:
- if random.choice([1, 1, 1]) or is_tips:
- font_weight = random.randint(200, 500)
- font_color = random.sample(font_colors, 1)[0]
- line_color = random.sample(font_colors, 1)[0]
- else:
- font_color = 'black'
- line_color = 'black'
- font_weight = random.randint(40, 100)
- if rotate:
- rotation = random.choice([random.randint(0, 70), random.randint(290, 360), 0])
- else:
- rotation = 0
- # print("font_color", font_color, char)
- # 写字
- plt.rcParams['font.weight'] = font_weight
- text = fig.text(0.5, 0.5, char, ha="center", va="center", size=55, color=font_color, rotation=rotation)
- # 字体有边框
- text.set_path_effects([path_effects.Stroke(linewidth=line_width, foreground=line_color), path_effects.Normal()])
- # 更新图
- # plt
- plt.pause(0.0000001)
- # fig = plt.gcf()
- # plt.plot()
- # plt.show()
- # plt.legend()
- # 申请缓存
- _buffer = io.BytesIO()
- fig.savefig(_buffer, format='png')
- _buffer.seek(0)
- image_pil = Image.open(_buffer)
- plt.cla()
- plt.clf()
- # plt.close('all')
- gc.collect()
- # 大小限制
- image_pil = image_pil.crop((105, 0, image_pil.size[0]-105, image_pil.size[1]-12))
- image_pil = image_pil.resize(char_shape)
- # image_pil.show()
- return image_pil
- def get_tips_image(tips_image_list, char_shape, image_shape):
- new_list = []
- for img in tips_image_list:
- # if random.choice([0, 0, 1]):
- # angle = random.choice([random.randint(0, 70), random.randint(290, 360)])
- # img = pil_rotate(img, angle, (255, 255, 255))
- new_list.append(img)
- tips_image_np = np.concatenate(new_list, axis=1)
- new_image = np.full((image_shape[0], image_shape[1], 3), 0, np.uint8)
- new_image[:tips_image_np.shape[0], :tips_image_np.shape[1], :] = tips_image_np
- position_list = []
- for i in range(len(new_list)):
- h = 0
- w = i*char_shape[1]
- position_list.append([h, w])
- if random.choice([0, 1]):
- new_image[np.all(new_image==0, axis=2)] = 255
- return new_image, position_list
- def get_image_roi(image_bg, image_fg, roi_h, roi_w):
- # h, w
- fg_h, fg_w = image_fg.shape[:2]
- bg_h, bg_w = image_bg.shape[:2]
- # roi取值范围
- roi = image_bg[roi_h:roi_h+fg_h, roi_w:roi_w+fg_w]
- # 获取bg中非fg字体部分的掩码,相当于排除fg的字体部分,只保留bg的除fg字体外的部分
- img_fg_gray = cv2.cvtColor(image_fg, cv2.COLOR_BGR2GRAY)
- ret, mask = cv2.threshold(img_fg_gray, 0, 255, cv2.THRESH_OTSU)
- bg_roi = cv2.bitwise_and(roi, roi, mask=mask)
- # 获取fg中字体部分的掩码,相当于排除fg中的白色背景,只保留fg的字体部分
- mask_inv = cv2.bitwise_not(mask)
- fg_roi = cv2.bitwise_and(image_fg, image_fg, mask=mask_inv)
- # 膨胀腐蚀去掉白色颗粒
- # kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
- # fg_roi = cv2.erode(fg_roi, kernel)
- # fg_roi = cv2.dilate(fg_roi, kernel)
- # bg的除字体外背景部分 + fg的字体部分
- image_roi = cv2.add(bg_roi, fg_roi)
- # 将roi部分放回原bg
- image_bg[roi_h:roi_h+fg_h, roi_w:roi_w+fg_w, :] = image_roi
- # cv2.imshow("image_fg", image_fg)
- # cv2.imshow("get_image_roi", image_bg)
- # cv2.waitKey(0)
- return image_bg
- def get_image_paste(image_bg, image_fg, roi_h, roi_w, stretch=False):
- fg_h, fg_w = image_fg.size[:2]
- if stretch:
- if random.choice([1, 1, 1]):
- threshold = random.randint(7, 13)
- # print(threshold)
- if random.choice([0, 1]):
- image_fg = image_fg.resize((fg_w, fg_h-threshold), Image.BICUBIC)
- else:
- image_fg = image_fg.resize((fg_w-threshold, fg_h), Image.BICUBIC)
- image_bg = cv2.cvtColor(image_bg, cv2.COLOR_BGR2BGRA)
- image_bg = np2pil_a(image_bg)
- # image_fg = np2pil_a(image_fg)
- image_bg.paste(image_fg, (roi_w, roi_h), image_fg)
- image_bg = pil2np(image_bg)
- # cv2.imshow("get_image_paste", image_bg)
- # cv2.waitKey(0)
- return image_bg
- def random_color(dims=3):
- color = [0]*dims
- find_flag = 0
- while not find_flag:
- for dim in range(dims):
- color[dim] = random.randint(0, 255)
- if color[dim] <= 125:
- find_flag = 1
- # RGB
- # color_list = [
- # [207, 91, 85],
- # [0, 201, 88],
- # [117, 74, 57],
- # [210, 210, 27],
- # [160, 157, 152],
- # [181, 210, 210],
- # [27, 112, 107],
- # [87, 26, 44],
- # [115, 19, 20],
- # [161, 210, 68],
- # [210, 108, 12],
- # [112, 9, 142],
- # [50, 41, 84],
- # [72, 52, 210],
- # [210, 177, 89],
- # [148, 200, 89],
- # [173, 116, 109],
- # [185, 185, 210],
- # [181, 7, 210],
- # [80, 210, 30],
- # [65, 72, 98],
- # [210, 123, 109],
- # [19, 64, 95],
- # [128, 21, 210],
- # [129, 137, 60]
- # ]
- # color = random.sample(color_list, 1)[0]
- return tuple(color)
- def create_noise(image_np):
- ic = ImageCaptcha()
- image_pil = np2pil(image_np)
- if random.choice([0, 1]):
- draw = ImageDraw.Draw(image_pil)
- for i in range(0, random.randint(100, 1000)):
- xy = (random.randrange(0, image_pil.size[0]), random.randrange(0, image_pil.size[1]))
- r = random.randint(0, 200)
- fill = (r, r, r)
- draw.point(xy, fill=fill)
- del draw
- for i in range(random.randint(2, 4)):
- image_pil = ic.create_noise_curve(image_pil, random_color())
- image_pil = ic.create_noise_dots(image_pil, random_color(), random.randint(3, 6), random.randint(30, 60))
- image_np = pil2np(image_pil)
- return image_np
- def create_noise_short_lines(image_np, line_num):
- image_np = copy.deepcopy(image_np)
- h, w = image_np.shape[:2]
- len_range = (5, int(w/2))
- for i in range(line_num):
- p1 = (random.randint(0, w), random.randint(0, h))
- find_flag = False
- while not find_flag:
- p2 = (random.randint(0, w), random.randint(0, h))
- p_len = math.sqrt(math.pow((p1[0]-p2[0]), 2) + math.pow((p1[1]-p2[1]), 2))
- if p_len < len_range[0] or p_len > len_range[1]:
- find_flag = True
- color = (random.randint(50, 200), random.randint(50, 200), random.randint(50, 200))
- thickness = random.randint(1, 1)
- cv2.line(image_np, p1, p2, color, thickness=thickness)
- return image_np
- def create_noise_point(image_np, point_num):
- image_np = copy.deepcopy(image_np)
- h, w = image_np.shape[:2]
- for i in range(point_num):
- p_center = (random.randint(1, w-2), random.randint(1, h-2))
- p_list = [(p_center[0]-1, p_center[1]),
- (p_center[0]+1, p_center[1]),
- p_center,
- (p_center[0], p_center[1]-1),
- (p_center[0], p_center[1]+1)]
- p_color = [random.randint(0, 130), random.randint(0, 130), random.randint(0, 130)]
- for p in p_list:
- if random.choice([0, 1]):
- if random.choice([0, 1]):
- _index = p_color.index(max(p_color))
- p_color[_index] = random.randint(200, 255)
- image_np[p[1], p[0], 0] = p_color[0]
- image_np[p[1], p[0], 1] = p_color[1]
- image_np[p[1], p[0], 2] = p_color[2]
- return image_np
- def get_iou(x1, y1, x2, y2, a1, b1, a2, b2):
- # 相交区域左上角横坐标
- ax = max(x1, a1)
- # 相交区域左上角纵坐标
- ay = max(y1, b1)
- # 相交区域右下角横坐标
- bx = min(x2, a2)
- # 相交区域右下角纵坐标
- by = min(y2, b2)
- area_n = (x2 - x1) * (y2 - y1)
- area_m = (a2 - a1) * (b2 - b1)
- w = max(0, bx - ax)
- h = max(0, by - ay)
- area_x = w * h
- return area_x / (area_n + area_m - area_x)
- def preprocess_true_boxes(true_boxes, input_shape, anchors, num_classes):
- """Preprocess true boxes to training input format
- Parameters
- ----------
- true_boxes: array, shape=(m, T, 5)
- Absolute x_min, y_min, x_max, y_max, class_id relative to input_shape.
- input_shape: array-like, hw, multiples of 32
- anchors: array, shape=(N, 2), wh
- num_classes: integer
- Returns
- -------
- y_true: list of array, shape like yolo_outputs, xywh are reletive value
- """
- # print(true_boxes.shape)
- # print(true_boxes[..., 4])
- # print(num_classes)
- try:
- sss = (true_boxes[..., 4] < num_classes).all()
- except:
- print(true_boxes[..., 4])
- assert (true_boxes[..., 4] < num_classes).all(), 'class id must be less than num_classes'
- # default setting
- num_layers = len(anchors)//3
- anchor_mask = [[6, 7, 8], [3, 4, 5], [0, 1, 2]] if num_layers == 3 else [[3, 4, 5], [1, 2, 3]]
- true_boxes = np.array(true_boxes, dtype='float32')
- input_shape = np.array(input_shape, dtype='int32')
- boxes_xy = (true_boxes[..., 0:2] + true_boxes[..., 2:4]) // 2
- boxes_wh = true_boxes[..., 2:4] - true_boxes[..., 0:2]
- true_boxes[..., 0:2] = boxes_xy/input_shape[::-1]
- true_boxes[..., 2:4] = boxes_wh/input_shape[::-1]
- m = true_boxes.shape[0]
- grid_shapes = [input_shape//{0: 32, 1: 16, 2: 8}[l] for l in range(num_layers)]
- y_true = [np.zeros((m, grid_shapes[l][0], grid_shapes[l][1],len(anchor_mask[l]), 5+num_classes),
- dtype='float32') for l in range(num_layers)]
- # Expand dim to apply broadcasting.
- anchors = np.expand_dims(anchors, 0)
- anchor_maxes = anchors / 2.
- anchor_mins = -anchor_maxes
- valid_mask = boxes_wh[..., 0] > 0
- for b in range(m):
- # Discard zero rows.
- wh = boxes_wh[b, valid_mask[b]]
- if len(wh) == 0:
- continue
- # Expand dim to apply broadcasting.
- wh = np.expand_dims(wh, -2)
- box_maxes = wh / 2.
- box_mins = -box_maxes
- intersect_mins = np.maximum(box_mins, anchor_mins)
- intersect_maxes = np.minimum(box_maxes, anchor_maxes)
- intersect_wh = np.maximum(intersect_maxes - intersect_mins, 0.)
- intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1]
- box_area = wh[..., 0] * wh[..., 1]
- anchor_area = anchors[..., 0] * anchors[..., 1]
- iou = intersect_area / (box_area + anchor_area - intersect_area)
- # Find best anchor for each true box
- best_anchor = np.argmax(iou, axis=-1)
- for t, n in enumerate(best_anchor):
- for l in range(num_layers):
- if n in anchor_mask[l]:
- i = np.floor(true_boxes[b,t,0]*grid_shapes[l][1]).astype('int32')
- j = np.floor(true_boxes[b,t,1]*grid_shapes[l][0]).astype('int32')
- k = anchor_mask[l].index(n)
- c = true_boxes[b, t, 4].astype('int32')
- y_true[l][b, j, i, k, 0:4] = true_boxes[b, t, 0:4]
- y_true[l][b, j, i, k, 4] = 1
- y_true[l][b, j, i, k, 5+c] = 1
- return y_true
- def get_puzzle(shape=(80, 80)):
- # 创建空图
- image_pil = Image.new('RGBA', (shape[1], shape[0]), (255, 255, 255, 0))
- draw = ImageDraw.Draw(image_pil)
- # 居中创建矩形
- rec_shape = (40, 40)
- left_up_point = [int((shape[0]-rec_shape[0])/2), int((shape[1]-rec_shape[1])/2)]
- right_down_point = [left_up_point[0]+rec_shape[0], left_up_point[1]+rec_shape[1]]
- # 透明度
- flag = random.choice([0, 1])
- if flag:
- alpha = random.randint(100, 150)
- else:
- alpha = random.randint(160, 255)
- # 背景色
- if flag:
- r = random.randint(0, 30)
- else:
- r = random.randint(100, 180)
- # r = random.randint(0, 255)
- fill_color = (r, r, r, alpha)
- # 边缘色
- if random.choice([0, 1, 1]):
- if flag:
- r = random.randint(140, 170)
- else:
- r = random.randint(70, 100)
- outline_color = (r, r, r)
- else:
- outline_color = (fill_color[0], fill_color[1], fill_color[2])
- draw.rectangle((left_up_point[1],
- left_up_point[0],
- left_up_point[1]+rec_shape[1],
- left_up_point[0]+rec_shape[0]),
- fill=fill_color,
- outline=outline_color)
- # 拼图的圆或半圆
- radius = random.randint(int(rec_shape[0] / 3 / 2), int(rec_shape[0] / 3 / 1.2))
- center_list = [[left_up_point[1], int((right_down_point[0]+left_up_point[0])/2), 1],
- [right_down_point[1], int((right_down_point[0]+left_up_point[0])/2), 1],
- [int((right_down_point[1]+left_up_point[1])/2), left_up_point[0], 0],
- [int((right_down_point[1]+left_up_point[1])/2), right_down_point[0], 0]
- ]
- circle_num = random.randint(1, 4)
- # print("circle_num", circle_num)
- center_list = random.sample(center_list, circle_num)
- min_w, min_h = left_up_point[1], left_up_point[0]
- max_w, max_h = right_down_point[1], right_down_point[0]
- for center in center_list:
- w, h = center[:2]
- is_width = center[2]
- # 判断长宽
- into_ratio = random.randint(int(1/2*radius), int(3/4*radius))
- if is_width:
- # 挑选圆是凸还是凹进去
- if random.choice([0, 1]):
- center = (center[0]+into_ratio, center[1])
- else:
- center = (center[0]-into_ratio, center[1])
- else:
- if random.choice([0, 1]):
- center = (center[0], center[1]+into_ratio)
- else:
- center = (center[0], center[1]-into_ratio)
- # 判断透明度
- color = fill_color
- if is_width:
- if left_up_point[1] <= center[0] <= right_down_point[1]:
- color = (0, 0, 0, 0)
- else:
- if left_up_point[0] <= center[1] <= right_down_point[0]:
- color = (0, 0, 0, 0)
- # print("center, color, alpha", center, color, alpha)
- draw.ellipse([(center[0]-radius, center[1]-radius),
- (center[0]+radius, center[1]+radius)],
- fill=color,
- outline=outline_color)
- # 修补内部圆的边缘颜色
- if color[3] == alpha:
- if is_width:
- if center[0] < w:
- draw.rectangle((w,
- h-radius,
- center[0]+radius,
- center[1]+radius),
- fill=fill_color)
- else:
- draw.rectangle((center[0]-radius,
- center[1]-radius,
- w,
- h+radius),
- fill=fill_color)
- else:
- if center[1] < h:
- draw.rectangle((w-radius,
- h,
- center[0]+radius,
- center[1]+radius),
- fill=fill_color)
- else:
- draw.rectangle((center[0]-radius,
- center[1]-radius,
- w+radius,
- h),
- fill=fill_color)
- # 修补外部圆的边缘颜色
- else:
- if is_width:
- if center[0] > w:
- draw.rectangle((center[0]-radius,
- center[1]-radius,
- w,
- h+radius),
- fill=(0, 0, 0, 0))
- else:
- draw.rectangle((w,
- h-radius,
- center[0]+radius,
- center[1]+radius),
- fill=(0, 0, 0, 0))
- else:
- if center[1] > h:
- draw.rectangle((center[0]-radius,
- center[1]-radius,
- w+radius,
- h),
- fill=(0, 0, 0, 0))
- else:
- draw.rectangle((w-radius,
- h,
- center[0]+radius,
- center[1]+radius),
- fill=(0, 0, 0, 0))
- # 新增面积
- if color[3] == alpha:
- if center[0]-radius <= min_w:
- min_w = center[0]-radius
- if center[0]+radius >= max_w:
- max_w = center[0]+radius
- if center[1]-radius <= min_h:
- min_h = center[1]-radius
- if center[1]+radius >= max_h:
- max_h = center[1]+radius
- image_pil = image_pil.crop([min_w, min_h, max_w+1, max_h+1])
- # image_pil.show("2")
- return image_pil
- def puzzle_on_image(image_np, puzzle_shape, image_shape):
- position_list = []
- # 获取拼图图片
- puzzle_image_pil = get_puzzle()
- puzzle_image_pil = puzzle_image_pil.resize(puzzle_shape)
- image_np = pil_resize(image_np, image_shape[0], image_shape[1])
- # h, w
- fg_w, fg_h = puzzle_image_pil.size[:2]
- bg_h, bg_w = image_np.shape[:2]
- # 拼图放置的位置
- position_h = random.randint(0, bg_h-fg_h)
- position_w = random.randint(0, bg_w-fg_w)
- position_list.append([position_h, position_w])
- # for p in position_list:
- # cv2.rectangle(image_np, (p[1], p[0]),
- # (p[1]+puzzle_shape[1], p[0]+puzzle_shape[0]),
- # (0, 0, 255), 1)
- # 拼图添加到背景图上
- image_np = get_image_paste(image_np, puzzle_image_pil, position_h, position_w)
- # cv2.imshow("puzzle_on_image", image_np)
- # cv2.waitKey(0)
- return image_np, position_list
- def get_drag_image(image_np):
- h, w = image_np.shape[:2]
- # 取一定高度图片
- clip_h = random.randint(int(1/4*h), int(3/4*h))
- image_clip = image_np[:clip_h, ...]
- # 将图片在一定宽度截断,重新拼接
- clip_w = random.randint(int(1/6*w), int(5/6*w))
- image_w1 = image_clip[:, :clip_w, ...]
- image_w2 = image_clip[:, clip_w:, ...]
- image_new = np.concatenate([image_w2, image_w1], axis=1)
- # 分割线
- clip_line = [(image_w2.shape[1], 0), (image_w2.shape[1], clip_h)]
- # show
- # print(clip_line)
- # cv2.line(image_new, clip_line[0], clip_line[1], (0, 0, 255), 2)
- # cv2.imshow("get_drag_image", image_new)
- # cv2.waitKey(0)
- return image_new, clip_line
- def get_real_data_puzzle(shape=(160, 256)):
- paths = glob("../data/detect2_real/*")
- i = 10000
- for p in paths:
- image = cv2.imread(p)
- image = pil_resize(image, shape[0], shape[1])
- cv2.imwrite("../data/detect2_real/"+str(i)+".jpg", image)
- i += 1
- image = image_enhance_distort(image)
- cv2.imwrite("../data/detect2_real/"+str(i)+".jpg", image)
- i += 1
- image = image_enhance_flip(image)
- cv2.imwrite("../data/detect2_real/"+str(i)+".jpg", image)
- i += 1
- def read_label_puzzle():
- paths = glob("../data/detect2_real/*.json")
- map_path = "../data/detect2_real/map.txt"
- with open(map_path, "a") as f:
- for p in paths:
- with open(p, "r") as fp:
- _dict = json.loads(fp.read())
- points = _dict.get("shapes")[0].get("points")
- image_path = _dict.get("imagePath")
- ps = [str(int(points[0][0])), str(int(points[0][1])),
- str(int(points[1][0])), str(int(points[1][1]))]
- p_str = ",".join(ps)
- f.write(image_path + " " + p_str + ",0" + "\n")
- def fix_map_txt():
- path = "../data/map.txt"
- with open(path, "r") as f:
- _list = f.readlines()
- with open("../data/map_new.txt", "w") as f:
- new_list = []
- for line in _list:
- ss = line.split(" ")
- ps = ss[-1][:-1].split(",")[:-1]
- if random.choice([0, 1, 1, 1]):
- pix = random.choice([1, 2, 2, 3, 3, 4, 4])
- for i in range(len(ps)):
- if i < 2:
- ps[i] = str(int(ps[i]) - pix)
- else:
- ps[i] = str(int(ps[i]) + pix)
- new_line = ss[0] + " " + ",".join(ps) + ",0\n"
- new_list.append(new_line)
- print("line", line)
- print("new_line", new_line)
- f.writelines(new_list)
- def get_char_map():
- path = "../data/phrase/phrase3.txt"
- with open(path, "r") as f:
- _list = f.readlines()
- path = "../data/phrase/phrase4.txt"
- with open(path, "r") as f:
- _list += f.readlines()
- path = "../data/phrase/phrase5.txt"
- with open(path, "r") as f:
- _list += f.readlines()
- _str = "".join(_list)
- _str = re.sub("\n", "", _str)
- _list = list(set([x+"\n" for x in _str]))
- _list.sort(key=lambda x: x)
- with open("../data/phrase/char.txt", "w") as f:
- f.writelines(_list)
- def resize_base_image():
- paths = glob("../data/base/*")
- for p in paths:
- _img = cv2.imread(p)
- h, w = _img.shape[:2]
- if h > 600 and h > w:
- best_h = 600
- best_w = int(best_h * w / h)
- _img = pil_resize(_img, best_h, best_w)
- cv2.imwrite(p, _img)
- elif w > 600 and w > h:
- best_w = 600
- best_h = int(best_w * h / w)
- _img = pil_resize(_img, best_h, best_w)
- cv2.imwrite(p, _img)
- def image_enhance_distort(image_np, hue=.3, sat=3., val=3.):
- """
- 数据增强:色彩扰动
- :return:
- """
- def rand(a=0, b=1):
- return np.random.rand()*(b-a) + a
- # cv2.imshow("distort_image1", image_np)
- hue = rand(-hue, 0)
- sat = rand(1, sat) if rand() < .2 else 1/rand(1, sat)
- val = rand(1, val) if rand() < .2 else 1/rand(1, val)
- # print(hue, sat, val)
- image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
- x = rgb_to_hsv(image_np/255.)
- x[..., 0] += hue
- x[..., 0][x[..., 0] > 1] -= 1
- x[..., 0][x[..., 0] < 0] += 1
- x[..., 1] *= sat
- x[..., 2] *= val
- x[x > 1] = 1
- x[x < 0] = 0
- image_np = hsv_to_rgb(x)
- image_np = cv2.cvtColor(np.uint8(image_np*255), cv2.COLOR_RGB2BGR)
- # cv2.imshow("distort_image2", image_np)
- # cv2.waitKey(0)
- return image_np
- def image_enhance_flip(image_np):
- """
- 数据增强:翻转
- :param image_np:
- :return:
- """
- # cv2.imshow("flip_image1", image_np)
- if random.choice([0, 1]):
- # 水平翻转
- image_np = cv2.flip(image_np, 1)
- else:
- # 垂直翻转
- image_np = cv2.flip(image_np, 0)
- # cv2.imshow("flip_image2", image_np)
- # cv2.waitKey(0)
- return image_np
- def image_enhance_blur(image_pil, blur_mode=2):
- """
- 数据增强:模糊
- :param image_pil:
- :param blur_mode: 1:标准模糊 2:高斯模糊 3:均值模糊
- :return:
- """
- # image_pil.show()
- if blur_mode == 1:
- image_pil = image_pil.filter(ImageFilter.BLUR)
- elif blur_mode == 2:
- image_pil = image_pil.filter(ImageFilter.GaussianBlur(1))
- elif blur_mode == 3:
- image_pil = image_pil.filter(ImageFilter.BoxBlur(10))
- # image_pil.show()
- return image_pil
- def image_enhance_color(image_pil, _range=(.5, 6.)):
- """
- 数据增强:色彩饱和度
- :param image_pil:
- :param _range:
- :return:
- """
- # image_pil.show()
- enh_col = ImageEnhance.Color(image_pil)
- color = random.uniform(_range[0], _range[1])
- # print("color", color)
- image_colored = enh_col.enhance(color)
- # image_colored.show()
- return image_colored
- def image_enhance_brightness(image_pil, _range=(.3, 2.)):
- """
- 数据增强:亮度
- :param image_pil:
- :param _range:
- :return:
- """
- enh_bri = ImageEnhance.Brightness(image_pil)
- brightness = random.uniform(_range[0], _range[1])
- # print("brightness", brightness)
- image_brightened = enh_bri.enhance(brightness)
- return image_brightened
- def image_enhance_contrast(image_pil, _range=(.3, 6.)):
- """
- 数据增强:对比度
- :param image_pil:
- :param _range:
- :return:
- """
- # image_pil.show()
- enh_con = ImageEnhance.Contrast(image_pil)
- contrast = random.uniform(_range[0], _range[1])
- # print("contrast", contrast)
- image_contrasted = enh_con.enhance(contrast)
- # image_contrasted.show()
- return image_contrasted
- def image_enhance_sharpness(image_pil, _range=(4., 8.)):
- """
- 数据增强:锐度
- :param image_pil:
- :param _range:
- :return:
- """
- # image_pil.show()
- enh_sha = ImageEnhance.Sharpness(image_pil)
- sharpness = random.uniform(_range[0], _range[1])
- # print("sharpness", sharpness)
- image_sharped = enh_sha.enhance(sharpness)
- # image_sharped.show()
- return image_sharped
- def image_enhance_wrap(image_pil, dx_factor=0.3, dy_factor=0.3):
- """图像扭曲"""
- width, height = image_pil.size
- dx = width * dx_factor
- dy = height * dy_factor
- x1 = int(random.uniform(-dx, dx))
- y1 = int(random.uniform(-dy, dy))
- x2 = int(random.uniform(-dx, dx))
- y2 = int(random.uniform(-dy, dy))
- warp_image = Image.new('RGB', (width + abs(x1) + abs(x2), height + abs(y1) + abs(y2)))
- warp_image.paste(image_pil, (abs(x1), abs(y1)))
- width2, height2 = warp_image.size
- warp_image = warp_image.transform((width, height),
- Image.QUAD,
- (x1, y1, -x1, height2 - y2, width2 + x2, height2 + y2, width2 - x2, -y1))
- warp_image.show()
- return warp_image
- def eight_neighbour(image_np, k=4):
- """
- 8邻域降噪
- :return:
- """
- def calculate_noise_count(img_obj, w, h):
- """
- 计算邻域非白色的个数
- """
- count = 0
- width, height = img_obj.shape
- for _w_ in [w - 1, w, w + 1]:
- for _h_ in [h - 1, h, h + 1]:
- if _w_ > width - 1:
- continue
- if _h_ > height - 1:
- continue
- if _w_ == w and _h_ == h:
- continue
- if img_obj[_w_, _h_] < 230: # 二值化的图片设置为255
- count += 1
- return count
- w, h = image_np.shape
- for _w in range(w):
- for _h in range(h):
- if _w == 0 or _h == 0:
- image_np[_w, _h] = 255
- continue
- # 计算邻域pixel值小于255的个数
- pixel = image_np[_w, _h]
- if pixel == 255:
- continue
- if calculate_noise_count(image_np, _w, _h) < k:
- image_np[_w, _h] = 255
- return image_np
- def connected_component(image_np):
- # 黑白转化
- image_np = 255 - image_np
- # 二值化
- ret, image_np = cv2.threshold(image_np, 60, 255, cv2.THRESH_BINARY)
- # 膨胀
- kernel = np.ones((2, 2), np.uint8)
- image_np = cv2.dilate(image_np, kernel, iterations=1)
- kernel = np.ones((2, 2), np.uint8)
- image_np = cv2.erode(image_np, kernel, iterations=1)
- return image_np
- w, h = binary_img.shape
- color = []
- color.append((0, 0, 0))
- img_color = np.zeros((w, h, 3), dtype=np.uint8)
- retval, labels, stats, centroids = cv2.connectedComponentsWithStats(binary_img)
- for num in range(1, retval):
- color_b = random.randint(0, 255)
- color_g = random.randint(0, 255)
- color_r = random.randint(0, 255)
- color.append((color_b, color_g, color_r))
- for x in range(w):
- for y in range(h):
- lable = labels[x, y]
- img_color[x, y, :] = color[int(lable)]
- # cv2.imshow("img_color", img_color)
- #
- # cv2.waitKey(0)
- return img_color
- def add_contrast(image_np):
- # cv2.imshow("image_np", image_np)
- img = image_np.astype(np.float32)
- bri_mean = np.mean(img)
- # a = np.arange(5, 16, 5) / 10
- # b = np.arange(-30, 31, 30)
- #
- # a_len = len(a)
- # b_len = len(b)
- # print(a_len, b_len)
- #
- # for i in range(a_len):
- # for j in range(b_len):
- # aa = a[i]
- # bb = b[j]
- # img_a = aa * (img-bri_mean) + bb + bri_mean
- # print(i, j, aa, bb)
- # img_a = np.clip(img_a, 0, 255).astype(np.uint8)
- # cv2.imshow("img_a", img_a)
- # cv2.waitKey(0)
- aa = 3
- bb = -50
- img_a = aa * (img-bri_mean) + bb + bri_mean
- img_a = np.clip(img_a, 0, 255).astype(np.uint8)
- # cv2.imshow("img_a", img_a)
- # cv2.waitKey(0)
- return img_a
- def get_image_legal(image_np):
- white = np.sum(image_np > 80)
- black = np.sum(image_np <= 80)
- if black > white:
- ratio = white / black
- else:
- ratio = black / white
- return ratio
- if __name__ == "__main__":
- generate_data_char()
- # for _p in glob("../data/test/FileInfo1021/*"):
- # add_contrast(cv2.imread(_p))
- # print(random.uniform(0.5, 0.5))
- # paths = glob('../data/equation/*')
- # for p in paths:
- # img = cv2.imread(p)
- # if img.shape[0] == 69 and img.shape[1] == 330:
- # print(p)
- # os.remove(p)
- # image_enhance_wrap(Image.open("D:/Project/captcha/data/equation/30_4_-_零.jpg"))
|