Spaces:
Sleeping
Sleeping
| # -*- coding: utf-8 -*- | |
| """ | |
| Spyder Editor | |
| Use pretrained model to recognize an image. | |
| """ | |
| import os | |
| import numpy as np | |
| import pandas as pd | |
| import PIL | |
| import tensorflow as tf | |
| from tensorflow.keras.applications import EfficientNetV2S | |
| from tensorflow.keras.applications import Xception | |
| from tensorflow.keras.applications import MobileNetV3Small | |
| """ | |
| model = ResNet50(weights='imagenet') | |
| img_path = 'elephant.jpg' | |
| img = keras.utils.load_img(img_path, target_size=(224, 224)) | |
| x = keras.utils.img_to_array(img) | |
| x = np.expand_dims(x, axis=0) | |
| x = preprocess_input(x) | |
| preds = model.predict(x) | |
| # decode the results into a list of tuples (class, description, probability) | |
| # (one such list for each sample in the batch) | |
| print('Predicted:', decode_predictions(preds, top=3)[0]) | |
| """ | |
| # Take a pretrained model, its preprocess input functions, prediction function, | |
| # the image to be predicted. | |
| # Return the decoded prediction. | |
| def pretrained_model(image, model_name, preprocess_fn, | |
| decoded_predictions_fn): | |
| # preprocess with the model's function | |
| x = preprocess_fn(image) | |
| # predictions | |
| pred = model_name.predict(x) | |
| # decode the results into a list of tuples (class, description, probability) | |
| decoded_pred = decoded_predictions_fn(pred, top=3) | |
| # Must put in the [0] first element of the | |
| # decoded predictions. | |
| pred_df = pd.DataFrame( | |
| decoded_pred[0], | |
| columns=["class", "description", "probability"] | |
| ) | |
| return pred_df | |
| # Load the specific image size for the model | |
| # path = filepath of the image | |
| # size = size to be loaded | |
| def load_and_resize_image(img_path, size=(224, 224, 3)): | |
| img = tf.keras.utils.load_img( | |
| path=img_path, target_size=size) | |
| x = tf.keras.utils.img_to_array(img) | |
| x = np.expand_dims(x, axis=0) | |
| print(f"Shape of x: {x.shape}") | |
| return x | |
| def predict(img_path): | |
| # Set the data directory | |
| data_dir = "images" | |
| image_list = ["cat.jpg", "mrt_train.jpg", | |
| "apples.jpg", "duck.jpg", "daisy.jpg"] | |
| # Xception: (299x299) | |
| # EfficientNetV2S: (384, 384) | |
| # MobileNetV2: (224, 224, 3) | |
| # Load the image | |
| # img_path = os.path.join(data_dir, image_list[4]) | |
| print(f"image path: {img_path}") | |
| # Efficient Net V2 Small | |
| img = load_and_resize_image( | |
| img_path, | |
| size=(384, 384)) | |
| eff_net = EfficientNetV2S(include_top=True, | |
| weights='imagenet', | |
| input_shape=(384, 384, 3)) | |
| pred_df_eff = pretrained_model( | |
| img, | |
| eff_net, | |
| tf.keras.applications.efficientnet_v2.preprocess_input, | |
| tf.keras.applications.efficientnet_v2.decode_predictions, | |
| ) | |
| print(pred_df_eff[["description", "probability"]]) | |
| # Mobile Net V3 Small | |
| img = load_and_resize_image(img_path, size=(224, 224)) | |
| mobile_net = MobileNetV3Small(include_top=True, | |
| weights="imagenet", | |
| input_shape=(224, 224, 3)) | |
| pred_df_mob = pretrained_model( | |
| img, | |
| eff_net, | |
| tf.keras.applications.mobilenet_v3.preprocess_input, | |
| tf.keras.applications.mobilenet_v3.decode_predictions, | |
| ) | |
| print(pred_df_mob[["description", "probability"]]) | |
| # Xception | |
| img = load_and_resize_image(img_path, size=(299, 299)) | |
| xcept = Xception(include_top=True, | |
| weights="imagenet", | |
| input_shape=(299, 299, 3)) | |
| pred_df_xcept = pretrained_model( | |
| img, | |
| xcept, | |
| tf.keras.applications.xception.preprocess_input, | |
| tf.keras.applications.xception.decode_predictions, | |
| ) | |
| print(pred_df_xcept[["description", "probability"]]) | |
| # Format the output for gradio label | |
| dict_eff = dict({ | |
| "Description": pred_df_eff["description"], | |
| "Probability": pred_df_eff["probability"] | |
| }) | |
| dict_mob = dict({ | |
| "Description": pred_df_mob["description"], | |
| "Probability": pred_df_mob["probability"] | |
| }) | |
| dict_xcept = dict({ | |
| "Description": pred_df_xcept["description"], | |
| "Probability": pred_df_xcept["probability"] | |
| }) | |
| return pred_df_eff[["description", "probability"]], \ | |
| pred_df_mob[["description", "probability"]], \ | |
| pred_df_xcept[["description", "probability"]] | |
| # Main | |
| if __name__ == "__main__": | |
| predict() | |