Blame for app.py

import torch
import torch.optim as optim
import lightning.pytorch as pl
from lightning.pytorch.tuner import Tuner

# import pytorch_lightning as pl
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
import cv2
import torch
from pytorch_grad_cam.utils.image import show_cam_on_image
import numpy as np
from albumentations.pytorch import ToTensorV2

from pytorchyolo.utils.utils import load_classes, rescale_boxes, print_environment_info

from utils_for_app import cells_to_bboxes,plot_image,YoloCAM,non_max_suppression
from yolov3 import YOLOv3
from loss import YoloLoss
from utils import LearningRateFinder
# Create your config module or import it from the existing config.py file.
import config
from main_yolov3_lightening import YOLOv3Lightning
import torch
import cv2
import numpy as np
import gradio as gr
import os

print("Loaded....weights 1")
model = YOLOv3Lightning()
print("Loaded....weights 2")
model.load_state_dict(torch.load("yolov3_608_ckpt_40.pth", map_location=torch.device('cpu')), strict=False)
print("Loaded....weights 3")
model.setup(stage="test")

print("Loaded....weights 4") 
IMAGE_SIZE = 608
transformsImG = A.Compose(
    [
        A.LongestMaxSize(max_size=IMAGE_SIZE),
        A.PadIfNeeded(
            min_height=IMAGE_SIZE, min_width=IMAGE_SIZE, border_mode=cv2.BORDER_CONSTANT
        ),
        A.Normalize(mean=[0, 0, 0], std=[1, 1, 1], max_pixel_value=255,),
        ToTensorV2(),
    ],
)

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,iou_threshold=0.5, threshold=0.4):
	
	
    transformed_image = transformsImG(image=image)["image"].unsqueeze(0)
    #output = model(transformed_image)
    
    img_size=608
    
    
    model.eval()  # Set model to evaluation mode


    nms_boxes = None

    # Get detections
    with torch.no_grad(): 
        output  = model(transformed_image)
        
        
        
        boxes = [[] for _ in range(1)]
    
        for i in range(3):
            print("output[i].shape", output[i].shape)
            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_threshold, threshold=threshold, box_format="midpoint",
        )   

        # BBOX = rescale_boxes(detect[0], img_size, image.shape[:2])
    

    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, (608, 608))
    img = np.float32(img) / 255
    cam_image = show_cam_on_image(img, img, use_rgb=True)


    return fig,cam_image


examples = [
    ["images/Grass_divot.jpg"],
    ["images/Grass_divot1.jpg"],
    ["images/Grass_divot2.jpg"],
    ["images/Grass_divot3.jpg"],
    ["images/Grass_divot4.jpg"]
]


icon_html = '<i class="fas fa-chart-bar"></i>'
title_with_icon = f"""
<div style="background-color: #f5f1f2; padding: 10px; display: flex; align-items: center;">
    {icon_html} <span style="margin-left: 10px;">Object Detection on Pascal VOC Dataset with YoloV3</span>
</div>
"""
description_with_icon = f"""
<div style="background-color: #f1f1f5; padding: 8px; display: flex; align-items: center;">
    {icon_html}
    <span style="margin-left: 8px;">
        <p><strong>PyTorch YOLOv3</strong></p>

</div>
"""


demo = gr.Interface(process_image_and_plot,
    inputs=[gr.Image(label="Input Image"),
            gr.Slider(0, 1, value=0.5, label="Intersection over Union (IOU) Threshold",info="Determines how much overlap between two boxes is allowed before they are considered redundant"),
            gr.Slider(0, 1, value=0.4, label="Threshold", info="It is used to filter out boxes with confidence scores below it. Higher value reduces the weaker classes"),],
    outputs=[
        gr.Plot(label="Output with Classes",), 
        gr.Image(shape=(32, 32), label="GradCAM Output"),
    ],
    title=title_with_icon,
    description=description_with_icon,
    examples=examples,
)
demo.launch()