initial project version!

tracking/utils/videot.py

@@ -0,0 +1,87 @@
+# -*- 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
+
+
+def video2imgs(videopath):
+# =============================================================================
+#     videopath：视频文件地址，在该地址的 "/file_imgs/" 文件加下存储视频帧图像
+# =============================================================================
+
+    path, filename = os.path.split(videopath)
+    file, ext = os.path.splitext(filename)
+    
+    savepath = os.path.join(path, "{}_imgs".format(file))
+    if not os.path.exists(savepath):
+        os.makedirs(savepath)
+
+    cap = cv2.VideoCapture(videopath)
+    k = 0
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        
+        k += 1
+        cv2.imwrite(os.path.join(savepath, "{}.png".format(k)), frame)
+        
+
+
+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()