Extracting data on classes detected

[kwayne-08]

Search before asking

• I have searched the HUB issues and discussions and found no similar questions.

Question

I am having difficulty extracting class information from inference on my custom model. When I run inference using the Ultralytics API I am able to use these python lines to request a response and convert the results to a useable form:

response = requests.post(url, headers=headers, data=data, files={"image": f})
response.raise_for_status()
json_str = json.dumps(response.json(), indent=2)

I can then iterate on the json_str to extract the name of classes detected in each image and how many occurances for each class name for each image and save it to a csv file.

However, when I run local inference using these lines, I am unable to parse the results to get the information I need:

results = subprocess.run(['yolo', 'task=detect', 'mode=predict', 'model=housing_model-2024-01-16.pt', 'conf=0.50',
                          'source=C:/Users/kwban/PycharmProjects/flask_website/uploads/resized'])

The type of the results variable is said to be <class 'subprocess.CompletedProcess'> but I have not found a way to parse it. The inference prints out these results:
image 1/3 C:\Users\kwban\PycharmProjects\flask_website\uploads\resized\p10001-s.jpg: 320x640 (no detections), 99.5ms
image 2/3 C:\Users\kwban\PycharmProjects\flask_website\uploads\resized\p10003-s.jpg: 320x640 1 ceilinglight, 55.0ms
image 3/3 C:\Users\kwban\PycharmProjects\flask_website\uploads\resized\p10004-s.jpg: 320x640 1 kitchenrange, 53.0ms

This information would be helpful, but I am unable to capture this information for a csv report. What is the solution for this?

Additional

No response

[kalenmike]

@kwayne-08 This is related to the ultralytics package. You can find some support in the predict docs.

Here is an example to get the boxes:

from ultralytics import YOLO

# Load a pretrained YOLOv8n model
model = YOLO('yolov8n.pt')

# Run inference on an image
results = model('bus.jpg')  # results list

# View results
for r in results:
    print(r.boxes)  # print the Boxes object containing the detection bounding boxes

[kwayne-08]

kalenmike,

Thanks for your quick response. However, I am getting bounding boxes, and I can get a full list of class names for the dataset. What I can't get is the class names for those objects detected in the image. It prints that information, but I'm unable to capture that info for use. For instance, even when I don't add a print function to the script, it will print "0: 320x640 1 ceilinglight, 1 door, 1 toilet, 1 window, 62.6ms." The information I need from this line is 1 ceilinglight, 1 door, 1 toilet, etc. But I can not find a way to capture it for use.

[UltralyticsAssistant]

Hello again @kwayne-08,

To extract the class names and their respective counts from the results after running local inference, you should interact with the results object directly rather than parsing the printed output. The results object contains a .pred attribute which holds the predictions in a tensor, where each detection is a row with the format [x1, y1, x2, y2, confidence, class].

You can iterate over the .pred tensor to count the occurrences of each class. Here's a high-level explanation of the steps you'd take:

1. Iterate over results.pred[0] to access the detections for the first image (assuming batch size of 1).
2. For each detection, extract the class index (the last value in the detection tensor).
3. Convert the class index to the class name using results.names, which maps indices to class names.
4. Use a dictionary to count the occurrences of each class name.
5. Once you have the counts, you can write them to a CSV file using Python's csv module.

Remember, you won't be able to capture the printed output directly, but by accessing the results object's attributes, you can programmatically obtain the information you need and format it as desired for your CSV report.

For more detailed guidance on working with the results object and accessing the predictions, please refer to the Ultralytics HUB Docs.

[kalenmike]

@kwayne-08 You have to implement some custom logic to solve that. Each box has its class id in box.cls these id's are the key position of the results.names dictionary. So you just need to fetch the class name as you loop through all the boxes. To get the total count of items you will just need to sum the results based on class index or name.

from ultralytics import YOLO

# Load a pretrained YOLOv8n model
model = YOLO('yolov8n.pt')

# Run inference on an image
results = model('bus.jpg')  # results list

r = results[0] # only one result as only one image was inferred

class_names = r.names

# View results
for b in r.boxes:
    print(f"Class name: {class_names[b.cls]}") 
    # Add your grouping and counting logic

[kwayne-08]

kalenmike,

Thank you for your responses. You have been a big help. Your latest example pointed me in the right direction, though I did find it necessary to modify the for loop in this way:

for b in r.boxes: clnm = str(b.cls[0]) clnm = int(clnm[-4:-2]) print(f"Class name: {class_names[clnm]}")

Your line print(f"Class name: {class_names[b.cls]}") was throwing an error since class_names[b.cls]} returned something like 'Tensor[32.0]' Therefore, I converted it to string, sliced it to get the class and entered that in the equation you sent. There is probably a better way of doing this, but it gets the result I was after.

Wayne

[UltralyticsAssistant]

@kwayne-08 hello Wayne,

I'm glad to hear that you've found a solution that works for you! It's great that you were able to adapt the example to fit your needs. Indeed, there might be more streamlined ways to extract the class index directly from the tensor, but what matters most is that you're able to retrieve the class names as required.

If you encounter any more issues or have further questions, feel free to reach out. The YOLO community and the Ultralytics team are always here to help. Keep up the good work! 🚀