Return the images that contain a defined 'class label'

Ultralytics HUB ⭐
1

I currently pull the pre-trained model from Roboflow and filter the dataset images in a directory to return the images that contain the user’s input ‘class’. My current app code is without Ultrayltics integration.

I need help changing the code so it uses the Yolov8n.pt and data.yaml files which contain the classes.
The goal of the app is to:

  1. Allow the user to enter a ‘class number’
  2. Give it a folder of images.
  3. Return all the images that contain the entered ‘class number’\

https://github.com/Yaandle/ImageFinder/blob/master/testing.py

1 Like
2

Hi! We are moving the Ultralytics community to Discord. To receive support for your questions, problems, issues, etc. please join us on the new server: Ultralytics

3

You can have an input variable that collects the user’s input

detections = sv.Detections.from_yolov8(result)

                detections = detections[detections.class_id == short_id]

                labels = [
                f"{model.model.names[class_id]}"
                for box, confidence, class_id in zip(detections.xyxy, detections.confidence, detections.class_id)
                ]

                image = box_annotator.annotate(
                    scene=image,
                    detections=detections,
                    labels=labels
                )

This worked for me. You can try it