register_coco_instances
and load_coco_json
require images with annotations to properly label the predicted objects.dataset_name = "test_data"
image_dir = "data/test"
coco_file = "data/test_annotations.json"
# Register dataset
# A COCO file is needed with image info, which I don't have
register_coco_instances(dataset_name , {}, coco_file, image_dir)
test_dict = load_coco_json(coco_file, image_dir, dataset_name=dataset_name )
metadata = MetadataCatalog.get(dataset_name)
# config details omitted for brevity
cfg = get_cfg()
predictor = DefaultPredictor(cfg)
# Make predictions for all images
for sample in test_dict:
image_filename = sample["file_name"]
img = cv2.imread(image_filename)
outputs = predictor(img)
# Display or save image with predictions to file
Here's a method to generate the image details from a directory of images and write it to an existing COCO JSON file:
from PIL import Image
def add_images_to_coco(image_dir, coco_filename):
image_filenames = list(Path(image_dir).glob('*.jpg'))
images = []
for i, image_filename in enumerate(image_filenames):
im = Image.open(image_filename)
width, height = im.size
image_details = {
"id": id + 1,
"height": height,
"width": width,
"file_name": str(image_filename.resolve()),
}
images.append(image_details)
# This will overwrite the image tags in the COCO JSON file
with open(coco_filename) as f:
data = json.load(f)
coco['images'] = images
with open(coco_filename, 'w') as coco_file:
json.dump(data, coco_file, indent = 4)
You'll need to create a baseline COCO JSON file with your categories if you don't already have one. It should look something like this:
{
"images": [ ],
"annotations": [ ],
"categories": [
{
"id": 1,
"name": "Car",
"supercategory": "Car"
},
{
"id": 2,
"name": "Person",
"supercategory": "Person"
}
]
}