data/convert_seg_dataset_to_cls.py

from data import ABDataset
from utils.common.data_record import read_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_cls(seg_imgs_path, seg_labels_path, target_cls_data_dir, 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 = {}
    
    for seg_img_path, seg_label_path in tqdm.tqdm(zip(seg_imgs_path, seg_labels_path), total=len(seg_imgs_path), 
                                                   dynamic_ncols=True, leave=False, desc=f'thread {thread_i}'):

        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)
            
        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()
                
                classes_imgs_id_map[class_idx] += 1

    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()
            

if __name__ == '__main__':
    # SuperviselyPerson
    # root_dir = '/data/zql/datasets/supervisely_person/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_dir = '/data/zql/datasets/supervisely_person_for_cls_task'
    # if os.path.exists(target_dir):
    #     shutil.rmtree(target_dir)
    # convert_seg_dataset_to_cls(
    #     seg_imgs_path=images_path,
    #     seg_labels_path=labels_path,
    #     target_cls_data_dir=target_dir,
    #     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_dir = '/data/zql/datasets/gta5_for_cls_task'
    # if os.path.exists(target_dir):
    #     shutil.rmtree(target_dir)
    
    # convert_seg_dataset_to_cls(
    #     seg_imgs_path=images_path,
    #     seg_labels_path=labels_path,
    #     target_cls_data_dir=target_dir,
    #     ignore_classes_idx=[],
    #     thread_i=0
    # )
    
    # 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_dir = '/data/zql/datasets/cityscapes_for_cls_task'
    # if os.path.exists(target_dir):
    #     shutil.rmtree(target_dir)
    # convert_seg_dataset_to_cls(
    #     seg_imgs_path=images_path,
    #     seg_labels_path=labels_path,
    #     target_cls_data_dir=target_dir,
    #     ignore_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_cls(
    #     seg_imgs_path=images_path,
    #     seg_labels_path=labels_path,
    #     target_cls_data_dir=target_dir,
    #     ignore_classes_idx=[1],
    #     # num_threads=8
    #     thread_i=1,
    #     min_img_size=224,
    #     label_after_hook=label_after_hook
    # )