Spaces:
Build error
Build error
File size: 4,036 Bytes
ade0106 bb029d9 ade0106 3a42afa ade0106 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 |
import gradio
import torch
import torchvision.transforms as T
import numpy as np
import matplotlib.pyplot as plt
import os
import random
import PIL.Image as Image
import time
from model import create_fasterrcnn_model
categories = [
{
"id": 0,
"name": "creatures",
"supercategory": "none"
},
{
"id": 1,
"name": "fish",
"supercategory": "creatures"
},
{
"id": 2,
"name": "jellyfish",
"supercategory": "creatures"
},
{
"id": 3,
"name": "penguin",
"supercategory": "creatures"
},
{
"id": 4,
"name": "puffin",
"supercategory": "creatures"
},
{
"id": 5,
"name": "shark",
"supercategory": "creatures"
},
{
"id": 6,
"name": "starfish",
"supercategory": "creatures"
},
{
"id": 7,
"name": "stingray",
"supercategory": "creatures"
}
]
# 1, Create title, description and article strings
title = "Ocean creatures detection Faster-R-CNN"
description = "A Faster-RCNN-ResNet-50 backbone feature extractor computer vision model to classify images of fish, penguin, shark, etc"
faster_rcnn = create_fasterrcnn_model(
num_classes=8, # len(class_names) would also work
)
# Load saved weights
faster_rcnn.load_state_dict(
torch.load(
f="./third_train.pth",
map_location=torch.device("cpu"), # load to CPU
)
)
import random
# Create predict function
def predict(img):
"""Transforms and performs a prediction on img and returns prediction and time taken.
"""
# Start the timer
start_time = time.time()
device = 'cpu'
transform = T.Compose([T.ToPILImage(),T.ToTensor()])
image_tensor = transform(img).to(device)
image_tensor = image_tensor.unsqueeze(0)
faster_rcnn.eval()
with torch.no_grad():
predictions = faster_rcnn(image_tensor)
pred_boxes = predictions[0]['boxes'].cpu().numpy()
pred_scores = predictions[0]['scores'].cpu().numpy()
pred_labels = predictions[0]['labels'].cpu().numpy()
label_names = [categories[label]['name'] for label in pred_labels]
fig, ax = plt.subplots(1)
ax.imshow(img)
for box, score, label_name in zip(pred_boxes, pred_scores, label_names):
if score > 0.5:
x1, y1, x2, y2 = box
w, h = x2 - x1, y2 - y1
rect = plt.Rectangle((x1, y1), w, h, fill=False, edgecolor='red', linewidth=2)
ax.add_patch(rect)
ax.text(x1, y1, f'{label_name}: {score:.2f}', fontsize=5, color='white', bbox=dict(facecolor='red', alpha=0.2))
# save the figure to an image file
random_name = str(random.randint(0,99))
img_path = f"./{random_name}.png"
fig.savefig(img_path)
# convert the figure to an image
fig.canvas.draw()
# Calculate the prediction time
pred_time = round(time.time() - start_time, 5)
# return the predicted label, the path to the saved image, and the prediction time
return img_path, str(pred_time)
### 4. Gradio app ###
# Get a list of all image file paths in the folder
example_list = [["examples/" + example] for example in os.listdir("examples")]
# Create the Gradio demo
demo = gradio.Interface(fn=predict, # mapping function from input to output
inputs=gradio.Image(type= "numpy"), # what are the inputs?
outputs=[gradio.outputs.Image(type= "filepath", label="Image with Bounding Boxes"),
gradio.outputs.Label(type="auto", label="Prediction Time")], # our fn has two outputs
# Create examples list from "examples/" directory
examples=example_list,
title=title,
description=description)
# Launch the demo!
demo.launch(debug =True) |