127 lines
3.2 KiB
Python
127 lines
3.2 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
Created on Sun Sep 29 08:59:21 2024
|
||
|
||
针对现场采集的视频，利用算法pipeline提取运动轨迹内的subimg，代替人工图像筛选、标注
|
||
|
||
@author: ym
|
||
"""
|
||
import os
|
||
import cv2
|
||
import pickle
|
||
from pathlib import Path
|
||
from track_reid import parse_opt, yolo_resnet_tracker
|
||
|
||
from tracking.dotrack.dotracks_back import doBackTracks
|
||
from tracking.dotrack.dotracks_front import doFrontTracks
|
||
|
||
IMGFORMATS = '.bmp', '.jpeg', '.jpg', 'png', 'tif', 'tiff', 'webp', 'pfm'
|
||
VIDFORMATS = '.avi', '.gif', '.m4v', '.mkv', '.mov', '.mp4', '.ts', '.wmv'
|
||
|
||
std_feature_path = r"\\192.168.1.28\share\测试_202406\contrast\std_features_2192_ft32vsft16"
|
||
|
||
|
||
opt = parse_opt()
|
||
optdict = vars(opt)
|
||
|
||
|
||
def pipeline(eventpath, savepath):
|
||
# eventpath = r"\\192.168.1.28\share\测试_202406\0918\images1\20240918-110822-1bc3902e-5a8e-4e23-8eca-fb3f02738551_6938314601726"
|
||
optdict["project"] = savepath
|
||
|
||
'''Yolo + Resnet + Tracker'''
|
||
optdict["source"] = eventpath
|
||
optdict["save_dir"] = savepath
|
||
optdict["nosave"] = False
|
||
|
||
tracksdict = yolo_resnet_tracker(**optdict)
|
||
|
||
bboxes = tracksdict['TrackBoxes']
|
||
|
||
basename = os.path.basename(eventpath)
|
||
base, ext = os.path.splitext(basename)
|
||
|
||
if base.find('front')>=0:
|
||
vts = doFrontTracks(bboxes, tracksdict)
|
||
vts.classify()
|
||
|
||
if base.find('back')>=0:
|
||
vts = doBackTracks(bboxes, tracksdict)
|
||
vts.classify()
|
||
|
||
tracks = [t for t in vts.tracks if t.cls>0 and not t.is_static()]
|
||
# tracks = [t for t in vts.tracks if t.cls>0]
|
||
|
||
for track in tracks:
|
||
# for track in vts.Residual:
|
||
for *xyxy, tid, conf, cls, fid, bid in track.boxes:
|
||
img = tracksdict[f'frame_{int(fid)}']["imgs"][int(bid)]
|
||
|
||
imgpth = savepath / Path(f'{base}_tid-{int(tid)}_fid-{int(fid)}_bid-{int(bid)}.jpg')
|
||
|
||
cv2.imwrite(imgpth, img)
|
||
|
||
return len(vts.Residual)
|
||
|
||
|
||
|
||
|
||
def main():
|
||
|
||
videopath = r"\\192.168.1.28\share\上海中环店采集视频\21-25\videos\1\back"
|
||
savepath = r"D:\contrast\barcodes"
|
||
|
||
vpaths = []
|
||
for root, dirs, files in os.walk(videopath):
|
||
vpth = [os.path.join(root, f) for f in files if os.path.splitext(f)[-1] in VIDFORMATS]
|
||
vpaths.extend(vpth)
|
||
|
||
manual_txt = os.path.join(savepath, 'manual_videos.txt')
|
||
file = open(manual_txt, 'a', encoding='utf-8')
|
||
|
||
|
||
|
||
manual = []
|
||
k = 0
|
||
for vpath in vpaths:
|
||
videoname = os.path.basename(vpath)
|
||
|
||
vname, ext = os.path.splitext(videoname)
|
||
barcode = videoname.split('_')[0]
|
||
|
||
subpath = os.path.join(savepath, barcode, vname)
|
||
subpath = Path(subpath)
|
||
if not subpath.exists():
|
||
subpath.mkdir(parents=True, exist_ok=True)
|
||
|
||
ntract = pipeline(vpath, subpath)
|
||
|
||
if ntract==0:
|
||
manual.append(vpath)
|
||
file.write(vpath)
|
||
file.write("\n")
|
||
|
||
print(f"{videoname} done!!!")
|
||
k += 1
|
||
if k==10:
|
||
break
|
||
|
||
file.close()
|
||
|
||
|
||
|
||
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|