File size: 8,439 Bytes
d4b77ac |
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 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 |
import os
import argparse
import shutil
from glob import glob
import numpy as np
from PIL import Image
from utils.logging_config import logger
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
'-v', '--video_dir',
type=str,
help="Video directory name"
)
parser.add_argument(
'-fl', '--flow_dir',
type=str,
help="Optical flow ground truth directory name"
)
parser.add_argument(
'-od', '--output_dir',
type=str,
help="Output directory name"
)
parser.add_argument(
'-o', '--output_filename',
type=str,
help="Output output filename"
)
args = parser.parse_args()
return args
def make_dirs(dir_name):
if not os.path.exists(dir_name):
os.makedirs(dir_name)
logger.info(f"Directory {dir_name} made")
ensure_dir = make_dirs
def make_dir_under_root(root_dir, name):
full_dir_name = os.path.join(root_dir, name)
make_dirs(full_dir_name)
return full_dir_name
def rm_dirs(dir_name, ignore_errors=False):
if os.path.exists(dir_name):
shutil.rmtree(dir_name, ignore_errors)
logger.info(f"Directory {dir_name} removed")
def read_dirnames_under_root(root_dir, skip_list=[]):
dirnames = [
name for i, name in enumerate(sorted(os.listdir(root_dir)))
if (os.path.isdir(os.path.join(root_dir, name))
and name not in skip_list
and i not in skip_list)
]
logger.info(f"Reading directories under {root_dir}, exclude {skip_list}, num: {len(dirnames)}")
return dirnames
def bbox_offset(bbox, location):
x0, y0 = location
(x1, y1), (x2, y2) = bbox
return ((x1 + x0, y1 + y0), (x2 + x0, y2 + y0))
def cover2_bbox(bbox1, bbox2):
x1 = min(bbox1[0][0], bbox2[0][0])
y1 = min(bbox1[0][1], bbox2[0][1])
x2 = max(bbox1[1][0], bbox2[1][0])
y2 = max(bbox1[1][1], bbox2[1][1])
return ((x1, y1), (x2, y2))
def extend_r_bbox(bbox, w, h, r):
(x1, y1), (x2, y2) = bbox
x1 = max(x1 - r, 0)
x2 = min(x2 + r, w)
y1 = max(y1 - r, 0)
y2 = min(y2 + r, h)
return ((x1, y1), (x2, y2))
def mean_squared_error(A, B):
return np.square(np.subtract(A, B)).mean()
def bboxes_to_mask(size, bboxes):
mask = Image.new("L", size, 255)
mask = np.array(mask)
for bbox in bboxes:
try:
(x1, y1), (x2, y2) = bbox
except Exception:
(x1, y1, x2, y2) = bbox
mask[y1:y2, x1:x2] = 0
mask = Image.fromarray(mask.astype("uint8"))
return mask
def get_extended_from_box(img_size, box, patch_size):
def _decide_patch_num(box_width, patch_size):
num = np.ceil(box_width / patch_size).astype(np.int)
if (num * patch_size - box_width) < (patch_size // 2):
num += 1
return num
x1, y1 = box[0]
x2, y2 = box[1]
new_box = (x1, y1, x2 - x1, y2 - y1)
box_x_start, box_y_start, box_x_size, box_y_size = new_box
patchN_x = _decide_patch_num(box_x_size, patch_size)
patchN_y = _decide_patch_num(box_y_size, patch_size)
extend_x = (patch_size * patchN_x - box_x_size) // 2
extend_y = (patch_size * patchN_y - box_y_size) // 2
img_x_size = img_size[0]
img_y_size = img_size[1]
x_start = max(0, box_x_start - extend_x)
x_end = min(box_x_start - extend_x + patchN_x * patch_size, img_x_size)
y_start = max(0, box_y_start - extend_y)
y_end = min(box_y_start - extend_y + patchN_y * patch_size, img_y_size)
x_start, y_start, x_end, y_end = int(x_start), int(y_start), int(x_end), int(y_end)
extented_box = ((x_start, y_start), (x_end, y_end))
return extented_box
# code modified from https://github.com/WonwoongCho/Generative-Inpainting-pytorch/blob/master/util.py
def spatial_discounting_mask(mask_width, mask_height, discounting_gamma):
"""Generate spatial discounting mask constant.
Spatial discounting mask is first introduced in publication:
Generative Image Inpainting with Contextual Attention, Yu et al.
Returns:
np.array: spatial discounting mask
"""
gamma = discounting_gamma
mask_values = np.ones((mask_width, mask_height), dtype=np.float32)
for i in range(mask_width):
for j in range(mask_height):
mask_values[i, j] = max(
gamma**min(i, mask_width - i),
gamma**min(j, mask_height - j))
return mask_values
def bboxes_to_discounting_loss_mask(img_size, bboxes, discounting_gamma=0.99):
mask = np.zeros(img_size, dtype=np.float32) + 0.5
for bbox in bboxes:
try:
(x1, y1), (x2, y2) = bbox
except Exception:
(x1, y1, x2, y2) = bbox
mask_width, mask_height = y2 - y1, x2 - x1
mask[y1:y2, x1:x2] = spatial_discounting_mask(mask_width, mask_height, discounting_gamma)
return mask
def find_proper_window(image_size, bbox_point):
'''
parameters:
image_size(2-tuple): (height, width)
bbox_point(2-2-tuple): (first_point, last_point)
return values:
window left-up point, (2-tuple)
window right-bottom point, (2-tuple)
'''
bbox_height = bbox_point[1][0] - bbox_point[0][0]
bbox_width = bbox_point[1][1] - bbox_point[0][1]
window_size = min(
max(bbox_height, bbox_width) * 2,
image_size[0], image_size[1]
)
# Limit min window size due to the requirement of VGG16
window_size = max(window_size, 32)
horizontal_span = window_size - (bbox_point[1][1] - bbox_point[0][1])
vertical_span = window_size - (bbox_point[1][0] - bbox_point[0][0])
top_bound, bottom_bound = bbox_point[0][0] - \
vertical_span // 2, bbox_point[1][0] + vertical_span // 2
left_bound, right_bound = bbox_point[0][1] - \
horizontal_span // 2, bbox_point[1][1] + horizontal_span // 2
if left_bound < 0:
right_bound += 0 - left_bound
left_bound += 0 - left_bound
elif right_bound > image_size[1]:
left_bound -= right_bound - image_size[1]
right_bound -= right_bound - image_size[1]
if top_bound < 0:
bottom_bound += 0 - top_bound
top_bound += 0 - top_bound
elif bottom_bound > image_size[0]:
top_bound -= bottom_bound - image_size[0]
bottom_bound -= bottom_bound - image_size[0]
return (top_bound, left_bound), (bottom_bound, right_bound)
def drawrect(drawcontext, xy, outline=None, width=0, partial=None):
(x1, y1), (x2, y2) = xy
if partial is None:
points = (x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1)
drawcontext.line(points, fill=outline, width=width)
else:
drawcontext.line([(x1, y1), (x1, y1 + partial)], fill=outline, width=width)
drawcontext.line([(x1 + partial, y1), (x1, y1)], fill=outline, width=width)
drawcontext.line([(x2, y1), (x2, y1 + partial)], fill=outline, width=width)
drawcontext.line([(x2, y1), (x2 - partial, y1)], fill=outline, width=width)
drawcontext.line([(x1, y2), (x1 + partial, y2)], fill=outline, width=width)
drawcontext.line([(x1, y2), (x1, y2 - partial)], fill=outline, width=width)
drawcontext.line([(x2 - partial, y2), (x2, y2)], fill=outline, width=width)
drawcontext.line([(x2, y2), (x2, y2 - partial)], fill=outline, width=width)
def get_everything_under(root_dir, pattern='*', only_dirs=False, only_files=False):
assert not(only_dirs and only_files), 'You will get nothnig '\
'when "only_dirs" and "only_files" are both set to True'
everything = sorted(glob(os.path.join(root_dir, pattern)))
if only_dirs:
everything = [f for f in everything if os.path.isdir(f)]
if only_files:
everything = [f for f in everything if os.path.isfile(f)]
return everything
def read_filenames_from_dir(dir_name, reader, max_length=None):
logger.debug(
f"{reader} reading files from {dir_name}")
filenames = []
for root, dirs, files in os.walk(dir_name):
assert len(dirs) == 0, f"There are direcories: {dirs} in {root}"
assert len(files) != 0, f"There are no files in {root}"
filenames = [os.path.join(root, name) for name in sorted(files)]
for name in filenames:
logger.debug(name)
if max_length is not None:
return filenames[:max_length]
return filenames
|