def infer_and_save_video(video_path, model, labels_dict, cfg, output_video='output.avi'):

    """视频推理"""

    image_widget = widgets.Image(format='jpeg', width=800, height=600)

    display(image_widget)

 

    # 读入视频

    cap = cv2.VideoCapture(video_path)

    fourcc = cv2.VideoWriter_fourcc(*'XVID')

    width = int(cap.get(3))

    height = int(cap.get(4))

    fps = cap.get(cv2.CAP_PROP_FPS)

    export_cap = cv2.VideoWriter(output_video, fourcc, fps, (width, height), True)

    while True:

        ret, img_frame = cap.read()

        if not ret:

            break

        # 对视频帧进行推理

        image_pred = infer_frame_with_vis(img_frame, model, labels_dict, cfg, bgr2rgb=True)

        export_cap.write(image_pred)  # 写入视频

        image_widget.value = img2bytes(image_pred)

    export_cap.release()

 

video_path = 'racing.mp4'

infer_and_save_video(video_path, model, labels_dict, cfg)

if not os.path.exists('infer_result'):

    os.makedirs('infer_result')

ans = {}

import ipywidgets as widgets

import numpy as np

def img2bytes(image):

    """将图片转换为字节码"""

    return bytes(cv2.imencode('.jpg', image)[1])

   

def find_camera_index():

    max_index_to_check = 10  # Maximum index to check for camera

 

    for index in range(max_index_to_check):

        cap = cv2.VideoCapture(index)

        if cap.read()[0]:

            cap.release()

            return index

 

    # If no camera is found

    raise ValueError("No camera found.")

 

# 获取摄像头

camera_index = find_camera_index()

cap = cv2.VideoCapture(camera_index)

# 初始化可视化对象

image_widget = widgets.Image(format='jpeg', width=1280, height=720)

display(image_widget)

while True:

    # 对摄像头每一帧进行推理和可视化

    _, img_src = cap.read()

    # image_pred = infer_frame_with_vis(img_frame, model, labels_dict, cfg)

    # basename = os.path.basename(IMAGE_PATH)

   

    image_h, image_w = img_src.shape[:2]

   

    det_input_tensor = det_model.preprocess(img_src)

   

    output = det_model.infer(det_input_tensor)

   

    bboxes = det_model.postprocess(output)

    # im = Image.open(IMAGE_PATH)

    im = Image.fromarray(img_src)

    draw = ImageDraw.Draw(im)

    ans['image_name'] = basename

    ans['result'] = []

    for bbox in bboxes:

        bbox_detail = {}

        x1 = int(bbox[0] / det_model.model_width * image_w)

        y1 = int(bbox[1] / det_model.model_height * image_h)

        x2 = int(bbox[2] / det_model.model_width * image_w)

        y2 = int(bbox[3] / det_model.model_height * image_h)

        draw.line([(x1, y1), (x1, y2), (x2, y2), (x2, y1), (x1, y1)], fill='red', width=2)

        bbox_detail['bbox'] = [x1, y1, x1, y2, x2, y2, x2, y1]

        res  = ','.join(map(str,bbox_detail['bbox']))

        print(f'det result: {res}')

        crop_img = img_src[y1:y2, x1:x2]

   

        rec_input_tensor = rec_model.preprocess(crop_img)

        rec_output = rec_model.infer(rec_input_tensor)

        bbox_detail['text'] = rec_model.postprocess(rec_output)

        print(f'rec result: {bbox_detail["text"]}')

        ans['result'].append(bbox_detail)

 

    image_widget.value = img2bytes(np.array(im))

cap.release()