osmr / imgclsmob

Sandbox for training deep learning networks
MIT License
2.95k stars 562 forks source link

incomplete example #50

Closed thunderbug1 closed 4 years ago

thunderbug1 commented 4 years ago

hi, thank you for the great repository. The problem that I came to is that I cannot find a list of classes for which the tf2 models of imagenet1k were trained for. Also the website of imagenet is not very helpfull in that regard.

For example, how could I load an image, classify it and print the class label? An example for how to do this would be very helpful to add.

osmr commented 4 years ago

See for example:

thunderbug1 commented 4 years ago

Thank you very much! It helped me a lot.

I made a small minimal example to classify an image, maybe we should add it to the examples:

from tf2cv.model_provider import get_model as tf2cv_get_model
import tensorflow as tf
import numpy as np

# load the model
net = tf2cv_get_model("resnet18", pretrained=True, data_format="channels_last")

# load image
import cv2
raw_image = cv2.imread("test_image.jpg")
img = cv2.resize(raw_image, dsize=(224, 224), interpolation=cv2.INTER_CUBIC)
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) # change color channel ordering

# show the image
import matplotlib.pyplot as plt
plt.imshow(np.squeeze(img))
plt.axis('off')

# loading could also be done with PIL where it is easier to display the image but it is a little bit slower
#from keras.preprocessing.image import load_img, img_to_array
#from PIL import Image
#raw_img = load_img("test_image.jpg") # load the image to classify
#img = raw_img.resize((224, 224), Image.BICUBIC) #  resize the image to the neuronal network input size
#img = img_to_array(img) # convert the image to an array
#img.show()

# prepare image for neural network
img = np.float32(img/255.0) # convert from uint8 to float32 and change value range
img = np.expand_dims(img, axis=0) # add a batch dimension for the network

y_net = net(img)  # make the prediction
pred_class = np.argmax(y_net) # find most probable class

# convert the predicted class index back to a name
from gluoncv.data import ImageNet1kAttr
attr = ImageNet1kAttr()
pred_name = attr.classes[pred_class]
print(pred_name)

# optionally: get the top 5 predictions
for pred_class in np.array(y_net).argsort()[0][-5:]:
    pred_name = attr.classes[pred_class]
    print(pred_name)
osmr commented 4 years ago

Unfortunately, the direct use of classification models trained on ImageNet-1K looks a bit more complicated. I will try to give a more correct script a little later.

thunderbug1 commented 4 years ago

oh yeah, you are probably right. I tried the code for some example images of the class "agama" and it worked fine, but for other classes it doesn't work at all.

osmr commented 4 years ago

Pay attention to an example https://github.com/osmr/imgclsmob/blob/master/examples/demo_tf2.py

thunderbug1 commented 4 years ago

Thank you very much! Nice demo ;)