|
|
|
""" |
|
Script for Gradio Application for Yolov3 model |
|
""" |
|
import numpy as np |
|
import gradio as gr |
|
import torch |
|
from yolo_model import YOLOv3 |
|
|
|
import config |
|
from inference import predict, YoloCAM |
|
|
|
|
|
model_path = 'epoch=39-step=41400.ckpt' |
|
model = YOLOv3.load_from_checkpoint(model_path, map_location=torch.device('cpu'), strict=False, num_classes=config.NUM_CLASSES) |
|
model.eval() |
|
|
|
cam = YoloCAM(model=model, target_layers=[model.layers[-2]], use_cuda=False) |
|
|
|
|
|
def inference( |
|
org_img: np.ndarray, |
|
iou_thresh: float, thresh: float, |
|
show_cam: str, |
|
transparency: float, |
|
): |
|
outputs = predict(cam, model, org_img, iou_thresh, thresh, show_cam, transparency) |
|
return outputs |
|
|
|
|
|
if __name__ == '__main__': |
|
title = "YoloV3 from Scratch on Pascal VOC Dataset with GradCAM" |
|
description = f"Pytorch Implemetation of YoloV3 trained from scratch on Pascal VOC dataset with GradCAM \n Classes in pascol voc: {', '.join(config.PASCAL_CLASSES)}" |
|
examples = [ |
|
["examples/dog.jpg", 0.5, 0.4, True, 0.5], |
|
["examples/cat.jpg", 0.6, 0.5, True, 0.5], |
|
["examples/car.jpg", 0.6, 0.5, True, 0.5], |
|
["examples/plane.jpg", 0.6, 0.5, True, 0.5], |
|
["examples/bird.jpg", 0.6, 0.5, True, 0.5], |
|
] |
|
|
|
demo = gr.Interface( |
|
inference, |
|
inputs=[ |
|
gr.Image(label="Input Image"), |
|
gr.Slider(0, 1, value=0.5, label="IOU Threshold"), |
|
gr.Slider(0, 1, value=0.4, label="Threshold"), |
|
gr.Checkbox(label="Show Grad Cam"), |
|
gr.Slider(0, 1, value=0.5, label="Opacity of GradCAM"), |
|
], |
|
outputs=[ |
|
gr.Gallery(rows=2, columns=1), |
|
], |
|
title=title, |
|
description=description, |
|
examples=examples, |
|
) |
|
gr.close_all() |
|
demo.launch() |
|
|