|
|
|
from model import create_model |
|
import pandas as pd |
|
import torch |
|
from typing import Tuple, Dict |
|
from timeit import default_timer as timer |
|
import gradio as gr |
|
import os |
|
import numpy as np |
|
|
|
|
|
labels_csv = pd.read_csv('./labels.csv') |
|
labels = labels_csv['breed'] |
|
labels = np.array(labels) |
|
unique_labels = np.unique(labels) |
|
|
|
unique_labels = [' '.join([word.capitalize() for word in label.split('_')]) for label in unique_labels] |
|
|
|
|
|
model, model_transforms = create_model(num_classes=len(unique_labels)) |
|
model = torch.compile(model) |
|
|
|
|
|
|
|
model.load_state_dict(torch.load(f='./convnext_model.pth', map_location='cpu',weights_only=True)) |
|
|
|
|
|
|
|
def predict(img) -> Tuple[Dict[str, float], str]: |
|
""" |
|
Predicts the class probabilities for a given image using a pre-trained model. |
|
|
|
Args: |
|
img: A PIL image to be predicted. |
|
|
|
Returns: |
|
A tuple containing: |
|
- A formatted string displaying class labels and their respective probabilities. |
|
- The time taken for inference in seconds as a string. |
|
""" |
|
|
|
start_time = timer() |
|
|
|
|
|
model.eval() |
|
with torch.inference_mode(): |
|
|
|
|
|
img = model_transforms(img).unsqueeze(dim=0) |
|
|
|
|
|
pred_logit = model(img) |
|
|
|
|
|
pred_prob = torch.softmax(pred_logit, dim=1) |
|
|
|
pred_label = torch.argmax(pred_prob, dim=1) |
|
|
|
|
|
prediction = unique_labels[pred_label] |
|
probabilities = {unique_labels[i]: pred_prob[0, i].item() for i in range(len(unique_labels))} |
|
|
|
|
|
end_time = timer() |
|
inference_time = end_time - start_time |
|
|
|
|
|
return probabilities, f"{inference_time:.4f} seconds" |
|
|
|
|
|
|
|
|
|
title = "Dogvision πΆ" |
|
description = "A [ConvNeXt Tiny](https://pytorch.org/vision/stable/models/generated/torchvision.models.convnext_tiny.html#torchvision.models.convnext_tiny) Computer Vision Model To Classify 120 Dog Breeds π© Ranging fro A Labrador π to A German Shepherd! πβπ¦Ί" |
|
article = "Created with π€ (and a mixture of mathematics, statistics, and tons of calculations π©π½βπ¬) by Arpit Vaghela [GitHub](https://github.com/magnifiques)" |
|
|
|
|
|
example_list = [["./examples/" + example] for example in os.listdir("examples")] |
|
|
|
demo = gr.Interface(fn=predict, |
|
inputs=gr.Image(type='pil'), |
|
outputs=[ |
|
gr.Label(num_top_classes=3, label="Top Predictions"), |
|
gr.Textbox(label="Prediction Time (s)") |
|
], |
|
examples=example_list, |
|
title=title, |
|
description=description, |
|
article=article) |
|
|
|
|
|
demo.launch(debug=False, |
|
share=True) |
|
|