Kurz: Spuštění modelu TensorFlow v Pythonu

Článek
09/01/2024

V tomto kurzu se dozvíte, jak použít exportovaný model TensorFlow místně ke klasifikaci obrázků.

Poznámka:

Tento kurz se vztahuje pouze na modely exportované z projektů klasifikace obrázků Obecné (kompaktní). Pokud jste exportovali jiné modely, navštivte prosím naše úložiště vzorového kódu.

Požadavky

Nainstalujte Python 2.7 nebo Python 3.6 nebo novější.
Nainstalovat pip.

Dále budete muset nainstalovat následující balíčky:

pip install tensorflow
pip install pillow
pip install numpy
pip install opencv-python

Načtení modelu a značek

Stažený soubor .zip z kroku exportu obsahuje soubor model.pb a soubor labels.txt. Tyto soubory představují natrénovaný model a popisky klasifikace. Prvním krokem je načíst model do projektu. Do nového skriptu Pythonu přidejte následující kód.

import tensorflow as tf
import os

graph_def = tf.compat.v1.GraphDef()
labels = []

# These are set to the default names from exported models, update as needed.
filename = "model.pb"
labels_filename = "labels.txt"

# Import the TF graph
with tf.io.gfile.GFile(filename, 'rb') as f:
    graph_def.ParseFromString(f.read())
    tf.import_graph_def(graph_def, name='')

# Create a list of labels.
with open(labels_filename, 'rt') as lf:
    for l in lf:
        labels.append(l.strip())

Příprava obrázku pro předpověď

Při přípravě obrázku na predikci je potřeba provést několik kroků. Tyto kroky napodobují manipulaci s obrázky prováděnou během trénování.

Otevření souboru a vytvoření obrázku v barevném prostoru BGR

from PIL import Image
import numpy as np
import cv2

# Load from a file
imageFile = "<path to your image file>"
image = Image.open(imageFile)

# Update orientation based on EXIF tags, if the file has orientation info.
image = update_orientation(image)

# Convert to OpenCV format
image = convert_to_opencv(image)

Pokud má obrázek rozměr větší než 1600 pixelů, zavolejte tuto metodu (definovanou později).
```
image = resize_down_to_1600_max_dim(image)
```

Oříznutí největšího středového čtverce

h, w = image.shape[:2]
min_dim = min(w,h)
max_square_image = crop_center(image, min_dim, min_dim)

Změna velikosti čtverce dolů na 256x256

augmented_image = resize_to_256_square(max_square_image)

Oříznutí středu na konkrétní velikost vstupu pro model

# Get the input size of the model
with tf.compat.v1.Session() as sess:
    input_tensor_shape = sess.graph.get_tensor_by_name('Placeholder:0').shape.as_list()
network_input_size = input_tensor_shape[1]

# Crop the center for the specified network_input_Size
augmented_image = crop_center(augmented_image, network_input_size, network_input_size)

Definování pomocných funkcí Ve výše uvedených krocích se používají následující pomocné funkce:

def convert_to_opencv(image):
    # RGB -> BGR conversion is performed as well.
    image = image.convert('RGB')
    r,g,b = np.array(image).T
    opencv_image = np.array([b,g,r]).transpose()
    return opencv_image

def crop_center(img,cropx,cropy):
    h, w = img.shape[:2]
    startx = w//2-(cropx//2)
    starty = h//2-(cropy//2)
    return img[starty:starty+cropy, startx:startx+cropx]

def resize_down_to_1600_max_dim(image):
    h, w = image.shape[:2]
    if (h < 1600 and w < 1600):
        return image

    new_size = (1600 * w // h, 1600) if (h > w) else (1600, 1600 * h // w)
    return cv2.resize(image, new_size, interpolation = cv2.INTER_LINEAR)

def resize_to_256_square(image):
    h, w = image.shape[:2]
    return cv2.resize(image, (256, 256), interpolation = cv2.INTER_LINEAR)

def update_orientation(image):
    exif_orientation_tag = 0x0112
    if hasattr(image, '_getexif'):
        exif = image._getexif()
        if (exif != None and exif_orientation_tag in exif):
            orientation = exif.get(exif_orientation_tag, 1)
            # orientation is 1 based, shift to zero based and flip/transpose based on 0-based values
            orientation -= 1
            if orientation >= 4:
                image = image.transpose(Image.TRANSPOSE)
            if orientation == 2 or orientation == 3 or orientation == 6 or orientation == 7:
                image = image.transpose(Image.FLIP_TOP_BOTTOM)
            if orientation == 1 or orientation == 2 or orientation == 5 or orientation == 6:
                image = image.transpose(Image.FLIP_LEFT_RIGHT)
    return image

Klasifikace obrázku

Jakmile je obrázek připravený jako tensor, můžeme ho odeslat prostřednictvím modelu pro predikci.

# These names are part of the model and cannot be changed.
output_layer = 'loss:0'
input_node = 'Placeholder:0'

with tf.compat.v1.Session() as sess:
    try:
        prob_tensor = sess.graph.get_tensor_by_name(output_layer)
        predictions = sess.run(prob_tensor, {input_node: [augmented_image] })
    except KeyError:
        print ("Couldn't find classification output layer: " + output_layer + ".")
        print ("Verify this a model exported from an Object Detection project.")
        exit(-1)

Zobrazení výsledků

Výsledky zpracování tensoru obrázku modelem pak bude potřeba namapovat zpět na popisky.

    # Print the highest probability label
    highest_probability_index = np.argmax(predictions)
    print('Classified as: ' + labels[highest_probability_index])
    print()

    # Or you can print out all of the results mapping labels to probabilities.
    label_index = 0
    for p in predictions:
        truncated_probablity = np.float64(np.round(p,8))
        print (labels[label_index], truncated_probablity)
        label_index += 1

Další kroky

Dále se dozvíte, jak model zabalit do mobilní aplikace:

Sdílet prostřednictvím