Vitis AI on Zynq Ultrascale+ MPSoC

    
import cv2
import numpy as np
import vart
import xir
import math
import time

def sigmoid(x):
        return 1 / (1 + np.exp(-x))

def calculate_iou(box1, box2):
    x1, y1, x2, y2 = box1
    x1_b, y1_b, x2_b, y2_b = box2
    inter_x1 = max(x1, x1_b)
    inter_y1 = max(y1, y1_b)
    inter_x2 = min(x2, x2_b)
    inter_y2 = min(y2, y2_b)
    inter_area = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)
    box1_area = (x2 - x1) * (y2 - y1)
    box2_area = (x2_b - x1_b) * (y2_b - y1_b)
    union_area = box1_area + box2_area - inter_area
    return inter_area / union_area if union_area > 0 else 0

def get_child_subgraph_dpu(graph):
        root_subgraph = graph.get_root_subgraph()
        child_subgraphs = root_subgraph.toposort_child_subgraph()
        return [cs for cs in child_subgraphs if cs.has_attr("device") and cs.get_attr("device").upper() == "DPU"]

def main():
    model_path = "/home/root/work/B1152/ofa_yolo_pruned_0_50_pt.xmodel"
    anchors = [
        [(10, 13), (16, 30), (33, 23)],
        [(30, 61), (62, 45), (59, 119)],
        [(116, 90), (156, 198), (373, 326)],
    ]
    num_classes = 80
    output_shapes = [(80, 80), (40, 40), (20, 20)]
    scale_factor = 0.00392156
    conf_threshold = 0.2
    nms_threshold = 0.2
    inputId = 0

    cam = cv2.VideoCapture(inputId)
    cam.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cam.set(cv2.CAP_PROP_FRAME_HEIGHT, 360)
    if not cam.isOpened():
        print("[ERROR] Failed to open camera ", inputId)
        exit()

    graph = xir.Graph.deserialize(model_path)
    subgraphs = get_child_subgraph_dpu(graph)
    assert len(subgraphs) == 1, "Expected exactly one DPU subgraph."
    dpu_runner = vart.Runner.create_runner(subgraphs[0], "run")

    input_tensors = dpu_runner.get_input_tensors()
    output_tensors = dpu_runner.get_output_tensors()

    fpga_input = [np.empty(input_tensors[0].dims, dtype=np.int8)]
    fpga_output = [np.empty(tuple(tensor.dims[1:]), dtype=np.int8, order="C") for tensor in output_tensors]

    fix_i = 2 ** input_tensors[0].get_attr("fix_point") * scale_factor
    fix_o = [1 / (2 ** tensor.get_attr("fix_point")) for tensor in output_tensors]

    frame_count = 0
    capture_time_total = 0
    preprocess_time_total = 0
    inference_time_total = 0
    postprocess_time_total = 0
    nms_time_total = 0
    draw_time_total = 0
    start_time = time.time()
    try:
        while True:
            capture_start = time.time()            
            ret, frame_raw = cam.read()
            capture_end = time.time()
            if not ret:
                break
            capture_time_total += (capture_end - capture_start)

            preprocess_start = time.time()
            frame = cv2.resize(frame_raw, (640, 640), interpolation=cv2.INTER_LINEAR)
            img_quantized = np.round(frame.astype(np.float32) * fix_i).astype(np.int8)
            preprocess_end = time.time()
            preprocess_time_total += (preprocess_end - preprocess_start)

            inference_start = time.time()
            fpga_input[0][0, ...] = img_quantized
            job_id = dpu_runner.execute_async(fpga_input, fpga_output)
            dpu_runner.wait(job_id)
            inference_end = time.time()
            inference_time_total += (inference_end - inference_start)

            postprocess_start = time.time()
            scaled_outputs = [
            (fpga_output[idx].astype(np.float32) * fix_o[idx]).reshape(tuple(tensor.dims))  
            for idx, tensor in enumerate(output_tensors)]

            reshaped_outputs = []
            for idx, scaled_output in enumerate(scaled_outputs):
                grid_height, grid_width = output_shapes[idx]
                expected_shape = (grid_height, grid_width, len(anchors[idx]), 5 + num_classes)
                reshaped_outputs.append(scaled_output.reshape(expected_shape))

            decoded_boxes = []
            for layer_idx, output in enumerate(reshaped_outputs):
                grid_height, grid_width, num_anchors, _ = output.shape
                anchors_layer = np.array(anchors[layer_idx]).reshape(-1, 2)
                tx, ty, tw, th, confidence = np.split(output[..., :5], 5, axis=-1)
                class_scores = output[..., 5:]
                confidence = sigmoid(confidence)
                class_scores = sigmoid(class_scores)
                for h in range(grid_height):
                    for w in range(grid_width):
                        for anchor_idx in range(num_anchors):
                            
                            tx_val = tx[h, w, anchor_idx]
                            ty_val = ty[h, w, anchor_idx]
                            tw_val = tw[h, w, anchor_idx]
                            th_val = th[h, w, anchor_idx]
                            conf_val = confidence[h, w, anchor_idx]

                            if conf_val <= conf_threshold:
                                continue

                            
                            bx = (w + sigmoid(tx_val)) / grid_width
                            by = (h + sigmoid(ty_val)) / grid_height
                            bw = np.exp(tw_val) * anchors_layer[anchor_idx, 0] / 640
                            bh = np.exp(th_val) * anchors_layer[anchor_idx, 1] / 640

                            
                            x_min = (bx - bw / 2) * 640
                            y_min = (by - bh / 2) * 640
                            x_max = (bx + bw / 2) * 640
                            y_max = (by + bh / 2) * 640

                            
                            class_probs = class_scores[h, w, anchor_idx, :]

                            
                            decoded_boxes.append({
                                "box": [x_min, y_min, x_max, y_max],
                                "confidence": conf_val,
                                "class_probs": class_probs,
                            })
            postprocess_end = time.time()
            postprocess_time_total += (postprocess_end - postprocess_start)
            nms_start = time.time()
            decoded_boxes.sort(key=lambda x: x["confidence"], reverse=True)
            nms_boxes = []
            while decoded_boxes:
                best_box = decoded_boxes.pop(0)
                nms_boxes.append(best_box)
                decoded_boxes = [
                    box for box in decoded_boxes if calculate_iou(best_box["box"], box["box"]) < nms_threshold
                ]
            nms_end = time.time()
            nms_time_total += (nms_end - nms_start)
            draw_start = time.time()
            image_height, image_width = frame_raw.shape[:2]
            scale_x = 640 / image_width
            scale_y = 640 / image_height
            pad_x = (640 - image_width * scale_x) / 2
            pad_y = (640 - image_height * scale_y) / 2
            for detection in nms_boxes:
                x_min, y_min, x_max, y_max = detection["box"]
                x_min = max(0, int((x_min - pad_x) / scale_x))
                y_min = max(0, int((y_min - pad_y) / scale_y))
                x_max = min(image_width, int((x_max - pad_x) / scale_x))
                y_max = min(image_height, int((y_max - pad_y) / scale_y))
                confidence = float(detection["confidence"])
                class_probs = detection["class_probs"]
                class_id = int(np.argmax(class_probs))
                label = f"Class {class_id} ({confidence:.2f})"
                cv2.rectangle(frame_raw, (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 2)
                cv2.putText(frame_raw, label, (int(x_min), int(y_min - 10)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
            draw_end = time.time()
            draw_time_total += (draw_end - draw_start)
            frame_count += 1
            current_time = time.time()
            elapsed_time = current_time - start_time
            fps = frame_count / elapsed_time
            cv2.putText(frame_raw, f"FPS: {fps:.2f}", (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (20, 20, 180), 2)
            cv2.imshow("Non-optimized ofa-yolo detector", frame_raw)
            key = cv2.waitKey(1) & 0xFF
            if key == 27:  # Press 'ESC' to exit
                break
    finally:
        cam.release()
        cv2.destroyAllWindows()
        print("\n[INFO] Average Timings per Frame (in ms):")
        print(f"[INFO] {'Frame Capture:':<20} {capture_time_total / frame_count * 1000:.2f} ms")
        print(f"[INFO] {'Preprocessing:':<20} {preprocess_time_total / frame_count * 1000:.2f} ms")
        print(f"[INFO] {'Run DPU:':<20} {inference_time_total / frame_count * 1000:.2f} ms")
        print(f"[INFO] {'Postprocessing:':<20} {postprocess_time_total / frame_count * 1000:.2f} ms")
        print(f"[INFO] {'NMS filtering:':<20} {nms_time_total / frame_count * 1000:.2f} ms")
        print(f"[INFO] {'Rendering/Draw:':<20} {draw_time_total / frame_count * 1000:.2f} ms")
        print(f"[INFO] {'Overall FPS:':<20} {fps:.2f}")

if __name__ == "__main__":
    main()