退购1.1定位算法

2023-08-10 12:25:23 +08:00
commit 11e12f1899
371 changed files with 46027 additions and 0 deletions
--- a/ultralytics/yolo/engine/ids.py
+++ b/ultralytics/yolo/engine/ids.py
@ -0,0 +1,111 @@
+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)