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