Руководство по классификации изображений для Python

Задача MediaPipe Image Classifier позволяет выполнять классификацию изображений. Эту задачу можно использовать, чтобы определить, что представляет собой изображение среди набора категорий, определенных во время обучения. Эти инструкции покажут вам, как использовать классификатор изображений с Python.

Вы можете увидеть эту задачу в действии, просмотрев веб-демо . Дополнительные сведения о возможностях, моделях и параметрах конфигурации этой задачи см. в разделе Обзор .

Пример кода

В примере кода классификатора изображений представлена ​​полная реализация этой задачи на Python. Этот код поможет вам протестировать эту задачу и приступить к созданию собственного классификатора изображений. Вы можете просматривать, запускать и редактировать пример кода классификатора изображений, используя только веб-браузер.

Если вы реализуете классификатор изображений для Raspberry Pi, обратитесь к примеру приложения Raspberry Pi .

Настраивать

В этом разделе описаны ключевые шаги по настройке среды разработки и проектов кода специально для использования Image Classifier. Общие сведения о настройке среды разработки для использования задач MediaPipe, включая требования к версии платформы, см. в руководстве по настройке Python .

Пакеты

Задача классификатора изображений — пакет pip mediapipe. Вы можете установить зависимость следующим образом:

$ python -m pip install mediapipe

Импорт

Импортируйте следующие классы, чтобы получить доступ к функциям задачи классификатора изображений:

import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision

Модель

Для задачи классификатора изображений MediaPipe требуется обученная модель, совместимая с этой задачей. Дополнительную информацию о доступных обученных моделях для Классификатора изображений см. в разделе «Модели » обзора задач.

Выберите и загрузите модель, а затем сохраните ее в локальном каталоге. Вы можете использовать рекомендованную модель EfficientNet-Lite0 .

model_path = '/absolute/path/to/efficientnet_lite0_int8_2.tflite'

Укажите путь к модели в параметре «Имя модели», как показано ниже:

base_options = BaseOptions(model_asset_path=model_path)

Создать задачу

Используйте функцию create_from_options для создания задачи. Функция create_from_options принимает параметры конфигурации, включая режим работы, локаль отображаемых имен, максимальное количество результатов, порог уверенности, список разрешенных категорий и список запрещенных. Дополнительные сведения о параметрах конфигурации см. в разделе Обзор конфигурации .

Задача «Классификатор изображений» поддерживает три типа входных данных: неподвижные изображения, видеофайлы и прямые видеопотоки. Выберите вкладку, соответствующую вашему типу входных данных, чтобы узнать, как создать задачу и выполнить вывод.

import mediapipe as mp

BaseOptions = mp.tasks.BaseOptions
ImageClassifier = mp.tasks.vision.ImageClassifier
ImageClassifierOptions = mp.tasks.vision.ImageClassifierOptions
VisionRunningMode = mp.tasks.vision.RunningMode

options = ImageClassifierOptions(
    base_options=BaseOptions(model_asset_path='/path/to/model.tflite'),
    max_results=5,
    running_mode=VisionRunningMode.IMAGE)

with ImageClassifier.create_from_options(options) as classifier:
  # The classifier is initialized. Use it here.
  # ...
    

import mediapipe as mp

BaseOptions = mp.tasks.BaseOptions
ImageClassifier = mp.tasks.vision.ImageClassifier
ImageClassifierOptions = mp.tasks.vision.ImageClassifierOptions
VisionRunningMode = mp.tasks.vision.RunningMode

options = ImageClassifierOptions(
    base_options=BaseOptions(model_asset_path='/path/to/model.tflite'),
    max_results=5,
    running_mode=VisionRunningMode.VIDEO)

with ImageClassifier.create_from_options(options) as classifier:
  # The classifier is initialized. Use it here.
  # ...
    

import mediapipe as mp

BaseOptions = mp.tasks.BaseOptions
ImageClassifierResult = mp.tasks.vision.ImageClassifier.ImageClassifierResult
ImageClassifier = mp.tasks.vision.ImageClassifier
ImageClassifierOptions = mp.tasks.vision.ImageClassifierOptions
VisionRunningMode = mp.tasks.vision.RunningMode

def print_result(result: ImageClassifierResult, output_image: mp.Image, timestamp_ms: int):
    print('ImageClassifierResult result: {}'.format(result))

options = ImageClassifierOptions(
    base_options=BaseOptions(model_asset_path='/path/to/model.tflite'),
    running_mode=VisionRunningMode.LIVE_STREAM,
    max_results=5,
    result_callback=print_result)

with ImageClassifier.create_from_options(options) as classifier:
  # The classifier is initialized. Use it here.
  # ...
    

Полный пример создания классификатора изображений для использования с изображением см. в примере кода .

Варианты конфигурации

Эта задача имеет следующие параметры конфигурации для приложений Python:

Подготовьте данные

Подготовьте входные данные в виде файла изображения или массива numpy, а затем преобразуйте их в объект mediapipe.Image . Если ваши входные данные представляют собой видеофайл или прямую трансляцию с веб-камеры, вы можете использовать внешнюю библиотеку, такую ​​как OpenCV, для загрузки входных кадров в виде числовых массивов.

Следующие примеры объясняют и показывают, как подготовить данные к обработке для каждого из доступных типов данных.

import mediapipe as mp

# Load the input image from an image file.
mp_image = mp.Image.create_from_file('/path/to/image')

# Load the input image from a numpy array.
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=numpy_image)
    

import mediapipe as mp

# Use OpenCV’s VideoCapture to load the input video.

# Load the frame rate of the video using OpenCV’s CV_CAP_PROP_FPS
# You’ll need it to calculate the timestamp for each frame.

# Loop through each frame in the video using VideoCapture#read()

# Convert the frame received from OpenCV to a MediaPipe’s Image object.
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=numpy_frame_from_opencv)
    

import mediapipe as mp

# Use OpenCV’s VideoCapture to start capturing from the webcam.

# Create a loop to read the latest frame from the camera using VideoCapture#read()

# Convert the frame received from OpenCV to a MediaPipe’s Image object.
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=numpy_frame_from_opencv)
    

Запустить задачу

Вы можете вызвать функцию классификации, соответствующую вашему режиму работы, чтобы инициировать выводы. API классификатора изображений вернет возможные категории объекта во входном изображении или фрейме.

# Perform image classification on the provided single image.
classification_result = classifier.classify(mp_image)
    

# Calculate the timestamp of the current frame
frame_timestamp_ms = 1000 * frame_index / video_file_fps

# Perform image classification on the video frame.
classification_result = classifier.classify_for_video(mp_image, frame_timestamp_ms)
    

# Send the latest frame to perform image classification.
# Results are sent to the `result_callback` provided in the `ImageClassifierOptions`.
classifier.classify_async(mp_image, frame_timestamp_ms)
    

Обратите внимание на следующее:

• При работе в режиме видео или режиме прямой трансляции необходимо также предоставить задаче «Классификатор изображений» временную метку входного кадра.
• При запуске изображения или видеомодели задача «Классификатор изображений» блокирует текущий поток до тех пор, пока он не завершит обработку входного изображения или кадра.
• При работе в режиме прямой трансляции задача классификатора изображений не блокирует текущий поток, а немедленно возвращается. Он будет вызывать свой прослушиватель результатов с результатом классификации каждый раз, когда завершает обработку входного кадра. Если функция classifyAsync вызывается, когда задача классификатора изображений занята обработкой другого кадра, задача игнорирует новый входной кадр.

Полный пример создания классификатора изображений для использования с изображением см. в примере кода .

Обработка и отображение результатов

После выполнения вывода задача «Классификатор изображений» возвращает объект ImageClassifierResult , который содержит список возможных категорий для объектов во входном изображении или кадре.

Ниже показан пример выходных данных этой задачи:

ImageClassifierResult:
 Classifications #0 (single classification head):
  head index: 0
  category #0:
   category name: "/m/01bwb9"
   display name: "Passer domesticus"
   score: 0.91406
   index: 671
  category #1:
   category name: "/m/01bwbt"
   display name: "Passer montanus"
   score: 0.00391
   index: 670

Этот результат был получен путем запуска Классификатора птиц на:

Пример кода классификатора изображений демонстрирует, как отображать результаты классификации, возвращенные из задачи. Подробности см. в примере кода .