# -*- 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()