animal-locator-drone / detection-service

YOLO detection service for classifying dogs
0 stars 0 forks source link

How to extract and process detections from Ultralytics YOLO #2

Closed RustyRaptor closed 7 months ago

RustyRaptor commented 7 months ago
  1. We need to be able to extract detections from the YOLO detector. Here is my current code.
  2. How do I extract the class names of detections in a stream?
  3. I cannot figure this out in the documentation except for the result.probs object but for some reason it is set to None in the results.
from ultralytics import YOLO

model = YOLO("yolov8n.pt")  # initialize model
results = model(source='./dogs.mp4', show=True, conf=0.4, stream=True)  # detect video stream

for result in results:
        boxes = result.boxes 
        names = result.names

        # Print out the class confidence and bounding box coordinates
        class_confidence = result.probs
        print("Class confidence: ", class_confidence)

        print("Boxes: ", boxes)
        print("Names: ", names)
        print("---------------------------------------------------")

Here is an example output for a frame detecting a person and a dog.

video 1/1 (frame 1724/2710) /home/ziadarafat/src/detection-service/dogs.mp4: 384x640 1 person, 1 dog, 130.0ms
Class confidence:  None
Boxes:  ultralytics.engine.results.Boxes object with attributes:

cls: tensor([ 0., 16.])
conf: tensor([0.8254, 0.7600])
data: tensor([[406.5982, 154.8007, 465.3931, 302.1017,   0.8254,   0.0000],
        [503.3336, 249.9007, 595.2296, 299.5467,   0.7600,  16.0000]])
id: None
is_track: False
orig_shape: (720, 1280)
shape: torch.Size([2, 6])
xywh: tensor([[435.9957, 228.4512,  58.7949, 147.3011],
        [549.2816, 274.7237,  91.8960,  49.6460]])
xywhn: tensor([[0.3406, 0.3173, 0.0459, 0.2046],
        [0.4291, 0.3816, 0.0718, 0.0690]])
xyxy: tensor([[406.5982, 154.8007, 465.3931, 302.1017],
        [503.3336, 249.9007, 595.2296, 299.5467]])
xyxyn: tensor([[0.3177, 0.2150, 0.3636, 0.4196],
        [0.3932, 0.3471, 0.4650, 0.4160]])
Names:  {0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane', 5: 'bus', 6: 'train', 7: 'truck', 8: 'boat', 9: 'traffic light', 10: 'fire hydrant', 11: 'stop sign', 12: 'parking meter', 13: 'bench', 14: 'bird', 15: 'cat', 16: 'dog', 17: 'horse', 18: 'sheep', 19: 'cow', 20: 'elephant', 21: 'bear', 22: 'zebra', 23: 'giraffe', 24: 'backpack', 25: 'umbrella', 26: 'handbag', 27: 'tie', 28: 'suitcase', 29: 'frisbee', 30: 'skis', 31: 'snowboard', 32: 'sports ball', 33: 'kite', 34: 'baseball bat', 35: 'baseball glove', 36: 'skateboard', 37: 'surfboard', 38: 'tennis racket', 39: 'bottle', 40: 'wine glass', 41: 'cup', 42: 'fork', 43: 'knife', 44: 'spoon', 45: 'bowl', 46: 'banana', 47: 'apple', 48: 'sandwich', 49: 'orange', 50: 'broccoli', 51: 'carrot', 52: 'hot dog', 53: 'pizza', 54: 'donut', 55: 'cake', 56: 'chair', 57: 'couch', 58: 'potted plant', 59: 'bed', 60: 'dining table', 61: 'toilet', 62: 'tv', 63: 'laptop', 64: 'mouse', 65: 'remote', 66: 'keyboard', 67: 'cell phone', 68: 'microwave', 69: 'oven', 70: 'toaster', 71: 'sink', 72: 'refrigerator', 73: 'book', 74: 'clock', 75: 'vase', 76: 'scissors', 77: 'teddy bear', 78: 'hair drier', 79: 'toothbrush'}
RustyRaptor commented 7 months ago

Ok I figured it out. It's quite simple it was just kinda weird because it's stored in the Boxes object and the variable names aren't very descriptive. Anyhow...

Boxes.cls is the class id which you can find in result.names Boxes.conf is the confidence level associated with said class.

Hence you can zip them to get a nice pairing of each. Additionally you can use the names field to fetch the actual class name.

from ultralytics import YOLO

model = YOLO("yolov8n.pt")  # initialize model
results = model(source='./dogs.mp4', show=True, conf=0.4, stream=True)  # detect video stream

for result in results:
        boxes = result.boxes
        confidences = list(zip(boxes.cls, boxes.conf))
        names = result.names

        if len(boxes) > 0:
                print("Class and Confidence", confidences)
                print(names)
                print("---------------------------------------------------")