How to select detections by index?

[lomnes-atlast-food]

Search before asking

• I have searched the Supervision issues and found no similar feature requests.

Question

I'd like to get a single detection by index but the way I've found to do it seems a little clunky. Is there a cleaner way to do this?

For example:

detections = sv.Detections.from_sam(sam_annotations)
big_detections = detections[detections.area > 1000] # filters detections by area
gets_single_detection = detections[np.array([0])] # gets the first detection
gives_error = detectons[0] # could return the first detection but throws an error
gives_error = detectons[[0,1,2]] # could return the first detection but throws an error

Question: Is there a different intended way to access detections by index?

If not, you could update the __getitem__ method of the Detections dataclass to allow passing ints, or lists of ints instead of only allowing np.array objects.

Suggested modification:

    def __getitem__(self, index: np.ndarray) -> Detections:
        index = np.asarray(index) # ADDED
        if isinstance(index, np.ndarray) and (
            index.dtype == bool or index.dtype == int
        ):
            return Detections(
                xyxy=self.xyxy[index],
                mask=self.mask[index] if self.mask is not None else None,
                confidence=self.confidence[index]
                if self.confidence is not None
                else None,
                class_id=self.class_id[index] if self.class_id is not None else None,
                tracker_id=self.tracker_id[index]
                if self.tracker_id is not None
                else None,
            )
        raise TypeError(
            f"Detections.__getitem__ not supported for index of type {type(index)}."
        )

Additional

No response

[github-actions]

Hello there, thank you for opening an Issue ! 🙏🏻 The team was notified and they will get back to you asap.

[SkalskiP]

Hi @lomnes-atlast-food 👋🏻! Thanks a lot for your interest in our project.

Let me start by saying that I 100% agree. Yesterday, I was working on a small project using supervision and had the same problem - there is no easy way to access detections by index. Let's add it!

Having said that, I'm just slightly concerned about API consistency. Here is what I think:

1. I love what we have here:

   detections[detections.area > 1000]

2. I love what you proposed here, and I really think we should add that:

   detectons[0]

3. Same goes for this. It is much easier to do detectons[[0,1,2]] and not detectons[np.array([0,1,2])]:

   detectons[[0,1,2]]

4. Problem I see is an inconsistency between __getitem__ and __iter__. I feel that the result type in those two cases below should be the same, but it is not.

   for detection in detection:
       type(detection) # -> tuple
   
   type(detection[0]) # -> sv.Detections

   I feel we need to agree on a consistent API at this point. Initially, I decided to return tuple in __iter__ to prevent creating objects that are going to be destroyed at the end of the for loop. I'm not sure what is the best approach here.

[lomnes-atlast-food]

Thanks for the feedback. Good observation about the type consistency. I also really like the filtering syntax!

An alternative would be a class method that returns the detections by index, or adding index as a parameter to the base class and use the same filtering notation.

detection = detections.by_index(0)
some_detections = detections.by_index([0,1,5])

Or

detection = detections[detections.index == 0]
detections[detections.index in [0,1,5]

Thanks so much for taking the time to review!

[SkalskiP]

I pinned 📌 the issue to give it a high priority. I don't have time to think about their API today. I have a lot of work to do. But I'll keep it in the back of my head, and I hope we will work on it next week 🔥 !

[SkalskiP]

Hi @lomnes-atlast-food 👋🏻! I just merged #118 into master branch. The code will be released with the supervision version 0.9.0.

import supervision as sv

detections = sv.Detections(...)

# select detections by index
first_detection = detections[0]

# select detections by list or array of indexes
few_detections = detections[[0, 2, 4]]

# select detections by index slice
first_detections = detections[:2]

I'm closing the issue for now. But if you notice any bugs or unexpected behavior, do not hesitate to reopen the issue.