Spaces:
Sleeping
Sleeping
| ### 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 foodvision_mini_examples_path.rglob("*.jpg")] | |
| # Create title, description and article | |
| title = "FoodVisionMini V0 🍕🍖🍣" | |
| description = "An <a href='https://pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b2.html#torchvision.models.efficientnet_b2'>EfficientNetB2</a> feature extractor computer vision model to classify images into pizza, steak or sushi<br>I have yet to improve it to label non-food images. Paciencia amigos" | |
| article = "Created at <a href='#'>09_pytorch_model_deploy.ipynb</a> 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(debug=False,# print errors locally in local colab debug=False | |
| 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*** | |