|
import torch |
|
import albumentations as A |
|
import cv2 |
|
|
|
from albumentations.pytorch import ToTensorV2 |
|
import numpy as np |
|
import config |
|
from PIL import Image |
|
from pytorch_grad_cam import GradCAM |
|
from pytorch_grad_cam.utils.image import show_cam_on_image |
|
from model import YOLOv3 |
|
import gradio as gr |
|
import os |
|
import matplotlib.pyplot as plt |
|
from utils import * |
|
|
|
|
|
|
|
|
|
model = YOLOv3(num_classes=config.NUM_CLASSES).to(config.DEVICE) |
|
model.load_state_dict(torch.load("custom_yolo_v3.pt", map_location=torch.device('cpu')), strict=False) |
|
|
|
IMAGE_SIZE = config.IMAGE_SIZE |
|
test_transforms = A.Compose( |
|
[ |
|
A.LongestMaxSize(max_size=IMAGE_SIZE), |
|
A.PadIfNeeded( |
|
min_height=IMAGE_SIZE, min_width=IMAGE_SIZE, border_mode=cv2.BORDER_CONSTANT |
|
), |
|
A.Normalize(mean=[0, 0, 0], std=[1, 1, 1], max_pixel_value=255,), |
|
ToTensorV2(), |
|
]) |
|
|
|
anchors = ( |
|
torch.tensor(config.ANCHORS) |
|
* torch.tensor(config.S).unsqueeze(1).unsqueeze(1).repeat(1,3,2) |
|
).to(config.DEVICE) |
|
|
|
def inference(input_img, transparency = 0.5, thresh=0.8, iou_thresh=0.3): |
|
|
|
x = test_transforms(image=input_img)['image'].unsqueeze(0).to(config.DEVICE) |
|
with torch.no_grad(): |
|
out = model(x) |
|
bboxes = [[] for _ in range(x.shape[0])] |
|
for i in range(3): |
|
batch_size, A, S, _, _ = out[i].shape |
|
anchor = anchors[i] |
|
boxes_scale_i = cells_to_bboxes( |
|
out[i], anchor, S=S, is_preds=True |
|
) |
|
for idx, (box) in enumerate(boxes_scale_i): |
|
bboxes[idx] += box |
|
|
|
model.train() |
|
|
|
nms_boxes = non_max_suppression( |
|
bboxes[0], iou_threshold=iou_thresh, threshold=thresh, box_format="midpoint", |
|
) |
|
visualization = plot_image(x[0].permute(1,2,0).detach().cpu(), nms_boxes) |
|
|
|
return visualization |
|
|
|
title = "Object Detection using YOLOv3" |
|
description = "A simple Gradio interface to show object detection on images" |
|
examples = [['./images/006294.jpg'], |
|
['./images/005898.jpg'], |
|
['./images/003785.jpg'], |
|
['./images/001624.jpg'], |
|
['./images/006796.jpg'], |
|
['./images/003388.jpg'], |
|
['./images/002216.jpg'], |
|
['./images/000341.jpg'], |
|
['./images/006818.jpg'], |
|
] |
|
|
|
|
|
demo = gr.Interface( |
|
inference, |
|
inputs = [gr.Image(label="Input Image", type='numpy'), |
|
gr.Slider(0, 1, value = 0.75, label="Threshold"), |
|
gr.Slider(0, 1, value = 0.75, label="IoU Threshold"), |
|
gr.Slider(0, 1, value = 0.8, label="Opacity of GradCAM")], |
|
outputs = [gr.Image(width=600, height=600,label="Output")], |
|
title = title, |
|
description = description, |
|
examples = examples, |
|
) |
|
demo.launch() |
