Files
ieemoo-ai-searchv2/utils/forsegmentation.py
2022-11-22 15:32:06 +08:00

122 lines
4.2 KiB
Python

import requests
from base64 import b64encode
from json import dumps
import os
import cv2
import numpy as np
import pdb
max_nu_area= 0
def get_object_mask(frame, ori_mask, mask_path, all_nu, result_path):
global max_nu_area
maskimg = cv2.imread(mask_path, 0)
kernel = np.ones((5, 5), np.uint8)
dst = ori_mask
dst = cv2.erode(dst, kernel)
dst = cv2.dilate(dst, kernel)
dst = cv2.medianBlur(dst,3)
if (cv2.__version__).split('.')[0] == '3':
_, contours, _ = cv2.findContours(dst, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
else:
contours, _ = cv2.findContours(dst, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
area,mask_area = 0,0
if len(contours) == 1 or len(contours) == 0:
return None
for contour in contours:
area_now = cv2.contourArea(contour)
dst = cv2.fillPoly(dst, [contour], (0))
if area_now > area:
area = area_now
Matrix = contour
max_area = area_now
if max_area > max_nu_area:
max_nu_area = max_area #维护最大的五个
flag = False
else:flag = True
(x, y, w, h) = cv2.boundingRect(Matrix)
dst = cv2.fillPoly(dst, [Matrix], (255))
coordination = [x, y, x + w, y + h]
if max_area/(w*h)<0.3:
return None
#print('masking', maskimg.shape)
#print('dst', dst.shape)
#pdb.set_trace()
mask_dst = cv2.bitwise_and(maskimg, dst)
if w < 350 or h <350:
return None
if (cv2.__version__).split('.')[0] == '3':
_, contours, _ = cv2.findContours(mask_dst, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
else:
contours, _ = cv2.findContours(mask_dst, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
for contour in contours:
mask_area_now = cv2.contourArea(contour)
if mask_area_now > mask_area:
mask_area = mask_area_now
proportion = float(mask_area/max_area)
if proportion<1.05 and proportion>0.5:
#print(coordination)
A,B,C = mask_dst, mask_dst, mask_dst
mask_dst = cv2.merge([A,B,C])
img = cv2.bitwise_and(mask_dst, frame)
img = img[coordination[1]:coordination[3],coordination[0]:coordination[2]]
frame = frame[coordination[1]:coordination[3],coordination[0]:coordination[2]]
#cv2.imshow('dst',img)
#cv2.waitKey(1000)
#print(all_nu)
#if all_nu>4: return 'True'
ratio = (w/h if h>w else h/w)
if ratio<0.5:
return None
if all_nu<5:
savenu = all_nu
elif all_nu>=5 and not(flag):
savenu = all_nu%5
print(savenu)
else:
return 'True'
cv2.imwrite('images/' + str(savenu)+'.jpg', img)
cv2.imwrite('images/' + 'ori'+ str(savenu)+'.jpg', frame)
#cv2.imwrite(os.sep.join([result_path, str(all_nu)+'.jpg']), img)
return 'con'
else:
return None
def get_object_location(file_test, mask_path,result_path):
cap = cv2.VideoCapture(file_test)
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
fgbg = cv2.createBackgroundSubtractorMOG2(detectShadows = False)#高斯混合模型为基础背景
nu= 10
working_nu = 0
while True:
ret, frame = cap.read()
if not ret:
break
frame = cv2.medianBlur(frame, ksize=3)
frame_motion = frame.copy()
fgmask = fgbg.apply(frame_motion)
mask= cv2.threshold(fgmask, 25, 255, cv2.THRESH_BINARY)[1] # 二值化
mask = cv2.dilate(mask, kernel, iterations=1)
if nu<=30:
res = get_object_mask(frame, mask, mask_path, working_nu, result_path)
#print(res)
if res=='True':
break
elif res == 'con':
working_nu+=1
else: continue
else:break
nu+=1
def analysis(file_test, mask_path, result_path):
#mask_path = 'mask.jpg'
#result_path = 'result'
if not (os.path.exists(result_path)):
os.mkdir(result_path)
get_object_location(file_test, mask_path, result_path)
if __name__ == '__main__':
mask_path = 'mask2.jpg'
result_path = 'images'
file_test = "ftp/anonymous/20210908-150524_e8e24395-fc7b-42f2-a50a-068c4ac73ee9_6921168509256.mp4"
analysis(file_test, mask_path, result_path)