diff --git a/data_utils/bboxCropper/README.md b/data_utils/bboxCropper/README.md new file mode 100644 index 0000000..911efd0 --- /dev/null +++ b/data_utils/bboxCropper/README.md @@ -0,0 +1,18 @@ +### Cropping tool for dataset images +Before running **bboxCropper.py**: update **config.cfg** with your data. + + +#### Inputs +Dataset, located at **DATASET_PATH**. + + +#### Outputs +- creates folders **CROPS_FOLDER_NAME**, **FRAMES_FOLDER_NAME**, **MASKS_FOLDER_NAME** in dataset folder. +- crops square of **CROP_SIZE*****CROP_SIZE** around bbox center on the image, with random shift, and saves it to **CROPS_FOLDER_NAME** folder. +- cuts out bbox content from this crop, saves it to **FRAMES_FOLDER_NAME** folder. +- creates binary mask of the size of the crop, with True pixels under bbox and all other zeros, saves it to **MASKS_FOLDER_NAME** folder. + +#### About masks +- by default, mask pixels are True in places where image inpainting required, others are False. For mask inversion: set INVERT_MASKS in cfg file. + +![Example](https://github.com/lacmus-foundation/lacmus/blob/master/data_utils/bboxCropper/screenshot.PNG) diff --git a/data_utils/bboxCropper/bboxCropper.py b/data_utils/bboxCropper/bboxCropper.py new file mode 100644 index 0000000..2fe7d3e --- /dev/null +++ b/data_utils/bboxCropper/bboxCropper.py @@ -0,0 +1,146 @@ +import cv2 +from pathlib import Path +import xml.etree.ElementTree as ET +import numpy as np +import os +import random + +def main(): + + # open CFG file, define paths and crops shapes + + config = open('config.cfg') + for line in config: + if line.startswith('CROP_SIZE'): + crop_size = int(line.split('=')[1].strip().replace('\n','')) + if line.startswith('DATASET_PATH'): + dataset_path = Path(line.split('=')[1].strip().replace('\n','')) + if line.startswith('CROPS_FOLDER'): + crops_folder = line.split('=')[1].strip().replace('\n','') + if line.startswith('FRAMES_FOLDER'): + frames_folder = line.split('=')[1].strip().replace('\n','') + if line.startswith('MASKS_FOLDER'): + masks_folder = line.split('=')[1].strip().replace('\n','') + if line.startswith('INVERT_MASKS'): + invert_masks = line.split('=')[1].strip().replace('\n','') + if invert_masks.lower() in ['false']: + invert_masks = False + else: + invert_masks = True + + images_folder = 'JPEGImages' + annotations_folder = 'Annotations' + config.close() + + print('Dataset location: ', dataset_path) + + # Create folders for outputs + if frames_folder not in os.listdir(dataset_path): + os.mkdir(Path(dataset_path, frames_folder)) + print('Created folder: ', Path(dataset_path, frames_folder)) + if crops_folder not in os.listdir(dataset_path): + os.mkdir(Path(dataset_path, crops_folder)) + print('Created folder: ', Path(dataset_path, crops_folder)) + if masks_folder not in os.listdir(dataset_path): + os.mkdir(Path(dataset_path, masks_folder)) + print('Created folder: ', Path(dataset_path, masks_folder)) + + print('Processing started...') + + # parse each annotation file + + n_files = len(os.listdir(Path(dataset_path, annotations_folder))) + passed_files=1 + + for filename in os.listdir(Path(dataset_path, annotations_folder)): + + if not filename.endswith('.xml'): continue + + fullname = Path(dataset_path, annotations_folder, filename) + tree = ET.parse(fullname) + root = tree.getroot() + bbox_num = 0 + + img = cv2.imread(str(Path(dataset_path, images_folder, filename[:-3]+'jpg'))) + + for rec in root: + + # get source image size + if rec.tag == 'size': + height = int(rec.findtext('height')) + width = int(rec.findtext('width')) + + # list all available bboxes + if rec.tag == 'object': + for box in rec: + if box.tag=='bndbox': + + # get initial bbox corners + ymin = int(box.findtext('ymin')) + ymax = int(box.findtext('ymax')) + xmin = int(box.findtext('xmin')) + xmax = int(box.findtext('xmax')) + + # calculate necessary padding to get crop of crop_size + padding_w = int((crop_size - (xmax - xmin))/2.) + padding_h = int((crop_size - (ymax - ymin))/2.) + + # get random shift within 25% of crop_size from bbox center + random_dx = int((random.random()-.5)*.5*crop_size) + random_dy = int((random.random()-.5)*.5*crop_size) + + # calculate crop corners + new_xmin = xmin - padding_w + random_dx + new_xmax = xmax + padding_w + random_dx + new_ymin = ymin - padding_h + random_dy + new_ymax = ymax + padding_h + random_dy + + # do not proceed if crop is outside of image + if (new_xmin<1 or new_xmax>width-1 or new_ymin<1 or new_ymax>height-1):continue + + dx = new_xmax - new_xmin + dy = new_ymax - new_ymin + + # correct crop corners to get exact crop_size + if dx