Spaces:
Runtime error
Runtime error
File size: 2,312 Bytes
e80e8f4 caa33e3 a741f39 caa33e3 e80e8f4 ef44018 e80e8f4 |
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 |
# AUTOGENERATED! DO NOT EDIT! File to edit: ../app.ipynb.
# %% auto 0
__all__ = ['device', 'model', 'MEAN', 'STD', 'transform', 'image', 'label', 'examples', 'intf', 'to_img', 'draw_image_with_bbox',
'localize_dog']
# %% ../app.ipynb 2
from models import Model
import torch
from torchvision import transforms
import gradio as gr
import numpy as np
from PIL import Image, ImageDraw
import cv2
# %% ../app.ipynb 3
device = "cuda" if torch.cuda.is_available() else "cpu"
model = Model()
model.load_state_dict(torch.load('model.pt', map_location=torch.device('cpu')))
model = model.to(device)
model.eval()
# %% ../app.ipynb 4
MEAN = [0.485, 0.456, 0.406]
STD = [0.229, 0.224, 0.225]
# %% ../app.ipynb 5
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(MEAN, STD),
])
# %% ../app.ipynb 6
def to_img(inp):
mean = np.array(MEAN)
std = np.array(STD)
inp = std * inp + mean
inp = np.clip(inp, 0, 1)
return inp * 255
def draw_image_with_bbox(im, shape, pred_bbox=None, pred_obj=1):
im = im.numpy().transpose((1, 2, 0))
im = cv2.resize(im, dsize=shape)
image_with_bbox = Image.fromarray(to_img(im).astype(np.uint8))
image_draw = ImageDraw.Draw(image_with_bbox)
xc, yc, w, h = pred_bbox
xmin = (xc - w / 2) * shape[0]
ymin = (yc - h / 2) * shape[1]
w = w * shape[0]
h = h * shape[1]
xmin, ymin, w, h = map(int, [xmin, ymin, w, h])
if pred_obj > 0.5:
image_draw.rectangle((max(xmin, 1), max(ymin, 1), min(xmin+w, shape[0] - 1), min(ymin+h, shape[1] - 1)), outline='red')
return image_with_bbox
# %% ../app.ipynb 7
def localize_dog(im):
shape = im.size[:2]
im = im.convert('RGB')
im = transform(im)
pred_label, pred_bbox = model(im.unsqueeze(0).to(device))
prediction = draw_image_with_bbox(im, shape, pred_bbox[0], pred_label[0])
return prediction
# %% ../app.ipynb 9
image = gr.inputs.Image(type="pil")
label = gr.outputs.Image(type="pil")
examples = ['1.jpg', '2.jpg', '3.jpg']
intf = gr.Interface(fn=localize_dog,
inputs=image,
outputs=label,
title='Dog localization',
examples=examples)
intf.launch()
|