class ImageProcessor(object): def __init__(self, path=config.video_path, show_img=True): self.pose_estimator = PoseEstimator() self.object_detector = ObjectDetectionYolo() self.BBV = BBoxVisualizer() self.IDV = IDVisualizer() self.object_tracker = ObjectTracker() self.video_path = path self.cap = cv2.VideoCapture(self.video_path) self.img = [] self.img_black = [] self.show_img = show_img self.locator = Locator([1, 2]) def process(self): cnt = 0 while True: ret, frame = self.cap.read() if ret: frame = cv2.resize(frame, config.frame_size) with torch.no_grad(): inps, orig_img, boxes, scores, pt1, pt2 = self.object_detector.process( frame) if boxes is not None: # cv2.imshow("bbox", self.BBV.visualize(boxes, copy.deepcopy(frame))) key_points, self.img, self.img_black = self.pose_estimator.process_img( inps, orig_img, boxes, scores, pt1, pt2) if len(key_points) > 0: id2ske_all, id2bbox_all = self.object_tracker.track( boxes, key_points) id2ske, id2bbox = self.locator.locate_user( id2ske_all, id2bbox_all) # process skeleton if self.show_img: cv2.imshow( "id_bbox_all", self.IDV.plot_bbox_id( id2bbox_all, copy.deepcopy(frame))) cv2.imshow( "id_ske_all", self.IDV.plot_skeleton_id( id2ske_all, copy.deepcopy(frame))) cv2.imshow( "id_bbox_located", self.IDV.plot_bbox_id( id2bbox, copy.deepcopy(frame))) cv2.imshow( "id_ske_located", self.IDV.plot_skeleton_id( id2ske, copy.deepcopy(frame))) self.__show_img() else: if self.show_img: self.__show_img() else: # cv2.imshow("bbox", frame) # cv2.imshow("id", frame) self.img, self.img_black = frame, frame if self.show_img: self.__show_img() cnt += 1 print(cnt) else: self.cap.release() cv2.destroyAllWindows() break def __show_img(self): cv2.imshow("result", self.img) cv2.moveWindow("result", 1200, 90) cv2.imshow("result_black", self.img_black) cv2.moveWindow("result_black", 1200, 540) cv2.waitKey(1) def process_single_person(self, frame): frame = cv2.resize(frame, config.frame_size) with torch.no_grad(): inps, orig_img, boxes, scores, pt1, pt2 = self.object_detector.process( frame) key_points, img, img_black = self.pose_estimator.process_img( inps, orig_img, boxes, scores, pt1, pt2) return key_points[0], img, img_black def process_multiple_person(self, frame): frame = cv2.resize(frame, config.frame_size) with torch.no_grad(): inps, orig_img, boxes, scores, pt1, pt2 = self.object_detector.process( frame) key_points, img, img_black = self.pose_estimator.process_img( inps, orig_img, boxes, scores, pt1, pt2) if len(key_points) > 0: id2ske_all, id2bbox_all = self.object_tracker.track( boxes, key_points) id2ske, id2bbox = self.locator.locate_user( id2ske_all, id2bbox_all) return id2ske, img, img_black
class HumanDetection: def __init__(self, resize_size, show_img=True): self.object_detector = ObjectDetectionYolo(cfg=config.yolo_cfg, weight=config.yolo_weight) self.object_tracker = ObjectTracker() self.pose_estimator = PoseEstimator(pose_cfg=config.pose_cfg, pose_weight=config.pose_weight) self.BBV = BBoxVisualizer() self.KPV = KeyPointVisualizer() self.IDV = IDVisualizer() self.boxes = tensor([]) self.boxes_scores = tensor([]) self.frame = np.array([]) self.id2bbox = {} self.CNN_model = CNNInference() self.kps = {} self.kps_score = {} self.show_img = show_img self.resize_size = resize_size def init(self): self.object_tracker.init_tracker() def clear_res(self): self.boxes = tensor([]) self.boxes_scores = tensor([]) self.frame = np.array([]) self.id2bbox = {} self.kps = {} self.kps_score = {} def visualize(self): img_black = np.full((self.resize_size[1], self.resize_size[0], 3), 0).astype(np.uint8) if config.plot_bbox and self.boxes is not None: self.BBV.visualize(self.boxes, self.frame) if config.plot_kps and self.kps is not []: self.KPV.vis_ske(self.frame, self.kps, self.kps_score) self.KPV.vis_ske_black(img_black, self.kps, self.kps_score) if config.plot_id and self.id2bbox is not None: self.IDV.plot_bbox_id(self.id2bbox, self.frame) self.IDV.plot_skeleton_id(self.kps, self.frame) return self.frame, img_black def process_img(self, frame, gray=False): self.clear_res() self.frame = frame with torch.no_grad(): if gray: gray_img = gray3D(copy.deepcopy(frame)) box_res = self.object_detector.process(gray_img) else: box_res = self.object_detector.process(frame) self.boxes, self.boxes_scores = self.object_detector.cut_box_score(box_res) if box_res is not None: self.id2bbox = self.object_tracker.track(box_res) self.id2bbox = eliminate_nan(self.id2bbox) boxes = self.object_tracker.id_and_box(self.id2bbox) inps, pt1, pt2 = crop_bbox(frame, boxes) if inps is not None: kps, kps_score, kps_id = self.pose_estimator.process_img(inps, boxes, pt1, pt2) self.kps, self.kps_score = self.object_tracker.match_kps(kps_id, kps, kps_score) return self.kps, self.id2bbox, self.kps_score def classify_whole(self, pred_img, show_img): pred = self.CNN_model.classify_whole(pred_img, show_img) return pred def classify(self, pred_img, show_img, id2bbox): preds = self.CNN_model.classify(pred_img, id2bbox) self.CNN_model.visualize(show_img, preds) return preds