ManjunathReddy/Yolo3_from_scratch

1from typing import List2import cv23import torch4import numpy as np5import src.config as config6from pytorch_grad_cam.utils.image import show_cam_on_image7 8from src.model_obj import Assignment139from src.utils import cells_to_bboxes, non_max_suppression, draw_predictions, YoloCAM10 11 12 13 14weights_path = "model_ass_13_up.ckpt"15model = Assignment13().load_from_checkpoint(weights_path,map_location=torch.device("cpu"))16model = model.model17#ckpt = torch.load(weights_path, map_location="cpu")18#model.load_state_dict(ckpt)19model.eval()20print("[x] Model Loaded..")21 22scaled_anchors = (23    torch.tensor(config.ANCHORS)24    * torch.tensor(config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)25).to(config.DEVICE)26 27cam = YoloCAM(model=model, target_layers=[model.layers[-2]], use_cuda=False)28 29def predict(image: np.ndarray, iou_thresh: float = 0.5, thresh: float = 0.4, show_cam: bool = False, transparency: float = 0.5) -> List[np.ndarray]:30    with torch.no_grad():31        transformed_image = config.transforms(image=image)["image"].unsqueeze(0)32        output = model(transformed_image)33        34        bboxes = [[] for _ in range(1)]35        for i in range(3):36            batch_size, A, S, _, _ = output[i].shape37            anchor = scaled_anchors[i]38            boxes_scale_i = cells_to_bboxes(39                output[i], anchor, S=S, is_preds=True40            )41            for idx, (box) in enumerate(boxes_scale_i):42                bboxes[idx] += box43 44    nms_boxes = non_max_suppression(45        bboxes[0], iou_threshold=iou_thresh, threshold=thresh, box_format="midpoint",46    )47    plot_img = draw_predictions(image, nms_boxes, class_labels=config.PASCAL_CLASSES)48    if not show_cam:49        return [plot_img]50    51    grayscale_cam = cam(transformed_image, scaled_anchors)[0, :, :]52    img = cv2.resize(image, (416, 416))53    img = np.float32(img) / 25554    cam_image = show_cam_on_image(img, grayscale_cam, use_rgb=True, image_weight=transparency)55    return [plot_img, cam_image]56 57 58