utils_data.py

Go to the documentation of this file.

import copy
import json
import os
import platform
import re
 
 
from smdataset.version_info import *
 
 
def check_json_format(path_format, json_file_name):
    """
    json 파일 이름이 해당 포맷에 부합하는지 확인하는 기능
    args:
        path_format : 확인하려는 포맷
        json_file_name : 확인하려는 파일 이름
    return (bool): 부합여부
    """
    try :
        for pf in path_format:
            filename_re = re.compile(pf)
            str_format = filename_re.search(json_file_name).group()
    except:
        return False
    
    return True
 
def check_json_version(json_file_name):
    """
    json_file_name 이름으로 버전을 확인하는 기능
    return (int): json 파일의 버전
    """
    for key, path_format in PATH_FORMAT.items():
        if check_json_format(path_format, json_file_name):
            return key        
    return None
 
def join(path_list):
    """
    os.path.join 기능 커스텀, 플랫폼에 따라 다르게 동작함
    args:
        path_list (list): 이어붙이려는 폴더 및 경로 리스트
    return (str) : path_list를 결합한 경로
    """
    if platform.system().lower() == 'windows':
        return "/".join(path_list)
    else:
        path_list_tmp = [path_list[0]]
        for p in path_list[1:]:
            p_list = p.split('/')
            if '' in p_list:
                p_list.remove('')
            path_list_tmp.extend(p_list)
        return os.path.join(*path_list_tmp)
 
def make_dataset_folder_path(json_file_name, ret_img_folder=False, json_version=None):
    """
    json파일 이름으로 json 파일이 저장된 폴더와 이미지 폴더를 만드는 기능
    args:
        json_file_name (str): json 파일 이름 
        ret_img_folder (bool): 이미지 폴더 반환 여부
        json_version (bool): json 파일 버전, 설정하지 않으면 json 파일 이름으로 json 파일 버전을 확인함
    return (str): json 파일이 저장된 폴더경로, 이미지 경로
    """
    if json_version is None:
        json_version = check_json_version(json_file_name)
        
    path_format = PATH_FORMAT[json_version]
    dataset_path_list = []
    for pf in path_format:
        filename_re = re.compile(pf)
        str_format = filename_re.search(json_file_name).group()
        dataset_path_list.append(str_format)
        
    json_path_list = dataset_path_list[:2]
    json_folder = join(json_path_list)
    image_folder = join(dataset_path_list)
    
    
    if ret_img_folder:
        return json_folder, image_folder
    else:
        return json_folder
    
def make_image_base_path(json_version, anno_path):
    """
    데이터셋이 저장된 기본 경로를 생성하는 기능
    ver 11
        json 파일 저장 경로 : base_path/annotation/P/C/json_name
        이미지 파일 저장 경로 : base_path/images/P/C/A/R/img_name
    args:
        json_version (int): json 파일 버전
        anno_path (str): json 파일 경로
    return:
        base_path : 데이터(json, 이미지)가 저장된 폴더
        image_base_path : 이미지가 저장된 기본 경로(base_path/images)
    """
    split_dataset = anno_path.split("/")
    json_file_name = split_dataset[-1]
    if json_version == 10:
        base_path_list = split_dataset[:-3]
        image_base_path_list = copy.deepcopy(base_path_list)
    elif json_version == 11:        
        base_path_list = split_dataset[:-4]
        image_base_path_list = copy.deepcopy(base_path_list)
        image_base_path_list.append('images')
 
    path_format = PATH_FORMAT[json_version]
    for pf in path_format:
        filename_re = re.compile(pf)
        str_format = filename_re.search(json_file_name).group()
        image_base_path_list.append(str_format)
    
    if platform.system().lower() == 'windows':
        image_base_path = "/".join(image_base_path_list)
        base_path = "/".join(base_path_list)
    else:
        image_base_path = os.path.join("/", *image_base_path_list)
        base_path = os.path.join("/", *base_path_list)
 
    return base_path, image_base_path
 
def make_image_folder(json_version, anno_path):
    """
    어노테이션 파일에 대응하는 이미지가 저장된 경로를 생성하는 기능
    image_base_path 뒤에 붙는 폴더
    args:
        json_version (int): json 파일 버전
        anno_path (str): json 파일 경로
    return:
        image_folder_path : 이미지가 저장된 폴더
    """
    split_dataset = anno_path.split("/")
    json_file_name = split_dataset[-1]
 
    image_folder_path_list = []
    path_format = PATH_FORMAT[json_version]
    for pf in path_format:
        filename_re = re.compile(pf)
        str_format = filename_re.search(json_file_name).group()
        image_folder_path_list.append(str_format)
 
    if platform.system().lower() == 'windows':
        image_folder_path = "/".join(image_folder_path_list)
    else:
        image_folder_path = os.path.join(*image_folder_path_list)
 
    return image_folder_path
 
 
def load_annotation(anno_path):
    """
    어노테이션 파일을 로드해서 기본 정보를 반환하는 기능
    args:
        anno_path (str): json 파일 경로
    return:
        json_version (int): json 파일 버전
        label_map_version  (int): 레이블맵 버전
        base_path (str): 데이터가 저장된 기본 경로
        image_base_path (str): 이미지가 저장된 기본 경로
        dataset (dict): 어노테이션 데이터
    """
    with open(anno_path, "r") as json_file:
        dataset = json.load(json_file)
 
    json_version = int(float(dataset["info"]["version"])*10)
    label_map_version = int(float(dataset["info"]["label_map_version"])*10)
    base_path, image_base_path = make_image_base_path(json_version, anno_path)
    dataset = check_and_init_dataset(dataset, CHECK_LIST[json_version]["keys"], CHECK_LIST[json_version]["default_value"])
 
    return json_version, label_map_version, base_path, image_base_path, dataset
 
def save_json(path, dataset):
    """
    어노테이션 데이터를 저장하는 기능
    args:
        path (str): 저장 경로
        dataset (dict): 저장하는 데이터
    """
    try:
        with open(path, "w") as json_file:
            json.dump(dataset, json_file)
    except Exception  as e:
        print(e)
 
def load_json(path):
    """
    어노테이션 데이터를 로드하는 기능
    args:
        path (str): 어노테이션 경로
    return (dict): 어노테이션 데이터
    """
    try:
        with open(path, "r") as json_file:
            data = json.load(json_file)
    except Exception  as e:
        data = None
 
    return data
    
def check_and_init_dataset(dataset, check_list, default_value):
    """
    어노테이션 데이터를 확인해서 필요한 정보를 초기화 하는 기능
    args:
        dataset (dict): 어노테이션 데이터
        check_list (list[str]): 확인하는 정보, 어노테이션 데이터에 해당 키가 없으면 default_value로 초기화 시켜줌
        default_value : 초기화 값
    return (dict): 어노테이션 데이터
    """
    for idx, anno in enumerate(dataset['annotations']):
        for check_key, init_val in zip(check_list, default_value):
            if not (check_key in anno):
                dataset['annotations'][idx][check_key] = init_val
    return dataset
 
def check_and_init_dictionary(dictionary, check_list, default_value): 
    """
    데이터를 확인해서 필요한 정보를 초기화 하는 기능
    args:
        dictionary (dict): 데이터
        check_list (list[str]): 확인하는 정보, 데이터에 해당 키가 없으면 default_value로 초기화 시켜줌
        default_value : 초기화 값
    return (dict): 데이터
    """   
    for key, val in zip(check_list, default_value):
        if not key in dictionary:
            dictionary[key] = val
    return dictionary
 
def make_dummy_annotation(anno_id, image_id, track_id=0):
    """
    더미데이터를 생성하는 기능
    args:
        anno_id (int): 어노테이션 아이디
        image_id (int): 이미지 아이디
        track_id (int): 추적 아이디
    return (dict): 더미 데이터
    """
    tmp_annotation = {
        "id": anno_id,
        "image_id": image_id,
        "category_id": 1,
        "iscrowd": 0,
        "keypoints": [
            [178.87325418994442, 168.6456005586591, 1.0],
            [197.7807262569836, 135.67562849162016, 1.0],
            [149.12639664804487, 140.08938547486025, 1.0],
            [241.84322625698292, 161.00872905027938, 1.0],
            [121.84322625698292, 161.00872905027938, 1.0],
            [218.3100558659221, 205.59811452513907, 1.0],
            [151.26187150837973, 209.51885474860308, 1.0],
            [274.60125698323964, 257.7213687150835, 1.0],
            [112.21508379888269, 268.0677374301673, 1.0],
            [320.4018854748606, 337.25139664804476, 1.0],
            [87.6002094972066, 346.36592178770945, 1.0],
            [246.41445530726241, 379.5006983240222, 1.0],
            [175.2807262569836, 380.26012569832403, 1.0],
            [283.3463687150843, 503.20949720670353, 1.0],
            [168.83449720670387, 514.8495111731842, 1.0],
            [316.25453910614533, 637.7988826815641, 1.0],
            [157.21508379888246, 641.4738128491614, 1.0],
        ],
        "area": 903482.25,
        "bbox": [32.3882681564246, 17.625698324022324, 391.0, 710.0],
        "ageclass": 0,
        "genderclass": 0,
        "upperclass": 0,
        "uppercolor": [255.0, 255.0, 255.0],
        "lowerclass": 0,
        "lowercolor": [255.0, 255.0, 255.0],
        "hatvisible": 0,
        "hatcolor": [-1.0, -1.0, -1.0],
        "bagvisible": 0,
        "bagcolor": [-1.0, -1.0, -1.0],
        "action_id": {'action': 0, 'pose':0, 'hand':0, 'foot':0},
        "mutual_action": 0,
        "mutual_action_target": [],
        "abnormal_id": 0,
        "dangerzone": [0, 0, 0, 0, 0, 0, 0, 0],
        "track_id": track_id,
        "ismodify": 0,
    }
    return tmp_annotation
 
def load_scenario(json_version, scenario_num):
    path = os.path.join(SCENARIO_FILE_BASE, SCENARIO_FILE[json_version])
    with open(path, "r", encoding="UTF8") as file:
        lines = file.readlines()
        lines = [line.rstrip() for line in lines]
 
    sceneLine = []
    for i in lines:
        tmp_i = i.split(".")
        if tmp_i[0] == scenario_num:
            sceneLine.append(tmp_i[1])
    return sceneLine
 
def load_label_map(label_map_ver, label_key):
    file_name = LABEL_FILE_INFO[label_map_ver][label_key]['file_name']
    splitStr = LABEL_FILE_INFO[label_map_ver][label_key]['keys']
    path = os.path.join(LABEL_BASE, file_name)
 
    labels = dict.fromkeys(splitStr)
    with open(path, "r", encoding='UTF-8') as txt_file:
        lines = txt_file.readlines()
 
        key = splitStr[0]
 
        for line in lines:
            line = line.strip()
            if len(line) < 1:
                continue
 
            if line in splitStr:
                key = line
                labels[key] = []
                continue
            
            labels[key].append(line)
 
    return labels, splitStr
 
 
def search_json_by_path_format(folder, data_root = None):
    if data_root is None:
        data_root = folder
        
    fileEx = '.json'
 
    json_list = []
    for path, dirs, files in os.walk(folder):
        if path.split('/')[-1].startswith('.'):
            continue
        json_list.extend([ file for file in files if file.endswith(fileEx)])
 
    json_list.sort()
 
    json_path = []
    for file in json_list:            
        json_path.append(join( [data_root, make_dataset_folder_path(file), file] ))
 
    return json_list, json_path