Spaces:
Runtime error
Runtime error
File size: 3,988 Bytes
7629b39 |
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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
import os
from matplotlib import pyplot as plt
import glob
import json
import numpy as np
from scipy.io import loadmat
from csv import DictReader
from collections import OrderedDict
from pycocotools.mask import decode as decode_RLE
import sys
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..'))
from configs.dataset_path_configs import IMG_V12_DIR, JSON_V12_DIR, STAN_V12_TRAIN_LIST_DIR, STAN_V12_VAL_LIST_DIR, STAN_V12_TEST_LIST_DIR
def get_img_dir(V12):
if V12:
return IMG_V12_DIR
else:
return IMG_DIR
def get_seg_from_entry(entry):
"""Given a .json entry, returns the binary mask as a numpy array"""
rle = {
"size": [entry['img_height'], entry['img_width']],
"counts": entry['seg']}
decoded = decode_RLE(rle)
return decoded
def full_animal_visible(seg_data):
if seg_data[0, :].sum() == 0 and seg_data[seg_data.shape[0]-1, :].sum() == 0 and seg_data[:, 0].sum() == 0 and seg_data[:, seg_data.shape[1]-1].sum() == 0:
return True
else:
return False
def load_train_and_test_lists(train_list_dir=None , test_list_dir=None):
""" returns sets containing names such as 'n02085620-Chihuahua/n02085620_5927.jpg' """
# train data
train_list_mat = loadmat(train_list_dir)
train_list = []
for ind in range(0, train_list_mat['file_list'].shape[0]):
name = train_list_mat['file_list'][ind, 0][0]
train_list.append(name)
# test data
test_list_mat = loadmat(test_list_dir)
test_list = []
for ind in range(0, test_list_mat['file_list'].shape[0]):
name = test_list_mat['file_list'][ind, 0][0]
test_list.append(name)
return train_list, test_list
def _filter_dict(t_list, j_dict, n_kp_min=4):
""" should only be used by load_stanext_json_as_dict() """
out_dict = {}
for sample in t_list:
if sample in j_dict.keys():
n_kp = np.asarray(j_dict[sample]['joints'])[:, 2].sum()
if n_kp >= n_kp_min:
out_dict[sample] = j_dict[sample]
return out_dict
def load_stanext_json_as_dict(split_train_test=True, V12=True):
# load json into memory
if V12:
with open(JSON_V12_DIR) as infile:
json_data = json.load(infile)
# with open(JSON_V12_DIR) as infile: json_data = json.load(infile, object_pairs_hook=OrderedDict)
else:
with open(JSON_DIR) as infile:
json_data = json.load(infile)
# convert json data to a dictionary of img_path : all_data, for easy lookup
json_dict = {i['img_path']: i for i in json_data}
if split_train_test:
if V12:
train_list_numbers = np.load(STAN_V12_TRAIN_LIST_DIR)
val_list_numbers = np.load(STAN_V12_VAL_LIST_DIR)
test_list_numbers = np.load(STAN_V12_TEST_LIST_DIR)
train_list = [json_data[i]['img_path'] for i in train_list_numbers]
val_list = [json_data[i]['img_path'] for i in val_list_numbers]
test_list = [json_data[i]['img_path'] for i in test_list_numbers]
train_dict = _filter_dict(train_list, json_dict, n_kp_min=4)
val_dict = _filter_dict(val_list, json_dict, n_kp_min=4)
test_dict = _filter_dict(test_list, json_dict, n_kp_min=4)
return train_dict, test_dict, val_dict
else:
train_list, test_list = load_train_and_test_lists(train_list_dir=STAN_ORIG_TRAIN_LIST_DIR , test_list_dir=STAN_ORIG_TEST_LIST_DIR)
train_dict = _filter_dict(train_list, json_dict)
test_dict = _filter_dict(test_list, json_dict)
return train_dict, test_dict, None
else:
return json_dict
def get_dog(json_dict, name, img_dir=None): # (json_dict, name, img_dir=IMG_DIR)
""" takes the name of a dog, and loads in all the relevant information as a dictionary:
dict_keys(['img_path', 'img_width', 'img_height', 'joints', 'img_bbox',
'is_multiple_dogs', 'seg', 'img_data', 'seg_data'])
img_bbox: [x0, y0, width, height] """
data = json_dict[name]
# load img
img_data = plt.imread(os.path.join(img_dir, data['img_path']))
# load seg
seg_data = get_seg_from_entry(data)
# add to output
data['img_data'] = img_data # 0 to 255
data['seg_data'] = seg_data # 0: bg, 1: fg
return data
|