更新 detacttracking

detecttracking/tracking/utils/videot.py

@@ -0,0 +1,111 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Wed Sep 20 14:10:09 2023
+
+@author: ym
+"""
+
+import numpy as np
+import os
+import cv2
+
+# from pathlib import Path
+# import math
+# import sys
+# from scipy.spatial.distance import cdist
+
+VideoFormat = ['.mp4', '.avi', '.ts']
+def video2imgs(videof, imgdir):
+    cap = cv2.VideoCapture(videof)
+    i = 0
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        imgp = os.path.join(imgdir, f"{i}.png")
+        i += 1
+        cv2.imwrite(imgp, frame)
+        
+        if i == 400:
+            break
+    cap.release()
+    
+    print(os.path.basename(videof) + f" haved resolved")
+
+def videosave(bboxes, videopath="100_1688009697927.mp4"):
+
+    cap = cv2.VideoCapture(videopath)
+    
+    fps = int(cap.get(cv2.CAP_PROP_FPS))  # integer required, floats produce error in MP4 codec
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    
+    ## =========================================== 在当前模块地址下存储图像和视频
+    path = os.path.split(os.path.realpath(__file__))[0]
+    _, filename = os.path.split(videopath)
+    file, ext = os.path.splitext(filename)
+    
+    ## ======================================================== 视频保存设置
+    fourcc = cv2.VideoWriter_fourcc(*'MP4V')
+    save_video_path = os.path.join(path, "{}_show_1.mp4".format(file))    
+    vid_writer = cv2.VideoWriter(save_video_path, fourcc, fps, (width, height))
+    
+    ## ======================================================== 图像保存路径设置
+    save_img_path = os.path.join(path, "{}_show".format(file))
+    if not os.path.exists(save_img_path):
+        os.makedirs(save_img_path)
+    
+    cout = 0
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if ret:
+            
+            idx = np.where(bboxes[:, 7] == cout)[0]
+            box = bboxes[idx, 0:4].astype(int)
+            for i in range(box.shape[0]):
+                x1, y1 = box[i, :2]
+                x2, y2 = box[i, 2:4]
+                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 125, 255), 2)
+                
+                
+            cv2.imwrite(os.path.join(save_img_path, "{}.png".format(cout)), frame)
+            vid_writer.write(frame)
+            
+            cout += 1
+        else:
+            print("end!!!!!!!!!!!!!!!!!!!")
+            break
+    
+    vid_writer.release()
+    cap.release()
+    
+def main():
+    videopath = r'D:\videos'
+    savepath = r'D:\videos'
+    # video2imgs(videopath, savepath)
+    k = 0
+    for filename in os.listdir(videopath):
+        # filename = "20240929-155533.ts"
+        
+        file, ext = os.path.splitext(filename)
+        if ext not in VideoFormat:
+            continue
+        
+        basename = os.path.basename(videopath)
+        imgbase = basename + '-&-' + file
+        imgdir = os.path.join(savepath, imgbase)
+        if not os.path.exists(imgdir):
+            os.mkdir(imgdir)
+
+        videof = os.path.join(videopath, filename)
+        video2imgs(videof, imgdir)
+        
+        # k += 1
+        # if k == 1:
+        #     break
+
+    
+    
+if __name__ == '__main__':
+    main()
+