Debug: refactor src

app.py

@@ -14,7 +14,7 @@ def process_image(image, yolo_versions=["yolov5"]):
     for yolo_version in yolo_versions:
         if yolo_version == "yolov5":
             result_images.append(xai_yolov5(image)) 
-        elif yolo_version == "yolov8n":
+        elif yolo_version == "yolov8s":
             result_images.append(xai_yolov8n(image))
         else:
             result_images.append((Image.fromarray(image), f"{yolo_version} not yet implemented."))
@@ -26,7 +26,7 @@ interface = gr.Interface(
     inputs=[
         gr.Image(type="pil", label="Upload an Image"),
         gr.CheckboxGroup(
-            choices=["yolov5", "yolov8n", "yolov10"],
+            choices=["yolov5", "yolov8s", "yolov10"],
             value=["yolov5"],  # Set default selection to YOLOv5
             label="Select Model(s)",
         )

requirements.txt

@@ -8,4 +8,3 @@ grad-cam==1.4.8
 gradio
 ultralytics
 torchcam
-YOLOv8-Explainer

yolov8.py

@@ -1,4 +1,3 @@
-from ultralytics import YOLO
 import torch
 import cv2
 import numpy as np
@@ -12,17 +11,18 @@ import gradio as gr
 COLORS = np.random.uniform(0, 255, size=(80, 3))
 
 def parse_detections(results):
+    detections = results.pandas().xyxy[0].to_dict()
     boxes, colors, names = [], [], []
-    for result in results:
-        # Accessing boxes directly from the result
-        for box in result.boxes:
-            xmin, ymin, xmax, ymax = box.xyxy[0].int().tolist()  # Convert to list of integers
-            category = int(box.cls[0].item())  # Class index
-            name = result.names[category]  # Get class name from names
-            boxes.append((xmin, ymin, xmax, ymax))
-            colors.append(COLORS[category])  # Ensure COLORS is defined elsewhere in your code
-            names.append(name)
-
+    for i in range(len(detections["xmin"])):
+        confidence = detections["confidence"][i]
+        if confidence < 0.2:
+            continue
+        xmin, ymin = int(detections["xmin"][i]), int(detections["ymin"][i])
+        xmax, ymax = int(detections["xmax"][i]), int(detections["ymax"][i])
+        name, category = detections["name"][i], int(detections["class"][i])
+        boxes.append((xmin, ymin, xmax, ymax))
+        colors.append(COLORS[category])
+        names.append(name)
     return boxes, colors, names
 
 def draw_detections(boxes, colors, names, img):
@@ -34,12 +34,15 @@ def draw_detections(boxes, colors, names, img):
                     lineType=cv2.LINE_AA)
     return img
 
-
 def generate_cam_image(model, target_layers, tensor, rgb_img, boxes):
     cam = EigenCAM(model, target_layers)
     grayscale_cam = cam(tensor)[0, :, :]
     img_float = np.float32(rgb_img) / 255
+
+    # Generate Grad-CAM
     cam_image = show_cam_on_image(img_float, grayscale_cam, use_rgb=True)
+
+    # Renormalize Grad-CAM inside bounding boxes
     renormalized_cam = np.zeros(grayscale_cam.shape, dtype=np.float32)
     for x1, y1, x2, y2 in boxes:
         renormalized_cam[y1:y2, x1:x2] = scale_cam_image(grayscale_cam[y1:y2, x1:x2].copy())
@@ -48,69 +51,28 @@ def generate_cam_image(model, target_layers, tensor, rgb_img, boxes):
 
     return cam_image, renormalized_cam_image
 
-
-def xai_yolov8nnnn(image):
-    model = YOLO('yolov8n.pt')  # Load YOLOv8n pre-trained weights
+def xai_yolov8s(image):
+    # Load YOLOv8 model
+    model = torch.hub.load('ultralytics/yolov8', 'yolov8s', pretrained=True)
     model.eval()
-
-    # Check if GPU is available and use it
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    model.to(device)
+    model.cpu()
 
     target_layers = [model.model.model[-2]]  # Grad-CAM target layer
-    
-    # Process the image through the model
-    results = model([image])
-    
-    # If results are a list, extract the first element (detected results)
-    if isinstance(results, list):
-        results = results[0]  # Extracting the first result (if list)
-    
-    # Ensure that outputs are in tensor form
-    logits = results.pred[0]  # Get the prediction tensor from the results
 
-    # Parse the detections
-    boxes, colors, names = parse_detections([results])  # Ensure results are passed as a list
+    # Run YOLO detection
+    results = model([image])
+    boxes, colors, names = parse_detections(results)
     detections_img = draw_detections(boxes, colors, names, image.copy())
-    
-    # Prepare image for Grad-CAM
+
+    # Prepare input tensor for Grad-CAM
     img_float = np.float32(image) / 255
     transform = transforms.ToTensor()
-    tensor = transform(img_float).unsqueeze(0).to(device)  # Ensure tensor is on the right device
-    
-    # Generate CAM images
+    tensor = transform(img_float).unsqueeze(0)
+
+    # Grad-CAM visualization
     cam_image, renormalized_cam_image = generate_cam_image(model, target_layers, tensor, image, boxes)
 
-    # Combine original image, CAM image, and renormalized CAM image
+    # Combine results
     final_image = np.hstack((image, cam_image, renormalized_cam_image))
-
-    # Return final image and a caption
-    caption = "Results using YOLOv8n"
+    caption = "Results using YOLOv8"
     return Image.fromarray(final_image), caption
-
-
-
-from YOLOv8_Explainer import yolov8_heatmap, display_images
-
-def xai_yolov8n(image):
-    model = yolov8_heatmap(
-        weight="yolov8n.pt", 
-            conf_threshold=0.4, 
-            device = "cpu", 
-            method = "EigenCAM", 
-            layer=[10, 12, 14, 16, 18, -3],
-            backward_type="all",
-            ratio=0.02,
-            show_box=True,
-            renormalize=False,
-    )
-
-    # Pass the NumPy array to the model
-    imagelist = model(image)  # Use the image array directly
-
-    # Display the resulting images
-    # Pass the NumPy array to the model
-    imagelist = model(image)  # Use the image array directly
-
-# Display the resulting images
-    print(imagelist)