diff --git a/deep_sort/deep/feature_extractor.py b/deep_sort/deep/feature_extractor.py index 5b54525..c7e48ec 100644 --- a/deep_sort/deep/feature_extractor.py +++ b/deep_sort/deep/feature_extractor.py @@ -52,8 +52,9 @@ class MyExtractor(object): def __init__(self, model_path, use_cuda=True): self.device = "cuda" if torch.cuda.is_available() and use_cuda else "cpu" self.net = RestNet18() - self.net.load_state_dict(torch.load(model_path)) - self.net = self.net.to(self.device) + self.net.load_state_dict(torch.load(model_path,map_location=torch.device(self.device))) + if self.device=="cuda": + self.net = self.net.to(self.device) logger = logging.getLogger("root.tracker") logger.info("Loading weights from {}... Done!".format(model_path)) self.raw_transformer = transforms.Compose([ diff --git a/example_reid.py b/example_reid.py index 6ca6d8a..75806d2 100644 --- a/example_reid.py +++ b/example_reid.py @@ -30,6 +30,20 @@ def opencv_box_plot(img:cv2.Mat,pred_img:np.array): 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 = {} @@ -49,8 +63,6 @@ def opencv_sort_plot(img:cv2.Mat,pred_yolo:np.array,pred_sort:np.array): 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 def extractidfeature(id_path:str,extractor): @@ -115,6 +127,7 @@ if __name__ == '__main__': 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)): @@ -139,17 +152,21 @@ if __name__ == '__main__': if term[0] == MOT_id: count+=1 if count > MAX_BUFFLEN//3: - cx = int((dpsort[i][0]+dpsort[i][2])/2) - cy = int((dpsort[i][1]+dpsort[i][3])/2) + 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)) - cv2.putText(img_drawing,f"{paths[index_d]}",(cx-10,cy-10), cv2.FONT_HERSHEY_PLAIN, 1, color, 2) + + # 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 - frame_t = opencv_sort_plot(img_drawing,detections,dpsort) - target.write(frame_t) + cv2.imshow("test",img_drawing) + cv2.waitKey(1) + target.write(img_drawing) frame_counter += 1 # save the video and end the program