Files
ieemoo-ai-conpurchase/utils/filter.py
2023-06-25 13:55:22 +08:00

56 lines
1.7 KiB
Python

import cv2
from floder.config import cfg
def filt(video_path):
#mask_path = '../../module/ieemoo-ai-search/model/now/ori_old.jpg'
maskimg = cv2.imread(cfg.maskImg)
fgbg = cv2.createBackgroundSubtractorMOG2(history=500, varThreshold=20, detectShadows = False)
capture = cv2.VideoCapture(video_path)
ret,frame = capture.read(0)
if frame.shape[0]<frame.shape[1]:
maskimg = cv2.imread(cfg.maskImg)
else:
maskimg = cv2.imread(cfg.maskImg1)
print('capture >>>> {}'.format(frame.shape))
imglist = []
re = False
nn = 0
while True:
ret,frame = capture.read()
nn += 1
#print('>>>>{}'.format(nn))
if not ret:break
if not re:
re = img_filter(frame, maskimg, fgbg)
else:
imglist.append(frame)
#cv2.imwrite('./imgs/tmpimgs/'+str(nn)+'.jpg', frame)
if len(imglist) > 30:
break
return imglist #-->list imgs
def img_filter(frame, maskimg, fgbg):
dic,dics = {},{}
iouArea = 0
height, width = frame.shape[:2]
frame = cv2.resize(frame, (int(width/2), int(height/2)), interpolation=cv2.INTER_CUBIC)
fgmask = fgbg.apply(frame)
draw1 = cv2.threshold(fgmask, 25, 255, cv2.THRESH_BINARY)[1]
draw1 = cv2.bitwise_and(maskimg[:, :, 0], draw1)
contours_m, hierarchy_m = cv2.findContours(draw1.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
for contour in contours_m:
dics[len(contour)] = contour
if len(dics.keys())>0:
cc = sorted(dics.keys())
iouArea = cv2.contourArea(dics[cc[-1]])
if iouArea>3000 and iouArea<50000:
return True
return False
if __name__ == '__main__':
videoName = 'filterImg.mp4'
filt(videoName)