import os import math import numpy as np import matplotlib.pyplot as plt from PIL import Image, ImageDraw import cv2 from tqdm import tqdm import torch import torchvision import torchvision.transforms as transforms from torchvision.transforms import ToTensor import timm from mmdet.apis import async_inference_detector, inference_detector from mmdet.apis.inference import init_detector from mmpose.apis import (inference_topdown, init_model) # dpt from torchvision.transforms import Compose # from dpt.models import DPTDepthModel # from dpt.midas_net import MidasNet_large # from dpt.transforms import Resize, NormalizeImage, PrepareForNet def load_combined_models(): # Pose estimation model pose_config_file = './mmpose/configs/body_2d_keypoint/topdown_heatmap/coco/td-hm_hrnet-w32_8xb64-210e_coco-256x192.py' pose_checkpoint_file = './pretrained/td-hm_hrnet-w32_8xb64-210e_coco-256x192-81c58e40_20220909.pth' pose_model = init_model(pose_config_file, pose_checkpoint_file, device='cuda:0') person_config_file = "./mmyolo/configs/yolov8/yolov8_x_mask-refine_syncbn_fast_8xb16-500e_coco.py" person_checkpoint = './pretrained/yolov8_x_mask-refine_syncbn_fast_8xb16-500e_coco_20230217_120411-079ca8d1.pth' person_model = init_detector(person_config_file, person_checkpoint, device='cuda:0') return pose_model, person_model def mmdet3x_convert_to_bboxes_mmdet(results, threshold, prefix = 'class'): boxes_list = [] scores_list = [] labels_list = [] confidence_score = results.pred_instances.scores.tolist() for i, conf in enumerate(confidence_score): if conf >= threshold: # print(conf) extracted_box = results.pred_instances.bboxes[i].cpu().tolist() extracted_label = results.pred_instances.labels[i].cpu().tolist() boxes_list.append([int(extracted_box[0]), int(extracted_box[1]), int(extracted_box[2]), int(extracted_box[3])]) scores_list.append(conf) labels_list.append('{}_{}'.format(prefix,extracted_label+1)) return boxes_list, labels_list, scores_list def depth_preprocess_image(image): # Resize the input image to dimensions divisible by 32 h, w, _ = image.shape new_h = h - (h % 32) new_w = w - (w % 32) resized_image = cv2.resize(image, (new_w, new_h)) input_transform = transforms.Compose( [ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ] ) return input_transform(resized_image).unsqueeze(0) def postprocess_depth(depth_tensor, original_size): depth_map = depth_tensor.squeeze().cpu().numpy() depth_map = cv2.resize(depth_map, original_size) depth_map = depth_map / depth_map.max() return depth_map def extract_depth_map(model, image): img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) input_tensor = depth_preprocess_image(img_rgb) with torch.no_grad(): depth_tensor = model(input_tensor) depth_map = postprocess_depth(depth_tensor, (img_rgb.shape[1], img_rgb.shape[0])) return depth_map def extract_ankles_points(keypoints, keypoints_confidence): keypoint_dict = {} keypoint_dict['left_ankle'] = [keypoints[15][0], keypoints[15][1], keypoints_confidence[15]] keypoint_dict['right_ankle'] = [keypoints[16][0], keypoints[16][1], keypoints_confidence[16]] return keypoint_dict def extract_keypoints(keypoints): left_shoulder= [keypoints[5][0], keypoints[5][1]] right_shoulder= [keypoints[6][0], keypoints[6][1]] left_hip= [keypoints[11][0], keypoints[11][1]] right_hip= [keypoints[12][0], keypoints[12][1]] return left_shoulder, right_shoulder, left_hip, right_hip def estimate_ankle_midpoint(keypoints): left_shoulder_coord, right_shoulder_coord, left_hip_coord, right_hip_coord = extract_keypoints(keypoints) # Calculate the midpoint of the shoulders and hips shoulder_midpoint =[int((left_shoulder_coord[0] + right_shoulder_coord[0]) / 2), int((left_shoulder_coord[1] + right_shoulder_coord[1]) / 2)] hip_midpoint = [int((left_hip_coord[0] + right_hip_coord[0]) / 2), int((left_hip_coord[1] + right_hip_coord[1]) / 2)] # Calculate the vector from the shoulder midpoint to the hip midpoint shoulder_to_hip_vector = [hip_midpoint[0] - shoulder_midpoint[0], hip_midpoint[1] - shoulder_midpoint[1]] # # Double the length of the vector to approximate the upper body's mirror image mirroring_fator= 1.2 mirrored_vector = [shoulder_to_hip_vector[0] * mirroring_fator, shoulder_to_hip_vector[1] * mirroring_fator] # Add the resulting vector to the hip midpoint to obtain the estimated ankle midpoint estimated_ankle_midpoint = [hip_midpoint[0] + mirrored_vector[0], hip_midpoint[1] + mirrored_vector[1]] # estimated_ankle_midpoint should be perpendicular with hip_midpoint estimated_ankle_midpoint =[hip_midpoint[0],estimated_ankle_midpoint[1]] return estimated_ankle_midpoint def is_point_inside_polygon(x, y, danger_points): n = len(danger_points) inside = False p1x, p1y = danger_points[0] for i in range(n + 1): p2x, p2y = danger_points[i % n] if y > min(p1y, p2y): if y <= max(p1y, p2y): if x <= max(p1x, p2x): if p1y != p2y: xinters = (y - p1y) * (p2x - p1x) / (p2y - p1y) + p1x if p1x == p2x or x <= xinters: inside = not inside p1x, p1y = p2x, p2y return inside def check_inside_polygon(worker_centers, danger_points): inside = [] for i, worker_center in enumerate(worker_centers): x, y = worker_center inside.append(is_point_inside_polygon(x, y, danger_points)) return inside def draw_polygon(frame, danger_points): img = Image.fromarray(frame) draw = ImageDraw.Draw(img) draw.polygon(danger_points, outline ="red", width = 5) return np.array(img) # DPT def preprocess_img_dpt(image_path): img = cv2.imread(image_path) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0 return img