185 lines
8.0 KiB
Python
185 lines
8.0 KiB
Python
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") |