### 1. Import and class names setup ### import gradio as gr import os from pathlib import Path import torch from model import create_effnetb2_model from time import perf_counter from typing import Tuple, Dict from PIL import Image import torchvision # Setup class names (hardcoded, these shall reside in a json file or sth like that...) class_names = ["pizza","steak","sushi"] ### 2. Model and transforms preparation ### effnetb2_model, effnetb2_transforms = create_effnetb2_model(num_classes=len(class_names)) # Load save weights effnetb2_model.load_state_dict(torch.load(f="09_pretrained_effnetb2_feature_extractor_pizza_steak_sushi_20percent.pth", map_location=torch.device("cpu"))) # map location to cpu is a must, as we have trained our model in the GPU ### 3. Predict function def predict(img) -> Tuple[Dict,float]: # Start a timer start_time = perf_counter() # Transform the input image for use with EffNetB2 effnetb2_transforms = torchvision.models.EfficientNet_B2_Weights.DEFAULT.transforms() img_tensor = effnetb2_transforms(img) # Put model in eval and inference effnetb2_model.eval() with torch.inference_mode(): y_logits = effnetb2_model(img_tensor.unsqueeze(dim=0)) y_pred_probs = torch.softmax(y_logits,dim=1) y_pred_probs_list = y_pred_probs.squeeze().tolist() # Creatae a prediction probability dictionary pred_prob_dict = {class_names[i]:float(prob) for i,prob in enumerate(y_pred_probs_list)} # End timer end_time = perf_counter() return pred_prob_dict, round(end_time-start_time,4) ### 4. Launch app import gradio as gr foodvision_mini_examples_path = "examples" example_list = [str(path) for path in Path(foodvision_mini_examples_path).rglob("*.jpg")] # Create title, description and article title = "FoodVisionMini V0 🍕🍖🍣" + os.environ["test_ando_secret"] description = "An EfficientNetB2 feature extractor computer vision model to classify images into pizza, steak or sushi
I have yet to improve it to label non-food images. Paciencia amigos" article = "Created at 09_pytorch_model_deploy.ipynb Google Colab notebook" # Create the Gradio demo demo = gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs=[gr.Label(num_top_classes=3, label="predictions"), gr.Number(label="Prediction time (s)")], examples=example_list, title=title, description=description, article=article) # Launch the demo demo.launch(share=True) # run on public url # *** IMPORTANTE: The Flag button of the interface will create a folder named "flagged" that will contain the images and predictions of those images that someone has Flagged***