Spaces:
Runtime error
Runtime error
# 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) | |