56 lines
1.5 KiB
Python
56 lines
1.5 KiB
Python
from AIDetector_pytorch import Detector
|
|
import imutils
|
|
import cv2
|
|
import os
|
|
os.environ['KMP_DUPLICATE_LIB_OK']='TRUE'
|
|
file = open('track_result.txt', 'w')
|
|
|
|
def main():
|
|
|
|
name = 'demo'
|
|
|
|
det = Detector()
|
|
# cap = cv2.VideoCapture('D:/TH/5_smoke/Yolov5-Deepsort-main/smogfire_video14.mp4')
|
|
cap = cv2.VideoCapture('/home/thsw/WJ/nyh/CODE/Yolov5-Deepsort-main/video/roadsign_orin.mp4')
|
|
fps = int(cap.get(5))
|
|
print('fps:', fps)
|
|
t = int(1000/fps)
|
|
|
|
videoWriter = None
|
|
frame_count = 0
|
|
|
|
while True:
|
|
|
|
# try:
|
|
_, im = cap.read()
|
|
if im is None:
|
|
break
|
|
frame_count += 1
|
|
file.write(str('第'+str(frame_count)+'帧\n*********************************************'))
|
|
result = det.feedCap(im)
|
|
result = result['frame']
|
|
result = imutils.resize(result, height=500)
|
|
if videoWriter is None:
|
|
fourcc = cv2.VideoWriter_fourcc(
|
|
'm', 'p', '4', 'v') # opencv3.0
|
|
videoWriter = cv2.VideoWriter(
|
|
'./detect_result/result_roadsign.mp4', fourcc, fps, (result.shape[1], result.shape[0]))
|
|
|
|
videoWriter.write(result)
|
|
# cv2.imshow(name, result)
|
|
# cv2.waitKey(t)
|
|
|
|
# if cv2.getWindowProperty(name, cv2.WND_PROP_AUTOSIZE) < 1:
|
|
# # 点x退出
|
|
# break
|
|
# except Exception as e:
|
|
# print(e)
|
|
# break
|
|
|
|
cap.release()
|
|
videoWriter.release()
|
|
# cv2.destroyAllWindows()
|
|
|
|
if __name__ == '__main__':
|
|
|
|
main() |