{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"os.chdir('../../')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## GPU 设置"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"GPUID='0'##调用GPU序号\n",
"os.environ[\"CUDA_VISIBLE_DEVICES\"] = GPUID"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"import tensorflow as tf\n",
"from glob import glob\n",
"from PIL import Image\n",
"import cv2\n",
"Input =tf.keras.layers.Input\n",
"Lambda = tf.keras.layers.Lambda\n",
"load_model = tf.keras.models.load_model\n",
"Model = tf.keras.models.Model\n",
"\n",
"from apphelper.image import get_box_spilt,read_voc_xml,resize_im,read_singLine_for_yolo\n",
"from text.keras_yolo3 import preprocess_true_boxes, yolo_text\n",
"from train.text.utils import get_random_data_ as get_random_data\n",
"\n",
"\n",
"def data_generator(roots, anchors, num_classes,splitW):\n",
" '''data generator for fit_generator\n",
" @@roots:jpg/png\n",
" '''\n",
" n = len(roots)\n",
" np.random.shuffle(roots)\n",
" scales = [416,608,608,608]##多尺度训练\n",
" i = 0\n",
" j = 0\n",
" m = len(scales)\n",
" while True:\n",
" root = roots[i]\n",
" i+=1\n",
" if i>=n:\n",
" i=0\n",
" scale = scales[j]\n",
" j+=1\n",
" if j>=m:\n",
" j=0\n",
" \n",
" xmlP = os.path.splitext(root)[0]+'.xml'\n",
" boxes = read_voc_xml(xmlP)\n",
" im = Image.open(root)\n",
" \n",
" w,h = resize_im(im.size[0],im.size[1], scale=scale, max_scale=None)\n",
" if max(w,h)>2048:\n",
" w,h = resize_im(im.size[0],im.size[1], scale=scale, max_scale=2048)\n",
" \n",
" input_shape = (h,w)\n",
" isRoate=True\n",
" rorateDegree=np.random.uniform(-5,5)\n",
" rorateDegree = 0\n",
" newBoxes,newIm = get_box_spilt(boxes, im, w,h,splitW=splitW, isRoate=isRoate, rorateDegree=rorateDegree)\n",
" newBoxes = np.array(newBoxes)\n",
" if len(newBoxes)==0:\n",
" continue\n",
" if np.random.randint(0,100)>70:\n",
" if np.random.randint(0,100)>50:\n",
" ##图像水平翻转\n",
" newBoxes[:,[0,2]] = w-newBoxes[:,[2,0]]\n",
" \n",
" im = Image.fromarray(cv2.flip(np.array(im),1))\n",
" else:\n",
" ##垂直翻转\n",
" newBoxes[:,[1,3]] = h-newBoxes[:,[3,1]]\n",
" \n",
" im = Image.fromarray(cv2.flip(np.array(im),0))\n",
" \n",
" \n",
" maxN = 128##随机选取128个box用于训练\n",
" image_data = []\n",
" box_data = []\n",
" \n",
" image, box = get_random_data(newIm,newBoxes, input_shape,max_boxes=maxN)\n",
" \n",
" image_data = np.array([image])\n",
" box_data = np.array([box])\n",
" y_true = preprocess_true_boxes(box_data, input_shape, anchors, num_classes)\n",
" yield [image_data, *y_true], [np.zeros(1)]*4\n",
" "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 加载训练数据集，标注XML软件参考https://github.com/cgvict/roLabelImg.git"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"val_split = 0.1\n",
"root='train/data/text/*/*.[j|p|J]*'\n",
"jpgPath = glob(root)\n",
"##剔除为标记的图像\n",
"delPaths = []\n",
"for p in jpgPath:\n",
" xmlP = os.path.splitext(p)[0]+'.xml'\n",
" if not os.path.exists(xmlP):\n",
" delPaths.append(p)\n",
" \n",
"print('total:',len(jpgPath))\n",
"jpgPath = list(set(jpgPath) - set(delPaths))\n",
"print('total:',len(jpgPath))\n",
"np.random.shuffle(jpgPath)\n",
"\n",
"\n",
"num_val = int(len(jpgPath)*val_split)\n",
"num_train = len(jpgPath) - num_val \n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 定义anchors及加载训练模型"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"## 计算训练集anchors\n",
"from train.text.gen_anchors import YOLO_Kmeans## anchors生产\n",
"splitW = 8##文本分割最小宽度\n",
"#cluster = YOLO_Kmeans(cluster_number=9, root=root, scales=[416, 512, 608, 608, 608, 768, 960, 1024], splitW=splitW)\n",
"#8,9, 8,18, 8,31, 8,59, 8,124, 8,351, 8,509, 8,605, 8,800\n",
"#print(cluster.anchors)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"## 数据事例\n",
"from apphelper.image import xy_rotate_box,box_rotate\n",
"def plot_boxes(img,angle, result,color=(0,0,0)):\n",
" tmp = np.array(img)\n",
" c = color\n",
" w,h = img.size\n",
" thick = int((h + w) / 300)\n",
" i = 0\n",
" if angle in [90,270]:\n",
" imgW,imgH = img.size[::-1]\n",
" \n",
" else:\n",
" imgW,imgH = img.size\n",
"\n",
" for line in result:\n",
" cx =line['cx']\n",
" cy = line['cy']\n",
" degree =line['angle']\n",
" w = line['w']\n",
" h = line['h']\n",
" x1,y1,x2,y2,x3,y3,x4,y4 = xy_rotate_box(cx, cy, w, h, degree)\n",
" x1,y1,x2,y2,x3,y3,x4,y4 = box_rotate([x1,y1,x2,y2,x3,y3,x4,y4],angle=(360-angle)%360,imgH=imgH,imgW=imgW)\n",
" cx =np.mean([x1,x2,x3,x4])\n",
" cy = np.mean([y1,y2,y3,y4])\n",
" cv2.line(tmp,(int(x1),int(y1)),(int(x2),int(y2)),c,1)\n",
" cv2.line(tmp,(int(x2),int(y2)),(int(x3),int(y3)),c,1)\n",
" cv2.line(tmp,(int(x3),int(y3)),(int(x4),int(y4)),c,1)\n",
" cv2.line(tmp,(int(x4),int(y4)),(int(x1),int(y1)),c,1)\n",
" mess=str(i)\n",
" cv2.putText(tmp, mess, (int(cx), int(cy)),0, 1e-3 * h, c, thick // 2)\n",
" i+=1\n",
" return Image.fromarray(tmp)\n",
"\n",
"def plot_box(img,boxes):\n",
" blue = (0, 0, 0) #18\n",
" tmp = np.copy(img)\n",
" for box in boxes:\n",
" cv2.rectangle(tmp, (int(box[0]),int(box[1])), (int(box[2]), int(box[3])), blue, 1) #19\n",
" \n",
" return Image.fromarray(tmp) \n",
"\n",
"def show(p,scale=608):\n",
" im = Image.open(p)\n",
" xmlP = p.replace('.jpg','.xml').replace('.png','.xml')\n",
" boxes = read_voc_xml(xmlP)\n",
" im = Image.open(p)\n",
" w,h = resize_im(im.size[0],im.size[1], scale=scale, max_scale=4096)\n",
" input_shape = (h,w)\n",
" isRoate=True\n",
" \n",
" rorateDegree=np.random.uniform(-5,5) \n",
" rorateDegree=0\n",
" newBoxes,newIm = get_box_spilt(boxes, im, sizeW=w, SizeH=h, splitW=splitW, isRoate=isRoate, rorateDegree=rorateDegree)\n",
" return plot_boxes(im,0, boxes,color=(0,0,0)),plot_box(newIm,newBoxes),newBoxes\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"a,b,newBoxes = show(jpgPath[9])\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"b"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#anchors = cluster.anchors\n",
"anchors = '8,9, 8,18, 8,31, 8,59, 8,124, 8,351, 8,509, 8,605, 8,800'\n",
"anchors = [float(x) for x in anchors.split(',')]\n",
"anchors = np.array(anchors).reshape(-1, 2)\n",
"num_anchors = len(anchors)\n",
"class_names = ['none','text',]##text \n",
"num_classes = len(class_names)\n",
"textModel = yolo_text(num_classes,anchors,train=True)\n",
"#textModel.load_weights('models/text.h5')##加载预训练模型权重\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"textModel.load_weights('models/text.h5')##加载预训练模型权重"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"trainLoad = data_generator(jpgPath[:num_train], anchors, num_classes,splitW)\n",
"testLoad = data_generator(jpgPath[num_train:], anchors, num_classes,splitW)\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"adam = tf.keras.optimizers.Adam(lr=0.0005)\n",
"textModel.compile(optimizer=adam, loss={'xy_loss':lambda y_true, y_pred:y_pred,\n",
" 'wh_loss':lambda y_true, y_pred:y_pred,\n",
" 'confidence_loss':lambda y_true, y_pred:y_pred,\n",
" 'class_loss':lambda y_true, y_pred:y_pred,\n",
" }\n",
" )\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"textModel.fit_generator(generator=trainLoad, \n",
" steps_per_epoch=num_train, \n",
" epochs=2,\n",
" verbose=2, \n",
" callbacks=None,\n",
" validation_data=testLoad, \n",
" validation_steps=num_val)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from text.keras_yolo3 import yolo_text,box_layer,K\n",
"from config import kerasTextModel,IMGSIZE,keras_anchors,class_names\n",
"from apphelper.image import resize_im,letterbox_image\n",
"from PIL import Image\n",
"import numpy as np\n",
"import tensorflow as tf\n",
"graph = tf.get_default_graph()##解决web.py 相关报错问题\n",
"\n",
"anchors = [float(x) for x in keras_anchors.split(',')]\n",
"anchors = np.array(anchors).reshape(-1, 2)\n",
"num_anchors = len(anchors)\n",
"\n",
"num_classes = len(class_names)\n",
"textModelTest = yolo_text(num_classes,anchors)\n",
"kerasTextModel = '/tmp/textModel.h5'\n",
"textModelTest.load_weights(kerasTextModel)\n",
"\n",
"\n",
"sess = K.get_session()\n",
"image_shape = K.placeholder(shape=(2, ))##图像原尺寸:h,w\n",
"input_shape = K.placeholder(shape=(2, ))##图像resize尺寸:h,w\n",
"box_score = box_layer([*textModelTest.output,image_shape,input_shape],anchors, num_classes)\n",
"\n",
"\n",
"\n",
"def text_detect(img,prob = 0.05):\n",
" im = Image.fromarray(img)\n",
" scale = IMGSIZE[0]\n",
" w,h = im.size\n",
" w_,h_ = resize_im(w,h, scale=scale, max_scale=2048)##短边固定为608,长边max_scale<4000\n",
" #boxed_image,f = letterbox_image(im, (w_,h_))\n",
" boxed_image = im.resize((w_,h_), Image.BICUBIC)\n",
" image_data = np.array(boxed_image, dtype='float32')\n",
" image_data /= 255.\n",
" image_data = np.expand_dims(image_data, 0) # Add batch dimension.\n",
" imgShape = np.array([[h,w]])\n",
" inputShape = np.array([[h_,w_]])\n",
" \n",
" \n",
" global graph\n",
" with graph.as_default():\n",
" ##定义 graph变量 解决web.py 相关报错问题\n",
" \"\"\"\n",
" pred = textModel.predict_on_batch([image_data,imgShape,inputShape])\n",
" box,scores = pred[:,:4],pred[:,-1]\n",
" \n",
" \"\"\"\n",
" box,scores = sess.run(\n",
" [box_score],\n",
" feed_dict={\n",
" textModelTest.input: image_data,\n",
" input_shape: [h_, w_],\n",
" image_shape: [h, w],\n",
" K.learning_phase(): 0\n",
" })[0]\n",
" \n",
"\n",
" keep = np.where(scores>prob)\n",
" \n",
" box[:, 0:4][box[:, 0:4]<0] = 0\n",
" box[:, 0][box[:, 0]>=w] = w-1\n",
" box[:, 1][box[:, 1]>=h] = h-1\n",
" box[:, 2][box[:, 2]>=w] = w-1\n",
" box[:, 3][box[:, 3]>=h] = h-1\n",
" box = box[keep[0]]\n",
"\n",
" scores = scores[keep[0]]\n",
" return box,scores\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"p='./train/text/26BB94CA21C11AB38BC5FC2E08D140CD.jpg'\n",
"IMGSIZE=416,416\n",
"img = np.array(Image.open(p))\n",
"box,scores = text_detect(img,prob = 0.01)\n",
"plot_box(img,box)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "chineseocr",
"language": "python",
"name": "chineseocr"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.6.2"
}
},
"nbformat": 4,
"nbformat_minor": 2
}