112 lines
4.3 KiB
Python
112 lines
4.3 KiB
Python
import cv2 as cv
|
|
# from segmentation import get_object_mask
|
|
import os, time
|
|
|
|
import numpy
|
|
|
|
|
|
def get_object_location(file_dir, mask_path, frame_path, result_path):
|
|
# cap = cv.VideoCapture(0)
|
|
# 设置变量
|
|
kernel = cv.getStructuringElement(cv.MORPH_ELLIPSE, (5, 5)) # 定义结构元素
|
|
# 背景差法
|
|
# fgbg = cv.bgsegm.createBackgroundSubtractorMOG()
|
|
# fgbg = cv.createBackgroundSubtractorMOG2(detectShadows = False)#高斯混合模型为基础背景
|
|
# fgbg = cv.bgsegm.createBackgroundSubtractorGMG(2)#结合静态背景图像估计和每个像素的贝叶斯分割
|
|
# fgbg = cv.createBackgroundSubtractorKNN()
|
|
for num, name in enumerate(os.listdir(file_dir)):
|
|
fgbg = cv.createBackgroundSubtractorMOG2(history=500, varThreshold=16, detectShadows=True)
|
|
file_test = os.sep.join([file_dir, name])
|
|
nu, nn = 0, 1
|
|
flag = False
|
|
T1 = time.time()
|
|
# 设置文件
|
|
cap = cv.VideoCapture(file_test)
|
|
while True:
|
|
# 读取一帧
|
|
ret, frame = cap.read()
|
|
# 如果视频结束，跳出循环
|
|
# if nn%2 == 0 or nn%3==0:
|
|
# nn += 1
|
|
# continue
|
|
nn += 1
|
|
# cv.imwrite(os.sep.join([frame_path, 'ori_' + str(nn) + '.jpg']), frame)
|
|
if (not ret):
|
|
break
|
|
if flag:
|
|
flag = False
|
|
print('flag change>>{}>>{}'.format(name, nn))
|
|
|
|
frame = cv.resize(frame, (512, 640), interpolation=cv.INTER_CUBIC)
|
|
cv.imwrite('images/' + str(nn) + '.jpg', frame)
|
|
frame = cv.medianBlur(frame, ksize=3)
|
|
frame_motion = frame.copy()
|
|
|
|
# 计算前景掩码
|
|
fgmask = fgbg.apply(frame_motion)
|
|
# cv.imwrite('fgmask'+'/fgmask_'+str(nn)+'.jpg',fgmask)
|
|
draw1 = cv.threshold(fgmask, 230, 255, cv.THRESH_BINARY)[1] # 二值化
|
|
draw1 = cv.erode(draw1, kernel, iterations=1)
|
|
draw1 = cv.dilate(draw1, kernel, iterations=1)
|
|
|
|
# cv.imwrite('frame'+'/'+str(nn)+'.jpg',draw1)
|
|
# cv.imshow(str(nn)+'.jpg', draw1)
|
|
# cv.waitKey()
|
|
|
|
if nn<100:
|
|
flag = check_tings(mask_path, draw1, nu)
|
|
|
|
# cv.imwrite(os.sep.join([frame_path, 'draw_' + str(nn) + '.jpg']), draw1)
|
|
# cv.imwrite(os.sep.join([frame_path, 'ori_' + str(nn) + '.jpg']), frame)
|
|
|
|
# cv.imread('mask', draw1)
|
|
# cv.waitKey(1)
|
|
if nu<=500:
|
|
# cv.imwrite('frame/frame_motion' + str(nu) + '.jpg', frame_motion)
|
|
# cv.imwrite(os.sep.join([frame_path, 'frame_motion'+str(nu) + '.jpg']), frame_motion)
|
|
# cv.imwrite('frame/draw_' + str(nu) + '.jpg', draw1)
|
|
draw1 = cv.erode(draw1, kernel)
|
|
draw1 = cv.dilate(draw1, kernel)
|
|
draw1 = cv.medianBlur(draw1, 3)
|
|
cv.imwrite(os.sep.join([frame_path, 'draw_'+str(nu) + '.jpg']), draw1)
|
|
|
|
else:break
|
|
nu+=1
|
|
T2 = time.time()
|
|
print('single video >>> {}-->{}-->{}-->{}'.format(name, nn, num, (T2 - T1)))
|
|
|
|
def check_tings(mask_path, img, nu):
|
|
dics = {}
|
|
mask_img = cv.imread(mask_path)
|
|
# print('mask_img',mask_img[:,:,0])
|
|
# cv.imwrite('D:/workspace/Track/yolov8_ultralytics/ultralytics/yolo/engine/draw1/' + str(nu) + '.jpg', img)
|
|
img = cv.bitwise_and(mask_img[:,:,0], img)
|
|
# cv.imshow('1.jpg', img)
|
|
# cv.waitKey()
|
|
contours_m, hierarchy_m = cv.findContours(img.copy(), cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
|
|
for contour in contours_m:
|
|
# print('contour', hierarchy_m)
|
|
dics[len(contour)] = contour
|
|
|
|
# print('dics',nu, dics)
|
|
if len(dics.keys()) > 0:
|
|
cc = sorted(dics.keys())
|
|
iouArea = cv.contourArea(dics[cc[-1]])
|
|
print('iouArea', nu, iouArea)
|
|
# if iouArea>10000 and iouArea<40000:
|
|
# if iouArea>1000 and iouArea<4000:
|
|
if iouArea>1000:
|
|
return True
|
|
else:
|
|
return False
|
|
else:
|
|
return False
|
|
|
|
if __name__ == '__main__':
|
|
# file_dir = "D:/Project/ieemoo/target-location/videos"
|
|
file_dir = "videos/"
|
|
mask_path = 'mask\lianhua_1.jpg'
|
|
frame_path = 'frame'
|
|
result_path = 'result'
|
|
get_object_location(file_dir, mask_path, frame_path, result_path)
|