Spaces:
Runtime error
Runtime error
File size: 3,237 Bytes
903f9d3 b15f6e1 903f9d3 f9c5d15 903f9d3 b15f6e1 903f9d3 b15f6e1 903f9d3 b15f6e1 903f9d3 b15f6e1 903f9d3 b15f6e1 903f9d3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
import gradio as gr
from model import Lightning_YOLO
import config
from utils import non_max_suppression, cells_to_bboxes, draw_bounding_boxes, get_annotations
import torch
scaled_anchors = config.scaled_anchors
model = Lightning_YOLO()
model.load_state_dict(torch.load("yolov3.pth", map_location="cpu"), strict=False)
model.eval()
def inference(image, threst = 0.5, iou_tresh = 0.5):
image_copy = image.copy()
transformed_image = config.transforms(image=image)["image"].unsqueeze(0)
output = model(transformed_image)
bboxes = [[] for _ in range(1)]
nms_boxes_output = []
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_tresh, threshold=threst, box_format="midpoint",
# )
for i in range(image.shape[0]):
nms_boxes = non_max_suppression(
bboxes[i], iou_threshold=iou_tresh, threshold=threst, box_format="midpoint",
)
nms_boxes_output.append(nms_boxes)
annotations = get_annotations(nms_boxes_output,config.IMAGE_SIZE,config.IMAGE_SIZE)
# plot_img = draw_bounding_boxes(image.copy(), nms_boxes, class_labels=config.PASCAL_CLASSES)
# return plot_img
return [image_copy, annotations]
def visualize_gradcam(image, target_layer=-5, 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.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
pass
window1 = gr.Interface(
inference,
inputs=[
gr.Image(label="Input Image"),
gr.Slider(0, 1, value=0.5, step=0.1, label="Threshold", info="Set Threshold value"),
gr.Slider(0, 1, value=0.5, step=0.1, label="IOU Threshold", info="Set IOU Threshold value"),
],
outputs=[
gr.Image(label="YoloV3 Object Detection"),
],
# examples=ex1,
)
window2 = gr.Interface(
visualize_gradcam,
inputs=[
gr.Image(label="Input Image"),
gr.Slider(-5, -2, value=-3, step=-1, label="Network Layer", info="GRAD-CAM Layer to visualize?"),
gr.Checkbox(label="GradCAM", value=True, info="Visualize Class Activation Maps ??"),
gr.Slider(0, 1, value=0.5, step=0.1, label="Transparency", info="Set Transparency of GRAD-CAMs"),
],
outputs=[
# gr.Image(label="Grad-CAM Visualization"),
gr.AnnotatedImage(label='BBox Prediction',
height=config.IMAGE_SIZE,
width=config.IMAGE_SIZE)
],
# examples=ex2,
)
app = gr.TabbedInterface([window1, window2], ["YOLO V3 Detection", "GradCAM Visualization"])
app.launch() |