Spaces:
Runtime error
Runtime error
File size: 1,715 Bytes
47cc4c7 |
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 |
# AUTOGENERATED! DO NOT EDIT! File to edit: ../app.ipynb.
# %% auto 0
__all__ = ['device', 'model', 'transforms', 'image', 'result', 'examples', 'intf', 'detect_stamps']
# %% ../app.ipynb 1
from model import YOLOStamp
from utils import *
import torch
import gradio as gr
import albumentations as A
from albumentations.pytorch.transforms import ToTensorV2
from PIL import Image
# %% ../app.ipynb 2
device = "cuda" if torch.cuda.is_available() else "cpu"
model = YOLOStamp()
model.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))
model = model.to(device)
model.eval()
# %% ../app.ipynb 3
transforms = A.Compose([
A.Resize(height=448, width=448),
A.Normalize(),
ToTensorV2(p=1.0),
])
# %% ../app.ipynb 7
def detect_stamps(image):
shape = image.size[:2]
image = image.convert('RGB')
image = np.array(image)
image = transforms(image=image)['image']
output = model(image.unsqueeze(0).to(device))[0]
boxes = output_tensor_to_boxes(output.detach().cpu())
boxes = nonmax_suppression(boxes)
img = image.permute(1, 2, 0).cpu().numpy()
img = visualize_bbox(img.copy(), boxes=boxes, draw_center=False)
img = cv2.resize(img, dsize=shape)
return Image.fromarray((255. * (img * np.array(STD) + np.array(MEAN))).astype(np.uint8))
# %% ../app.ipynb 9
image = gr.inputs.Image(type="pil")
result = gr.outputs.Image(type="pil")
examples = ['./examples/1.jpg', './examples/2.jpg', './examples/3.jpg']
intf = gr.Interface(fn=detect_stamps,
inputs=image,
outputs=result,
title='Stamp detection',
examples=examples)
intf.launch(inline=False)
|