Files
detecttracking/pipeline_extract_subimg.py
2024-10-07 17:27:00 +08:00

125 lines
3.1 KiB
Python

# -*- coding: utf-8 -*-
"""
Created on Sun Sep 29 08:59:21 2024
@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()