Spaces:
Runtime error
Runtime error
File size: 2,818 Bytes
5d3cdc9 | 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 | import torch
import torch.optim as optim
import lightning.pytorch as pl
from lightning.pytorch.tuner import Tuner
from tqdm import tqdm
from torch.optim.lr_scheduler import OneCycleLR
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import albumentations as A
from pytorch_grad_cam.utils.image import show_cam_on_image
from albumentations.pytorch import ToTensorV2
import config
from yolo_lightning import YOLOv3Lightning
import torch
import cv2
import numpy as np
import gradio as gr
import os
from utils_app import *
model = YOLOv3Lightning (config)
model.load_state_dict(torch.load("custom_yolo_model.pth", map_location=torch.device('cpu')), strict=False)
model.setup(stage="test")
IMAGE_SIZE = 416
ANCHORS = [
[(0.28, 0.22), (0.38, 0.48), (0.9, 0.78)],
[(0.07, 0.15), (0.15, 0.11), (0.14, 0.29)],
[(0.02, 0.03), (0.04, 0.07), (0.08, 0.06)],
] # Note these have been rescaled to be between [0, 1]
S = [IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8]
scaled_anchors = (
torch.tensor(config.ANCHORS)
* torch.tensor(config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)
)
def process_image_and_plot(image, model, scaled_anchors):
transformed_image = config.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=0.5, threshold=0.4, box_format="midpoint",
)
fig = plot_image(transformed_image[0].permute(1, 2, 0), nms_boxes)
cam = YoloCAM(model=model, target_layers=[model.model.layers[-2]], use_cuda=False)
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)
return fig,cam_image
examples = [
["images/2012_004288.jpg"],
["images/2012_004314.jpg"],
["images/car.jpg"],
]
def processed_image(image):
figure,gradcam = process_image_and_plot(image, model, scaled_anchors)
return figure,gradcam
title = "YoloV3 on Pascal VOC Dataset (GradCAM)"
description = f"Pytorch Implemetation of YoloV3 trained from scratch on Pascal VOC dataset with GradCAM \n Class in pascol voc: {', '.join(config.PASCAL_CLASSES)}"
demo = gr.Interface(processed_image,
inputs=[
gr.Image(label="Input Image"),
],
outputs=[gr.Plot(),gr.Image(shape=(32, 32), label="GradCAM Prediction")],
title=title,
description=description,
examples=examples,
)
demo.launch() |