ieemoo-ai-conpurchase/network/vanalysis_video.py

import argparse
import glob, cv2, os, pdb, time, sys
sys.path.append('utils')
sys.path.append('network/core')
import numpy as np
import torch
from network.core.raft import RAFT
from network.core.utils import flow_viz
from network.core.utils.utils import InputPadder
from floder.config import cfg
#from utils.retrieval_feature import AntiFraudFeatureDataset

DEVICE = 'cuda'
pre_area = 0
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))  # 定义膨胀结构元素
kernel1 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2, 2))  # 定义腐蚀结构元素

def load_image(imfile):
    #img = np.array(Image.open(imfile)).astype(np.uint8)
    img = np.array(imfile).astype(np.uint8)
    img = torch.from_numpy(img).permute(2, 0, 1).float()
    return img[None].to(DEVICE)

def viz(img, flo):
    img = img[0].permute(1,2,0).cpu().numpy()
    flo = flo[0].permute(1,2,0).cpu().numpy()
    flo = flow_viz.flow_to_image(flo)
    return flo

def raft_init_model(args):
    model = torch.nn.DataParallel(RAFT(args))
    model.load_state_dict(torch.load(args.model))
    model = model.module
    model.to(DEVICE)
    model.eval()
    return model

def vanalysis(model, imgsList):
    imfile1, imfile2 = None,None
    re = []
    with torch.no_grad():
        for nn, frame in enumerate(imgsList):
            #print('frame {}>>>{}'.format(nn, type(frame)))
            if frame.shape[0]<frame.shape[1]:#1024x1280
                MASKIMG = cfg.maskAreaImg
            else:#1280x1024
                MASKIMG = cfg.maskAreaImg1
            frame_show = frame
            height, width = frame.shape[:2]
            frame = cv2.GaussianBlur(frame,(5,5),0)
            frame = cv2.resize(frame, (int(width/2), int(height/2)), interpolation=cv2.INTER_CUBIC)
            if nn == 0:
                imfile1 = frame
                #nu += 1
                continue
            else:
                imfile2 = frame
            image1 = load_image(imfile1)
            image2 = load_image(imfile2)

            padder = InputPadder(image1.shape)
            image1, image2 = padder.pad(image1, image2)

            flow_low, flow_up = model(image1, image2, iters=2, test_mode=True)
            flo = viz(image1, flow_up)
            result, coordination = get_target(flo, imfile1, MASKIMG)
            #print('result >>> {} coordination >>>{}'.format(result, coordination))
            if not result is  None:
                re.append(result)
                #cv2.imwrite('./imgs/tmpimgs/'+str(nn)+'.jpg', result)
            imfile1 = imfile2
        return re[:10] # ——> list

#def get_target(path, img, ori_img, nu, ori_mask, MASKIMG):
def get_target(img, ori_img, MASKIMG):
    global pre_area
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    ret, mask = cv2.threshold(img, 249, 255, cv2.THRESH_BINARY)
    mask_max_area, mask_max_contour = 0, 0
    mask = cv2.bitwise_not(mask)
    mask_image = np.zeros((ori_img.shape[0], ori_img.shape[1], 1), np.uint8)
    if (cv2.__version__).split('.')[0] == '3':
        _, contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    else:
        contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    if len(contours)>100:
        return None, '_'
    for contour in contours:
        mask_area_now = cv2.contourArea(contour)
        if mask_area_now > mask_max_area:
            mask_max_area = mask_area_now
            mask_max_contour = contour
    if mask_max_area == 0 :return None, '_'  #mask_max_area 目标位的面积
    (x, y, w, h) = cv2.boundingRect(mask_max_contour)
    if (w*h)/(img.shape[0]*img.shape[1])>0.80:
        return None, '_'
    if min(w,h) <100 or max(w,h)>1000:
        return None, '_'
    coordination = [x, y, x + w, y + h]
    mask_image = cv2.fillPoly(mask_image, [mask_max_contour], (255))
    if pre_area==0:
        pre_area = mask_max_area
        return None, '_'
    else:
        if abs(mask_max_area-pre_area)/pre_area > 0.4:
            pre_area = mask_max_area
            #print('abs:',abs(mask_max_area-pre_area)/pre_area)
            return None, '_'
        else:
            pre_area = mask_max_area
    A,B,C = mask_image, mask_image, mask_image
    mask_image = cv2.merge([A,B,C])

    #该方法去除框外干扰
    if not get_iou_ratio(mask_image, MASKIMG):
        return None, '_'

    show = cv2.bitwise_and(ori_img, mask_image)
    show = ori_img[coordination[1]:coordination[3], coordination[0]:coordination[2]]
    return show, coordination

def get_iou_ratio(oimg, MASKIMG):
    mimg = cv2.imread(MASKIMG)
    iimg = cv2.bitwise_and(oimg, mimg)
    iimgarea = get_area(iimg)
    oimgarea = get_area(oimg)
    if iimgarea/oimgarea < 0.1:
        return False
    else: return True

def get_area(img):
    kernel = np.ones((3, 3), dtype=np.uint8)
    img = cv2.dilate(img, kernel, 1)
    img = cv2.erode(img, kernel, 1)
    maxcontour, nu = 0,0
    contours, _ = cv2.findContours(img[:,:,1] ,cv2.RETR_TREE , cv2.CHAIN_APPROX_NONE)
    if len(contours) == 0:
        return 0
    for i in range(len(contours)):
        if maxcontour < len(contours[i]):
            maxcontour = len(contours[i])
            nu = i
    area = cv2.contourArea(contours[nu])
    return area

if __name__ == '__main__':
    model = raft_init_model()
    from utils.tools import createNet
    import pdb
    #uuid_barcode = '6907992825762'
    imgslist = []
    for imgname in os.listdir('test_imgs'):
        imgslist.append(cv2.imread(os.sep.join(['test_imgs', imgname])))
    pdb.set_trace()
    analysis = vanalysis(model=model, video_path=imgslist)
#    analysis_video(model, video_path, result_path)