master
/ text / .ipynb_checkpoints / opencv_dnn_detect-checkpoint.py

opencv_dnn_detect-checkpoint.py @c1cd888 — raw · history · blame 

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
from config import yoloCfg,yoloWeights,opencvFlag
from config import AngleModelPb,AngleModelPbtxt
from config import IMGSIZE
from PIL import Image
import numpy as np
import cv2

if opencvFlag=='keras':
    ##转换为tf模型，以便GPU调用
    import tensorflow as tf
    from tensorflow.python.platform import gfile
    config = tf.ConfigProto(allow_soft_placement=True)
    sess = tf.Session(config=config)
    with gfile.FastGFile(AngleModelPb, 'rb') as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())
            sess.graph.as_default()
            tf.import_graph_def(graph_def, name='')
    inputImg =  sess.graph.get_tensor_by_name('input_1:0')
    predictions = sess.graph.get_tensor_by_name('predictions/Softmax:0')
    keep_prob = tf.placeholder(tf.float32)

    
else:
   angleNet = cv2.dnn.readNetFromTensorflow(AngleModelPb,AngleModelPbtxt)##dnn 文字方向检测
textNet  = cv2.dnn.readNetFromDarknet(yoloCfg,yoloWeights)##文字定位
def text_detect(img):
    thresh=0
    h,w = img.shape[:2]
    
    inputBlob = cv2.dnn.blobFromImage(img, scalefactor=0.00390625, size=IMGSIZE,swapRB=True ,crop=False);
    textNet.setInput(inputBlob)
    pred = textNet.forward()
    cx = pred[:,0]*w
    cy = pred[:,1]*h
    xmin = cx - pred[:,2]*w/2
    xmax = cx + pred[:,2]*w/2
    ymin = cy - pred[:,3]*h/2
    ymax = cy + pred[:,3]*h/2
    scores = pred[:,4]
    indx = np.where(scores>thresh)[0]
    scores = scores[indx]
    boxes = np.array(list(zip(xmin[indx],ymin[indx],xmax[indx],ymax[indx])))
    return boxes,scores


def angle_detect_dnn(img,adjust=True):
    """
    文字方向检测
    """
    h,w = img.shape[:2]
    ROTATE = [0,90,180,270]
    if adjust:
       thesh = 0.05
       xmin,ymin,xmax,ymax = int(thesh*w),int(thesh*h),w-int(thesh*w),h-int(thesh*h)
       img = img[ymin:ymax,xmin:xmax]##剪切图片边缘
    
    
    inputBlob = cv2.dnn.blobFromImage(img, 
                                      scalefactor=1.0, 
                                      size=(224, 224),
                                      swapRB=True ,
                                      mean=[103.939,116.779,123.68],crop=False);
    angleNet.setInput(inputBlob)
    pred = angleNet.forward()
    index = np.argmax(pred,axis=1)[0]
    return ROTATE[index]


def angle_detect_tf(img,adjust=True):
    """
    文字方向检测
    """
    h,w = img.shape[:2]
    ROTATE = [0,90,180,270]
    if adjust:
       thesh = 0.05
       xmin,ymin,xmax,ymax = int(thesh*w),int(thesh*h),w-int(thesh*w),h-int(thesh*h)
       img = img[ymin:ymax,xmin:xmax]##剪切图片边缘
    img = cv2.resize(img,(224,224))
    img = img[..., ::-1].astype(np.float32)
        
    img[..., 0] -= 103.939
    img[..., 1] -= 116.779
    img[..., 2] -= 123.68
    img          = np.array([img])
        
    
    
    out = sess.run(predictions, feed_dict={inputImg: img,
                                              keep_prob: 0
                                             })
    

    index = np.argmax(out,axis=1)[0]
    return ROTATE[index]

def angle_detect(img,adjust=True):
    """
    文字方向检测
    """
    if opencvFlag=='keras':
        return angle_detect_tf(img,adjust=adjust)
    else:
        return angle_detect_dnn(img,adjust=adjust)