vvd2003 commited on
Commit
cd17245
β€’
1 Parent(s): 214b796

Delete app.py

Browse files
Files changed (1) hide show
  1. app.py +0 -75
app.py DELETED
@@ -1,75 +0,0 @@
1
- ### 1. Imports and class names setup ###
2
- import gradio as gr
3
- import os
4
- import torch
5
-
6
- from class_names import class_names
7
- from model import Load_model
8
- from timeit import default_timer as timer
9
- from typing import Tuple, Dict
10
-
11
-
12
- ### 1. Model and transforms preparation ###
13
-
14
- # Create model and transform
15
- model, transforms = Load_model()
16
-
17
- # Load saved weights
18
- def load_checkpoint(checkpoint_file, model, device='cpu'):
19
- print("=> Loading checkpoint")
20
- checkpoint = torch.load(checkpoint_file, map_location=device)
21
- model.load_state_dict(checkpoint["state_dict"])
22
- load_checkpoint('model_checkpoint.pt', model)
23
-
24
- ### 2. Predict function ###
25
-
26
- # Create predict function
27
- def predict(img) -> Tuple[Dict, float]:
28
- """Transforms and performs a prediction on img and returns prediction and time taken.
29
- """
30
- # Start the timer
31
- start_time = timer()
32
-
33
- # Transform the target image and add a batch dimension
34
- img = transforms(img).unsqueeze(0)
35
-
36
- # Put model into evaluation mode and turn on inference mode
37
- model.eval()
38
- with torch.inference_mode():
39
- # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
40
- pred_probs = torch.softmax(model(img), dim=1)
41
-
42
- # Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
43
- pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
44
-
45
- # Calculate the prediction time
46
- pred_time = round(timer() - start_time, 5)
47
-
48
- # Return the prediction dictionary and prediction time
49
- return pred_labels_and_probs, pred_time
50
-
51
-
52
- ### 3. Gradio app ###
53
-
54
- # Create title, description and article strings
55
- title = "BirdVision 500 πŸ¦…πŸ¦†πŸ¦πŸ•ŠπŸ¦€πŸ¦’πŸ¦œ"
56
- description = "A model based on YoLov8 classification 500 birds."
57
- article = "Created on [GITHUB](https://github.com/vvduc1803/Yolov8_cls)."
58
-
59
- # Create examples list from "examples/" directory
60
- example_list = [["examples/" + example] for example in os.listdir("examples")]
61
-
62
- # Create the Gradio demo
63
- demo = gr.Interface(fn=predict, # mapping function from input to output
64
- inputs=gr.Image(type="pil"), # what are the inputs?
65
- outputs=[gr.Label(num_top_classes=10, label="Predictions"), # what are the outputs?
66
- gr.Number(label="Prediction time (s)")],
67
- # our fn has two outputs, therefore we have two outputs
68
- # Create examples list from "examples/" directory
69
- examples=example_list,
70
- title=title,
71
- description=description,
72
- article=article)
73
-
74
- # Launch the demo!
75
- demo.launch()