first commit

network/vanalysis_video.py

@@ -0,0 +1,156 @@
+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)