Skip to content

Conversation

capjamesg
Copy link
Collaborator

This PR adds a new annotator for annotating classification results. This PR also adds a new data loader that enables you to load CLIP classifications into a sv.Classifications object.

Testing code

import supervision as sv
import clip
from PIL import Image
import torch
import cv2

SOURCE_IMAGE_PATH = "cafe.jpg"

image = Image.open(SOURCE_IMAGE_PATH)
model, preprocess = clip.load('ViT-B/16')

image = preprocess(image).unsqueeze(0).to("cpu")

text = clip.tokenize(["a piano", "a cafe"])

with torch.no_grad():
    logits_per_image, logits_per_text = model(image, text)

    classifications = sv.Classifications.from_clip(logits_per_image)

    # should show two results, indicating all predictions are loaded into the object
    print(classifications.get_top_k(3))

annotator = sv.ClassificationAnnotator()

image = annotator.annotate(
    scene=cv2.imread(SOURCE_IMAGE_PATH),
    classifications=classifications,
    labels=["a piano", "a cafe"]
)

sv.plot_image(image)

Here is an example of the annotator in use:

Screenshot 2023-10-13 at 23 49 36

Here is the raw image on which inference was run with CLIP in the code above:

cafe

@capjamesg capjamesg self-assigned this Oct 13, 2023
@capjamesg capjamesg requested a review from SkalskiP October 18, 2023 15:08
@capjamesg
Copy link
Collaborator Author

@capjamesg capjamesg changed the title Add sv.ClassificationAnnotator and sv.Classifications.from_clip Add sv.Classifications.from_clip data loader Nov 27, 2023
@SkalskiP
Copy link
Collaborator

Awesome @capjamesg 🔥 Merging!

@SkalskiP SkalskiP merged commit b63457e into develop Nov 27, 2023
@SkalskiP SkalskiP deleted the add-classification-annotator branch January 2, 2024 18:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants