Debug: refactor src

yolov8.py

@@ -8,21 +8,19 @@ from pytorch_grad_cam.utils.image import show_cam_on_image, scale_cam_image
 import gradio as gr
 from ultralytics import YOLO
 
-# Global Color Palette
-COLORS = np.random.uniform(0, 255, size=(80, 3))
 
-def parse_detections(results):
-    detections = results.pandas().xyxy[0].to_dict()
+COLORS = np.random.uniform(0, 255, size=(80, 3))
+def parse_detections(detections, model):
     boxes, colors, names = [], [], []
-    for i in range(len(detections["xmin"])):
-        confidence = detections["confidence"][i]
+    for detection in detections.boxes:
+        xmin, ymin, xmax, ymax = map(int, detection.xyxy[0].tolist())
+        confidence = detection.conf.item()
         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])
+        class_id = int(detection.cls.item())
+        name = model.names[class_id]
         boxes.append((xmin, ymin, xmax, ymax))
-        colors.append(COLORS[category])
+        colors.append(COLORS[class_id])
         names.append(name)
     return boxes, colors, names
 
@@ -53,27 +51,18 @@ def generate_cam_image(model, target_layers, tensor, rgb_img, boxes):
     return cam_image, renormalized_cam_image
 
 def xai_yolov8s(image):
-    # Load YOLOv8 model
-    model = YOLO('yolov8s.pt')
+    model = YOLO('yolov8s.pt')  # Ensure the model weights are available
     model.eval()
-    model.cpu()
-
-    target_layers = [model.model.model[-2]]  # Grad-CAM target layer
-
-    # Run YOLO detection
-    results = model([image])
-    boxes, colors, names = parse_detections(results)
+    results = model(image)
+    detections = results[0]
+    
+    boxes, colors, names = parse_detections(detections, model)
     detections_img = draw_detections(boxes, colors, names, image.copy())
-
-    # Prepare input tensor for Grad-CAM
     img_float = np.float32(image) / 255
     transform = transforms.ToTensor()
     tensor = transform(img_float).unsqueeze(0)
-
-    # Grad-CAM visualization
-    cam_image, renormalized_cam_image = generate_cam_image(model, target_layers, tensor, image, boxes)
-
-    # Combine results
+    target_layers = [model.model[-2]]  # Adjust according to YOLOv8 architecture
+    cam_image, renormalized_cam_image = generate_cam_image(model.model, target_layers, tensor, image, boxes)
     final_image = np.hstack((image, cam_image, renormalized_cam_image))
     caption = "Results using YOLOv8"
     return Image.fromarray(final_image), caption