Spaces:
Sleeping
Sleeping
Commit
·
5f61be0
1
Parent(s):
6e5e82f
Fix: version dependency
Browse files
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
|
37 |
-
xmin, ymin, xmax, ymax
|
|
|
|
|
38 |
if confidence > 0.2:
|
39 |
boxes.append((xmin, ymin, xmax, ymax))
|
40 |
-
colors.append(COLORS[
|
41 |
-
names.append(result.names[
|
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
|
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
|