IBM_AI/vi_server/scale_test.py

#!/usr/bin/env python
# -*- coding=utf-8 -*-
# coding=utf-8
import sys
reload(sys)
sys.setdefaultencoding('utf8')
import os
import numpy as np
import cv2
import time
from log import *
from skimage import measure, color


def get_image_box2(img,lower_color,upper_color,img_width,median=3):
    img_hsv = cv2.inRange(cv2.cvtColor(img,cv2.COLOR_BGR2HSV), lower_color, upper_color)
    img_hsv=cv2.medianBlur(img_hsv, median)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (median, median))
    img_hsv = cv2.dilate(img_hsv, kernel)
    img_hsv = cv2.erode(img_hsv, kernel)

    min_length=15
    # labels = measure.label(img_hsv)
    # props = measure.regionprop(labels)
    # bbox=None
    # for ia in range(len(props)):
        
    #     if  props[ia].bbox[3]-props[ia].bbox[1]<img_width/4 and props[ia].bbox[3]-props[ia].bbox[1]>min_length and (bbox is None or props[ia].bbox[3]-props[ia].bbox[1]>bbox[3]-bbox[1] ):
    #         #print("props[ia].bbox:",props[ia].bbox," ",props[ia].bbox[3]-props[ia].bbox[1])
    #         bbox=props[ia].bbox
    #     if not bbox is None:
    #         cv2.rectangle(img_hsv,(props[ia].bbox[1],props[ia].bbox[0]),(props[ia].bbox[3],props[ia].bbox[2]),(255,255,255),2) 
    if lower_color[2]==0:
        cv2.imwrite(r"e:\b2\hsv4.jpg",img_hsv)
    edgeThresh =70
    img2 = cv2.Canny(img_hsv,edgeThresh,edgeThresh*3,apertureSize =3)
    if lower_color[2]==0:
        cv2.imwrite(r"e:\b2\canny.jpg",img2)
    lines=cv2.HoughLinesP(img_hsv, 1, np.pi/180, 120,None,min_length,median)
    if not lines is None:
        cv2.imwrite(r"e:\b2\img2"+str(lower_color[2])+".jpg",img2)
        maxline=None
        max_verticals=[]
        for line in lines:
            x1, y1, x2, y2=line[0]
            if x2==x1 and abs(y1-y2)>min_length*2:
                if len(max_verticals)<2:
                    max_verticals.append((x1,y1,x2,y2))
                    # cv2.line(img_hsv, (x1, y2-50), (x2,y1+50), (255, 255, 255), 5)
                else:
                    for index in range(len(max_verticals)):
                        if abs(y1-y2)>abs(max_verticals[index][1]-max_verticals[index][3]):
                            max_verticals[index]=(x1,y1,x2,y2)
                            break
        offset=3
        # print("max_verticals:",max_verticals)
        for line in lines:
            x1, y1, x2, y2=line[0]
            print("line[0]:[",lower_color[2],"]",line[0])
            cv2.line(img_hsv, (x1, y1), (x2, y2), (125, 124, 124), 5)
            cv2.imwrite(r"e:\b2\hsv"+str(lower_color[2])+".jpg",img_hsv)
            if x2-x1>min_length and (maxline==None or x2-x1>maxline[2]-maxline[0]):
                # cv2.line(img_hsv, (x1-50, y1+3), (x2+50, y2+3), (255, 255, 255), 2)
                # cv2.imwrite(r"e:\b2\hsv"+str(lower_color[2])+".jpg",img_hsv)
                if len(max_verticals)==0 :
                    maxline=(x1,y1,x2,y2)
                else:
                    is_offset=False
                    for index in range(len(max_verticals)):
                        if abs(max_verticals[index][0]-x1)<=offset or abs(max_verticals[index][2]-x2)<=offset:
                            is_offset=True
                            break
                    if not is_offset:
                        maxline=(x1,y1,x2,y2)
        return maxline
    else:
        return None

def get_image_box(img,lower_color,upper_color,img_width,median=3):
    img_hsv = cv2.inRange(cv2.cvtColor(img,cv2.COLOR_BGR2HSV), lower_color, upper_color)
    img_hsv=cv2.medianBlur(img_hsv, median)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (median, median))
    img_hsv = cv2.dilate(img_hsv, kernel)
    img_hsv = cv2.erode(img_hsv, kernel)

    min_length=15
    edgeThresh =70
    img2 = cv2.Canny(img_hsv,edgeThresh,edgeThresh*3,apertureSize =3)
    # if lower_color[2]==0:
    #     cv2.imwrite(r"e:\b2\canny.jpg",img2)
    
    contours, hierarchy= cv2.findContours(img2,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
    max_rect=None
    for i in range(len(contours)):
        rect=cv2.boundingRect(contours[i])
        if rect[2]>min_length and (max_rect==None or rect[2]>max_rect[2]):
            max_rect=rect
    if not max_rect is None:
        for i in range(len(contours)):
            rect=cv2.boundingRect(contours[i])
            if is_in_rect(rect,max_rect):
                if rect[2]>max_rect[2]/2:
                    max_rect=rect
    
    #print("max_rect:",max_rect)
    offset=4
    return max_rect if max_rect is None else (max_rect[0]+offset,max_rect[1]+max_rect[3],max_rect[0]+max_rect[2]-offset,max_rect[1]+max_rect[3])
    
def is_in_rect(rect_small,rect_large):
    return rect_small[0]>rect_large[0] and rect_small[0]+rect_small[2]<rect_large[0]+rect_large[2] and rect_small[1]>rect_large[1] and rect_small[1]+rect_small[3]<rect_large[1]+rect_large[3]

def get_scale_rect(img,deep_height=300):
    height=img.shape[0]
    width=img.shape[1]
    print("width:",width,", height:",height)
    #color
    color_tuples=[
        #red
        (np.array([0, 127, 128]) ,np.array([10,  255, 255]),5),
        #green
        (np.array([60, 76, 46]) ,np.array([100, 255, 255]),5)]
    black_tuples=[
        #black
        (np.array([0, 0, 0]) ,np.array([180,  255, 46]),1)
        #white
        ,(np.array([0, 0, 221]) ,np.array([180,  30, 255]),3)
    ]
    banner_imgs=[img[height-deep_height:,]
                   ,img[:deep_height,]
                  ]
    deeps=[height-deep_height,0]

    for color_tuple in color_tuples:
        i=0
        for banner_img in banner_imgs:
            bbox=get_image_box(banner_img,color_tuple[0],color_tuple[1],width,color_tuple[2])
            if  not bbox is None:
                #print("banner_img:",banner_img,", color_tuple:",color_tuple)
                return (bbox[0],bbox[1],bbox[2],bbox[3],deeps[i])
            i+=1
    scale=None
    for color_tuple in black_tuples:
        i=0
        for banner_img in banner_imgs:
            bbox=get_image_box(banner_img,color_tuple[0],color_tuple[1],width,color_tuple[2])
            if  not bbox is None:
                # if not scale is None:
                #          print("scale: bbox[2]:", bbox[2],", bbox[0]:",bbox[0],", scale[2]:",scale[2],",scale[0],",scale[0],", bbox[2]-bbox[0]:",bbox[2]-bbox[0],", scale[2]-scale[0]:",scale[2]-scale[0])
                if scale is None or bbox[2]-bbox[0]>scale[2]-scale[0]:
                    #print("black bbox:",bbox,", ",bbox[2]-bbox[0],", deep:",deeps[i],", color_tuple:",color_tuple)
                    scale= (bbox[0],bbox[1],bbox[2],bbox[3],deeps[i])
                    return scale
            i+=1
    return scale
def get_scale(img):
    scale=get_scale_rect(img)
    return 0 if scale is None else scale[2]-scale[0]

def test():
    jpg_file=r"e:\b2\M_B-Other-C-500X-2019121601_23b0f2d2.jpg"
    jpg_file=r"e:\b2\55MnB.jpg"
    jpg_file=r"e:\b2\20200617160856.jpg"
    jpg_file=r"e:\b2\03116102_A2.jpg"
    jpg_file=r"e:\b2\20#rz.jpg"
    jpg_file=r"e:\b2\20200515140955.jpg"
    jpg_file=r"e:\b2\6587a18d99434adb9dd1b5350bc52fd.jpg"
    jpg_file=r"e:\b2\20200617161139.jpg"
    jpg_file=r"e:\b2\4341.JPG"
    jpg_file=r"e:\b2\SNAP-111916-0055.jpg"
    jpg_file=r"e:\b2\3_100X.jpg"
    jpg_file=r"e:\b2\W_200X.jpg"
    jpg_file=r"e:\b2\1x500.jpg"
    ##jpg_file=r"e:\b2\30CrMoA500.jpg"
    #jpg_file=r"e:\b2\100-2.JPG"
    #jpg_file=r"e:\b2\625100X.JPG"
    #jpg_file=r"e:\b2\1000X-5.jpg"
    #jpg_file=unicode(r"E:\ibm\样品收集\照片\20200709\各类图片\轴承钢碳化物网状 带状\碳化物网状500x.JPG")


    imgcolor=cv2.imdecode(np.fromfile(jpg_file, dtype=np.uint8),1)
    maxrect=get_scale_rect(imgcolor)
    height=imgcolor.shape[0]
    width=imgcolor.shape[1]
    if maxrect:
        cv2.line(imgcolor, (maxrect[0], maxrect[3]+maxrect[4]), (maxrect[2], maxrect[3]+maxrect[4]), (255, 0, 0), 5)
    cv2.namedWindow("enhanced",0)
    cv2.resizeWindow("enhanced",width/2,height/2)
    cv2.imshow("enhanced",imgcolor)
    cv2.waitKey(0)

if __name__ == '__main__':
    start = time.time()
    log_info("scale test start...")
    if len(sys.argv)<2:
        log_error("missing parameter")
        sys.exit()
    jpg_file=sys.argv[1].decode('gbk');
    img=cv2.imdecode(np.fromfile(jpg_file, dtype=np.uint8), 1)
    log_info("jpg_file:"+jpg_file)
    length=get_scale(img)
    log_result(length)
    log_info("{0} scale length finished, result:{1} consume:{2}".format(jpg_file,length,str(time.time()-start)))