BhumikaMak commited on
Commit
5f61be0
·
1 Parent(s): 6e5e82f

Fix: version dependency

Browse files
Files changed (1) hide show
  1. app.py +7 -5
app.py CHANGED
@@ -33,12 +33,14 @@ def parse_detections(results, yolo_version):
33
  elif yolo_version == "yolov8":
34
  # For YOLOv8
35
  for result in results:
36
- for box in result.boxes.data.tolist(): # Access box data directly
37
- xmin, ymin, xmax, ymax, confidence, class_id = box
 
 
38
  if confidence > 0.2:
39
  boxes.append((xmin, ymin, xmax, ymax))
40
- colors.append(COLORS[int(class_id) % len(COLORS)])
41
- names.append(result.names[int(class_id)])
42
 
43
  return boxes, colors, names
44
 
@@ -87,7 +89,7 @@ def process_image(image, yolo_versions=["yolov5"]):
87
  # Parse detections using updated function
88
  boxes, colors, names = parse_detections(results[0], yolo_version)
89
 
90
- detections_img = draw_detections(boxes.copy(), colors.copy(), names.copy(), rgb_img.copy())
91
 
92
  # Grad-CAM visualization
93
  target_layers = [model.model.model[-1]] # Use last layer as target layer for Grad-CAM
 
33
  elif yolo_version == "yolov8":
34
  # For YOLOv8
35
  for result in results:
36
+ for box in result.boxes:
37
+ xmin, ymin, xmax, ymax = box.xyxy.tolist()
38
+ confidence = box.conf.item()
39
+ class_id = int(box.cls.item())
40
  if confidence > 0.2:
41
  boxes.append((xmin, ymin, xmax, ymax))
42
+ colors.append(COLORS[class_id % len(COLORS)])
43
+ names.append(result.names[class_id])
44
 
45
  return boxes, colors, names
46
 
 
89
  # Parse detections using updated function
90
  boxes, colors, names = parse_detections(results[0], yolo_version)
91
 
92
+ detections_img = draw_detections(boxes, colors.copy(), names.copy(), rgb_img.copy())
93
 
94
  # Grad-CAM visualization
95
  target_layers = [model.model.model[-1]] # Use last layer as target layer for Grad-CAM