使用 TensorFlow Hub 進行圖片分類

在這個 Colab 中,您將試用多個來自 TensorFlow Hub 的圖片分類模型,並決定哪個模型最適合您的使用案例。

由於 TF Hub 鼓勵操作圖片的模型採用一致的輸入慣例,因此您可以輕鬆試用不同的架構,找出最符合您需求的架構。

在 TensorFlow.org 上檢視 在 Google Colab 中執行 在 GitHub 上檢視 下載筆記本 查看 TF Hub 模型
import tensorflow as tf
import tensorflow_hub as hub

import requests
from PIL import Image
from io import BytesIO

import matplotlib.pyplot as plt
import numpy as np

用於載入圖片的輔助函式 (隱藏)

選取圖片分類模型。之後,系統會設定一些內部變數,並下載標籤檔案以供使用。

模型之間存在一些技術差異,例如不同的輸入大小、模型大小、準確度和推論時間。您可以在此處變更您使用的模型,直到找到最適合您使用案例的模型為止。

為了您的方便,系統會印出模型的控制代碼 (網址)。如需關於每個模型的詳細文件,請參閱該處。

選取圖片分類模型

Selected model: efficientnetv2-s : https://tfhub.dev/google/imagenet/efficientnet_v2_imagenet1k_s/classification/2
Images will be converted to 384x384

您可以選取下列其中一張圖片,或使用您自己的圖片。請記住,模型的輸入大小各不相同,其中有些模型使用動態輸入大小 (可在未縮放的圖片上進行推論)。因此,load_image 方法會將圖片重新縮放為預期格式。

選取輸入圖片

2024-03-09 13:46:31.989759: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:282] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected

png

現在已選取模型,使用 TensorFlow Hub 載入模型非常簡單。

這也會使用隨機輸入呼叫模型,作為「暖機」執行。後續呼叫通常會快得多,您可以將此與下方的延遲時間進行比較。

classifier = hub.load(model_handle)

input_shape = image.shape
warmup_input = tf.random.uniform(input_shape, 0, 1.0)
%time warmup_logits = classifier(warmup_input).numpy()
CPU times: user 2.12 s, sys: 188 ms, total: 2.31 s
Wall time: 1.36 s

一切都已準備就緒,可以進行推論。您可以在此處查看模型針對選取圖片輸出的前 5 個結果。

# Run model on image
%time probabilities = tf.nn.softmax(classifier(image)).numpy()

top_5 = tf.argsort(probabilities, axis=-1, direction="DESCENDING")[0][:5].numpy()
np_classes = np.array(classes)

# Some models include an additional 'background' class in the predictions, so
# we must account for this when reading the class labels.
includes_background_class = probabilities.shape[1] == 1001

for i, item in enumerate(top_5):
  class_index = item if includes_background_class else item + 1
  line = f'({i+1}) {class_index:4} - {classes[class_index]}: {probabilities[0][top_5][i]}'
  print(line)

show_image(image, '')
CPU times: user 911 ms, sys: 103 ms, total: 1.01 s
Wall time: 80.1 ms
(1)   35 - blowing glass: 0.7747849822044373
(2)   34 - blowing bubble gum: 0.1064402163028717
(3)   37 - blowing nose: 0.00587468734011054
(4)  148 - drinking shots: 0.0025945263914763927
(5)   36 - blowing leaves: 0.002559840213507414

png

瞭解詳情

如果您想瞭解詳情,並嘗試如何使用這些模型進行移轉學習,您可以試用本教學課程:用於圖片分類的移轉學習

如果您想查看更多圖片模型,可以查看 tfhub.dev 上的模型