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_seg_dataset_to_det(seg_imgs_path, seg_labels_path, root_dir, target_coco_ann_path, ignore_classes_idx, thread_i, min_img_size=224, label_after_hook=lambda x: x):
"""
Reference: https://blog.csdn.net/lizaijinsheng/article/details/119889946
NOTE:
Background class should not be considered.
However, if a seg dataset has only one valid class, so that the generated cls dataset also has only one class and
the cls accuracy will be 100% forever. But we do not use the generated cls dataset alone, so it is ok.
"""
assert len(seg_imgs_path) == len(seg_labels_path)
classes_imgs_id_map = {}
coco_ann = {
'categories': [],
"type": "instances",
'images': [],
'annotations': []
}
image_id = 0
ann_id = 0
pbar = tqdm.tqdm(zip(seg_imgs_path, seg_labels_path), total=len(seg_imgs_path),
dynamic_ncols=True, leave=False, desc=f'thread {thread_i}')
for seg_img_path, seg_label_path in pbar:
try:
seg_img = Image.open(seg_img_path)
seg_label = Image.open(seg_label_path).convert('L')
seg_label = np.array(seg_label)
seg_label = label_after_hook(seg_label)
except Exception as e:
print(e)
print(f'file {seg_img_path} error, skip')
exit()
# seg_img = Image.open(seg_img_path)
# seg_label = Image.open(seg_label_path).convert('L')
# seg_label = np.array(seg_label)
image_coco_info = {'file_name': os.path.relpath(seg_img_path, root_dir), 'height': seg_img.height, 'width': seg_img.width,
'id':image_id}
image_id += 1
coco_ann['images'] += [image_coco_info]
this_img_classes = set(seg_label.reshape(-1).tolist())
# print(this_img_classes)
for class_idx in this_img_classes:
if class_idx in ignore_classes_idx:
continue
if class_idx not in classes_imgs_id_map.keys():
classes_imgs_id_map[class_idx] = 0
mask = np.zeros((seg_label.shape[0], seg_label.shape[1]), dtype=np.uint8)
mask[seg_label == class_idx] = 1
mask_without_small = morphology.remove_small_objects(mask, min_size=10, connectivity=2)
label_image = measure.label(mask_without_small)
for region in measure.regionprops(label_image):
bbox = region.bbox # (top, left, bottom, right)
bbox = [bbox[1], bbox[0], bbox[3], bbox[2]] # (left, top, right, bottom)
width, height = bbox[2] - bbox[0], bbox[3] - bbox[1]
if width < min_img_size or height < min_img_size:
continue
# target_cropped_img_path = os.path.join(target_cls_data_dir, str(class_idx),
# f'{classes_imgs_id_map[class_idx]}.{seg_img_path.split(".")[-1]}')
# ensure_dir(target_cropped_img_path)
# seg_img.crop(bbox).save(target_cropped_img_path)
# print(target_cropped_img_path)
# exit()
ann_coco_info = {'area': width*height, 'iscrowd': 0, 'image_id':
image_id - 1, 'bbox': [bbox[0], bbox[1], width, height],
'category_id': class_idx,
'id': ann_id, 'ignore': 0,
'segmentation': []}
ann_id += 1
coco_ann['annotations'] += [ann_coco_info]
classes_imgs_id_map[class_idx] += 1
pbar.set_description(f'# ann: {ann_id}')
coco_ann['categories'] = [
{'id': ci, 'name': f'class_{c}_in_seg'} for ci, c in enumerate(classes_imgs_id_map.keys())
]
c_to_ci_map = {c: ci for ci, c in enumerate(classes_imgs_id_map.keys())}
for ann in coco_ann['annotations']:
ann['category_id'] = c_to_ci_map[
ann['category_id']
]
write_json(target_coco_ann_path, coco_ann, indent=0, backup=True)
write_json(os.path.join(root_dir, 'coco_ann.json'), coco_ann, indent=0, backup=True)
num_cls_imgs = 0
for k, v in classes_imgs_id_map.items():
# print(f'# class {k}: {v + 1}')
num_cls_imgs += v
# print(f'total: {num_cls_imgs}')
return classes_imgs_id_map
from concurrent.futures import ThreadPoolExecutor
# def convert_seg_dataset_to_cls_multi_thread(seg_imgs_path, seg_labels_path, target_cls_data_dir, ignore_classes_idx, num_threads):
# if os.path.exists(target_cls_data_dir):
# shutil.rmtree(target_cls_data_dir)
# assert len(seg_imgs_path) == len(seg_labels_path)
# n = len(seg_imgs_path) // num_threads
# pool = ThreadPoolExecutor(max_workers=num_threads)
# # threads = []
# futures = []
# for thread_i in range(num_threads):
# # thread = threading.Thread(target=convert_seg_dataset_to_cls,
# # args=(seg_imgs_path[thread_i * n: (thread_i + 1) * n],
# # seg_labels_path[thread_i * n: (thread_i + 1) * n],
# # target_cls_data_dir, ignore_classes_idx))
# # threads += [thread]
# future = pool.submit(convert_seg_dataset_to_cls, *(seg_imgs_path[thread_i * n: (thread_i + 1) * n],
# seg_labels_path[thread_i * n: (thread_i + 1) * n],
# target_cls_data_dir, ignore_classes_idx, thread_i))
# futures += [future]
# futures += [
# pool.submit(convert_seg_dataset_to_cls, *(seg_imgs_path[(thread_i + 1) * n: ],
# seg_labels_path[(thread_i + 1) * n: ],
# target_cls_data_dir, ignore_classes_idx, thread_i))
# ]
# for f in futures:
# f.done()
# res = []
# for f in futures:
# res += [f.result()]
# print(res[-1])
# res_dist = {}
# for r in res:
# for k, v in r.items():
# if k in res_dist.keys():
# res_dist[k] += v
# else:
# res_dist[k] = v
# print('results:')
# print(res_dist)
# pool.shutdown()
# import random
# def random_crop_aug(target_dir):
# for class_dir in os.listdir(target_dir):
# class_dir = os.path.join(target_dir, class_dir)
# for img_path in os.listdir(class_dir):
# img_path = os.path.join(class_dir, img_path)
# img = Image.open(img_path)
# w, h = img.width, img.height
# for ri in range(5):
# img.crop(
# [
# random.randint(0, w // 5),
# random.randint(0, h // 5),
# random.randint(w // 5 * 4, w),
# random.randint(h // 5 * 4, h)
# ]
# ).save(
# os.path.join(os.path.dirname(img_path), f'randaug_{ri}_' + os.path.basename(img_path))
# )
# # print(img_path)
# # exit()
def post_ignore_classes(coco_ann_json_path):
# from data.datasets.object_detection.yolox_data_util.api import remap_dataset
# remap_dataset(coco_ann_json_path, [], {})
pass
if __name__ == '__main__':
# SuperviselyPerson
# root_dir = '/data/zql/datasets/supervisely_person_full_20230635/Supervisely Person Dataset'
# images_path, labels_path = [], []
# for p in os.listdir(root_dir):
# if p.startswith('ds'):
# p1 = os.path.join(root_dir, p, 'img')
# images_path += [(p, os.path.join(p1, n)) for n in os.listdir(p1)]
# for dsi, img_p in images_path:
# target_p = os.path.join(root_dir, p, dsi, img_p.split('/')[-1])
# labels_path += [target_p]
# images_path = [i[1] for i in images_path]
# target_coco_ann_path = '/data/zql/datasets/supervisely_person_for_det_task/coco_ann.json'
# if os.path.exists(target_coco_ann_path):
# os.remove(target_coco_ann_path)
# convert_seg_dataset_to_det(
# seg_imgs_path=images_path,
# seg_labels_path=labels_path,
# root_dir=root_dir,
# target_coco_ann_path=target_coco_ann_path,
# ignore_classes_idx=[0, 2],
# # num_threads=8
# thread_i=0
# )
# random_crop_aug('/data/zql/datasets/supervisely_person_for_cls_task')
# GTA5
# root_dir = '/data/zql/datasets/GTA-ls-copy/GTA5'
# images_path, labels_path = [], []
# for p in os.listdir(os.path.join(root_dir, 'images')):
# p = os.path.join(root_dir, 'images', p)
# if not p.endswith('png'):
# continue
# images_path += [p]
# labels_path += [p.replace('images', 'labels_gt')]
# target_coco_ann_path = '/data/zql/datasets/gta5_for_det_task/coco_ann.json'
# if os.path.exists(target_coco_ann_path):
# os.remove(target_coco_ann_path)
# """
# [
# 'road', 'sidewalk', 'building', 'wall',
# 'fence', 'pole', 'light', 'sign',
# 'vegetation', 'terrain', 'sky', 'people', # person
# 'rider', 'car', 'truck', 'bus', 'train',
# 'motocycle', 'bicycle'
# ]
# """
# need_classes_idx = [13, 15]
# convert_seg_dataset_to_det(
# seg_imgs_path=images_path,
# seg_labels_path=labels_path,
# root_dir=root_dir,
# target_coco_ann_path=target_coco_ann_path,
# ignore_classes_idx=[i for i in range(20) if i not in need_classes_idx],
# thread_i=0
# )
# from data.datasets.object_detection.yolox_data_util.api import remap_dataset
# new_coco_ann_json_path = remap_dataset('/data/zql/datasets/GTA-ls-copy/GTA5/coco_ann.json', [-1], {0: 0, 1:-1, 2:-1, 3: 1, 4:-1, 5:-1})
# print(new_coco_ann_json_path)
# cityscapes
# root_dir = '/data/zql/datasets/cityscape/'
# def _get_target_suffix(mode: str, target_type: str) -> str:
# if target_type == 'instance':
# return '{}_instanceIds.png'.format(mode)
# elif target_type == 'semantic':
# return '{}_labelIds.png'.format(mode)
# elif target_type == 'color':
# return '{}_color.png'.format(mode)
# else:
# return '{}_polygons.json'.format(mode)
# images_path, labels_path = [], []
# split = 'train'
# images_dir = os.path.join(root_dir, 'leftImg8bit', split)
# targets_dir = os.path.join(root_dir, 'gtFine', split)
# for city in os.listdir(images_dir):
# img_dir = os.path.join(images_dir, city)
# target_dir = os.path.join(targets_dir, city)
# for file_name in os.listdir(img_dir):
# target_types = []
# for t in ['semantic']:
# target_name = '{}_{}'.format(file_name.split('_leftImg8bit')[0],
# _get_target_suffix('gtFine', t))
# target_types.append(os.path.join(target_dir, target_name))
# images_path.append(os.path.join(img_dir, file_name))
# labels_path.append(target_types[0])
# # print(images_path[0: 5], '\n', labels_path[0: 5])
# target_coco_ann_path = '/data/zql/datasets/cityscape/coco_ann.json'
# # if os.path.exists(target_dir):
# # shutil.rmtree(target_dir)
# need_classes_idx = [26, 28]
# convert_seg_dataset_to_det(
# seg_imgs_path=images_path,
# seg_labels_path=labels_path,
# root_dir=root_dir,
# target_coco_ann_path=target_coco_ann_path,
# ignore_classes_idx=[i for i in range(80) if i not in need_classes_idx],
# # num_threads=8
# thread_i=0
# )
# import shutil
# ignore_target_dir = '/data/zql/datasets/cityscapes_for_cls_task_ignored'
# ignore_label = 255
# raw_idx_map_in_y_transform = {-1: ignore_label, 0: ignore_label, 1: ignore_label, 2: ignore_label,
# 3: ignore_label, 4: ignore_label, 5: ignore_label, 6: ignore_label,
# 7: 0, 8: 1, 9: ignore_label, 10: ignore_label, 11: 2, 12: 3, 13: 4,
# 14: ignore_label, 15: ignore_label, 16: ignore_label, 17: 5,
# 18: ignore_label, 19: 6, 20: 7, 21: 8, 22: 9, 23: 10, 24: 11, 25: 12, 26: 13, 27: 14,
# 28: 15, 29: ignore_label, 30: ignore_label, 31: 16, 32: 17, 33: 18}
# ignore_classes_idx = [k for k, v in raw_idx_map_in_y_transform.items() if v == ignore_label]
# ignore_classes_idx = sorted(ignore_classes_idx)
# for class_dir in os.listdir(target_dir):
# if int(class_dir) in ignore_classes_idx:
# continue
# shutil.move(
# os.path.join(target_dir, class_dir),
# os.path.join(ignore_target_dir, class_dir)
# )
# else:
# shutil.move(
# os.path.join(target_dir, class_dir),
# os.path.join(target_dir, str(raw_idx_map_in_y_transform[int(class_dir)]))
# )
# continue
# print(class_dir)
# exit()
# baidu person
# root_dir = '/data/zql/datasets/baidu_person/clean_images/'
# images_path, labels_path = [], []
# for p in os.listdir(os.path.join(root_dir, 'images')):
# images_path += [os.path.join(root_dir, 'images', p)]
# labels_path += [os.path.join(root_dir, 'profiles', p.split('.')[0] + '-profile.jpg')]
# target_dir = '/data/zql/datasets/baiduperson_for_cls_task'
# # if os.path.exists(target_dir):
# # shutil.rmtree(target_dir)
# def label_after_hook(x):
# x[x > 1] = 1
# return x
# convert_seg_dataset_to_det(
# seg_imgs_path=images_path,
# seg_labels_path=labels_path,
# root_dir=root_dir,
# target_coco_ann_path='/data/zql/datasets/baidu_person/clean_images/coco_ann_zql.json',
# ignore_classes_idx=[1],
# # num_threads=8
# thread_i=1,
# min_img_size=224,
# label_after_hook=label_after_hook
# )
# from data.visualize import visualize_classes_in_object_detection
# from data import get_dataset
# d = get_dataset('CityscapesDet', '/data/zql/datasets/cityscape/', 'val', None, [], None)
# visualize_classes_in_object_detection(d, {'car': 0, 'bus': 1}, {}, 'debug.png')
# d = get_dataset('GTA5Det', '/data/zql/datasets/GTA-ls-copy/GTA5', 'val', None, [], None)
# visualize_classes_in_object_detection(d, {'car': 0, 'bus': 1}, {}, 'debug.png')
# d = get_dataset('BaiduPersonDet', '/data/zql/datasets/baidu_person/clean_images/', 'val', None, [], None)
# visualize_classes_in_object_detection(d, {'person': 0}, {}, 'debug.png')