Hi, I'm trying to make an object counter with a yolo v8 model that I trained localy but when I try to adapt your code I get this error: IndexError: too many indices for tensor of dimension 2, on this line of code: if prediction['class'] in target_classes: inside the count_objects function.
def count_objects(predictions, target_classes):
"""
Helper method to count the number of objects in an image for a given class
:param predictions: predictions returned from calling the predict method
:param target_class: str, target class for object count
:return: dictionary with target class and total count of occurrences in image
"""
object_counts = {x:0 for x in target_classes}
for prediction in predictions:
if prediction['class'] in target_classes:
object_counts[prediction['class']] += 1
elif prediction['class'] not in target_classes:
object_counts[prediction['class']] = 1
present_objects = object_counts.copy()
for i in object_counts:
if object_counts[i] < 1:
present_objects.pop(i)
return present_objects
predictions = model.predict(source = "folderImages", conf = 0.6)
class_counts= count_objects(predictions, model.names)
print('/n', "Class Counts:", '/n')
print(class_counts)