The roboflow-computer-vision-utilities from roboflow

roboflow-computer-vision-utilities's People

Contributors

Stargazers

Watchers

roboflow-computer-vision-utilities's Issues

Problem with object counting

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.

Here is my code:

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)

Recommend Projects