pepeline_extract_subimg.py: extract subimg

pipeline_extract_subimg.py

@@ -0,0 +1,125 @@
+# -*- 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()
+    
+    
+    
+    
+    
+    
+    
+    
+    
+