inference_axmodel.py

#!/usr/bin/env python3
import argparse
import json
import os
import os.path as osp
import cv2
import numpy as np
import axengine as axe
from collections import defaultdict
from tqdm import tqdm


def parse_args():
    parser = argparse.ArgumentParser(description='BEVFormer AXEngine Inference from Extracted Data')
    parser.add_argument('model', help='AXModel path')
    parser.add_argument('config_json', help='JSON config file path')
    parser.add_argument('data_dir', help='extracted data directory (extracted_data)')
    parser.add_argument('--output-dir', default='./inference_results_extracted', help='output directory')
    parser.add_argument('--score-thr', type=float, default=0.1, help='score threshold')
    parser.add_argument('--fps', type=int, default=3, help='video fps')
    parser.add_argument('--start-scene', type=int, default=0, help='start scene index')
    parser.add_argument('--end-scene', type=int, default=None, help='end scene index (None for all)')
    return parser.parse_args()


def load_axmodel(axmodel_path):
    """Load AXModel"""
    # 尝试使用 AxEngineExecutionProvider 而不是 AXCLRTExecutionProvider
    providers = ['AxEngineExecutionProvider']
    session = axe.InferenceSession(axmodel_path, providers=providers)
    return session


def load_config_from_json(config_path):
    """Load configuration from JSON file"""
    with open(config_path, 'r') as f:
        config = json.load(f)
    return config


def preprocess_image(img_path, img_norm_cfg, target_size=(480, 800)):
    """Preprocess image: load, resize, normalize
    
    Args:
        img_path: path to image file
        img_norm_cfg: normalization config with 'mean', 'std', 'to_rgb'
        target_size: (H, W) target size
    
    Returns:
        img: (C, H, W) normalized numpy array, float32
    """
    # Load image
    img = cv2.imread(img_path)
    if img is None:
        raise ValueError(f"Cannot load image: {img_path}")
    
    # Convert BGR to RGB if needed
    if img_norm_cfg.get('to_rgb', True):
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    
    # Resize if needed
    if img.shape[:2] != target_size:
        img = cv2.resize(img, (target_size[1], target_size[0]))  # (W, H)
    
    # Convert to float and normalize
    img = img.astype(np.float32)
    mean = np.array(img_norm_cfg.get('mean', [123.675, 116.28, 103.53]), dtype=np.float32)
    std = np.array(img_norm_cfg.get('std', [58.395, 57.12, 57.375]), dtype=np.float32)

    img = (img - mean) / std
    img = img.transpose(2, 0, 1)  # (H, W, C) -> (C, H, W)
    
    return img


def load_data(data_dir, scene_name, frame_idx):
    """Load data
    
    Args:
        data_dir: data directory path
        scene_name: scene name (scene token)
        frame_idx: frame index (sample index)
    
    Returns:
        img: (1, N, C, H, W) numpy array
        lidar2img: (1, N, 4, 4) numpy array
        can_bus: (1, 18) numpy array
        meta: dict with metadata
    """
    scene_dir = osp.join(data_dir, scene_name)
    
    # Load meta
    meta_path = osp.join(scene_dir, f'meta_{frame_idx:06d}.json')
    with open(meta_path, 'r') as f:
        meta = json.load(f)
    
    # Get normalization config
    img_norm_cfg = meta.get('img_norm_cfg', {
        'mean': [123.675, 116.28, 103.53],
        'std': [58.395, 57.12, 57.375],
        'to_rgb': True
    })
    
    # Get image shape
    img_shape = meta.get('img_shape', [[480, 800, 3]] * 6)
    target_size = (img_shape[0][0], img_shape[0][1])  # (H, W)
    
    # Load images for all cameras
    num_cams = meta.get('num_cams', 6)
    imgs = []
    for cam_idx in range(num_cams):
        img_path = osp.join(scene_dir, f'cam_{cam_idx:02d}_{frame_idx:06d}.png')
        img = preprocess_image(img_path, img_norm_cfg, target_size)
        imgs.append(img)
    
    # Stack images: (N, C, H, W) -> (1, N, C, H, W)
    img = np.stack(imgs, axis=0)  # (N, C, H, W)
    img = img[np.newaxis, ...]  # (1, N, C, H, W)
    
    # Load lidar2img: (N, 4, 4) -> (1, N, 4, 4)
    lidar2img = np.array(meta['lidar2img'], dtype=np.float32)  # (N, 4, 4)
    lidar2img = lidar2img[np.newaxis, ...]  # (1, N, 4, 4)
    
    # Load can_bus: (18,) -> (1, 18)
    can_bus = np.array(meta['can_bus'], dtype=np.float32)  # (18,)
    can_bus = can_bus[np.newaxis, ...]  # (1, 18)
    
    return img, lidar2img, can_bus, meta

CLASS_COLORS = {
    0: (0, 255, 0), 1: (255, 255, 0), 2: (0, 0, 255), 3: (0, 165, 255),
    4: (255, 0, 255), 5: (0, 255, 255), 6: (128, 0, 128), 7: (255, 165, 0),
    8: (0, 0, 255), 9: (128, 128, 128),
}


def denormalize_bbox_np(normalized_bboxes, pc_range):
    """Denormalize bbox using numpy only"""
    # rotation
    rot_sine = normalized_bboxes[..., 6:7]
    rot_cosine = normalized_bboxes[..., 7:8]
    
    rot = np.arctan2(rot_sine, rot_cosine)
    
    # center in the bev
    cx = normalized_bboxes[..., 0:1]
    cy = normalized_bboxes[..., 1:2]
    cz = normalized_bboxes[..., 4:5]
    
    # size
    w = normalized_bboxes[..., 2:3]
    l = normalized_bboxes[..., 3:4]
    h = normalized_bboxes[..., 5:6]
    
    w = np.exp(w)
    l = np.exp(l)
    h = np.exp(h)
    
    if normalized_bboxes.shape[-1] > 8:
        # velocity
        vx = normalized_bboxes[:, 8:9]
        vy = normalized_bboxes[:, 9:10]
        denormalized_bboxes = np.concatenate([cx, cy, cz, w, l, h, rot, vx, vy], axis=-1)
    else:
        denormalized_bboxes = np.concatenate([cx, cy, cz, w, l, h, rot], axis=-1)
    return denormalized_bboxes

def decode_bboxes_custom_np(all_cls_scores, all_bbox_preds, pc_range, post_center_range, max_num=100, score_threshold=None, num_classes=10):
    """Custom bbox decode function"""
    # Use output from the last decoder layer
    all_cls_scores = all_cls_scores[-1]  # (bs, num_query, num_classes)
    all_bbox_preds = all_bbox_preds[-1]  # (bs, num_query, 10)
    
    batch_size = all_cls_scores.shape[0]
    predictions_list = []
    
    for i in range(batch_size):
        cls_scores = all_cls_scores[i]  # (num_query, num_classes)
        bbox_preds = all_bbox_preds[i]  # (num_query, 10)
        
        # Apply sigmoid
        cls_scores = 1.0 / (1.0 + np.exp(-cls_scores))
        
        # TopK selection
        cls_scores_flat = cls_scores.reshape(-1)
        topk_indices = np.argsort(cls_scores_flat)[::-1][:max_num]
        scores = cls_scores_flat[topk_indices]
        labels = topk_indices % num_classes
        bbox_index = topk_indices // num_classes
        bbox_preds = bbox_preds[bbox_index]
        
        # Denormalize bbox
        final_box_preds = denormalize_bbox_np(bbox_preds, pc_range)  # (max_num, 9)
        final_scores = scores
        final_preds = labels
        
        # Apply score threshold
        if score_threshold is not None:
            thresh_mask = final_scores > score_threshold
            tmp_score = score_threshold
            while thresh_mask.sum() == 0:
                tmp_score *= 0.9
                if tmp_score < 0.01:
                    thresh_mask = np.ones(len(final_scores), dtype=bool)
                    break
                thresh_mask = final_scores >= tmp_score
        else:
            thresh_mask = np.ones(len(final_scores), dtype=bool)
        
        # Apply post processing range filtering
        if post_center_range is not None:
            post_center_range_arr = np.array(post_center_range)
            mask = (final_box_preds[..., :3] >= post_center_range_arr[:3]).all(1)
            mask &= (final_box_preds[..., :3] <= post_center_range_arr[3:]).all(1)
            mask &= thresh_mask
            
            boxes3d = final_box_preds[mask]
            scores = final_scores[mask]
            labels = final_preds[mask]
        else:
            boxes3d = final_box_preds[thresh_mask]
            scores = final_scores[thresh_mask]
            labels = final_preds[thresh_mask]
        
        predictions_list.append({
            'bboxes': boxes3d,
            'scores': scores,
            'labels': labels
        })
    
    return predictions_list


def get_bboxes_custom_np(preds_dicts, pc_range, post_center_range, max_num=100, score_threshold=None, num_classes=10):
    """Custom get_bboxes function"""
    # Decode bounding boxes
    preds_list = decode_bboxes_custom_np(
        preds_dicts['all_cls_scores'],
        preds_dicts['all_bbox_preds'],
        pc_range,
        post_center_range,
        max_num,
        score_threshold,
        num_classes
    )
    
    num_samples = len(preds_list)
    ret_list = []
    
    for i in range(num_samples):
        preds = preds_list[i]
        bboxes = preds['bboxes']
        
        if len(bboxes) == 0:
            ret_list.append((
                np.zeros((0, 9), dtype=np.float32),
                np.zeros((0,), dtype=np.float32),
                np.zeros((0,), dtype=np.int64)
            ))
            continue
        
        # Adjust z coordinate: convert center z to bottom center z
        bboxes[:, 2] = bboxes[:, 2] - bboxes[:, 5] * 0.5
        
        # Shrink box dimensions: multiply w, l, h by 0.9 to fix oversized boxes
        bboxes[:, 3:6] = bboxes[:, 3:6] * 0.9
        
        scores = preds['scores']
        labels = preds['labels']
        
        ret_list.append((bboxes, scores, labels))
    
    return ret_list


def format_bbox_result_np(bboxes, scores, labels):
    return {
        'boxes_3d': bboxes,
        'scores_3d': scores,
        'labels_3d': labels
    }


def rotation_3d_in_axis_np(points, angles, axis=2):
    """Rotate points by angles according to axis"""
    rot_sin = np.sin(angles)
    rot_cos = np.cos(angles)
    ones = np.ones_like(rot_cos)
    zeros = np.zeros_like(rot_cos)
    
    if axis == 2 or axis == -1:
        # Rotate around z-axis
        # Build rotation matrix: (N, 3, 3)
        N = len(angles)
        rot_mat = np.zeros((N, 3, 3), dtype=points.dtype)
        rot_mat[:, 0, 0] = rot_cos
        rot_mat[:, 0, 1] = -rot_sin
        rot_mat[:, 0, 2] = zeros
        rot_mat[:, 1, 0] = rot_sin
        rot_mat[:, 1, 1] = rot_cos
        rot_mat[:, 1, 2] = zeros
        rot_mat[:, 2, 0] = zeros
        rot_mat[:, 2, 1] = zeros
        rot_mat[:, 2, 2] = ones
        
        # Rotation: (N, M, 3) @ (N, 3, 3) -> (N, M, 3)
        return np.einsum('aij,ajk->aik', points, rot_mat)
    else:
        raise ValueError(f'Only axis=2 (z-axis) is supported for LiDAR boxes')


def compute_bbox_corners_np(bboxes):
    """Compute 8 corners of 3D bbox"""
    if len(bboxes) == 0:
        return np.zeros((0, 8, 3), dtype=np.float32)
    
    dtype = bboxes.dtype
    
    # Extract bbox parameters
    centers = bboxes[:, :3]  # (N, 3) [x, y, z] - the bottom center
    w = bboxes[:, 3:4]  # width (y direction)
    l = bboxes[:, 4:5]  # length (x direction)
    h = bboxes[:, 5:6]  # height (z direction)
    dims = np.concatenate([l, w, h], axis=1)  # (N, 3) [x_size, y_size, z_size] = [l, w, h]
    yaws = bboxes[:, 6]  # (N,) yaw angle
    
    # Fix: offset yaw by -80 degrees
    yaws = yaws - (np.pi / 2.0 - np.pi / 18.0)
    
    # Generate corners
    corners_norm = np.stack(np.unravel_index(np.arange(8), [2] * 3), axis=1).astype(dtype)
    
    # Rearrange to [0, 1, 3, 2, 4, 5, 7, 6]
    corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]]
    
    # Use relative origin [0.5, 0.5, 0] (bottom center)
    corners_norm = corners_norm - np.array([0.5, 0.5, 0], dtype=dtype)
    
    # Scale corners: dims is [x_size, y_size, z_size]
    corners = dims[:, np.newaxis, :] * corners_norm[np.newaxis, :, :]  # (N, 8, 3)
    
    # Rotate around z-axis
    corners = rotation_3d_in_axis_np(corners, yaws, axis=2)
    
    # Translate to center point
    corners += centers[:, np.newaxis, :]
    
    return corners


def draw_bbox3d_on_img_custom_np(bboxes, raw_img, lidar2img_rt, color=(0, 255, 0), thickness=2):
    """Custom 3D bbox drawing"""
    img = raw_img.copy()
    
    if len(bboxes) == 0:
        return img
    
    if not isinstance(bboxes, np.ndarray):
        bboxes = np.array(bboxes)
    if not isinstance(lidar2img_rt, np.ndarray):
        lidar2img_rt = np.array(lidar2img_rt)
    
    lidar2img_rt = lidar2img_rt.reshape(4, 4)
    
    # Compute corners
    corners_3d = compute_bbox_corners_np(bboxes)  # (N, 8, 3)
    
    num_bbox = corners_3d.shape[0]
    
    # Project to 2D
    corners_3d_flat = corners_3d.reshape(-1, 3)  # (N*8, 3)
    ones = np.ones((corners_3d_flat.shape[0], 1), dtype=np.float32)
    pts_4d = np.concatenate([corners_3d_flat, ones], axis=-1)  # (N*8, 4)
    
    # Project
    pts_2d = pts_4d @ lidar2img_rt.T  # (N*8, 4)
    
    # Perspective division
    pts_2d[:, 2] = np.clip(pts_2d[:, 2], a_min=1e-5, a_max=1e5)
    pts_2d[:, 0] /= pts_2d[:, 2]
    pts_2d[:, 1] /= pts_2d[:, 2]
    
    imgfov_pts_2d = pts_2d[:, :2].reshape(num_bbox, 8, 2)
    
    line_indices = ((0, 1), (0, 3), (0, 4), (1, 2), (1, 5), (3, 2), (3, 7),
                    (4, 5), (4, 7), (2, 6), (5, 6), (6, 7))
    
    for i in range(num_bbox):
        corners = imgfov_pts_2d[i].astype(np.int32)
        for start, end in line_indices:
            pt1 = (int(corners[start, 0]), int(corners[start, 1]))
            pt2 = (int(corners[end, 0]), int(corners[end, 1]))
            # Check if points are within image range
            h, w = img.shape[:2]
            if (0 <= pt1[0] < w and 0 <= pt1[1] < h) or (0 <= pt2[0] < w and 0 <= pt2[1] < h):
                cv2.line(img, pt1, pt2, color, thickness, cv2.LINE_AA)
    
    return img.astype(np.uint8)


def post_process_outputs_np(all_cls_scores, all_bbox_preds, config, score_thr=0.1):
    bbox_coder = config['model']['bbox_coder']
    pc_range = bbox_coder['pc_range']
    post_center_range = bbox_coder['post_center_range']
    max_num = bbox_coder['max_num']
    score_threshold = bbox_coder.get('score_threshold', None)
    num_classes = bbox_coder['num_classes']
    
    preds_dicts = {
        'all_cls_scores': all_cls_scores,
        'all_bbox_preds': all_bbox_preds
    }
    
    bbox_list = get_bboxes_custom_np(
        preds_dicts, pc_range, post_center_range,
        max_num, score_threshold, num_classes
    )
    
    results = []
    for bboxes, scores, labels in bbox_list:
        # Set class score thresholds
        class_score_thrs = {
            0: 0.3,  # Car
            1: 0.3,  # Truck
            2: 0.3,  # Construction vehicle
            3: 0.3,  # Bus
            4: 0.3,  # Trailer
            5: 0.3,  # Barrier
            6: 0.3,  # Motorcycle
            7: 0.3,  # Bicycle
            8: 0.3, # Pedestrian
            9: 0.3,  # Traffic cone
        }
        default_thr = score_thr
        
        keep_indices = []
        for i in range(len(scores)):
            cls_id = int(labels[i])
            thr = class_score_thrs.get(cls_id, default_thr)
            if scores[i] > thr:
                keep_indices.append(i)
        
        if len(keep_indices) == 0:
            results.append(format_bbox_result_np(
                np.zeros((0, 9), dtype=np.float32),
                np.zeros((0,), dtype=np.float32),
                np.zeros((0,), dtype=np.int64)
            ))
            continue
        
        keep_indices = np.array(keep_indices, dtype=np.int64)
        bboxes = bboxes[keep_indices]
        scores = scores[keep_indices]
        labels = labels[keep_indices]
        
        # Circle NMS
        dist_thrs = {
            0: 2.0, 1: 3.0, 2: 2.5, 3: 4.0, 4: 3.0,
            5: 1.0, 6: 1.5, 7: 1.0, 8: 0.5, 9: 0.3,
        }
        
        if len(scores) > 0:
            keep_nms = circle_nms_np(bboxes, scores, labels, dist_thrs)
            if len(keep_nms) > 0:
                bboxes = bboxes[keep_nms]
                scores = scores[keep_nms]
                labels = labels[keep_nms]
            else:
                results.append(format_bbox_result_np(
                    np.zeros((0, 9), dtype=np.float32),
                    np.zeros((0,), dtype=np.float32),
                    np.zeros((0,), dtype=np.int64)
                ))
                continue
        
        results.append(format_bbox_result_np(bboxes, scores, labels))
    
    return results


def circle_nms_np(bboxes, scores, labels, dist_thrs):
    if len(bboxes) == 0:
        return np.array([], dtype=np.int64)
    
    keep = []
    order = np.argsort(scores)[::-1]
    bboxes = bboxes[order]
    scores = scores[order]
    labels = labels[order]
    
    pts = bboxes[:, :2]
    labels_np = labels
    
    suppressed = np.zeros(len(bboxes), dtype=bool)
    
    for i in range(len(bboxes)):
        if suppressed[i]:
            continue
        keep.append(order[i])
        
        curr_cls = int(labels_np[i])
        radius = dist_thrs.get(curr_cls, 1.0)
        
        if i + 1 < len(bboxes):
            dists = np.linalg.norm(pts[i+1:] - pts[i], axis=1)
            idx_to_suppress = np.where(
                (dists < radius) & (labels_np[i+1:] == curr_cls)
            )[0]
            suppressed[i+1:][idx_to_suppress] = True
    
    return np.array(keep, dtype=np.int64)


def denormalize_img_np(img_array, img_norm_cfg):
    """Denormalize image array (C, H, W) to (H, W, C) BGR"""
    mean = np.array(img_norm_cfg.get('mean', [123.675, 116.28, 103.53]))
    std = np.array(img_norm_cfg.get('std', [58.395, 57.12, 57.375]))
    
    # (C, H, W) RGB -> (H, W, C) RGB
    if img_array.ndim == 3:
        img = img_array.transpose(1, 2, 0)
    else:
        img = img_array
    img = (img * std + mean)
    img = np.clip(img, 0, 255).astype(np.uint8)
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    return img


def draw_bev_map(bboxes, labels, scores, pc_range, bev_size=(800, 800), score_thr=0.1):
    """Draw BEV (Bird's Eye View) map with detections
    
    Args:
        bboxes: (N, 9) numpy array, format: [x, y, z, w, l, h, yaw, vx, vy]
        labels: (N,) numpy array, class labels
        scores: (N,) numpy array, detection scores
        pc_range: [x_min, y_min, z_min, x_max, y_max, z_max]
        bev_size: (width, height) of BEV image
        score_thr: score threshold
    
    Returns:
        bev_img: (H, W, 3) numpy array, BEV visualization
    """
    bev_w, bev_h = bev_size # BEV image size
    bev_img = np.ones((bev_h, bev_w, 3), dtype=np.uint8) * 255  # White background
    
    # Draw grid
    x_min, y_min, z_min, x_max, y_max, z_max = pc_range
    x_range = x_max - x_min
    y_range = y_max - y_min
    
    # Draw grid lines
    grid_color = (200, 200, 200)  # Light gray grid lines
    for i in range(-5, 6):
        x = x_min + (i + 5) * x_range / 10
        y = y_min + (i + 5) * y_range / 10
        # Vertical lines (y direction in LiDAR -> x direction in image)
        img_x = int((y - y_min) / y_range * bev_w)
        if 0 <= img_x < bev_w:
            cv2.line(bev_img, (img_x, 0), (img_x, bev_h), grid_color, 1)
        # Horizontal lines (x direction in LiDAR -> y direction in image, flipped)
        img_y = int((x_max - x) / x_range * bev_h)
        if 0 <= img_y < bev_h:
            cv2.line(bev_img, (0, img_y), (bev_w, img_y), grid_color, 1)
    
    # Draw center lines (ego vehicle position) - darker on white background
    center_x = int((0 - y_min) / y_range * bev_w)
    center_y = int((x_max - 0) / x_range * bev_h)
    cv2.line(bev_img, (center_x, 0), (center_x, bev_h), (150, 150, 150), 2)
    cv2.line(bev_img, (0, center_y), (bev_w, center_y), (150, 150, 150), 2)
    
    
    ego_length_px = 30  # pixels (representing ~4.5m, along x-axis rightward)
    ego_width_px = 12   # pixels (representing ~1.8m, along y-axis downward)
    
    ego_corners_local = np.array([
        [ego_length_px//2, -ego_width_px//2],   # front-top (head)
        [ego_length_px//2, ego_width_px//2],    # front-bottom
        [-ego_length_px//2, ego_width_px//2],   # back-bottom
        [-ego_length_px//2, -ego_width_px//2],  # back-top
    ], dtype=np.float32)
    

    rotation_angle_90 = np.pi / 2  # 90 degrees in radians
    cos_rot_90 = np.cos(rotation_angle_90)
    sin_rot_90 = np.sin(rotation_angle_90)
    rot_mat_90 = np.array([[cos_rot_90, -sin_rot_90], [sin_rot_90, cos_rot_90]])
    
    ego_corners_rotated_90 = ego_corners_local @ rot_mat_90.T
    
    ego_corners_rotated = ego_corners_rotated_90 @ rot_mat_90.T
    
    # Translate to image coordinates (center position)
    ego_corners = []
    for corner in ego_corners_rotated:
        corner_img_x = int(center_x + corner[0])
        corner_img_y = int(center_y + corner[1])
        ego_corners.append([corner_img_x, corner_img_y])
    ego_corners = np.array(ego_corners, dtype=np.int32)
    
    # Draw filled rectangle
    cv2.fillPoly(bev_img, [ego_corners], (0, 0, 255))  # Red filled
    cv2.polylines(bev_img, [ego_corners], True, (0, 0, 0), 2)  # Black outline
    

    arrow_length = ego_length_px // 2
    initial_direction = np.array([1.0, 0.0])  
    arrow_dir_rotated_90 = initial_direction @ rot_mat_90.T
    arrow_dir_rotated = arrow_dir_rotated_90 @ rot_mat_90.T
    arrow_end_x = int(center_x + arrow_length * arrow_dir_rotated[0])
    arrow_end_y = int(center_y + arrow_length * arrow_dir_rotated[1])
    cv2.arrowedLine(bev_img, (center_x, center_y), (arrow_end_x, arrow_end_y),
                   (0, 0, 0), 3, tipLength=0.3)  # Black arrow
    
    if len(bboxes) == 0:
        return bev_img
    
    if score_thr > 0:
        mask = scores > score_thr
        bboxes = bboxes[mask]
        labels = labels[mask]
        scores = scores[mask]
    
    if len(bboxes) == 0:
        return bev_img
    
    default_color = (255, 255, 255)
    

    for i in range(len(bboxes)):
        box = bboxes[i]
        label = int(labels[i])
        score = float(scores[i])
        color = CLASS_COLORS.get(label, default_color)
        
        x, y, z = box[0], box[1], box[2]  # center position
        w, l, h = box[3], box[4], box[5]  # width, length, height
        yaw = box[6]  # yaw angle
        
        yaw = yaw - np.pi / 2.0  # Subtract 90 degrees (counterclockwise)
        
        # Convert to image coordinates
        # Note: In LiDAR coordinate, x is forward, y is left, z is up
        # In BEV image (top-down view):
        #   - x (forward) -> image y (downward, flipped)
        #   - y (left) -> image x (rightward)
        # So: img_x = (y - y_min) / y_range * bev_w
        #     img_y = (x_max - x) / x_range * bev_h  (flip x to get top-down view)
        img_x = int((y - y_min) / y_range * bev_w)
        img_y = int((x_max - x) / x_range * bev_h)  # Flip x for top-down view
        
        # Skip if outside image
        if not (0 <= img_x < bev_w and 0 <= img_y < bev_h):
            continue
        
        # Calculate box dimensions in image space
        box_w_px = int(w / x_range * bev_w)
        box_l_px = int(l / y_range * bev_h)
        
        # Draw rotated rectangle
        # Calculate 4 corners of the box in LiDAR coordinates
        cos_yaw = np.cos(yaw)
        sin_yaw = np.sin(yaw)
        
        # Box corners relative to center (in LiDAR frame: x forward, y left)
        corners_local = np.array([
            [l/2, w/2],   # front-right
            [l/2, -w/2],  # front-left
            [-l/2, -w/2], # back-left
            [-l/2, w/2]   # back-right
        ])
        
        # Rotate corners
        rot_mat = np.array([[cos_yaw, -sin_yaw], [sin_yaw, cos_yaw]])
        corners_rotated = corners_local @ rot_mat.T
        
        # Translate to world coordinates and convert to image space
        corners_img = []
        for corner in corners_rotated:
            corner_x = x + corner[0]  # x in LiDAR (forward)
            corner_y = y + corner[1]  # y in LiDAR (left)
            corner_img_x = int((corner_y - y_min) / y_range * bev_w)  # y -> img_x
            corner_img_y = int((x_max - corner_x) / x_range * bev_h)  # x -> img_y (flipped)
            corners_img.append([corner_img_x, corner_img_y])
        
        corners_img = np.array(corners_img, dtype=np.int32)
        
        # Draw filled polygon (semi-transparent on white background)
        overlay = bev_img.copy()
        cv2.fillPoly(overlay, [corners_img], color)
        cv2.addWeighted(overlay, 0.5, bev_img, 0.5, 0, bev_img)
        # Draw outline (black on white background)
        cv2.polylines(bev_img, [corners_img], True, (0, 0, 0), 2)
        
        # Draw direction arrow (forward direction) - black on white
        # In LiDAR: forward is +x, left is +y
        # In BEV image: x -> img_y (flipped), y -> img_x
        # So rotation: img_x += sin(yaw) * length, img_y -= cos(yaw) * length
        arrow_length = max(box_l_px // 2, 10)
        arrow_end_x = int(img_x + arrow_length * sin_yaw)   # y component -> img_x
        arrow_end_y = int(img_y - arrow_length * cos_yaw)  # x component -> img_y (flipped)
        cv2.arrowedLine(bev_img, (img_x, img_y), (arrow_end_x, arrow_end_y),
                       (0, 0, 0), 2, tipLength=0.3)  # Black arrow
        
        # Draw center point
        cv2.circle(bev_img, (img_x, img_y), 3, (0, 0, 0), -1)  # Black center point
    
    # Rotate BEV map counterclockwise by 90 degrees (map only, not text)
    center = (bev_w // 2, bev_h // 2)
    rotation_matrix = cv2.getRotationMatrix2D(center, 90, 1.0)  # 90 degrees counterclockwise
    bev_img = cv2.warpAffine(bev_img, rotation_matrix, (bev_w, bev_h), borderValue=(255, 255, 255))
    
    # Flip horizontally to fix mirror effect
    bev_img = cv2.flip(bev_img, 1)  # 1 for horizontal flip
    
    text = 'BEV Map'
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 1
    thickness = 2
    (text_width, text_height), baseline = cv2.getTextSize(text, font, font_scale, thickness)
    text_x = bev_w - text_width - 10  
    text_y = text_height + 10  
    cv2.putText(bev_img, text, (text_x, text_y), font, font_scale, (0, 0, 0), thickness)
    
    return bev_img


def visualize_results_np(img, result, lidar2img, img_norm_cfg, class_names, score_thr=0.3, pc_range=None):
    num_cams = img.shape[1] if img.ndim == 5 else 1
    raw_imgs = [denormalize_img_np(img[0, cam_idx], img_norm_cfg) for cam_idx in range(num_cams)]
    boxes_3d = result.get('boxes_3d')
    scores_3d = result.get('scores_3d')
    labels_3d = result.get('labels_3d')
    vis_imgs = []
    boxes_3d_for_bev = labels_3d_for_bev = scores_3d_for_bev = None

    if boxes_3d is not None and len(boxes_3d) > 0:
        mask = (scores_3d > score_thr) if (score_thr > 0 and scores_3d is not None) else np.ones_like(scores_3d, dtype=bool)
        if np.any(mask):
            boxes_3d = boxes_3d[mask]
            scores_3d = scores_3d[mask]
            labels_3d = labels_3d[mask]
        boxes_3d_for_bev = boxes_3d.copy()
        labels_3d_for_bev = labels_3d.copy()
        scores_3d_for_bev = scores_3d.copy()
        for cam_idx, vis_img in enumerate(raw_imgs):
            vis_img = vis_img.copy()
            if lidar2img.shape[1] > cam_idx:
                cam_lidar2img = lidar2img[0, cam_idx]
                for box, label in zip(boxes_3d, labels_3d):
                    color = CLASS_COLORS.get(int(label), (255, 255, 255))
                    try:
                        vis_img = draw_bbox3d_on_img_custom_np(box[None], vis_img, cam_lidar2img, color=color, thickness=2)
                    except Exception:
                        pass
            vis_imgs.append(vis_img)
    else:
        vis_imgs = raw_imgs

    if pc_range is None:
        pc_range = [-51.2, -51.2, -5.0, 51.2, 51.2, 3.0]

    if boxes_3d_for_bev is not None and len(boxes_3d_for_bev) > 0:
        bev_size = (vis_imgs[0].shape[0], vis_imgs[0].shape[0]) if vis_imgs else (800, 800)
        bev_img = draw_bev_map(boxes_3d_for_bev, labels_3d_for_bev, scores_3d_for_bev, pc_range, bev_size=bev_size, score_thr=score_thr)
    else:
        bev_size = (vis_imgs[0].shape[0], vis_imgs[0].shape[0]) if vis_imgs else (800, 800)
        bev_img = np.full((bev_size[1], bev_size[0], 3), 255, np.uint8)
        cv2.putText(bev_img, 'BEV Map (No Detections)', (10, bev_size[1]//2), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)

    if len(vis_imgs) == 6:
        target_height = max(img.shape[0] for img in vis_imgs)
        resized_imgs = [img if img.shape[0] == target_height else cv2.resize(img, (int(img.shape[1] * target_height / img.shape[0]), target_height)) for img in vis_imgs]

        reordered_imgs = [
            resized_imgs[2], resized_imgs[0], resized_imgs[1],  
            cv2.flip(resized_imgs[4], 1), cv2.flip(resized_imgs[3], 1), cv2.flip(resized_imgs[5], 1)  
        ]
        top_row = np.hstack(reordered_imgs[:3])
        bottom_row = np.hstack(reordered_imgs[3:])
        left_side = np.vstack([top_row, bottom_row])
        bev_img = cv2.resize(bev_img, (int(bev_img.shape[1] * left_side.shape[0] / bev_img.shape[0]), left_side.shape[0]))
        vis_img = np.hstack([left_side, bev_img])
    elif len(vis_imgs) > 1:
        target_height = max(img.shape[0] for img in vis_imgs)
        resized_imgs = [img if img.shape[0] == target_height else cv2.resize(img, (int(img.shape[1] * target_height / img.shape[0]), target_height)) for img in vis_imgs]
        if bev_img.shape[0] != target_height:
            bev_img = cv2.resize(bev_img, (int(bev_img.shape[1] * target_height / bev_img.shape[0]), target_height))
        vis_img = np.hstack([np.hstack(resized_imgs), bev_img])
    else:
        cam_img = vis_imgs[0] if vis_imgs else bev_img
        if bev_img.shape[0] != cam_img.shape[0]:
            bev_img = cv2.resize(bev_img, (int(bev_img.shape[1] * cam_img.shape[0] / bev_img.shape[0]), cam_img.shape[0]))
        vis_img = np.hstack([cam_img, bev_img]) if vis_imgs else bev_img

    return vis_img


def create_video_from_images(image_dir, output_video_path, fps=3):
    import subprocess
    
    image_files = sorted([f for f in os.listdir(image_dir) if f.endswith(('.png', '.jpg', '.jpeg'))])
    if len(image_files) == 0:
        return
    
    first_img = cv2.imread(osp.join(image_dir, image_files[0]))
    if first_img is None:
        return
    
    height, width = first_img.shape[:2]
    
    max_width, max_height = 1920, 1080
    if width > max_width or height > max_height:
        scale = min(max_width / width, max_height / height)
        width, height = int(width * scale), int(height * scale)
    
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    video_writer = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
    if not video_writer.isOpened():
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        video_writer = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
    
    for img_file in tqdm(image_files, desc=f"Creating video: {osp.basename(output_video_path)}"):
        img_path = osp.join(image_dir, img_file)
        img = cv2.imread(img_path)
        if img is not None:
            if img.shape[:2] != (height, width):
                img = cv2.resize(img, (width, height))
            video_writer.write(img)
    
    video_writer.release()

def main():
    args = parse_args()
    
    # Load configuration from JSON
    config = load_config_from_json(args.config_json)
    
    # Create output directory
    os.makedirs(args.output_dir, exist_ok=True)
    
    # Load AXModel
    ax_session = load_axmodel(args.model)
    
    # Get model parameters from config
    transformer_cfg = config['model']['transformer']
    bev_h = transformer_cfg['bev_h']
    bev_w = transformer_cfg['bev_w']
    embed_dims = transformer_cfg['embed_dims']
    
    # Load scene index
    scene_index_path = osp.join(args.data_dir, 'scene_index.json')
    with open(scene_index_path, 'r') as f:
        scene_index_data = json.load(f)
    
    scenes_dict = scene_index_data['scenes']
    scene_names = list(scenes_dict.keys())
    
    end_scene = args.end_scene if args.end_scene is not None else len(scene_names)
    end_scene = min(end_scene, len(scene_names))
    
    prev_frame_info = {
        'prev_bev': None,
        'scene_token': None,
        'prev_pos': np.zeros(3, dtype=np.float32),
        'prev_angle': 0.0,
    }
    
    scene_results = defaultdict(list)
    
    # Process all scenes
    for scene_idx in range(args.start_scene, end_scene):
        scene_name = scene_names[scene_idx]
        scene_info = scenes_dict[scene_name]
        sample_indices = scene_info['samples']
        num_frames = len(sample_indices)
        
        print(f"Processing scene {scene_idx+1}/{len(scene_names)}: {scene_name} ({num_frames} frames)")
        
        # Reset prev_bev for new scene
        if scene_name != prev_frame_info['scene_token']:
            prev_frame_info['prev_bev'] = None
            prev_frame_info['prev_pos'] = np.zeros(3, dtype=np.float32)
            prev_frame_info['prev_angle'] = 0.0
        
        prev_frame_info['scene_token'] = scene_name
        
        # Process all frames in this scene
        for local_idx, frame_idx in enumerate(tqdm(sample_indices, desc=f"Scene {scene_name}")):
            # Load data
            img, lidar2img, can_bus, meta = load_data(args.data_dir, scene_name, frame_idx)
            
            # Process can_bus (compute delta)
            curr_can_bus_np = can_bus[0]  # (18,)
            
            tmp_pos = curr_can_bus_np[:3].copy()
            tmp_angle = curr_can_bus_np[-1]
            
            delta_can_bus_np = curr_can_bus_np.copy()
            
            if prev_frame_info['prev_bev'] is not None and prev_frame_info['scene_token'] == scene_name:
                delta_can_bus_np[:3] -= prev_frame_info['prev_pos']
                delta_can_bus_np[-1] -= prev_frame_info['prev_angle']
            else:
                delta_can_bus_np[:3] = 0.0
                delta_can_bus_np[-1] = 0.0
            
            prev_frame_info['prev_pos'] = tmp_pos
            prev_frame_info['prev_angle'] = tmp_angle
            
            # Prepare prev_bev
            prev_bev_input = next((inp for inp in ax_session.get_inputs() if inp.name == 'prev_bev'), None)
            expected_shape = (bev_h * bev_w, 1, embed_dims)
            if prev_bev_input is not None:
                expected_shape = list(prev_bev_input.shape)
                for i, dim in enumerate(expected_shape):
                    if isinstance(dim, str) or dim < 0:
                        expected_shape[i] = (bev_h * bev_w, 1, embed_dims)[i] if i < 3 else 1
                expected_shape = tuple(expected_shape)
            
            if prev_frame_info['prev_bev'] is None:
                prev_bev = np.zeros(expected_shape, dtype=np.float32)
            else:
                prev_bev = prev_frame_info['prev_bev']
                if prev_bev.shape != expected_shape and len(prev_bev.shape) == 3:
                    prev_bev = prev_bev.reshape(expected_shape)
            
            # Prepare AXEngine inputs
            img_np = img.astype(np.float32)
            lidar2img_np = lidar2img.astype(np.float32)
            can_bus_np = delta_can_bus_np.reshape(1, -1).astype(np.float32)
            
            input_names = [inp.name for inp in ax_session.get_inputs()]
            ax_inputs = {}
            for name in input_names:
                if name == 'img':
                    ax_inputs['img'] = img_np
                elif name == 'can_bus':
                    ax_inputs['can_bus'] = can_bus_np
                elif name == 'lidar2img':
                    ax_inputs['lidar2img'] = lidar2img_np
                elif name == 'prev_bev':
                    ax_inputs['prev_bev'] = prev_bev
            
            # Run inference
            ax_outputs = ax_session.run(None, ax_inputs)
            bev_embed, all_cls_scores, all_bbox_preds = ax_outputs
            
            prev_frame_info['prev_bev'] = bev_embed
            
            # Post-process
            results = post_process_outputs_np(
                all_cls_scores, all_bbox_preds, config, args.score_thr
            )
            
            # Visualize
            img_norm_cfg = config['img_norm']
            class_names = config['dataset']['class_names']
            pc_range = config['model']['bbox_coder']['pc_range']
            vis_img = visualize_results_np(
                img, results[0], lidar2img, img_norm_cfg, class_names, args.score_thr, pc_range=pc_range
            )
            
            scene_results[scene_name].append({
                'frame_idx': local_idx,
                'result': results[0],
                'vis_img': vis_img,
                'meta': meta
            })
    
    # Save results
    for scene_name, frames in tqdm(scene_results.items(), desc="Save scene results"):
        scene_dir = osp.join(args.output_dir, scene_name)
        os.makedirs(scene_dir, exist_ok=True)
        images_dir = osp.join(scene_dir, 'images')
        os.makedirs(images_dir, exist_ok=True)
        
        for local_idx, frame_data in enumerate(frames):
            vis_img = frame_data['vis_img']
            
            if vis_img is None:
                continue
            
            if not isinstance(vis_img, np.ndarray):
                vis_img = np.array(vis_img)
            
            if vis_img.dtype != np.uint8:
                vis_img = (vis_img * 255).astype(np.uint8) if vis_img.max() <= 1.0 else vis_img.astype(np.uint8)
            
            if len(vis_img.shape) == 3 and vis_img.shape[0] in (1, 3):
                vis_img = vis_img.transpose(1, 2, 0)
            
            if vis_img.shape[0] > 0 and vis_img.shape[1] > 0:
                cv2.imwrite(osp.join(images_dir, f'frame_{local_idx:06d}.png'), vis_img)
        
        create_video_from_images(images_dir, osp.join(scene_dir, f'{scene_name}_result.mp4'), args.fps)
        print(f"✓ Scene {scene_name}: {len(frames)} frames, video: {osp.join(scene_dir, f'{scene_name}_result.mp4')}")


if __name__ == '__main__':
    main()