detectron2-pipeline
detectron2-pipeline copied to clipboard
question: is there an easy way to get only class labels out for each image?
Hi, very nice work!
Can i ask if there is an easy way to get only class labels out for each image?
For my project i don't need visualizations, instead i'd love a dataframe with class labels.
Thank you
Hi Krzysiek,
I suppose that all you need is a detection mode (--config-file configs/COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml
) to detect classes on the frames.
Then to extract class ids you can use:
predictions = data["predictions"]
instances = instances.to(self.cpu_device)
print(predictions.pred_classes)
Looking at the detectron2/utils/video_visualizer.py
is the best way to know how to extract needed information from the model predictions.
Labels Can be shown
-
By configuring the metadata of the model by the following line of code
MetadataCatalog.get(cfg.DATASETS.TRAIN[0]).thing_classes = ['class_1', 'class_2', class_3', 'and so on']
-
By adding the following argument Visualizer object
"MetadataCatalog.get(cfg.DATASETS.TRAIN[0])"
The final line of code will be
MetadataCatalog.get(cfg.DATASETS.TRAIN[0]).thing_classes = ['class_1', 'class_2', class_3', 'and so on']
v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1)