File size: 1,933 Bytes
b84549f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
from data import ABDataset
from utils.common.data_record import read_json, write_json
from PIL import Image
import os
from utils.common.file import ensure_dir
import numpy as np
from itertools import groupby
from skimage import morphology, measure
from PIL import Image
from scipy import misc
import tqdm
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
import shutil
def convert_det_dataset_to_det(coco_ann_json_path, data_dir, target_data_dir, min_img_size=224):
coco_ann = read_json(coco_ann_json_path)
img_id_to_path = {}
for img in coco_ann['images']:
img_id_to_path[img['id']] = os.path.join(data_dir, img['file_name'])
classes_imgs_id_map = {}
for ann in tqdm.tqdm(coco_ann['annotations'], total=len(coco_ann['annotations']), dynamic_ncols=True):
img_id = ann['image_id']
img_path = img_id_to_path[img_id]
img = Image.open(img_path)
bbox = ann['bbox']
if bbox[2] < min_img_size or bbox[3] < min_img_size:
continue
bbox = [bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]]
class_idx = str(ann['category_id'])
if class_idx not in classes_imgs_id_map.keys():
classes_imgs_id_map[class_idx] = 0
target_cropped_img_path = os.path.join(target_data_dir, class_idx,
f'{classes_imgs_id_map[class_idx]}.{img_path.split(".")[-1]}')
classes_imgs_id_map[class_idx] += 1
ensure_dir(target_cropped_img_path)
img.crop(bbox).save(target_cropped_img_path)
if __name__ == '__main__':
convert_det_dataset_to_det(
coco_ann_json_path='/data/zql/datasets/coco2017/train2017/coco_ann.json',
data_dir='/data/zql/datasets/coco2017/train2017',
target_data_dir='/data/zql/datasets/coco2017_for_cls_task',
min_img_size=224
) |