open-mmlab / mmdetection

OpenMMLab Detection Toolbox and Benchmark
https://mmdetection.readthedocs.io
Apache License 2.0
29.21k stars 9.4k forks source link

Save only detected images #11719

Open Syed05 opened 4 months ago

Syed05 commented 4 months ago

I run the following script to perform inference on multiple images using a trained model (trained on custom data). I want to save only those images in which objects are detected.

Below is my inference code:

import os
from mmdet.apis import init_detector, inference_detector
import mmcv
import cv2
from mmdet.registry import VISUALIZERS

config_file = 'configs/_base_/custom2.py'
checkpoint_file = 'work_dirs/custom2/epoch_300.pth'

model = init_detector(config_file, checkpoint_file)

image_folder = 'images'
output_folder = 'outputs'
image_paths = [os.path.join(image_folder, img_name) for img_name in os.listdir(image_folder)]

visualizer = VISUALIZERS.build(model.cfg.visualizer)
visualizer.dataset_meta = model.dataset_meta

for img_path in image_paths:
        image = mmcv.imread(img_path)
        result = inference_detector(model, image)
        os.makedirs(output_folder, exist_ok=True)
        output_file = os.path.join(output_folder, os.path.basename(img_path))

    visualizer.add_datasample(
        'result',
        cv2.cvtColor(image, cv2.COLOR_BGR2RGB),
        data_sample=result,
        draw_gt=False,
        wait_time=0,
        pred_score_thr=0.3,
        out_file=output_file
    )