def _get_mmcv_home():
mmcv_home = os.path.expanduser(
os.getenv(
ENV_MMCV_HOME,
os.path.join(os.getenv(ENV_XDG_CACHE_HOME, DEFAULT_CACHE_DIR),
'mmcv')))
mkdir_or_exist(mmcv_home)
return mmcv_home
def main():
args = parse_args()
img_path = args.img_path
save_path = args.save_path
for i in range(1, 21):
print(i)
_img_path = os.path.join(img_path, f'{i:03d}')
_save_path = os.path.join(save_path, f'{i:03d}')
mkdir_or_exist(_save_path)
imagenames = os.listdir(_img_path)
for imagename in imagenames:
if int(imagename.split('.')[0]) % 10 == 0:
new_filename_path = os.path.join(_save_path, imagename)
# print(new_filename_path)
old_filename_path = os.path.join(_img_path, imagename)
shutil.copy(old_filename_path, new_filename_path)
def imwrite(img, file_path, params=None, auto_mkdir=True):
"""Write image to file.
Args:
img (ndarray): Image array to be written.
file_path (str): Image file path.
params (None or list): Same as opencv's :func:`imwrite` interface.
auto_mkdir (bool): If the parent folder of `file_path` does not exist,
whether to create it automatically.
Returns:
bool: Successful or not.
"""
if auto_mkdir:
dir_name = osp.abspath(osp.dirname(file_path))
mkdir_or_exist(dir_name)
return cv2.imwrite(file_path, img, params)
def cvt2frames(self,
frame_dir,
size=None,
step=1,
start=0,
end=None,
filename_tmpl='{:06d}.jpg'):
"""This function overwrites original ``mmcv.VideoReader.cvt2frame``.
This function converts a video to frames and save the frames in a
perticular pattern.
Args:
frame_dir (str): Output root directory to store the frames.
size (None or tuple <width, height>): frame size, if it is tuple
of target width and height, the frame will be resized.
If height is -1 or width is -1, the ratio will be kept.
step (int): frequency of saving frames.
start (int): start index.
end (int or None): end index.
filename_tmpl (str or callable): filename template. It should be
a formated string or some callable function/object that
accept an index as input parameter.
"""
assert isinstance(start, int) and start >= 0
assert end is None or isinstance(end, int)
assert isinstance(step, int) and step >= 1
_callable_fn_tmpl = callable(filename_tmpl)
if size is not None:
size = self._adjust_size(size)
mkdir_or_exist(frame_dir)
if end is None:
end = self.frame_cnt
else:
end = min(end, self.frame_cnt)
for i in range(start, end, step):
img = self.get_frame(i)
if size:
img = cv2.resize(img, size)
if _callable_fn_tmpl:
fn = filename_tmpl(i)
else:
fn = filename_tmpl.format(i)
cv2.imwrite(osp.join(frame_dir, fn), img)
def cvt2frames(self,
frame_dir,
file_start=0,
filename_tmpl='{:06d}.jpg',
start=0,
max_num=0,
show_progress=True):
"""Convert a video to frame images.
Args:
frame_dir (str): Output directory to store all the frame images.
file_start (int): Filenames will start from the specified number.
filename_tmpl (str): Filename template with the index as the
placeholder.
start (int): The starting frame index.
max_num (int): Maximum number of frames to be written.
show_progress (bool): Whether to show a progress bar.
"""
mkdir_or_exist(frame_dir)
if max_num == 0:
task_num = self.frame_cnt - start # task_num = (total frame) - (start frame)
else:
task_num = min(self.frame_cnt - start,
max_num) # task_num = min(74000 - 0, 600)
if task_num <= 0:
raise ValueError('start must be less than total frame number')
if start > 0:
self._set_real_position(start)
def write_frame(file_idx):
img = self.read()
filename = osp.join(frame_dir, filename_tmpl.format(file_idx))
cv2.imwrite(filename, img) # filename= 저장될 파일명, img= 저장할 이미지
if show_progress:
track_progress(write_frame, range(file_start,
file_start + task_num))
else:
for i in range(task_num):
img = self.read()
if img is None:
break
filename = osp.join(frame_dir,
filename_tmpl.format(i + file_start))
cv2.imwrite(filename, img)
def imshow_rbboxes(img_or_path,
rbboxes,
labels=None,
scores=None,
score_threshold=0.0,
colors='red',
show_label=False,
show_score=False,
thickness=3,
show=True,
win_name='',
wait_time=0,
out_file=None,
return_img=False):
""" Draw oriented bounding boxes on image
Args:
img (str or ndarray): The image to be displayed.
rbboxes (list or ndarray): A ndarray of shape (N, 8)
labels (list or ndarray): A ndarray of shape (N, 1)
scores (list or ndarray): A ndarray of shape (N, 1)
"""
if is_str(img_or_path):
img = cv2.imread(img_or_path)
else:
img = img_or_path
if rbboxes == []:
return
if isinstance(rbboxes, list):
rbboxes = np.array(rbboxes)
if rbboxes.shape[1] == 5:
rbboxes_ = []
for rbbox in rbboxes:
rbboxes_.append(thetaobb2pointobb(rbbox))
rbboxes = np.array(rbboxes_)
if rbboxes.ndim == 1:
rbboxes = np.array([rbboxes])
if labels is None:
labels_vis = np.array(['ins'] * rbboxes.shape[0])
else:
labels_vis = np.array(labels)
if labels_vis.ndim == 0:
labels_vis = np.array([labels_vis])
if scores is None:
scores_vis = np.array([1.0] * rbboxes.shape[0])
else:
scores_vis = np.array(scores)
if scores_vis.ndim == 0:
scores_vis = np.array([scores_vis])
if labels is None:
colors = dict()
colors[colors] = color_val(colors)
else:
max_label = max(labels)
colors = [color_val(_) for _ in range(max_label + 1)]
for rbbox, label, score in zip(rbboxes, labels_vis, scores_vis):
if score < score_threshold:
continue
if len(rbbox) == 5:
rbbox = np.array(thetaobb2pointobb(rbbox))
rbbox = rbbox.astype(np.int32)
cx = np.mean(rbbox[::2])
cy = np.mean(rbbox[1::2])
current_color = colors[label]
for idx in range(-1, 3, 1):
cv2.line(
img, (int(rbbox[idx * 2]), int(rbbox[idx * 2 + 1])),
(int(rbbox[(idx + 1) * 2]), int(rbbox[(idx + 1) * 2 + 1])),
current_color,
thickness=thickness)
if show_label:
cv2.putText(img,
label, (cx, cy),
cv2.FONT_HERSHEY_COMPLEX_SMALL,
fontScale=1.0,
color=current_color,
thickness=2,
lineType=8)
if show_score:
cv2.putText(img,
"{:.2f}".format(score), (cx, cy),
cv2.FONT_HERSHEY_COMPLEX_SMALL,
fontScale=1.0,
color=current_color,
thickness=2,
lineType=8)
if show:
cv2.namedWindow(win_name, cv2.WINDOW_NORMAL)
cv2.imshow(win_name, img)
cv2.waitKey(wait_time)
if out_file is not None:
dir_name = osp.abspath(osp.dirname(out_file))
mkdir_or_exist(dir_name)
cv2.imwrite(out_file, img)
if return_img:
return img
def imshow_bboxes(img_or_path,
bboxes,
labels=None,
scores=None,
score_threshold=0.0,
colors='red',
show_label=False,
show_score=False,
thickness=3,
show=True,
win_name='',
wait_time=0,
out_file=None,
origin_file=None,
return_img=False):
""" Draw horizontal bounding boxes on image
Args:
img (str or ndarray): The image to be displayed.
bboxes (list or ndarray): A ndarray of shape (N, 4)
labels (list or ndarray): A ndarray of shape (N, 1)
scores (list or ndarray): A ndarray of shape (N, 1)
"""
if is_str(img_or_path):
img = cv2.imread(img_or_path)
img_origin = img.copy()
else:
img = img_or_path
img_origin = img.copy()
if len(bboxes) == []:
return
if isinstance(bboxes, list):
bboxes = np.array(bboxes)
if bboxes.ndim == 1:
bboxes = np.array([bboxes])
if labels is None:
labels_vis = np.array(['ins'] * bboxes.shape[0])
else:
labels_vis = np.array(labels)
if labels_vis.ndim == 0:
labels_vis = np.array([labels_vis])
if scores is None:
scores_vis = np.array([1.0] * bboxes.shape[0])
else:
scores_vis = np.array(scores)
if scores_vis.ndim == 0:
scores_vis = np.array([scores_vis])
if labels is None:
colors = dict()
colors[colors] = color_val(colors)
else:
max_label = max(labels)
if max_label > 20:
max_label = max_label % 20
colors = [color_val(_) for _ in range(max_label + 1)]
for bbox, label, score in zip(bboxes, labels_vis, scores_vis):
if score < score_threshold:
continue
bbox = bbox.astype(np.int32)
xmin, ymin, xmax, ymax = bbox
current_color = colors[label % 20]
cv2.rectangle(img, (xmin, ymin), (xmax, ymax),
color=current_color,
thickness=thickness)
if show_label:
cv2.putText(img,
label, (xmin, ymin - 5),
cv2.FONT_HERSHEY_COMPLEX_SMALL,
fontScale=1.0,
color=current_color,
thickness=2,
lineType=8)
if show_score:
cv2.putText(img,
"{:.2f}".format(score), (xmin, ymin - 5),
cv2.FONT_HERSHEY_COMPLEX_SMALL,
fontScale=1.0,
color=current_color,
thickness=2,
lineType=8)
if show:
cv2.imshow(win_name, img)
cv2.waitKey(wait_time)
if out_file is not None:
dir_name = osp.abspath(osp.dirname(out_file))
mkdir_or_exist(dir_name)
cv2.imwrite(out_file, img)
if origin_file is not None:
dir_name = osp.abspath(osp.dirname(origin_file))
mkdir_or_exist(dir_name)
cv2.imwrite(origin_file, img_origin)
if return_img:
return img
def _plt_show_wrong_tracks(img,
bboxes,
ids,
error_types,
thickness=0.1,
font_scale=3,
text_width=8,
text_height=13,
show=False,
wait_time=100,
out_file=None):
"""Show the wrong tracks with matplotlib.
Args:
img (str or ndarray): The image to be displayed.
bboxes (ndarray): A ndarray of shape (k, 5).
ids (ndarray): A ndarray of shape (k, ).
error_types (ndarray): A ndarray of shape (k, ), where 0 denotes
false positives, 1 denotes false negative and 2 denotes ID switch.
thickness (float, optional): Thickness of lines.
Defaults to 0.1.
font_scale (float, optional): Font scale to draw id and score.
Defaults to 3.
text_width (int, optional): Width to draw id and score.
Defaults to 8.
text_height (int, optional): Height to draw id and score.
Defaults to 13.
show (bool, optional): Whether to show the image on the fly.
Defaults to False.
wait_time (int, optional): Value of waitKey param.
Defaults to 100.
out_file (str, optional): The filename to write the image.
Defaults to None.
Returns:
ndarray: Original image.
"""
assert bboxes.ndim == 2, \
f' bboxes ndim should be 2, but its ndim is {bboxes.ndim}.'
assert ids.ndim == 1, \
f' ids ndim should be 1, but its ndim is {ids.ndim}.'
assert error_types.ndim == 1, \
f' error_types ndim should be 1, but its ndim is {error_types.ndim}.'
assert bboxes.shape[0] == ids.shape[0], \
'bboxes.shape[0] and ids.shape[0] should have the same length.'
assert bboxes.shape[1] == 5, \
f' bboxes.shape[1] should be 5, but its {bboxes.shape[1]}.'
bbox_colors = sns.color_palette()
# red, yellow, blue
bbox_colors = [bbox_colors[3], bbox_colors[1], bbox_colors[0]]
if isinstance(img, str):
img = plt.imread(img)
else:
assert img.ndim == 3
img = mmcv.bgr2rgb(img)
img_shape = img.shape
bboxes[:, 0::2] = np.clip(bboxes[:, 0::2], 0, img_shape[1])
bboxes[:, 1::2] = np.clip(bboxes[:, 1::2], 0, img_shape[0])
plt.imshow(img)
plt.gca().set_axis_off()
plt.autoscale(False)
plt.subplots_adjust(top=1,
bottom=0,
right=1,
left=0,
hspace=None,
wspace=None)
plt.margins(0, 0)
plt.gca().xaxis.set_major_locator(plt.NullLocator())
plt.gca().yaxis.set_major_locator(plt.NullLocator())
plt.rcParams['figure.figsize'] = img_shape[1], img_shape[0]
for bbox, error_type, id in zip(bboxes, error_types, ids):
x1, y1, x2, y2, score = bbox
w, h = int(x2 - x1), int(y2 - y1)
left_top = (int(x1), int(y1))
# bbox
plt.gca().add_patch(
Rectangle(left_top,
w,
h,
thickness,
edgecolor=bbox_colors[error_type],
facecolor='none'))
# FN does not have id and score
if error_type == 1:
continue
# score
text = '{:.02f}'.format(score)
width = len(text) * text_width
plt.gca().add_patch(
Rectangle((left_top[0], left_top[1]),
width,
text_height,
thickness,
edgecolor=bbox_colors[error_type],
facecolor=bbox_colors[error_type]))
plt.text(left_top[0],
left_top[1] + text_height + 2,
text,
fontsize=font_scale)
# id
text = str(id)
width = len(text) * text_width
plt.gca().add_patch(
Rectangle((left_top[0], left_top[1] + text_height + 1),
width,
text_height,
thickness,
edgecolor=bbox_colors[error_type],
facecolor=bbox_colors[error_type]))
plt.text(left_top[0],
left_top[1] + 2 * (text_height + 1),
text,
fontsize=font_scale)
if out_file is not None:
mkdir_or_exist(osp.abspath(osp.dirname(out_file)))
plt.savefig(out_file, dpi=300, bbox_inches='tight', pad_inches=0.0)
if show:
plt.draw()
plt.pause(wait_time / 1000.)
plt.clf()
return img
print('opened:', video_reader.opened, 'resolution',
video_reader.resolution, 'fps:', video_reader.fps, 'total frames:',
video_reader.frame_cnt)
print('\n')
save_path = osp.join(camvid_path, video[:-4])
if osp.exists(save_path):
shutil.rmtree(save_path)
video_reader.cvt2frames(
frame_dir=save_path,
file_start=file_start[i],
filename_tmpl=filename_tmpl[i],
start=start[i]) # extract frames from the given video
#
first_frame = [0, 930, 390, 6690] # first frame to extract
mkdir_or_exist(osp.join(camvid_path, 'per_30'))
for j, video in enumerate(video_names):
for i in range(20):
save_path = osp.join(camvid_path, video[:-4])
img_name = filename_tmpl[j].format(first_frame[j] + i * 30)
shutil.copy(osp.join(save_path, img_name),
osp.join(camvid_path, 'per_30'))
camvid_path = "/data/datasets/video_ss/camvid/raw"
splits = ['train', 'val', 'test']
for split in splits:
f_name = osp.join(camvid_path, f'{split}.txt')
with open(f_name, 'w') as f:
f_list = os.listdir(osp.join(camvid_path, split))
for i in f_list:
def imshow_bboxes(img_or_path,
bboxes,
labels=None,
scores=None,
score_threshold=0.0,
show_label=False,
show_score=False,
thickness=2,
show=False,
win_name='',
wait_time=0,
out_file=None,
origin_file=None,
return_img=False):
""" Draw horizontal bounding boxes on image
Args:
img (str or ndarray): The image to be displayed.
bboxes (list or ndarray): A ndarray of shape (N, 4)
labels (dict): {"category": idx}
scores (list or ndarray): A ndarray of shape (N, 1)
"""
if is_str(img_or_path):
img = cv2.imread(img_or_path)
img_origin = img.copy()
else:
img = img_or_path
img_origin = img.copy()
if len(bboxes) == 0:
return
if isinstance(bboxes, list):
bboxes = np.array(bboxes)
if bboxes.ndim == 1:
bboxes = np.array([bboxes])
if labels is None:
labels_vis = np.array(['ins'] * bboxes.shape[0])
else:
labels_vis = [list(label.keys())[0] for label in labels]
class_idxs = [list(label.values())[0] for label in labels]
max_label_idx = max(class_idxs)
if max_label_idx > 20:
max_label_idx = max_label_idx % 20
color_dict = {
list(label.keys())[0]: color_val(list(label.values())[0])
for label in labels
}
if scores is None:
scores_vis = np.array([1.0] * bboxes.shape[0])
else:
scores_vis = np.array(scores)
if scores_vis.ndim == 0:
scores_vis = np.array([scores_vis])
for bbox, label, score in zip(bboxes, labels_vis, scores_vis):
if score < score_threshold:
continue
bbox = bbox.astype(np.int32)
xmin, ymin, xmax, ymax = bbox
if labels is None:
current_color = color_val()
else:
current_color = color_dict[label]
cv2.rectangle(img, (xmin, ymin), (xmax, ymax),
color=current_color,
thickness=thickness)
if show_label:
cv2.putText(img,
label, (xmin, ymin - 5),
cv2.FONT_HERSHEY_COMPLEX_SMALL,
fontScale=1.0,
color=current_color,
thickness=2,
lineType=8)
if show_score:
cv2.putText(img,
"{:.2f}".format(score), (xmin, ymin - 5),
cv2.FONT_HERSHEY_COMPLEX_SMALL,
fontScale=1.0,
color=current_color,
thickness=2,
lineType=8)
if show:
# cv2.namedWindow(win_name, cv2.WINDOW_NORMAL)
# cv2.resizeWindow(win_name, 360, 360)
cv2.imshow(win_name, img)
cv2.waitKey(wait_time)
if out_file is not None:
dir_name = osp.abspath(osp.dirname(out_file))
mkdir_or_exist(dir_name)
cv2.imwrite(out_file, img)
if origin_file is not None:
dir_name = osp.abspath(osp.dirname(origin_file))
mkdir_or_exist(dir_name)
cv2.imwrite(origin_file, img_origin)
if return_img:
return img
filepath_train_w = '/root/cwt1/ntire2021/work_dirs/edvr_g8_600k_large_s2_l_compress/results_600k_train_w/'
filepath_train_wh = '/root/cwt1/ntire2021/work_dirs/edvr_g8_600k_large_s2_l_compress/results_600k_train_wh/'
filepath_merge = '/root/cwt1/ntire2021/work_dirs/edvr_g8_600k_large_s2_l_compress/results_600k_train_merge/'
filepath_gt = '/root/cwt1/ntire2021/data/video_compress_track2/images/train_raw/'
frame_names = range(1, 21)
count = 0
psnr_mean = []
for frame in frame_names:
print("frame:", frame)
_train_path = os.path.join(filepath_train, f'{frame:03d}')
_train_path_h = os.path.join(filepath_train_h, f'{frame:03d}')
_train_path_w = os.path.join(filepath_train_w, f'{frame:03d}')
_train_path_wh = os.path.join(filepath_train_wh, f'{frame:03d}')
_gt_path = os.path.join(filepath_gt, f'{frame:03d}')
_save_path = os.path.join(filepath_merge, f'{frame:03d}')
mkdir_or_exist(_save_path)
imagenames = os.listdir(_train_path)
for imagename in imagenames:
count += 1
if count % 100 == 0:
print(count)
# if count > 10:
# break
img_train = cv2.imread(os.path.join(_train_path, imagename))
img_train_h = cv2.imread(os.path.join(_train_path_h, imagename))
img_train_w = cv2.imread(os.path.join(_train_path_w, imagename))
img_train_wh = cv2.imread(os.path.join(_train_path_wh, imagename))
img_train = img_train.astype(np.float32)
