forked from wffjwbbf/ComDesignProject
84 lines
3.5 KiB
Python
84 lines
3.5 KiB
Python
import torch
|
|
from torchvision import transforms
|
|
import numpy as np
|
|
import cv2
|
|
|
|
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 opencv_box_plot(img:cv2.Mat,pred_img:np.array):
|
|
pred_img = pred_img.astype(np.uint)
|
|
print(type(img))
|
|
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)
|
|
print(type(img))
|
|
cv2.imshow("test",img)
|
|
cv2.waitKey(0)
|
|
cv2.imwrite("result.png",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)
|
|
cv2.imshow("test",img)
|
|
cv2.waitKey(1)
|
|
return img
|
|
|
|
|
|
|
|
# example
|
|
if __name__ == '__main__':
|
|
dsort_path = "./deep_sort/deep/checkpoint/market_agw_R50.pth"
|
|
my_option = param.Parameters()
|
|
model = YOLOv7(my_option) # load model in head
|
|
deepsort = dsort.DeepSort(model_path= dsort_path,model_config=None,use_cuda=(torch.device("cuda:0") == my_option.device))
|
|
with torch.no_grad():
|
|
sources = cv2.VideoCapture("test.mp4")
|
|
target = cv2.VideoWriter("output.mp4",cv2.VideoWriter_fourcc('M', 'P', '4', '2'),24,(640,640))
|
|
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)
|
|
result = model.detect(img_tensor) # get the sequence of result
|
|
result = result[0].detach().cpu().numpy() # the single img is index 0
|
|
bbox_xywhs = []
|
|
confs = []
|
|
for xyxycc in result:
|
|
xywh = deepsort._xyxy_to_xywh(xyxycc[0:4])
|
|
conf = xyxycc[4]
|
|
clas = int(xyxycc[5])
|
|
bbox_xywhs.append(xywh[:])
|
|
confs.append(conf)
|
|
dpsort = deepsort.update(bbox_xywh=np.array(bbox_xywhs),confidences=np.array(confs),ori_img=np.array(img))
|
|
print(f"result:{result}")
|
|
print(f"dpsort:{dpsort}")
|
|
frame_t = opencv_sort_plot(np.array(img),result,dpsort)
|
|
target.write(frame_t)
|
|
target.release()
|
|
# img_src = cv2.imread("test.png") # re-read for imshow
|
|
# img_src = cv2.resize(img_src,(640,360))
|
|
# opencv_box_plot(img_src,result) # show the result in picture |