# coding:utf-8 import os import glob import json import re from copy import deepcopy import math import cv2 import copy import matplotlib.cm as cm import numpy as np from pyquaternion import Quaternion import random from multiprocessing import Pool, cpu_count import traceback # 前11个颜色对应车辆类别，颜色区别明显 # det_cls_index: car, suv, pickup, van, bus, medium_car, truck, tanker, large_truck, construction_vehicle, special_vehicle colors = [(255, 0, 0), # 0: car - 红色 (0, 255, 0), # 1: suv - 绿色 (0, 0, 255), # 2: pickup - 蓝色 (255, 255, 0), # 3: van - 黄色 (0, 255, 255), # 4: bus - 青色 (255, 0, 255), # 5: medium_car - 洋红色 (255, 165, 0), # 6: truck - 橙色 (128, 0, 255), # 7: tanker - 紫色 (50, 205, 50), # 8: large_truck - 酸橙绿 (255, 105, 180), # 9: construction_vehicle - 热粉色 (0, 191, 255), # 10: special_vehicle - 深天蓝 (240, 32, 160), (128, 128, 128), (204, 161, 141),(101, 134, 179),(89, 109, 61),(141, 137, 194),(132, 228, 208), (249, 118, 35),(80, 91, 182),(222, 91, 125),(91, 63, 123),(173, 232, 91), (255, 164, 26),(44, 56, 142),(74, 148, 81),(179, 42, 50),(250, 226, 21),(191, 81, 160), (6, 142, 172),(252, 252, 252),(230, 230, 230),(200, 200, 200),(143, 143, 142),(100, 100, 100),(50, 50, 50), (64, 224, 208), (255, 140, 0), (186, 85, 211), (255, 20, 147), (72, 209, 204), (218, 112, 214), (255, 215, 0), (127, 255, 212), (220, 20, 60), (46, 139, 87), (255, 99, 71)] # 扩充到40个颜色 subcls = { 'kNegative' : 0, # 背景 'kBus' : 1, # 大巴 'kCar' : 2, # 小轿车,suv 'kMiniBus' : 3, # 面包车 'kBucketTruck' : 4, # 斗卡 'kContainerTruck' : 5, # 箱卡 'kTricycle' : 6, # 三轮车 'kTanker' : 7, # 油罐车,晒水车（车身带有圆形,椭圆形,半圆形的罐） 'kCementTankTruck' : 8, # 水泥罐车 'kPickup' : 9, # 皮卡 'kSedimentTruck' : 10, # 渣土车 'kIveco' : 11, # 依维柯 'kSpecialCar' : 12, # 异型车 'kCityAuto' : 13, # 市政车 'kVehicleUnknown' : 14, # 未知车辆 'kWrecker' : 15, # 小型拖车，清障车 'kFlatTransporter' : 16,# 大型平板拖车 'kTractorHead' : 17, # 卡车车头,牵引车车头 'kSpecialFlatTransporter' : 18, # 大型特种拖车（相对于大型平板拖车会多挡板） 'kCarCarrier' : 19, # 轿运车 # 盖布车（车身上的货物全部被布盖着，同时对车辆尾部进行了遮挡） 'kTruckWithTrap' : 20, 'kEngineeringVehicle' : 21, # 工程车 'kSweeper' : 22, # 扫地车 'kRoadServiceCar' : 23, # 道路维修车 'kSanitationTruck' : 24, # 环卫车 'kGarbageTruck' : 25, # 垃圾车 } det_cls_index = [ "car", "suv", "pickup", "van", "bus", "medium_car", "truck", "tanker", "large_truck", "construction_vehicle", "special_vehicle", "unknown", 'pedestrian', 'bicycle', "bicyclist", "motorcycle", "motorcyclist", "tricycle", "tricyclist", 'traffic_light_bbox', 'traffic_light_bulb', 'traffic_sign', 'animal', 'movable_object', 'warning_triangle', 'traffic_cone', 'water_barrier', 'crash_barrel', 'movable_barrier', 'bollard', 'sphere_bollard', 'cube_bollard', 'cylinder_bollard', 'construction_barrier', 'other_barrier', 'road_barrier_unknown', "wheel", "plate", "face" ] def readcalib(calibpath): calib_params, cam_intrinsic, distortion_param, cam_rectify, _, _ = build_camera_config(calibpath) Vehiclepos = np.array(calib_params['pos']) VehiclePitch, VehicleYaw, VehicleRoll = calib_params['pitch'],calib_params['yaw'],calib_params['roll'] # camera_yaw = calib_params['yaw'] # camera_pitch = calib_params['pitch'] cu = calib_params['cu'] cv = calib_params['cv'] focal_u = calib_params['focal_u'] focal_v = calib_params['focal_v'] vanish_x = cu + focal_u * np.tan(VehicleYaw * np.pi/180) vanish_y = cv - focal_v * np.tan(VehiclePitch * np.pi/180) centerx ,centery = vanish_x,vanish_y + 0 roi0 = [0,int(centery - 880 / 2.0),1920,int(centery + 880 / 2.0)] roi1 = [int(centerx - 768 / 2.0),int(centery - 352 / 2.0),int(centerx + 768 / 2.0),int(centery + 352 / 2.0)] # opt.roi[0] = [0,120,1920,1080] # opt.roisize[0] = [704,352] # opt.roisize[1] = [704,352] # centerx ,centery = vanish_x,vanish_y + 0 # roi_rect = [int(centerx - 704 / 2.0), int(centery - 352 / 2.0), int(centerx + 704/2.0), int(centery + 352/2.0)] # opt.roi[1] = roi_rect return cam_intrinsic,distortion_param, cam_rectify,Vehiclepos,VehiclePitch,VehicleYaw, VehicleRoll def build_camera_config(calib_file): calib_params = json.load(open(calib_file, "r")) cam_intrinsic = np.array([calib_params['focal_u'], 0, calib_params['cu'], 0, calib_params['focal_v'], calib_params['cv'], 0, 0, 1], dtype=np.float64).reshape(3, 3) distortion_param = np.array(calib_params['distort_coeffs']) image_w = int(calib_params.get('image_width') or round(calib_params['cu'] * 2)) image_h = int(calib_params.get('image_height') or round(calib_params['cv'] * 2)) if len(distortion_param) == 5: cam_rectify, _ = cv2.getOptimalNewCameraMatrix(cam_intrinsic, distortion_param, (image_w, image_h), 1, (image_w, image_h)) elif len(distortion_param) == 4: cam_rectify = cv2.fisheye.estimateNewCameraMatrixForUndistortRectify(cam_intrinsic, distortion_param, [image_w, image_h], np.eye(3), balance=1) else: raise ValueError(f"Unsupported distort_coeffs length {len(distortion_param)} in {calib_file}") return calib_params, cam_intrinsic, distortion_param, cam_rectify, image_w, image_h def infer_camera_name_from_path(path): matches = re.findall(r'camera\d+', os.path.normpath(path), flags=re.IGNORECASE) return matches[-1].lower() if matches else None def resolve_camera_calib_path(calib_root, reference_path=''): if os.path.isfile(calib_root): return calib_root camera_name = infer_camera_name_from_path(reference_path) or infer_camera_name_from_path(calib_root) candidate_dirs = [calib_root] l2_calib_dir = os.path.join(calib_root, 'L2_calib') if os.path.isdir(l2_calib_dir): candidate_dirs.insert(0, l2_calib_dir) if camera_name: for candidate_dir in candidate_dirs: candidate_path = os.path.join(candidate_dir, f'{camera_name}.json') if os.path.exists(candidate_path): return candidate_path candidate_paths = [] for candidate_dir in candidate_dirs: candidate_paths.extend(sorted(glob.glob(os.path.join(candidate_dir, 'camera*.json')))) if len(candidate_paths) == 1: return candidate_paths[0] if not candidate_paths: raise FileNotFoundError(f'No calibration file found under {calib_root}') raise ValueError( f'Unable to determine calibration file for {reference_path or calib_root}. ' f'Candidates: {candidate_paths}' ) def resolve_image_path(label_path): image_base = label_path.replace('labels', 'images')[0:-4] for ext in ('.png', '.jpg', '.jpeg'): image_path = image_base + ext if os.path.exists(image_path): return image_path raise FileNotFoundError(f'No image file found for label {label_path}. Tried {image_base}.png/.jpg/.jpeg') def _project_label_3d_bbox(co, cam_intrinsic, distortion_param): if len(co) <= 11: return None x3d_ori, y3d_ori, z3d_ori, l3d, h3d, w3d, rotation_y = map(float, co[5:12]) if z3d_ori <= 0: return None corners = cam_corners_front_rear( np.array([x3d_ori, y3d_ori, z3d_ori, l3d, h3d, w3d, rotation_y]), 'whole', 0, ) pts2d = [] for x, y, z in corners: if z <= 0: return None if len(distortion_param) == 5: image_coord = xyz_xy(x / z, y / z, distortion_param, cam_intrinsic) else: image_coord = cam_img_kb(x, y, z, cam_intrinsic, distortion_param) pts2d.append([float(image_coord[0]), float(image_coord[1])]) pts2d = np.array(pts2d, dtype=np.float64) return np.array( [pts2d[:, 0].min(), pts2d[:, 1].min(), pts2d[:, 0].max(), pts2d[:, 1].max()], dtype=np.float64, ) def _build_label_box(co, canvas_w, canvas_h): cx = float(co[1]) * canvas_w cy = float(co[2]) * canvas_h box_w = float(co[3]) * canvas_w box_h = float(co[4]) * canvas_h return np.array([cx - 0.5 * box_w, cy - 0.5 * box_h, cx + 0.5 * box_w, cy + 0.5 * box_h], dtype=np.float64) def resolve_label_canvas_size(calib_path, image_w, image_h, label_temp_all, cam_intrinsic, distortion_param): camera_name = infer_camera_name_from_path(calib_path) if camera_name != 'camera2' or (image_w, image_h) != (1920, 1280): return image_w, image_h candidate_sizes = [(image_w, image_h), (3840, 2160)] candidate_scores = {candidate: [] for candidate in candidate_sizes} valid_count = 0 for co in label_temp_all: projected_box = _project_label_3d_bbox(co, cam_intrinsic, distortion_param) if projected_box is None: continue projected_center = np.array( [(projected_box[0] + projected_box[2]) * 0.5, (projected_box[1] + projected_box[3]) * 0.5], dtype=np.float64, ) for candidate in candidate_sizes: label_box = _build_label_box(co, candidate[0], candidate[1]) label_center = np.array( [(label_box[0] + label_box[2]) * 0.5, (label_box[1] + label_box[3]) * 0.5], dtype=np.float64, ) center_distance = np.linalg.norm(label_center - projected_center) size_distance = np.linalg.norm((label_box[2:] - label_box[:2]) - (projected_box[2:] - projected_box[:2])) candidate_scores[candidate].append(center_distance + 0.5 * size_distance) valid_count += 1 if valid_count >= 8: break if not valid_count: return image_w, image_h best_candidate = min( candidate_sizes, key=lambda candidate: float(np.mean(candidate_scores[candidate])) if candidate_scores[candidate] else float('inf'), ) return best_candidate def plot_box(x, img, color=None, label=None, line_thickness=None): # Plots one bounding box on image img tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thickness color = color or [random.randint(0, 255) for _ in range(3)] c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3])) cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA) if label: tf = max(tl - 1, 1) # font thickness t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0] c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3 cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA) def append_class_legend(img, legend_items): if not legend_items: return img font = cv2.FONT_HERSHEY_SIMPLEX font_scale = 0.6 thickness = 1 margin_x = 20 margin_y = 16 swatch_size = 18 text_gap = 8 item_gap = 24 row_gap = 12 title = '2D legend' title_size = cv2.getTextSize(title, font, font_scale, thickness)[0] text_height = cv2.getTextSize('Ag', font, font_scale, thickness)[0][1] item_height = max(swatch_size, text_height) row_width_limit = img.shape[1] - margin_x layout = [] x = margin_x row = 0 max_row = 0 for label_cls, label_name, color in legend_items: item_text = f'{label_cls}: {label_name}' text_size = cv2.getTextSize(item_text, font, font_scale, thickness)[0] item_width = swatch_size + text_gap + text_size[0] if x + item_width > row_width_limit and x > margin_x: row += 1 x = margin_x layout.append((x, row, item_text, color)) x += item_width + item_gap max_row = row legend_height = margin_y * 3 + title_size[1] + (max_row + 1) * item_height + max_row * row_gap legend_img = np.full((legend_height, img.shape[1], 3), 255, dtype=img.dtype) cv2.line(legend_img, (0, 0), (img.shape[1] - 1, 0), (200, 200, 200), 1, cv2.LINE_AA) cv2.putText(legend_img, title, (margin_x, margin_y + title_size[1]), font, font_scale, (0, 0, 0), thickness, cv2.LINE_AA) base_y = margin_y * 2 + title_size[1] for x, row, item_text, color in layout: row_y = base_y + row * (item_height + row_gap) swatch_top = row_y swatch_bottom = row_y + swatch_size cv2.rectangle(legend_img, (x, swatch_top), (x + swatch_size, swatch_bottom), color, -1, cv2.LINE_AA) cv2.rectangle(legend_img, (x, swatch_top), (x + swatch_size, swatch_bottom), (0, 0, 0), 1, cv2.LINE_AA) cv2.putText(legend_img, item_text, (x + swatch_size + text_gap, row_y + swatch_size - 3), font, font_scale, (0, 0, 0), thickness, cv2.LINE_AA) return cv2.vconcat([img, legend_img]) T1Q2Dlabel = {'vehicle': 0, 'pedestrian':1, 'bicycle':2, 'rider':3, 'roadblock':4,'cone': 4, 'crashbarrels' :4, 'plasticpile':4, 'bumpingpost':4, 'head':5, 'tsr':6, 'guideboard':7, 'plate':8, 'wheel':9, 'tl_border':10,'tl_wick':11,'tl_num':12} def detectmap(path): # labeldict = {0:'vehicle', 1:'pedestrian', 2:'bicycle', 3:'cone', 4:'warning', 5:'head', 6:'tsr', 7:'guideboard', 8:'plate', 9:'wheel', 10:'tl_border',11:'tl_wick',12:'tl_num'} labeldict = {0:'vehicle', 1:'pedestrian', 2:'bicycle', 3:'rider', 4:'roadblock', 5:'head', 6:'tsr', 7:'guideboard', 8:'plate', 9:'wheel', 10:'tl_border',11:'tl_wick',12:'tl_num'} textlists = open(path,'r').readlines() label_maps = {} for text in textlists: data = text.strip('\n').split(' ') filename = data[0] box = [float(data[2]), float(data[3]), float(data[4]),float(data[5])] # label = labeldict.get(int(data[1])) label = str(data[1]) # ignore = int(data[6]) ignore = 1 if filename not in label_maps: label_maps[filename] = [] label_maps[filename].append((label,box,ignore)) return label_maps def get_rot_mat_and_trans(calib_dict): su = np.sin(calib_dict['roll'] * np.pi / 180) cu = np.cos(calib_dict['roll'] * np.pi / 180) sv = np.sin(calib_dict['pitch'] * np.pi / 180) cv = np.cos(calib_dict['pitch'] * np.pi / 180) sw = np.sin(calib_dict['yaw'] * np.pi / 180) cw = np.cos(calib_dict['yaw'] * np.pi / 180) rot = np.array([[cv * cw, su * sv * cw - cu * sw, su * sw + cu * sv * cw], [cv * sw, cu * cw + su * sv * sw, cu * sv * sw - su * cw], [-sv, su * cv, cu * cv]]) rot = np.transpose(rot) rot_yxz = np.array([[0, -1, 0], [0, 0, -1], [1, 0, 0]]) rot_ego2cam = np.matmul(rot_yxz, rot) # rot_cam2ego = np.transpose(rot_ego2cam) # equals np.linalg.inv pos = np.array(calib_dict['pos']) trans_ego2cam = -1 * np.matmul(rot_ego2cam, pos) ego_to_cam = np.zeros((4, 4)) ego_to_cam[:3, :3] = rot_ego2cam ego_to_cam[:3, 3] = trans_ego2cam ego_to_cam[3, 3] = 1 # cam_to_ego = np.linalg.inv(ego_to_cam) return ego_to_cam def Vehicle2OrigImg(pts3d, ego2cam, distort_param, K, origW, origH): pts3d_cam = np.dot(ego2cam[:3, :3],[pts3d[0],pts3d[1],pts3d[2]]) + ego2cam[:3, 3] # pts2d_cam = cam_img_kb(pts3d_cam[0],pts3d_cam[1],pts3d_cam[2],K,distort_param) if len(distort_param)==5: pts2d_cam = xyz_xy(pts3d_cam[0] /pts3d_cam[2], pts3d_cam[1] / pts3d_cam[2], distort_param, K) else: pts2d_cam = cam_img_kb(pts3d_cam[0],pts3d_cam[1],pts3d_cam[2],K,distort_param) pts2d_cam = pts2d_cam[0:2] pts2d_cam[0] = np.clip(pts2d_cam[0],0,origW-1) pts2d_cam[1] = np.clip(pts2d_cam[1],0,origH-1) return pts2d_cam, pts3d_cam def deal_cutoff(cor,bbox_cam3d,cen2d,cen2d_di,alpha,lidar_to_cam,K,distortion_param,K_new,image_w,image_h,ydis_dict): cor = np.array(cor) # a = cor[:,0] * K_new[0][0] # b = -cor[:,2]*K_new[0][2] # c = cor[:,2]*(image_w-1-K_new[0][2]) x_left = cor[:,0]>=-cor[:,2]*K_new[0][2]/K_new[0][0] # (8,) x_right = cor[:,0]<=cor[:,2]*(image_w-1-K_new[0][2])/K_new[0][0] if (x_left==True).sum()==(x_right==True).sum() and (x_right==True).sum()<8: if bbox_cam3d[0]<=-bbox_cam3d[2]*(image_w-1-K_new[0][2])/K_new[0][0]: ####右边界线 is_left = True # assert bbox_cam3d[0] < -bbox_cam3d[2]*K_new[0][2]/K_new[0][0] else: is_left = False # assert bbox_cam3d[0]>-bbox_cam3d[2]*(image_w-1-K_new[0][2])/K_new[0][0] else: is_left = (x_left==True).sum()<(x_right==True).sum() # 左边界被切 have_in = ((x_left & x_right)[[0,1,4,5]]==True).sum() pts2ds_c,_ = cam2pixel(np.array([bbox_cam3d[0], bbox_cam3d[1], bbox_cam3d[2]]), K_new) condition_c = np.logical_and((pts2ds_c[1] <= image_h-1) & (pts2ds_c[1] >= 0),(pts2ds_c[0] <= image_w-1) & (pts2ds_c[0] >= 0)) cutoff = 0 if have_in ==2 and condition_c == False: cutoff = 1 y_dis = max(0,min(image_h-1,round(K_new[1][2]+K_new[1][1]*(bbox_cam3d[1]/bbox_cam3d[2])))) if is_left: x_dis = ydis_dict[0][y_dis] else: x_dis = ydis_dict[1][y_dis] z_new = (bbox_cam3d[2]*K_new[0][0]+bbox_cam3d[0]*K_new[0][0]*np.tan(bbox_cam3d[-1]))/((x_dis-K_new[0][2])*np.tan(bbox_cam3d[-1])+K_new[0][0]) x_new = z_new*(x_dis-K_new[0][2])/K_new[0][0] # cen2d = xyz_xy(x_new/z_new,bbox_cam3d[1]/z_new,distortion_param,K) cen2d = cam_img_kb(x_new,bbox_cam3d[1],z_new,K,distortion_param) x_nc,y_nc = float(cen2d[0]),float(cen2d[1]) if x_nc<0 or x_nc>image_w-1 or y_nc<0 or y_nc>image_h-1: x_nc,y_nc = max(0,min(image_w-1,x_nc)),max(0,min(image_h-1,y_nc)) img_corners = np.array([x_nc,y_nc]).reshape((-1, 2)) x_new,y_new = img_cam_kb(img_corners,K,distortion_param) x_new,y_new = x_new*z_new,y_new*z_new if y_new!=bbox_cam3d[1]: bbox_cam3d[4] = bbox_cam3d[4]-(y_new-bbox_cam3d[1])*2 bbox_cam3d[1] = y_new is_l = np.logical_or((x_left & x_right)[[0,1]].all(),(x_left & x_right)[[4,5]].all()) see_radio = [1,1] if is_l: l = (bbox_cam3d[3]/2 - np.sqrt(np.sum(np.square(np.array([x_new,bbox_cam3d[1],z_new],dtype=np.float32) - np.array(bbox_cam3d[0:3],dtype=np.float32)))))*2 assert l 0.1: # cen2d = xyz_xy(new_box[0]/new_box[2],new_box[1]/new_box[2],distortion_param, K) # 中心点（x,y,z）;底部中心点(x,y+l/2,z) 1933.3505753743707, 593.9974302913548 # cen2d_di = xyz_xy(new_box[0]/new_box[2],(new_box[1]+0.5*new_box[4])/new_box[2],distortion_param,K) cen2d = cam_img_kb(new_box[0],new_box[1],new_box[2],K,distortion_param) # 中心点（x,y,z）;底部中心点(x,y+l/2,z) 1933.3505753743707, 593.9974302913548 cen2d_di = cam_img_kb(new_box[0],(new_box[1]+0.5*new_box[4]),new_box[2],K,distortion_param) # assert new_box[2]>0 alpha = -np.arctan2(new_box[0],new_box[2])+new_box[-1] # -torch.atan2(torch.tensor(loc[0]), torch.tensor(loc[2]))+torch.tensor(rot[0]) if alpha<-np.pi: alpha+=2*np.pi if alpha>np.pi: alpha-=2*np.pi bbox_cam3d = new_box return bbox_cam3d, cen2d,cen2d_di, alpha, cutoff def rotation_3d_in_axis(points, angles, axis=0): ####需要特别注意旋转的轴和旋转矩阵正逆问题 以下为逆时针旋转矩阵 rot_sin = np.sin(angles) rot_cos = np.cos(angles) ones = np.ones_like(rot_cos) zeros = np.zeros_like(rot_cos) if axis == 1: ####Y rot_mat = np.stack([ np.stack([rot_cos, zeros, -rot_sin]), np.stack([zeros, ones, zeros]), np.stack([rot_sin, zeros, rot_cos]) ]) elif axis == 2: #####Z rot_mat = np.stack([ np.stack([rot_cos, rot_sin, zeros]), np.stack([-rot_sin, rot_cos, zeros]), np.stack([zeros, zeros, ones]) ]) elif axis == 0: #####X rot_mat = np.stack([ np.stack([ones, zeros, zeros]), np.stack([zeros, rot_cos, rot_sin]), np.stack([zeros, -rot_sin, rot_cos]) ]) else: raise ValueError(f'axis should in range [0, 1, 2], got {axis}') return np.dot(points, rot_mat) def create_corner(object = [0,0,0,0,0,0,0], type= 'cam',dtype=np.float32): dims = np.array([object[3],object[4],object[5]]) ###lwh # torch.Size([1, 3]) corners_norm =np.stack(np.unravel_index(np.arange(8), [2] * 3), axis=1) corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]] # use relative origin [0.5, 1, 0.5] corners_norm = corners_norm -[0.5, 0.5, 0.5] corners = dims.reshape([-1, 3]) * corners_norm.reshape([8, 3]) # rotate around y axisa corners = rotation_3d_in_axis(corners, object[6], axis=2) ###### corners += np.array(object[:3]).reshape(1, 3) return corners def cam_corners_front_rear(pred3d, facetype,pitch = 0): dims = pred3d[3:6] corners_norm = np.stack(np.unravel_index(np.arange(8), [2] * 3), axis=1) corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]] # use relative origin [0.5, 1, 0.5] if facetype == 'tail': # front corners_norm = corners_norm - [0, 0.5, 0.5] elif facetype == 'front': # tail corners_norm = corners_norm - [1, 0.5, 0.5] elif facetype == 'right': # left corners_norm = corners_norm - [0.5, 0.5, 0] elif facetype == 'left': # right corners_norm = corners_norm - [0.5, 0.5, 1] elif facetype == 'whole': # center corners_norm = corners_norm - [0.5, 0.5, 0.5] else: raise AssertionError('Non valid face type') corners = dims.reshape([1, 3]) * corners_norm.reshape([8, 3]) # rotate around y axis corners = rotation_3d_in_axis(corners, pred3d[6], axis=1) #### corners = rotation_3d_in_axis(corners, pitch * np.pi / 180.0, axis=0) #### corners += pred3d[:3].reshape(1, 3) return corners def ious_function(boxes1, boxes2, iom = False): # boxes1-(4,),mx4 boxes1 = np.array(boxes1) boxes2 = np.array(boxes2) boxes1_area = (boxes1[..., 2] - boxes1[..., 0]) * (boxes1[..., 3] - boxes1[..., 1]) boxes2_area = (boxes2[..., 2] - boxes2[..., 0]) * (boxes2[..., 3] - boxes2[..., 1]) left_up = np.maximum(boxes1[..., :2], boxes2[..., :2]) right_down = np.minimum(boxes1[..., 2:], boxes2[..., 2:]) inter_section = np.maximum(right_down - left_up, 0.0) inter_area = inter_section[..., 0] * inter_section[..., 1] if iom: union_area = boxes1_area else: union_area = boxes1_area + boxes2_area - inter_area ious = np.maximum(1.0 * inter_area / union_area, np.finfo(np.float32).eps) return ious def get_point_xyz(xyz1,xyz2,num=1): # num=1为黑色 (x1,y1,z1),(x2,y2,z2) = xyz1,xyz2 dis = np.array([abs(x1-x2),abs(y1-y2),abs(z1-z2)]) dis_idx = dis.argmax() res = [] if dis_idx==0: for x in np.concatenate((np.arange(x1,x2,0.01) ,np.arange(x1,x2,-0.01))): res.append([x,(x-x1)*(y2-y1)/(x2-x1)+y1,(x-x1)*(z2-z1)/(x2-x1)+z1,num]) elif dis_idx==1: for y in np.concatenate((np.arange(y1,y2,0.01) ,np.arange(y1,y2,-0.01))): res.append([(y-y1)*(x2-x1)/(y2-y1)+x1,y,(y-y1)*(z2-z1)/(y2-y1)+z1,num]) else: for z in np.concatenate((np.arange(z1,z2,0.01) ,np.arange(z1,z2,-0.01))): res.append([(z-z1)*(x2-x1)/(z2-z1)+x1,(z-z1)*(y2-y1)/(z2-z1)+y1,z,num]) return res def cam_img_kb(x,y,z,cam_intrinsic,distort_param): # objp=np.array([x/z,y/z,1]).reshape(1,-1,3) # rvec = np.array([[[0., 0., 0.]]]) # tvec = np.array([[[0., 0., 0.]]]) # image_coord, _ = cv2.fisheye.projectPoints(objp, rvec, tvec,cam_intrinsic, distort_param) # image_coord = image_coord.reshape(-1) # return image_coord camxyz = np.array([x/z, y/z]).reshape(-1,2) a = camxyz[...,0] b = camxyz[...,1] r = np.sqrt(a * a + b * b) theta = np.arctan(r) k1 = distort_param[0] k2 = distort_param[1] k3 = distort_param[2] k4 = distort_param[3] thetaD = theta * (1 + k1 * theta**2 + k2 * theta**4 + k3 * theta**6 + k4 * theta**8) xp = thetaD / r * a yp = thetaD / r * b fx = cam_intrinsic[0,0] fy = cam_intrinsic[1,1] cx = cam_intrinsic[0,2] cy = cam_intrinsic[1,2] imgx = (fx * xp + cx) imgy = (fy * yp + cy) # mask_2d = (imgx>=0)&(imgx<3840)&(imgy>=0)&(imgy<2160) pix_point = np.round(np.concatenate((imgx.reshape((-1,1)),imgy.reshape((-1,1))),axis=-1)) # pix_point = pix_point[mask_2d].astype(np.int32) pix_point = pix_point.astype(np.int32) return pix_point.reshape(-1) def img_cam_kb(img_corners,cam_intrinsic,distort_param): # nx2 ones_col = np.ones((img_corners.shape[0], 1), dtype=np.float32) img_corners = np.concatenate((img_corners.astype(np.float32),ones_col),axis=-1) # （n,3） k1 = distort_param[0] k2 = distort_param[1] k3 = distort_param[2] k4 = distort_param[3] points = np.dot(np.linalg.inv(cam_intrinsic),img_corners.T).T # (2073600, 3) xp = points[:,0] # n yp = points[:,1] # n thetaD = np.sqrt(xp*xp+yp*yp) cam_theta,femask = [],[] theta = thetaD for i in range(10): ftheta = theta * (1 + k1 * theta**2 + k2 * theta**4 + k3 * theta**6 + k4 * theta**8) - thetaD ftheta_dao = 1 + 3 * k1 * theta**2 + 5 * k2 * theta**4 + 7 * k3 * theta**6 + 9 * k4 * theta**8 theta = theta - ftheta/ftheta_dao if abs(ftheta)<0.0000001: break # print(ftheta) r = np.tan(theta) ################################### normx = xp * r / thetaD # 4,H,W normy = yp * r / thetaD ones_col = np.ones((normy.shape[0], 1), dtype=np.float32) x3dy3d = np.concatenate((normx.reshape((-1,1)),normy.reshape((-1,1)),ones_col),axis=-1) return x3dy3d[0,0], x3dy3d[0,1] # def img_cam_kb(img_corners,cam_intrinsic,distort_param): # nx2 # ones_col = np.ones((img_corners.shape[0], 1), dtype=np.float32) # img_corners = np.concatenate((img_corners.astype(np.float32),ones_col),axis=-1) # （n,3） # k1 = distort_param[0] # k2 = distort_param[1] # k3 = distort_param[2] # k4 = distort_param[3] # points = np.dot(np.linalg.inv(cam_intrinsic),img_corners.T).T # (2073600, 3) # xp = points[:,0] # n # yp = points[:,1] # n # thetaD = np.sqrt(xp*xp+yp*yp) # cam_theta,femask = [],[] # # for d in tqdm(thetaD): # for d in thetaD: # coeff = np.array([-d,1,0,k1,0,k2,0,k3,0,k4],dtype=np.float64) # k4*x**9+k3*x**7+k2*x**5+k1*x**3+x-d=0 # success,roots = cv2.solvePoly(coeff) # roots = roots.squeeze(1) # roots = np.array([root[0] for root in roots if abs(root[1]) < 1e-10]) # 找到唯一实数根解 # assert len(roots)==1,'come with many theta' # if roots[0]>=0 and roots[0] <=np.pi/2: # femask.append(True) # else: # femask.append(False) # cam_theta.append(roots[0]) # r = np.tan(cam_theta) # mask3d = np.array(femask) # ################################### # normx = xp * r / thetaD # 4,H,W # normy = yp * r / thetaD # ones_col = np.ones((normy.shape[0], 1), dtype=np.float32) # x3dy3d = np.concatenate((normx.reshape((-1,1)),normy.reshape((-1,1)),ones_col),axis=-1) # x3dy3d = x3dy3d[mask3d] # return x3dy3d[0,0], x3dy3d[0,1] def cam2RectifyImg(x,y, distort_param, K): # R = np.array([[0, -1, 0],[-1, 0, 0],[0, 0, 1]]) # pts3d_cam = np.dot(lidar_to_cam[:3, :3],[pts3d[0],pts3d[1],pts3d[2]]) + lidar_to_cam[:3, 3] # pts3d_cam = lidar_to_cam[:3, :3] @ (pts3d.T) + lidar_to_cam[:3, 3].reshape(3, 1) # 1x3 pts2d_cam = K @ np.array([x,y,1]) pts2d_cam[2] = np.clip(pts2d_cam[2],1,500) pts2d_cam = (pts2d_cam[:2] / pts2d_cam[2]).T pts2d_cam = pts2d_cam[0:2] image_h, image_w = 1080,1920 pts2d_cam[0] = np.clip(pts2d_cam[0],0,image_w-1) pts2d_cam[1] = np.clip(pts2d_cam[1],0,image_h-1) return pts2d_cam def xyz_xy(x,y,distort_param,K): r = np.sqrt(x * x + y * y) r_2 = r * r r_4 = r_2 * r_2 r_6 = r_4 * r_2 m_distort_coeffs = distort_param coeffs = np.zeros(12,dtype=np.float32) for i in range(min(m_distort_coeffs.shape[0],12)): coeffs[i] = m_distort_coeffs[i] radial_ratio = (1 + coeffs[0] * r_2 + coeffs[1] * r_4 + coeffs[4] * r_6) / (1 + coeffs[5] * r_2 + coeffs[6] * r_4 + coeffs[7] * r_6) res_x = x * radial_ratio + 2 * coeffs[2] * x * y + coeffs[3] * (r_2 + 2 * pow(x, 2)) + coeffs[8] * r_2 + coeffs[9] * r_4 res_y = y * radial_ratio + 2 * coeffs[2] * (r_2 + 2 * pow(y, 2)) + coeffs[3] * x * y + coeffs[10] * r_2 + coeffs[11] * r_4 image_coord = np.dot(K,np.array([res_x,res_y,1])).T return image_coord[:2] def xy_xyz(x,y,m_distort_coeffs,k): x = (x - k[0][2])/k[0][0] y = (y - k[1][2])/k[1][1] coeffs = np.zeros(12,dtype=np.float32) for i in range(min(m_distort_coeffs.shape[0],12)): coeffs[i] = m_distort_coeffs[i] iter_num = 5 x0,y0 = x,y for i in range(iter_num): r = np.sqrt(x * x + y * y) r_2 = r * r r_4 = r_2 * r_2 r_6 = r_4 * r_2 inv_radial_ratio = (1 + coeffs[5] * r_2 + coeffs[6] * r_4 + coeffs[7] * r_6) / (1 + coeffs[0] * r_2 + coeffs[1] * r_4 + coeffs[4] * r_6) delta_x = 2 * coeffs[2] * x * y + coeffs[3] * (r_2 + 2 * pow(x, 2)) + coeffs[8] * r_2 + coeffs[9] * r_4 delta_y = 2 * coeffs[2] * (r_2 + 2 * pow(y, 2)) + coeffs[3] * x * y + coeffs[10] * r_2 + coeffs[11] * r_4 x = (x0 - delta_x) * inv_radial_ratio y = (y0 - delta_y) * inv_radial_ratio # x = x*k[0][0]+k[0][2] # y = y*k[1][1]+k[1][2] return x,y def get_cutoff_corners(cam3d_corners,pts2ds,cam_intrinsic,distortion_param,cam_rectify,imgW,imgH): xyz_list = [] cut_pts2d = np.zeros_like(pts2ds) cam3d_corners_update = np.zeros_like(cam3d_corners) for (i,j) in [(0, 1), (1, 2), (2, 3), (3, 0),(0,4), (1,5), (2,6),(3,7)]: res = get_point_xyz(cam3d_corners[i],cam3d_corners[j],0) xyz_list.extend(res) if (i,j) in [(0,4), (1,5), (2,6),(3,7)]: point = np.array(res) pointx_in = np.logical_and(point[:,0]>=-point[:,2]*cam_rectify[0][2]/cam_rectify[0][0],point[:,0]<=point[:,2]*(imgW-1-cam_rectify[0][2])/cam_rectify[0][0]) pointy_in = np.logical_and(point[:,1]>=-point[:,2]*cam_rectify[1][2]/cam_rectify[1][1],point[:,1]<=point[:,2]*(imgH-1-cam_rectify[1][2])/cam_rectify[1][1]) pointxy_in = np.logical_and(pointx_in,pointy_in) t_point = point[pointxy_in] if t_point.shape[0] == 0: if len(distortion_param)==5: origimg0 = xyz_xy(point[0,0] /point[0,2], point[0,1] / point[0,2], distortion_param, cam_intrinsic) else: origimg0 = cam_img_kb(point[0,0] , point[0,1] , point[0,2],cam_intrinsic,distortion_param) cut_pts2d[i] = origimg0[:2] cut_pts2d[j] = origimg0[:2] continue if len(distortion_param)==5: point_origimg0 = xyz_xy(t_point[0,0] /t_point[0,2], t_point[0,1] / t_point[0,2], distortion_param, cam_intrinsic) point_origimg1 = xyz_xy(t_point[-1,0] /t_point[-1,2], t_point[-1,1] / t_point[-1,2], distortion_param, cam_intrinsic) else: point_origimg0 = cam_img_kb(t_point[0,0] , t_point[0,1] , t_point[0,2],cam_intrinsic,distortion_param) point_origimg1 = cam_img_kb(t_point[-1,0] , t_point[-1,1] , t_point[-1,2],cam_intrinsic,distortion_param) cut_pts2d[i] = point_origimg0[:2] cut_pts2d[j] = point_origimg1[:2] cam3d_corners_update[i] = t_point[0,:3] cam3d_corners_update[j] = t_point[-1,:3] return cut_pts2d ,cam3d_corners_update def cam2pixel(pts3d_cam, K): pts2d_cam = K @ pts3d_cam Z = pts2d_cam[2] pts2d_cam = (pts2d_cam[:2] / Z).T return pts2d_cam, Z def nms_area(boxes,lables,area_thresh=0.9): if len(boxes) < 2: discard = [] return discard # for i in range(boxes.shape[0]): # if boxes[i][0] < 40 or boxes[i][3] > 3840 - 40: # discard = [] # return discard rect_boxes = np.array(boxes) x1 = rect_boxes[:,0] y1 = rect_boxes[:,1] x2 = rect_boxes[:,2] y2 = rect_boxes[:,3] areas = (y2-y1+1)*(x2-x1+1) discard = [] index = areas.argsort()[::-1] ##降序 for idx in index[1:]: x11=np.maximum(x1[index[0]],x1[idx]) y11=np.maximum(y1[index[0]],y1[idx]) x22=np.minimum(x2[index[0]],x2[idx]) y22=np.minimum(y2[index[0]],y2[idx]) w=np.maximum(0,x22-x11+1) h=np.maximum(0,y22-y11+1) overlaps = w*h iou=overlaps /areas[idx] if iou > area_thresh: if lables[index[0]] == lables[idx]: discard.append(idx) elif lables[index[0]] in ['Pedestrian','Rider'] and lables[idx] in ['Pedestrian','Rider']: discard.append(idx) elif lables[index[0]] in ['Bicycle','Rider'] and lables[idx] in ['Bicycle','Rider']: discard.append(idx) return discard def Lidar2OrigImg(pts3d, lidar_to_cam, distort_param, K, origW, origH): pts3d_cam = np.dot(lidar_to_cam[:3, :3],[pts3d[0],pts3d[1],pts3d[2]]) + lidar_to_cam[:3, 3] # pts2d_cam = cam_img_kb(pts3d_cam[0],pts3d_cam[1],pts3d_cam[2],K,distort_param) if len(distort_param)==5: pts2d_cam = xyz_xy(pts3d_cam[0] /pts3d_cam[2], pts3d_cam[1] / pts3d_cam[2], distort_param, K) else: pts2d_cam = cam_img_kb(pts3d_cam[0],pts3d_cam[1],pts3d_cam[2],K,distort_param) pts2d_cam = pts2d_cam[0:2] pts2d_cam[0] = np.clip(pts2d_cam[0],0,origW-1) pts2d_cam[1] = np.clip(pts2d_cam[1],0,origH-1) return pts2d_cam, pts3d_cam def Lidar2RectifyImg(pts3d, lidar_to_cam, K,origW, origH): # R = np.array([[0, -1, 0],[-1, 0, 0],[0, 0, 1]]) pts3d_cam = np.dot(lidar_to_cam[:3, :3],[pts3d[0],pts3d[1],pts3d[2]]) + lidar_to_cam[:3, 3] # pts3d_cam = lidar_to_cam[:3, :3] @ (pts3d.T) + lidar_to_cam[:3, 3].reshape(3, 1) # 1x3 pts2d_cam = K @ pts3d_cam pts2d_cam[2] = np.clip(pts2d_cam[2],1,500) pts2d_cam = (pts2d_cam[:2] / pts2d_cam[2]).T pts2d_cam = pts2d_cam[0:2] image_h, image_w = origH,origW pts2d_cam[0] = np.clip(pts2d_cam[0],0,image_w-1) pts2d_cam[1] = np.clip(pts2d_cam[1],0,image_h-1) return pts2d_cam, pts3d_cam def Calculate_vector_angle(vector_a ,vector_b): dot_product = np.dot(vector_a, vector_b) # 计算两个向量的范数（大小） norm_a = np.linalg.norm(vector_a) norm_b = np.linalg.norm(vector_b) # 计算夹角的余弦值 cos_angle = dot_product / (norm_a * norm_b) # 使用 arccos 计算角度，得到的结果是弧度 angle_radians = np.arccos(cos_angle) # 将弧度转换为度 angle_degrees = np.degrees(angle_radians) return angle_degrees def angle2score(angle): angle = (angle - 90)/90 angle = min(max(0,angle),1) angle = np.sqrt(angle) return angle def drawPointBox(img, imgW, imgH, rect_corners, colors, thickness=1): line_indices = ((4, 5), (5, 6), (6, 7), (7, 4), # 尾部 (0,4), (1,5), (2,6),(3,7), (0, 1), (1, 2), (2, 3), (3, 0), (0, 2), (1, 3)) # 头部(1, 3), border_x = imgW #### border_y = imgH ##### for i in range(len(rect_corners)): if int(rect_corners[i][0]) < 0 or int(rect_corners[i][0]) > border_x or int(rect_corners[i][1]) < 0 or int(rect_corners[i][1]) > border_y: continue # cv2.circle(img,(int(rect_corners[i][0]),int(rect_corners[i][1])),5,(255,0,0),2,cv2.LINE_8,0) # cv2.putText(img, str(i), (int(rect_corners[i][0]+5),int(rect_corners[i][1] - 5)), 1, 1, (0, 0, 255), lineType=cv2.LINE_AA) corners = rect_corners.astype(np.int32) for start, end in line_indices: try: if (start,end) in [(0, 1), (1, 2), (2, 3), (3, 0),(0, 2), (1, 3)] and corners.shape[0]==8: if int(corners[start, 0]) < 5 or int(corners[start, 0]) > (border_x - 5) or int(corners[start, 1]) < 5 or int(corners[start, 1]) > (border_y - 5): continue if int(corners[end, 0]) < 5 or int(corners[end, 0]) > (border_x - 5) or int(corners[end, 1]) < 5 or int(corners[end, 1]) > (border_y - 5): continue cv2.line(img, (corners[start, 0], corners[start, 1]),(corners[end, 0], corners[end, 1]), colors[1], thickness*2,cv2.LINE_AA) else: cv2.line(img, (corners[start, 0], corners[start, 1]),(corners[end, 0], corners[end, 1]), colors[0], thickness,cv2.LINE_AA) except: pass return img def showFace2Dbox(img,label_temp_all, cam_intrinsic, distortion_param, cam_rectify,origW,origH,pitch,label_canvas_w=None,label_canvas_h=None): if label_canvas_w is None: label_canvas_w = origW if label_canvas_h is None: label_canvas_h = origH legend_items = [] legend_seen = set() for co in label_temp_all: try: label_det_cls,cx,cy,w,h =co[0], float(co[1])*label_canvas_w,float(co[2])*label_canvas_h,float(co[3])*label_canvas_w,float(co[4])*label_canvas_h box2d= [cx-0.5*w,cy-0.5*h,cx+0.5*w,cy+0.5*h] # 检查 label_det_cls 是否在 det_cls_index 中 if label_det_cls not in det_cls_index: print(f"Warning: label_det_cls '{label_det_cls}' not in det_cls_index, skipping") continue label_cls = det_cls_index.index(label_det_cls) # 检查 label_cls 是否超出 colors 范围 if label_cls >= len(colors): print(f"Warning: label_cls {label_cls} exceeds colors length {len(colors)}, using modulo") label_cls = label_cls % len(colors) color = colors[label_cls] if label_cls < 19: box_label = label_det_cls line_thickness = 2 else: box_label = None line_thickness = 1 plot_box(box2d, img, label=box_label, color=color, line_thickness=line_thickness) legend_key = (label_cls, label_det_cls) if legend_key not in legend_seen: legend_seen.add(legend_key) legend_items.append((label_cls, label_det_cls, color)) except Exception as e: print(f"Error processing label: {e}, co={co}") continue return append_class_legend(img, legend_items) def showFace3Dbox(img,label_temp_all, cam_intrinsic, distortion_param, cam_rectify,origW,origH,pitch): for co in label_temp_all: # try: # label_det_cls,cx,cy,w,h =co[0], float(co[1])*origW,float(co[2])*origH,float(co[3])*origW,float(co[4])*origH # box2d= [cx-0.5*w,cy-0.5*h,cx+0.5*w,cy+0.5*h] # # 检查 label_det_cls 是否在 det_cls_index 中 # if label_det_cls not in det_cls_index: # print(f"Warning: label_det_cls '{label_det_cls}' not in det_cls_index, skipping") # continue # label_cls = det_cls_index.index(label_det_cls) # # 检查 label_cls 是否超出 colors 范围 # if label_cls >= len(colors): # print(f"Warning: label_cls {label_cls} exceeds colors length {len(colors)}, using modulo") # label_cls = label_cls % len(colors) # plot_box(box2d, img, label=str('%s'%label_cls), color=colors[label_cls], line_thickness=1) # except Exception as e: # print(f"Error processing label: {e}, co={co}") # continue # if int(label_det_cls) < 11: # show_label_color = 0 # elif int(label_det_cls) > 11 : # show_label_color = int(label_det_cls) - 11 # plot_box(box2d, img, label=str('%s'%label_det_cls), color=colors[show_label_color], line_thickness=1) # # box2d_1=[[box2d[0],box2d[1]],[box2d[2],box2d[3]]] # # box2d_Rectify = cv2.undistortPoints(np.array(box2d_1).reshape(2,1,2), cam_intrinsic, distortion_param, None, cam_rectify).reshape((4,)) # # plot_box(box2d_Rectify, img, label=str('%s'%label_det_cls), color=colors[label_det_cls], line_thickness=1) pts2ds = [] if len(co)>7: x3d_ori,y3d_ori,z3d_ori,l3d,h3d,w3d,rotation_y = float(co[5]),float(co[6]),float(co[7]),float(co[8]),float(co[9]),float(co[10]),float(co[11]) color = (0,0,0) if rotation_y <= 0 else (0,255,255) corners = cam_corners_front_rear(np.array([x3d_ori,y3d_ori,z3d_ori,l3d,h3d,w3d,rotation_y]),'whole',pitch) pts2ds = [] for i in range(corners.shape[0]): x,y,z = corners[i] if len(distortion_param)==5: image_coord = xyz_xy(x/z,y/z,distortion_param,cam_intrinsic) # image_coord = cam2RectifyImg(x/z,y/z,distortion_param,cam_rectify) else: image_coord = cam_img_kb(x,y,z,cam_intrinsic,distortion_param) pts2ds.append([image_coord[0],image_coord[1]]) #####deal cutoff#### x_left = corners[:,0]>=-corners[:,2]*cam_rectify[0][2]/cam_rectify[0][0] # (8,) x_right = corners[:,0]<=corners[:,2]*(origW-1-cam_rectify[0][2])/cam_rectify[0][0] have_in = ((x_left & x_right)[[0,1,4,5]]==True).sum() # po2ds = None if have_in >= 1 and have_in < 4: cut_pts2d,cam3d_corners_update= get_cutoff_corners(corners,pts2ds,cam_intrinsic,distortion_param,cam_rectify,origW,origH) # cut_pts2d,cut_rect,cut_rectify_box3d, po2ds= get_cutoff_corners(corners,pts2ds_new,cam_intrinsic,distortion_param,cam_rectify,imgW,imgH) if cut_pts2d is not None: pts2ds = cut_pts2d sqe = [6,7,4,5,2,3,0,1] pts2ds = np.array(pts2ds)[sqe] img = drawPointBox(img, origW, origH,np.array(pts2ds), colors=[(0, 0, 255), color], thickness=1) if len(co)==51: anchortype_index = ['front','tail','left','right'] value_3d_4face = co[18:50] for jj in range(len(value_3d_4face)): value_3d_4face[jj] = float(value_3d_4face[jj]) value_3d_4face = np.reshape(value_3d_4face,(4,-1)) for anchortype in range(len(value_3d_4face)): x3d,y3d,z3d,alpha,xc,yc,score,is_good = value_3d_4face[anchortype] xc,yc = max(0,min(origW-1,float(xc)*origW)),max(0,min(origH-1,float(yc)*origH)) if len(distortion_param)==5: x3d_new ,y3d_new = xy_xyz(xc,yc,distortion_param,cam_intrinsic) # x3d_new ,y3d_new = cam2RectifyImg(x/z,y/z,distortion_param,cam_rectify) else: x3d_new ,y3d_new = img_cam_kb(np.array([xc,yc]).reshape((-1, 2)),cam_intrinsic,distortion_param) infer_x3d = x3d_new * z3d infer_y3d = y3d_new * z3d rotation_y = alpha + np.arctan2(infer_x3d, z3d) if is_good == -1: continue if rotation_y<0: color = (0,0,0) else: color = (0,255,255) corners = cam_corners_front_rear(np.array([infer_x3d,infer_y3d,z3d, l3d,h3d,w3d, rotation_y]),anchortype_index[anchortype],pitch) pts2ds_new = [] for i in range(corners.shape[0]): x,y,z = corners[i] if len(distortion_param)==5: image_coord = xyz_xy(x/z,y/z,distortion_param,cam_intrinsic) # image_coord = cam2RectifyImg(x/z,y/z,distortion_param,cam_rectify) else: image_coord = cam_img_kb(x,y,z,cam_intrinsic,distortion_param) pts2ds_new.append([image_coord[0],image_coord[1]]) #####deal cutoff#### x_left = corners[:,0]>=-corners[:,2]*cam_rectify[0][2]/cam_rectify[0][0] # (8,) x_right = corners[:,0]<=corners[:,2]*(origW-1-cam_rectify[0][2])/cam_rectify[0][0] have_in = ((x_left & x_right)[[0,1,4,5]]==True).sum() # po2ds = None if have_in >= 1 and have_in < 4: cut_pts2d,cam3d_corners_update= get_cutoff_corners(corners,pts2ds_new,cam_intrinsic,distortion_param,cam_rectify,origW,origH) if cut_pts2d is not None: pts2ds_new = cut_pts2d sqe = [6,7,4,5,2,3,0,1] pts2ds = np.array(pts2ds_new)[sqe] img = drawPointBox(img, origW, origH,np.array(pts2ds), colors=[(0, 0, 255),color], thickness=1) if len(distortion_param)==5: [xc,yc] = xyz_xy(infer_x3d/z3d,infer_y3d/z3d,distortion_param,cam_intrinsic) # [xc,yc] = cam2RectifyImg(infer_x3d/z3d,infer_y3d/z3d,distortion_param,cam_rectify) else: [xc,yc] = cam_img_kb(infer_x3d,infer_y3d,z3d,cam_intrinsic,distortion_param) img = cv2.circle(img, (int(xc),int(yc)), 4, colors[int(anchortype)],-1) return img def mono3dlabel_trans2train(calib_path,anno): det_cls_index = [ "car", "suv", "pickup", "van", "bus", "medium_car", "truck", "tanker", "large_truck", "construction_vehicle", "special_vehicle", "unknown", 'pedestrian', 'bicycle', "bicyclist", "motorcycle", "motorcyclist", "tricycle", "tricyclist", 'traffic_light_bbox', 'traffic_light_bulb', 'traffic_sign', 'animal', 'movable_object', 'warning_triangle', 'traffic_cone', 'water_barrier', 'crash_barrel', 'movable_barrier', 'bollard', 'sphere_bollard', 'cube_bollard', 'cylinder_bollard', 'construction_barrier', 'other_barrier', 'road_barrier_unknown', "wheel", "plate", "face" ] if os.path.exists(anno): text_basename = os.path.basename(anno) print(text_basename) text_split = text_basename.strip('_').split('_') calib_file = resolve_camera_calib_path(calib_path, anno) calib_param_vehicle, cam_intrinsic, distortion_param, cam_rectify, origW, origH = build_camera_config(calib_file) ego_to_cam = get_rot_mat_and_trans(calib_param_vehicle) objs = json.load(open(anno))['asso_list'] label_temp_all = [] for obj in objs: label_temp = [] label_2d = obj['camera_mea'] confidence = label_2d['confidence'] try: label_radar = obj['radar_mea'] except: label_radar=None label_lidar = obj['lidar_mea'] # label_det_cls = det_cls_map[label_2d['cls']] ############## label_det_cls,label_det_subcls = label_2d['cls'],label_2d['subcls'] if label_det_cls=='vehicle': label_det_cls=label_det_subcls label_det_cls = det_cls_index.index(label_det_cls) ############### if label_det_cls == -1: continue # sublabel = subcls.get(label_2d['subcls']) ###'kTricycle' : 6 # 三轮车 box2d = [label_2d['det_pt_x'],label_2d['det_pt_y'],label_2d['det_pt_width'],label_2d['det_pt_height']] box2d = [box2d[0] , box2d[1] , box2d[0] + box2d[2], box2d[1] + box2d[3]] label_temp=[label_det_cls,min((box2d[0]+box2d[2])/(2*origW),1), min((box2d[1]+box2d[3])/(2*origH),1),min((box2d[2]-box2d[0])/origW,1), min((box2d[3]-box2d[1])/origH,1)] # if label_radar is not None: # radar_conf = label_radar['asso_conf'] # if radar_conf < 0.5: # continue # radar_x3d ,radar_y3d, radar_z3d= label_radar['x'],label_radar['y'],label_radar['z'] ###rear plane # radarpoint = [radar_x3d ,radar_y3d, radar_z3d] # radarpts2d , radarpts3d = Lidar2OrigImg(radarpoint,lidar_to_cam,distortion_param,cam_intrinsic, origW, origH) # cv2.circle(img,(int(radarpts2d[0]),int(radarpts2d[1])),5,(255,0,0),2,cv2.LINE_8,0) # cv2.putText(img, str('%.2f'%radarpts3d[2]), (int(radarpts2d[0]) + 5,int(radarpts2d[1]) - 5), 1, 1, (0, 0, 255)) if (label_lidar is not None) and label_det_cls<=18: obj_lidar = [float(label_lidar[1]),float(label_lidar[2]),float(label_lidar[3]),float(label_lidar[4]),float(label_lidar[5]),float(label_lidar[6]),float(label_lidar[7])] l3d,w3d,h3d = obj_lidar[3],obj_lidar[4],obj_lidar[5] singleLW = int(float(label_lidar[11])) yaw_lidar = obj_lidar[6] Rz = np.array([ [np.cos(yaw_lidar), -np.sin(yaw_lidar), 0], [np.sin(yaw_lidar), np.cos(yaw_lidar), 0], [ 0, 0, 1.0], ]) #############lidar Z轴对应 cam Y轴 顺时针 RotateZ =np.dot(ego_to_cam[:3, :3],Rz) ##lidar yaw 累加旋转矩阵 orientation = Quaternion(matrix=RotateZ) rotation_y = -orientation.yaw_pitch_roll[0] # convert the rot to our cam coordinate campts2d , campts3d = Vehicle2OrigImg(np.array([obj_lidar[0],obj_lidar[1],obj_lidar[2]]), ego_to_cam, distortion_param, cam_intrinsic, origW, origH) campts2d_di, campts3d_di = Vehicle2OrigImg(np.array([obj_lidar[0],obj_lidar[1],obj_lidar[2] - 0.5 * h3d]),ego_to_cam, distortion_param, cam_intrinsic, origW, origH) # campts2d , campts3d = Lidar2OrigImg(np.array([obj_lidar[0],obj_lidar[1],obj_lidar[2]]), lidar_to_cam, distortion_param, cam_intrinsic, origW, origH) # campts2d_di, campts3d_di = Lidar2OrigImg(np.array([obj_lidar[0],obj_lidar[1],obj_lidar[2] - 0.5 * h3d]),lidar_to_cam, distortion_param, cam_intrinsic, origW, origH) box3d_cam = [campts3d[0], campts3d[1], campts3d[2],l3d,h3d,w3d,rotation_y] alpha = -np.arctan2(campts3d[0],campts3d[2]) + rotation_y if alpha<-np.pi: alpha+=2*np.pi if alpha>np.pi: alpha-=2*np.pi x3d_ori,y3d_ori,z3d_ori,l3d,h3d,w3d,rotation_y,alpha,singleLW = campts3d[0], campts3d[1], campts3d[2],l3d,h3d,w3d,rotation_y,alpha,singleLW label_temp = label_temp + [x3d_ori,y3d_ori,z3d_ori,l3d,h3d,w3d,rotation_y,campts2d[0]/origW,campts2d[1]/origH,campts2d_di[0]/origW,campts2d_di[1]/origH,alpha] corners = cam_corners_front_rear(np.array(box3d_cam),'whole', pitch=0) if label_det_cls > 11: label_temp = label_temp + [-1] + [confidence] else: xyz3d_0 = np.mean(corners[4:8,:],axis=0) ### Front_point xyz3d_1 = np.mean(corners[0:4,:],axis=0) ### Back_point xyz3d_2 = np.mean((corners[1,:],corners[2,:],corners[5,:],corners[6,:]),axis=0) ### Left_point xyz3d_3 = np.mean((corners[0,:],corners[4,:],corners[3,:],corners[7,:]),axis=0) ### Right_point xyz3d_4face = [xyz3d_0,xyz3d_1,xyz3d_2,xyz3d_3] value_3d_4face = [] x_left = corners[:,0]>=-corners[:,2]*cam_rectify[0][2]/cam_rectify[0][0] # (8,) x_right = corners[:,0]<=corners[:,2]*(origW-1-cam_rectify[0][2])/cam_rectify[0][0] x_down = corners[:,1]<=corners[:,2]*(origH-1-cam_rectify[1][2])/cam_rectify[1][1] have_in = ((x_left & x_right)[[0,1,4,5]]==True).sum() ##################### x_front_1 = corners[4:8,0]>=-corners[4:8,2]*cam_rectify[0][2]/cam_rectify[0][0] x_front_2 = corners[4:8,0]<=corners[4:8,2]*(origW-1-cam_rectify[0][2])/cam_rectify[0][0] have_in_front = ((x_front_1 & x_front_2)[[0,1,2,3]]==True).sum() y_front = corners[4:8,1]<=corners[4:8,2]*(origH-5-cam_rectify[1][2])/cam_rectify[1][1] y_front_center = xyz3d_0[1]<=xyz3d_0[2]*(origH-5-cam_rectify[1][2])/cam_rectify[1][1] x_tail_1 = corners[0:4,0]>=-corners[0:4,2]*cam_rectify[0][2]/cam_rectify[0][0] x_tail_2 = corners[0:4,0]<=corners[0:4,2]*(origW-1-cam_rectify[0][2])/cam_rectify[0][0] y_tail = corners[0:4,1]<=corners[0:4,2]*(origH-5-cam_rectify[1][2])/cam_rectify[1][1] y_tail_center = xyz3d_1[1]<=xyz3d_1[2]*(origH-5-cam_rectify[1][2])/cam_rectify[1][1] have_in_tail = ((x_tail_1 & x_tail_2)[[0,1,2,3]]==True).sum() ##################### anchorpoint = [x3d_ori, y3d_ori, z3d_ori] if have_in >= 1 and have_in < 4: ######切车 if y_front_center.sum() == 1 or y_tail_center.sum() == 1: if have_in_tail < 4 and have_in_front==4: anchorpoint = xyz3d_0 alpha = -np.arctan2(anchorpoint[0],anchorpoint[2]) + rotation_y if alpha<-np.pi: alpha+=2*np.pi if alpha>np.pi: alpha-=2*np.pi score = 1 if len(distortion_param)==5: [xc,yc] = xyz_xy(anchorpoint[0]/anchorpoint[2],anchorpoint[1]/anchorpoint[2],distortion_param,cam_intrinsic) else: [xc,yc] = cam_img_kb(anchorpoint[0],anchorpoint[1],anchorpoint[2],cam_intrinsic,distortion_param) xc,yc = xc/origW,yc/origH value_3d_face = [anchorpoint[0],anchorpoint[1],anchorpoint[2],alpha,xc,yc,score,1] value_3d_4face.append(value_3d_face) value_3d_4face.append([-1,-1,-1,-1,-1,-1,0,-1]) value_3d_4face.append([-1,-1,-1,-1,-1,-1,0,-1]) value_3d_4face.append([-1,-1,-1,-1,-1,-1,0,-1]) elif have_in_tail == 4 and have_in_front<4: anchorpoint = xyz3d_1 alpha = -np.arctan2(anchorpoint[0],anchorpoint[2]) + rotation_y if alpha<-np.pi: alpha+=2*np.pi if alpha>np.pi: alpha-=2*np.pi score = 1 if len(distortion_param)==5: [xc,yc] = xyz_xy(anchorpoint[0]/anchorpoint[2],anchorpoint[1]/anchorpoint[2],distortion_param,cam_intrinsic) else: [xc,yc] = cam_img_kb(anchorpoint[0],anchorpoint[1],anchorpoint[2],cam_intrinsic,distortion_param) xc,yc = xc/origW,yc/origH value_3d_face = [anchorpoint[0],anchorpoint[1],anchorpoint[2],alpha,xc,yc,score,1] value_3d_4face.append([-1,-1,-1,-1,-1,-1,0,-1]) value_3d_4face.append(value_3d_face) value_3d_4face.append([-1,-1,-1,-1,-1,-1,0,-1]) value_3d_4face.append([-1,-1,-1,-1,-1,-1,0,-1]) else: if 'G1M2' in text_basename and (x_down.sum() < 8 and rotation_y < 0 and y_tail.sum() < 4 and y_front_center.sum() == 1): ####G1M2车尾底边被截断### anchorpoint = xyz3d_0 alpha = -np.arctan2(anchorpoint[0],anchorpoint[2]) + rotation_y if alpha<-np.pi: alpha+=2*np.pi if alpha>np.pi: alpha-=2*np.pi score = 1 if len(distortion_param)==5: [xc,yc] = xyz_xy(anchorpoint[0]/anchorpoint[2],anchorpoint[1]/anchorpoint[2],distortion_param,cam_intrinsic) else: [xc,yc] = cam_img_kb(anchorpoint[0],anchorpoint[1],anchorpoint[2],cam_intrinsic,distortion_param) xc,yc = xc/origW,yc/origH value_3d_face = [anchorpoint[0],anchorpoint[1],anchorpoint[2],alpha,xc,yc,score,1] value_3d_4face.append(value_3d_face) value_3d_4face.append([-1,-1,-1,-1,-1,-1,0,-1]) value_3d_4face.append([-1,-1,-1,-1,-1,-1,0,-1]) value_3d_4face.append([-1,-1,-1,-1,-1,-1,0,-1]) elif 'G1M2' in text_basename and (x_down.sum() < 8 and rotation_y > 0 and y_front.sum() == 4 and y_tail_center.sum() == 1): ####G1M2车头底边被截断### anchorpoint = xyz3d_1 alpha = -np.arctan2(anchorpoint[0],anchorpoint[2]) + rotation_y if alpha<-np.pi: alpha+=2*np.pi if alpha>np.pi: alpha-=2*np.pi score = 1 if len(distortion_param)==5: [xc,yc] = xyz_xy(anchorpoint[0]/anchorpoint[2],anchorpoint[1]/anchorpoint[2],distortion_param,cam_intrinsic) else: [xc,yc] = cam_img_kb(anchorpoint[0],anchorpoint[1],anchorpoint[2],cam_intrinsic,distortion_param) xc,yc = xc/origW,yc/origH value_3d_face = [anchorpoint[0],anchorpoint[1],anchorpoint[2],alpha,xc,yc,score,1] value_3d_4face.append([-1,-1,-1,-1,-1,-1,0,-1]) value_3d_4face.append(value_3d_face) value_3d_4face.append([-1,-1,-1,-1,-1,-1,0,-1]) value_3d_4face.append([-1,-1,-1,-1,-1,-1,0,-1]) elif (rotation_y < 0 and y_front_center.sum() == 0) or (rotation_y > 0 and y_tail_center.sum() == 0): value_3d_4face = [] else: for xyz3d in xyz3d_4face: angle = Calculate_vector_angle(np.array([xyz3d[0] - x3d_ori, xyz3d[2] - z3d_ori]), np.array([xyz3d[0], xyz3d[2]])) score = angle2score(angle) alpha = -np.arctan2(xyz3d[0],xyz3d[2]) + rotation_y if alpha<-np.pi: alpha+=2*np.pi if alpha>np.pi: alpha-=2*np.pi if len(distortion_param)==5: [xc,yc] = xyz_xy(xyz3d[0]/xyz3d[2],xyz3d[1]/xyz3d[2],distortion_param,cam_intrinsic) else: [xc,yc] = cam_img_kb(xyz3d[0],xyz3d[1],xyz3d[2],cam_intrinsic,distortion_param) xc,yc = xc/origW,yc/origH value_3d_face = [xyz3d[0],xyz3d[1],xyz3d[2],alpha,xc,yc,score,1] value_3d_4face.append(value_3d_face) value_3d_4face = np.array(value_3d_4face) if singleLW == 0: #####L W Believable value_3d_4face = value_3d_4face elif singleLW == 1: #####L Believable value_3d_4face[0:2,7] = -1 elif singleLW == 2: #####W Believable value_3d_4face[2:4,7] = -1 label_temp.append(singleLW) side_diff = np.min((abs(0-box2d[0]),abs(origW-box2d[2]))) value_3d_4face = np.array(value_3d_4face) ################################ if 'G1M2' in text_basename or 'D4Q' in text_basename: side_diff_threshold = 12 elif 'G1M3' in text_basename: side_diff_threshold = 70 if len(value_3d_4face)!=0: if have_in >= 1 and have_in < 4: # if np.sum(value_3d_4face[:,-1])==-2 and np.sum(value_3d_4face[:,-2]) == 1: if side_diff>side_diff_threshold and abs(rotation_y) >= 0.785 and abs(rotation_y)<= 2.355: ##45 label_temp = label_temp[0:6] label_temp[5]=-1 # print(annotations) elif not(np.sum(value_3d_4face[:,-1])==-2 and np.sum(value_3d_4face[:,-2]) == 1) : if side_diff6: for value_3d_face in value_3d_4face: for value_3d in value_3d_face: label_temp.append(value_3d) label_temp.append(confidence) else: label_temp.append(-1) label_temp.append(confidence) label_temp_all.append(label_temp) return label_temp_all def select_texts_for_vis(textlists, select_mode, stride_step=1, sample_num_per_uuid=5): if select_mode == 'all': return textlists if select_mode == 'stride': return textlists[0::stride_step] if select_mode == 'per_uuid_random': uuid_to_texts = {} for text in textlists: uuid_dir = os.path.dirname(os.path.dirname(text)) uuid_to_texts.setdefault(uuid_dir, []).append(text) selected_texts = [] for uuid_dir in sorted(uuid_to_texts.keys()): uuid_texts = sorted(uuid_to_texts[uuid_dir]) sample_count = min(sample_num_per_uuid, len(uuid_texts)) if sample_count == len(uuid_texts): sampled_texts = uuid_texts else: sampled_texts = sorted(random.sample(uuid_texts, sample_count)) selected_texts.extend(sampled_texts) return selected_texts raise ValueError(f'Unsupported select_mode: {select_mode}') def process_single_label(args): """Worker function for processing a single label file""" text, showpath, origW, origH, count, save_mode, save_scale = args try: label_temp_all = [] textname = os.path.basename(text) group_id = os.path.basename(os.path.dirname(os.path.dirname(text))) print(count, textname) path, name = os.path.split(text) calib_path = resolve_camera_calib_path(path.replace('labels','calib'), text) cam_intrinsic, distortion_param, cam_rectify, Vehiclepos, VehiclePitch, VehicleYaw, VehicleRoll = readcalib(calib_path) context = open(text,'r').readlines() for ll in context: data = ll.strip('\n').split(' ') label_temp = [] for i in range(len(data)): label_temp.append(data[i]) label_temp_all.append(label_temp) img_path = resolve_image_path(text) img = cv2.imread(img_path) if img is None: print(f"Error: Failed to read image from {img_path}") return False origH, origW = img.shape[:2] label_canvas_w, label_canvas_h = resolve_label_canvas_size( calib_path, origW, origH, label_temp_all, cam_intrinsic, distortion_param, ) print(f"Image loaded successfully: shape={img.shape}") img2d = showFace2Dbox( img.copy(), label_temp_all, cam_intrinsic, distortion_param, cam_rectify, origW, origH, pitch=0, label_canvas_w=label_canvas_w, label_canvas_h=label_canvas_h, ) img3d = showFace3Dbox(img.copy(), label_temp_all, cam_intrinsic, distortion_param, cam_rectify, origW, origH, pitch=0) if img2d is None or img3d is None: print(f"Error: showFace2Dbox or showFace3Dbox returned None") return False image_name = os.path.basename(text)[0:-4] + '.jpg' save_targets = [] if save_mode == 'concat': img_show = cv2.vconcat([img2d, img3d]) if save_scale != 1.0: img_show = cv2.resize(img_show, (0,0), fx=save_scale, fy=save_scale, interpolation=cv2.INTER_LINEAR) save_targets.append((os.path.join(showpath, 'concat', group_id, image_name), img_show)) elif save_mode == 'overlay': img_overlay = showFace3Dbox(img2d.copy(), label_temp_all, cam_intrinsic, distortion_param, cam_rectify, origW, origH, pitch=0) if img_overlay is None: print("Error: showFace3Dbox returned None in overlay mode") return False if save_scale != 1.0: img_overlay = cv2.resize(img_overlay, (0,0), fx=save_scale, fy=save_scale, interpolation=cv2.INTER_LINEAR) save_targets.append((os.path.join(showpath, 'overlay', group_id, image_name), img_overlay)) elif save_mode == 'split': if save_scale != 1.0: img2d = cv2.resize(img2d, (0,0), fx=save_scale, fy=save_scale, interpolation=cv2.INTER_LINEAR) img3d = cv2.resize(img3d, (0,0), fx=save_scale, fy=save_scale, interpolation=cv2.INTER_LINEAR) save_targets.append((os.path.join(showpath, '2D', group_id, image_name), img2d)) save_targets.append((os.path.join(showpath, '3D', group_id, image_name), img3d)) else: print(f"Error: unsupported save_mode {save_mode}") return False for output_path, output_img in save_targets: os.makedirs(os.path.dirname(output_path), exist_ok=True) result = cv2.imwrite(output_path, output_img) if not result: print(f"Error: cv2.imwrite failed for {output_path}") return False # 验证文件是否成功写入且大小不为0 if os.path.exists(output_path): file_size = os.path.getsize(output_path) print(f"Successfully saved: {output_path}, size={file_size} bytes") if file_size == 0: print(f"Warning: Output file size is 0!") return False else: print(f"Error: Output file does not exist: {output_path}") return False return True except Exception as e: print(f"Error processing {text}: {str(e)}") traceback.print_exc() return False if __name__ == '__main__': ####GT inpath = '/mnt/mono3d/ydong_data/Mono3d_camera2_gt_by_group_id' random.seed(2100) showpath_root = '/data1/dongying/Mono3d/D4Q2/data_visualization_for_check_camera2' save_mode = 'split' # 'concat': 2D/3D上下拼接, 'overlay': 2D/3D叠加到同一张图, 'split': 分开保存到2D/3D目录 save_scale = 0.5 select_mode = 'stride' ###'all' 使用全部label 'stride'按步长抽取 'per_uuid_random',按每个 uuid 单独随机抽样.随机选 sample_num_per_uuid 个 label stride_step = 3 ##3 sample_num_per_uuid = 5 group_batch_size = 1000 origH, origW = 1280, 1920 #2160, 3840 num_processes = 24 # test # group_batch_size = 10 # num_processes = 4 if save_mode not in ['concat', 'overlay', 'split']: raise ValueError(f'Unsupported save_mode: {save_mode}') os.makedirs(showpath_root, exist_ok=True) group_dirs = sorted( group_dir for group_dir in glob.glob(os.path.join(inpath, '*')) if os.path.isdir(group_dir) ) # group_dirs = group_dirs[:20] # group_dirs = ['/mnt/mono3d/ydong_data/Mono3d_camera2_gt_by_group_id/019b45ab-8a37-7ae9-a448-6a8bc34f84b5'] total_groups = len(group_dirs) if total_groups == 0: raise FileNotFoundError(f'No group directories found under {inpath}') total_batches = math.ceil(total_groups / group_batch_size) total_success_count = 0 total_task_count = 0 for batch_idx, start_idx in enumerate(range(0, total_groups, group_batch_size), 1): batch_group_dirs = group_dirs[start_idx:start_idx + group_batch_size] batch_showpath = os.path.join(showpath_root, f'batch_{batch_idx:03d}') os.makedirs(batch_showpath, exist_ok=True) args_list = [] count = 0 for group_dir in batch_group_dirs: textlists = sorted(glob.glob(os.path.join(group_dir, 'labels', '*.txt'))) if not textlists: print(f"Skip group without labels: {group_dir}") continue selected_texts = select_texts_for_vis( textlists, select_mode, stride_step=stride_step, sample_num_per_uuid=sample_num_per_uuid, ) for text in selected_texts: count += 1 args_list.append((text, batch_showpath, origW, origH, count, save_mode, save_scale)) if not args_list: print(f"Batch {batch_idx}/{total_batches}: no labels selected, skip") continue print( f"Starting batch {batch_idx}/{total_batches} with {len(batch_group_dirs)} groups, " f"{len(args_list)} files, save root: {batch_showpath}" ) with Pool(processes=num_processes) as pool: results = pool.map(process_single_label, args_list) success_count = sum(results) total_success_count += success_count total_task_count += len(args_list) print( f"Completed batch {batch_idx}/{total_batches}: " f"{success_count}/{len(args_list)} files processed successfully" ) print(f"All batches completed: {total_success_count}/{total_task_count} files processed successfully")