ComDesignProject/example_reid.py

import torch
from torchvision import transforms
import numpy as np
import cv2
import os

import param
from detect import YOLOv7
from deep_sort import deep_sort as dsort

def get_img_test(raw_img:np.array):
    raw_transform = transforms.Compose([transforms.ToPILImage(),
                                        transforms.Resize((360,640)),
                                        transforms.Pad((0,(640-360)//2)),])
    return raw_transform(raw_img)

def reverse_box_get_from_yolo(detection):
    x1 = int(detection[0]*1920/640)
    x2 = int(detection[2]*1920/640)
    y1 = int((detection[1]-140)*1080/360)
    y2 = int((detection[3]-140)*1080/360)
    return x1,y1,x2,y2

def opencv_box_plot(img:cv2.Mat,pred_img:np.array):
    pred_img = pred_img.astype(np.uint)
    for info in pred_img:
        x1,y1,x2,y2 = info[0],info[1],info[2],info[3]
        cv2.rectangle(img,(x1,y1-140),(x2,y2-140),(255,0,0),2)

    cv2.imshow("test",img)
    cv2.waitKey(0)
    cv2.imwrite("result.png",img)
    
def opencv_match_pointer_plot(img:cv2.Mat,start_base:tuple,text:str,color:tuple):
    SIZE = 50
    brush = [0,0]
    brush[0]=start_base[0]
    brush[1]=start_base[1]
    brush_push = [brush[0]+SIZE,brush[1]-SIZE]
    cv2.line(img,brush,brush_push,color,1)
    # brush=brush_push
    # brush_push[0]+=SIZE*2
    # brush_push[1]+=0
    # cv2.line(img,brush,brush_push,color,1)    
    cv2.putText(img,text,brush_push, cv2.FONT_HERSHEY_PLAIN, 1, color, 2)
    return img

routine = {}
    
def opencv_sort_plot(img:cv2.Mat,pred_yolo:np.array,pred_sort:np.array):
    for info in pred_yolo:
        x1,y1,x2,y2 = int(info[0]),int(info[1]),int(info[2]),int(info[3])
        cv2.rectangle(img,(x1,y1),(x2,y2),(0,255,0),1)    
    for info in pred_sort:
        x1,y1,x2,y2,id = info[0],info[1],info[2],info[3],info[4]
        cxy = (int((x1+x2)/2),int((y1+y2)/2))
        color = (int(id%3*100),int(id%4*75),int(id%5*50))
        cv2.putText(img,str(id),cxy, cv2.FONT_HERSHEY_PLAIN, 1.0, color, 2)
        if routine.get(id) is None:
            routine[id] = [cxy]
        else:
            for i in range(1,len(routine[id])):
                cv2.line(img,routine[id][i-1],routine[id][i],color,1)
            cv2.line(img,routine[id][-1],cxy,color,1)
            routine[id].append(cxy)
    return img

def extractidfeature(id_path:str,extractor):
    img_id = cv2.imread(id_path)
    tensor_id = extractor([img_id])
    return tensor_id

# example        
if __name__ == '__main__':
    MAX_BUFFLEN        = 10
    DISTANCE_THRESHOLD = 0.4
    SAVE_TXT_FLAG = True
    
    my_option = param.Parameters()
    
    print(torch.__version__)
    print(my_option.device)
    
    query_path = f"./mydataset/query/cam{my_option.query_index}/"                   # query from .jpg photos
    gallary_path = f"mydataset/video/cam{my_option.gallary_index}.mp4"              # source from mp4 via yolo
    output_video_path = f"output/output{my_option.gallary_index}_{my_option.query_index}.mp4"    # output file path
    output_txt_path = f"output/rank10_detection_{my_option.gallary_index}_{my_option.query_index}.txt"
    
    query_features = []
    paths = os.listdir(query_path)
    query_match_buff = [[] for p in paths]
    query_matched = [-1 for p in paths]
    
    # load detection model and deepsort model
    model = YOLOv7(my_option)         
    deepsort = dsort.DeepSort(model_path=my_option.weights_reid,model_config="self",use_cuda=(torch.device("cuda:0") == my_option.device))
    
    # get query feature id
    for p in paths:
        query_features.append(extractidfeature(query_path+p,deepsort.extractor))
        
    with torch.no_grad():
        sources = cv2.VideoCapture(gallary_path)
        target = cv2.VideoWriter(output_video_path,cv2.VideoWriter_fourcc('m', 'p', '4', 'v'),24,(1920,1080))
        frame_counter = 0
        while True:
            ret,frame = sources.read()
            if ret is False:
                break
            
            img    = get_img_test(frame)                        # get image as torch.Tensor with size of [1,1,640,640]
            img_tensor = transforms.ToTensor()(img).unsqueeze(dim=0).to(my_option.device)     
            detections = model.detect(img_tensor)               # get the sequence of result   
            detections = detections[0].detach().cpu().numpy()   # the single img is index 0
            
            for detection in detections:
                detection[0],detection[1],detection[2],detection[3] = reverse_box_get_from_yolo(detection[0:4])
            
            bbox_xywhs = []
            confs      = []
            for xyxycc in detections: 
                xywh   =  deepsort._xyxy_to_xywh(xyxycc[0:4])
                conf   =  xyxycc[4]
                clas   =  int(xyxycc[5])
                bbox_xywhs.append(xywh[:])
                confs.append(conf)
            dpsort, features = deepsort.update(bbox_xywh=np.array(bbox_xywhs),confidences=np.array(confs),ori_img=np.array(frame))
            
            img_drawing = np.array(frame)
            img_drawing = opencv_sort_plot(img_drawing,detections,dpsort)
            
            # ReID: compute features
            for i in range(len(dpsort)):
                track = dpsort[i]
                MOT_id = track[-1]
                feature_distances = []
                for j in range(len(query_features)):
                    feature_distance = np.mean(np.power(features[i]-query_features[j],2))
                    feature_distances.append(feature_distance)
                    
                # ReID: rank distances for every feature   
                feature_distances_sorted = sorted(feature_distances)
                for distance in feature_distances_sorted:
                    if distance < DISTANCE_THRESHOLD:
                        index_d = feature_distances.index(distance)
                        query_match_buff[index_d].append((MOT_id,distance,(frame_counter,track[0],track[1],track[2],track[3]))) 
                        query_match_buff[index_d].sort(key=lambda x:x[1],reverse=True)
                        if len(query_match_buff[index_d]) > MAX_BUFFLEN:
                            query_match_buff[index_d].pop()
                            count = 0
                            for term in query_match_buff[index_d]:
                                if term[0] == MOT_id:
                                    count+=1
                            if count > MAX_BUFFLEN//3:
                                cx = int((track[0]+track[2])/2)
                                cy = int((track[1]+track[3])/2)
                                color = (int(index_d%2*100),int(index_d%3*75),int(index_d%4*50))
                                
                                # img_drawing=opencv_match_pointer_plot(img_drawing,(cx,track[1]-10),f"{paths[index_d]}",color)
                                # cv2.rectangle(img_drawing,(track[i][0],track[i][1]),(track[i][2],track[i][3]),color,1)                
                                cv2.putText(img_drawing,f"q_{paths[index_d]}".split(".")[0],(track[0],track[1]), cv2.FONT_HERSHEY_PLAIN, 1, color, 2)
                            break
                    else:
                        break
            
            # Final plot and show video
            cv2.imshow("test",img_drawing)
            cv2.waitKey(1)
            target.write(img_drawing)
            frame_counter += 1
        
        # save the video and end the program
        target.release()
        
        # rank10 output as txt
        if SAVE_TXT_FLAG is True:
            with open(output_txt_path,mode="w+") as f:
                result = []
                for i in range(len(query_match_buff)):
                    for raw_info in query_match_buff[i]:
                        detect_info = raw_info[-1]
                        result.append((detect_info[0],i,detect_info[1],detect_info[2],detect_info[3],detect_info[4]))
                        result.sort(key=lambda x:x[0])
                for info in result:
                    f.write(str(info[0])+","+str(info[1])+","+str(info[2])+","+str(info[3])+","+str(info[4])+","+str(info[5])+"\n")