OpenVINO™ による人物追跡#

import platform 

%pip install -q "openvino-dev>=2024.0.0" 
%pip install -q opencv-python requests scipy tqdm 

if platform.system() != "Windows":
     %pip install -q "matplotlib>=3.4" 
else:
     %pip install -q "matplotlib>=3.4,<3.7"

import collections 
from pathlib import Path 
import time 

import numpy as np 
import cv2 
from IPython import display 
import matplotlib.pyplot as plt 
import openvino as ov

# ローカルモジュールをインポート 

if not Path("./notebook_utils.py").exists():
    # `notebook_utils` モジュールを取得 
    import requests 

    r = requests.get( 

url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py", 
    ) 

    open("notebook_utils.py", "w").write(r.text) 

import notebook_utils as utils 
from deepsort_utils.tracker import Tracker 
from deepsort_utils.nn_matching import NearestNeighborDistanceMetric 
from deepsort_utils.detection import ( 
    Detection, 
    compute_color_for_labels, 
    xywh_to_xyxy, 
    xywh_to_tlwh, 
    tlwh_to_xyxy, 
)

# モデルがダウンロードされるディレクトリー 
base_model_dir = "model" 
precision = "FP16" 
# The name of the model from Open Model Zoo 
detection_model_name = "person-detection-0202" 

download_command = ( 
    f"omz_downloader " f"--name {detection_model_name} " f"--precisions {precision} " f"--output_dir {base_model_dir} " f"--cache_dir {base_model_dir}" 
) 
!  $download_command 

detection_model_path = 
f"model/intel/{detection_model_name}/{precision}/{detection_model_name}.xml" 

reidentification_model_name = "person-reidentification-retail-0287" 

download_command = ( 
    f"omz_downloader " f"--name {reidentification_model_name} " f"--precisions {precision} " f"--output_dir {base_model_dir} " f"--cache_dir {base_model_dir}" 
) 
!  $download_command 

reidentification_model_path = 
f"model/intel/{reidentification_model_name}/{precision}/{reidentification_model_name}.xml"

core = ov.Core() 

class Model: 
    """ 
    This class represents a OpenVINO model object.
    """ 

    def __init__(self, model_path, batchsize=1, device="AUTO"): 
        """ 
        Initialize the model object 

        Parameters 
        ---------- 
        model_path: path of inference model 
        batchsize: batch size of input data 
        device: device used to run inference 
        """ 
        self.model = core.read_model(model=model_path) 
        self.input_layer = self.model.input(0) 
        self.input_shape = self.input_layer.shape 
        self.height = self.input_shape[2] 
        self.width = self.input_shape[3] 

        for layer in self.model.inputs: 
            input_shape = layer.partial_shape 
            input_shape[0] = batchsize 
            self.model.reshape({layer: input_shape}) 
        self.compiled_model = core.compile_model(model=self.model, device_name=device) 
        self.output_layer = self.compiled_model.output(0) 

def predict(self, input): 
    """ 
    Run inference 

    Parameters 
    ---------- 
    input: array of input data 
    """ 
    result = self.compiled_model(input)[self.output_layer] 
    return result

def preprocess(frame, height, width): 
    """ 
    Preprocess a single image 

    Parameters 
    ---------- 
    frame: input frame 
    height: height of model input data 
    width: width of model input data 
    """ 
    resized_image = cv2.resize(frame, (width, height)) 
    resized_image = resized_image.transpose((2, 0, 1)) 
    input_image = np.expand_dims(resized_image, axis=0).astype(np.float32) 
    return input_image 

def batch_preprocess(img_crops, height, width): 
    """ 
    Preprocess batched images 

    Parameters 
    ---------- 
    img_crops: batched input images 
    height: height of model input data 
    width: width of model input data 
    """ 
    img_batch = np.concatenate([preprocess(img, height, width) for img in img_crops], axis=0) 
    return img_batch 

def process_results(h, w, results, thresh=0.5): 
    """ 
    postprocess detection results 

    Parameters 
    ---------- 
    h, w: original height and width of input image 
    results: raw detection network output 
    thresh: threshold for low confidence filtering 
    """ 
    # 'results' 変数は [1, 1, N, 7] テンソルです 
    detections = results.reshape(-1, 7) 
    boxes = [] 
    labels = [] 
    scores = [] 
    for i, detection in enumerate(detections):
        _, label, score, xmin, ymin, xmax, ymax = detection 
        # 検出されたオブジェクトをフィルター処理 
        if score > thresh:
            # 正規化された座標 [0,1] のボックスからピクセル座標のボックスを作成 
            boxes.append( 
                [ 
                    (xmin + xmax) / 2 * w, 
                    (ymin + ymax) / 2 * h, 
                    (xmax - xmin) * w, 
                    (ymax - ymin) * h, 
                ] 
            ) 
            labels.append(int(label)) 
            scores.append(float(score)) 

    if len(boxes) == 0: 
        boxes = np.array([]).reshape(0, 4) 
        scores = np.array([]) 
        labels = np.array([]) 
    return np.array(boxes), np.array(scores), np.array(labels) 

def draw_boxes(img, bbox, identities=None): 
    """ 
    Draw bounding box in original image 

    Parameters 
    ---------- 
    img: original image 
    bbox: coordinate of bounding box 
    identities: identities IDs 
    """ 
    for i, box in enumerate(bbox): 
        x1, y1, x2, y2 = [int(i) for i in box] 
        # ボックステキストとバー 
        id = int(identities[i]) if identities is not None else 0 
        color = compute_color_for_labels(id) 
        label = "{}{:d}".format("", id) 
        t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 2, 2)[0] 
        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2) 
        cv2.rectangle(img, (x1, y1), (x1 + t_size[0] + 3, y1 + t_size[1] + 4), color, -1) 
        cv2.putText( 
            img, 
            label, 
            (x1, y1 + t_size[1] + 4), 
            cv2.FONT_HERSHEY_PLAIN, 
            1.6, 
            [255, 255, 255], 
            2, 
        ) 
    return img 

def cosin_metric(x1, x2): 
    """ 
    Calculate the consin distance of two vector 

    Parameters 
    ---------- 
    x1, x2: input vectors 
    """ 
    return np.dot(x1, x2) / (np.linalg.norm(x1) * np.linalg.norm(x2))

# 人物追跡を実行するメイン処理関数 
def run_person_tracking(source=0, flip=False, use_popup=False, skip_first_frames=0): 
    """ 
    Main function to run the person tracking: 
    1. Create a video player to play with target fps (utils.VideoPlayer). 
    2. Prepare a set of frames for person tracking. 
    3. Run AI inference for person tracking. 
    4. Visualize the results. 
    Parameters:
    ---------- 
        source: The webcam number to feed the video stream with primary webcam set to "0", or the video path. 
        flip: To be used by VideoPlayer function for flipping capture image. 
        use_popup: False for showing encoded frames over this notebook, True for creating a popup window. 
        skip_first_frames: Number of frames to skip at the beginning of the video.
    """ 
    player = None 
    try:
        # ターゲット fps で再生するビデオプレーヤーを作成 
        player = utils.VideoPlayer( 
            source=source, 
            size=(700, 450), 
            flip=flip, 
            fps=24, 
            skip_first_frames=skip_first_frames, 
        ) 
        # キャプチャーを開始 
        player.start() 
        if use_popup: 
            title = "Press ESC to Exit" 
            cv2.namedWindow(winname=title, flags=cv2.WINDOW_GUI_NORMAL | cv2.WINDOW_AUTOSIZE) 

        processing_times = collections.deque() 
        while True:            # フレームをグラブ 
            frame = player.next() 
            if frame is None: 
                print("Source ended") 
                break 
            # フレームがフル HD より大きい場合は、サイズを縮小してパフォーマンスを向上させます
            # ニューラル・ネットワークの入力に合わせて画像のサイズを変更し、暗さを変更 
            h, w = frame.shape[:2] 
            input_image = preprocess(frame, detector.height, detector.width) 

            # 処理時間を測定 
            start_time = time.time() 
            # Get the results. 
            output = detector.predict(input_image) 
            stop_time = time.time() 
            processing_times.append(stop_time - start_time) 
            if len(processing_times) > 200: 
                processing_times.popleft() 

            _, f_width = frame.shape[:2] 
            # 平均処理時間 [ms] 
            processing_time = np.mean(processing_times) * 1100 
            fps = 1000 / processing_time 

            # 検出結果からポーズを取得 
            bbox_xywh, score, label = process_results(h, w, results=output)
 
            img_crops = [] 
            for box in bbox_xywh: 
                x1, y1, x2, y2 = xywh_to_xyxy(box, h, w) 
                img = frame[y1:y2, x1:x2] 
                img_crops.append(img) 

            # 各人の再識別機能を取得 
            if img_crops:
                # 前処理 
                img_batch = batch_preprocess(img_crops, extractor.height, extractor.width) 
                features = extractor.predict(img_batch) 
            else: 
                features = np.array([]) 

            # 検出結果と再識別結果をまとめる 
            bbox_tlwh = xywh_to_tlwh(bbox_xywh) 
            detections = [Detection(bbox_tlwh[i], features[i]) for i in range(features.shape[0])] 

            # 追跡対象の位置を予測 
            tracker.predict() 

            # トラッカーを更新 
            tracker.update(detections) 

            # bbox ID を更新 
            outputs = [] 
            for track in tracker.tracks: 
                if not track.is_confirmed() or track.time_since_update > 1: 
                    continue 
                box = track.to_tlwh() 
                x1, y1, x2, y2 = tlwh_to_xyxy(box, h, w) 
                track_id = track.track_id 
                outputs.append(np.array([x1, y1, x2, y2, track_id], dtype=np.int32)) 
            if len(outputs) > 0: 
                outputs = np.stack(outputs, axis=0) 

            # 視覚化のための描画ボックス 
            if len(outputs) > 0: 
                bbox_tlwh = [] 
                bbox_xyxy = outputs[:, :4] 
                identities = outputs[:, -1] 
                frame = draw_boxes(frame, bbox_xyxy, identities) 

            cv2.putText( 
                img=frame, 
                text=f"Inference time: {processing_time:.1f}ms ({fps:.1f} FPS)", 
                org=(20, 40), 
                fontFace=cv2.FONT_HERSHEY_COMPLEX, 
                fontScale=f_width / 1000, 
                color=(0, 0, 255), 
                thickness=1, 
                lineType=cv2.LINE_AA, 
            ) 

            if use_popup: 
                cv2.imshow(winname=title, mat=frame) 
                key = cv2.waitKey(1) 
                # escape = 27 
                if key == 27: 
                    break 
            else: 
                # numpy 配列を jpg にエンコード
                _, encoded_img = cv2.imencode(ext=".jpg", img=frame, params=[cv2.IMWRITE_JPEG_QUALITY, 100]) 
                # IPython イメージを作成 
                i = display.Image(data=encoded_img) 
                # このノートブックに画像を表示 
                display.clear_output(wait=True) 
                display.display(i) 

# ctrl-c 
except KeyboardInterrupt: 
    print("Interrupted") 
# 異なるエラー 
except RuntimeError as e: 
    print(e) 
finally: 
    if player is not None: 
        # キャプチャーを停止 
        Player.stop().stop() 
    if use_popup: 
        cv2.destroyAllWindows()