import shutil
import sys
from pathlib import Path
import numpy as np
from mmdet.apis import inference_detector, init_detector, show_result_pyplot
[docs]class Detector:
def __init__(self, config_file, checkpoint_file, device="cuda:0", cfg_options=None) -> None:
"""Instantiate an object detector.
Args:
config_file (str): Config file path.
checkpoint_file (str): Checkpoint file path.
device (str, optional): Device used for calculating. Defaults to "cuda:0".
cfg_options (dict, optional): Options to override some settings. Defaults to None.
"""
self.model = init_detector(config_file, checkpoint_file, device, cfg_options)
[docs] def test_images(self, images_or_dir, score_thr=0.3, out_dir=None):
"""Inference images with the detector.
Args:
images_or_dir (str or list[str]): Image files.
"""
if isinstance(images_or_dir, str):
image_paths = [str(f) for f in Path(images_or_dir).glob("*.jpg")]
else:
image_paths = [str(f) for f in images_or_dir]
assert isinstance(image_paths, list)
if out_dir is not None:
out_dir = Path(out_dir)
shutil.rmtree(out_dir, ignore_errors=True)
(out_dir / "images").mkdir(parents=True, exist_ok=False)
results = []
for image_path in image_paths:
out_file = None
if out_dir is not None:
out_file = str(out_dir / "images" / Path(image_path).name)
result = self.test_image(image_path, score_thr=score_thr, out_file=out_file)
results.append(result)
return results
[docs] def test_image(self, image_path, score_thr=0.3, out_file=None):
"""Inference image with the detector.
Args:
img (str): Image file.
"""
bbox_result = inference_detector(self.model, image_path)
if out_file is not None:
show_result_pyplot(self.model, image_path, bbox_result, score_thr=score_thr, out_file=out_file)
return bbox_result, image_path
[docs] def to_file(self, results, txt_file):
"""Format results.
Such as:
filepath,height,width,x1,y1,x2,y2,confidence,label,x1,y1,x2,y2,confidence,label
Args:
results (list[tuple]): The results.
"""
class_names = self.model.CLASSES
lines = []
for bbox_result, image_path in results:
bboxes = np.vstack(bbox_result)
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(bbox_result)
]
labels = np.concatenate(labels)
line = [image_path, "0", "0"]
for bbox, label in zip(bboxes, labels):
line.extend([f"{v:.2f}" for v in bbox[:4]])
line.extend([f"{bbox[4]:.4f}", class_names[label]])
lines.append(",".join(line))
text = "\n".join(lines)
Path(txt_file).parent.mkdir(parents=True, exist_ok=True)
with open(txt_file, "w") as f:
f.write(text)
[docs]def func(root, config_file, checkpoint_file, images_dir, score_thr, out_dir, txt_file):
root = Path(root)
config_file = str(root / config_file)
checkpoint_file = str(root / checkpoint_file)
executor = Detector(config_file, checkpoint_file)
results = executor.test_images(images_dir, score_thr, out_dir)
if txt_file is not None:
if out_dir is not None:
txt_file = Path(out_dir) / txt_file
executor.to_file(results, txt_file)
[docs]def parse_args(args=None):
from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
parser.add_argument("root", type=str,
help="the base dir")
parser.add_argument("config_file", type=str,
help="config file path")
parser.add_argument("checkpoint_file", type=str,
help="checkpoint file path")
parser.add_argument("images_dir", type=str,
help="a directory of images")
parser.add_argument("-s", dest="score_thr", type=float, default=0.05,
help="minimum score of bboxes to be shown")
parser.add_argument("-o", dest="out_dir", type=str, default=None,
help="draw boxes on the image")
parser.add_argument("-f", dest="txt_file", type=str, default="predictions.txt",
help="format and save results to a txt file")
args = parser.parse_args(args=args)
return vars(args)
[docs]def main(args=None):
kwargs = parse_args(args)
print(f"{__file__}: {kwargs}")
print(func(**kwargs))
return 0
if __name__ == "__main__":
sys.exit(main())