|
import gradio as gr |
|
import numpy as np |
|
import config |
|
from utils import * |
|
from pytorch_grad_cam.utils.image import show_cam_on_image |
|
from model import YOLOv3 |
|
|
|
ex1 = [[f'examples/{i}.jpg'] for i in range(1,8)] |
|
ex2 = [[f'examples/{i}.jpg'] for i in range(8,15)] |
|
|
|
scaled_anchors = config.scaled_anchors |
|
|
|
model = YOLOv3() |
|
model.load_state_dict(torch.load("S13_model.pth", map_location="cpu"), strict=False) |
|
model.eval() |
|
|
|
@torch.inference_mode() |
|
def YoloV3_classifier(image, thresh=0.5,iou_thresh=0.5): |
|
transformed_image = config.test_transforms(image=image)["image"].unsqueeze(0) |
|
output = model(transformed_image) |
|
|
|
bboxes = [[] for _ in range(1)] |
|
for i in range(3): |
|
batch_size, A, S, _, _ = output[i].shape |
|
anchor = scaled_anchors[i] |
|
boxes_scale_i = cells_to_bboxes( |
|
output[i], anchor, S=S, is_preds=True |
|
) |
|
for idx, (box) in enumerate(boxes_scale_i): |
|
bboxes[idx] += box |
|
|
|
nms_boxes = non_max_suppression( |
|
bboxes[0], iou_threshold=iou_thresh, threshold=thresh, box_format="midpoint", |
|
) |
|
plot_img = draw_bounding_boxes(image.copy(), nms_boxes, class_labels=config.PASCAL_CLASSES) |
|
|
|
return plot_img |
|
|
|
window1 = gr.Interface( |
|
YoloV3_classifier, |
|
inputs=[ |
|
gr.Image(label="Input Image"), |
|
gr.Slider(0, 1, value=0.5, step=0.1, label="Threshold"), |
|
gr.Slider(0, 1, value=0.5, step=0.1, label="IOU Threshold"), |
|
], |
|
outputs=[ |
|
gr.Image(label="YoloV3 Object Detection"), |
|
], |
|
examples=ex1, |
|
) |
|
|
|
|
|
def visualize_gradCam(image, target_layer=-14, show_cam=True, transparency=0.5): |
|
|
|
if show_cam: |
|
cam = YoloCAM(model=model, target_layers=[model.layers[target_layer]], use_cuda=False) |
|
transformed_image = config.test_transforms(image=image)["image"].unsqueeze(0) |
|
grayscale_cam = cam(transformed_image, scaled_anchors)[0, :, :] |
|
img = cv2.resize(image, (416, 416)) |
|
img = np.float32(img) / 255 |
|
cam_image = show_cam_on_image(img, grayscale_cam, use_rgb=True, image_weight=transparency) |
|
else: |
|
cam_image = image |
|
|
|
return cam_image |
|
|
|
window2 = gr.Interface( |
|
visualize_gradCam, |
|
inputs=[ |
|
gr.Image(label="Input Image"), |
|
gr.Radio(choices=[-11,-14,-17],value=-14, label="Network Layer"), |
|
gr.Checkbox(label="GradCAM", value=True), |
|
gr.Slider(0, 1, value=0.5, step=0.1, label="Transparency"), |
|
], |
|
outputs=[ |
|
gr.Image(label="Grad-CAM Visualization"), |
|
], |
|
examples=ex2, |
|
) |
|
|
|
app = gr.TabbedInterface([window1, window2], ["YOLO V3 Detection", "GradCAM Visualization"]) |
|
app.launch() |
|
|