app.py

# 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)