Last active
November 14, 2019 00:53
-
-
Save asanakoy/7a2a5b820473eaddcc56bf221e3ca0c9 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import os | |
from os.path import join | |
import json | |
import cv2 | |
import numpy as np | |
from tqdm import tqdm | |
import json | |
from pycocotools import mask as maskUtils | |
from scipy import ndimage | |
from PIL import Image | |
import PIL.ExifTags | |
def rotate90_dp_coordinated(angle, coords): | |
if angle == 90: | |
factor = 1 | |
elif angle in [-180, 180]: | |
factor = 2 | |
elif angle in [-90, 270]: | |
factor = 3 | |
else: | |
raise ValueError(f'unknown angle {angle}') | |
x, y = coords[:, 0], coords[:, 1] | |
width = height = 256 | |
if factor == 1: | |
x, y = y, (width - 1) - x | |
elif factor == 2: | |
x, y = (width - 1) - x, (height - 1) - y | |
elif factor == 3: | |
x, y = (height - 1) - y, x | |
coords[:, 0] = x | |
coords[:, 1] = y | |
return coords | |
def rotate_annotation(ann, angle=90): | |
ann = ann.copy() | |
mask = maskUtils.decode(ann['segmentation']) | |
print('mask shape befor rot:', mask.shape) | |
mask = ndimage.rotate(mask, angle) | |
print('mask shape after rot:', mask.shape) | |
rle = maskUtils.encode(np.asfortranarray(mask)) | |
# decode str from bytes | |
rle['counts'] = rle['counts'].decode("utf-8") | |
ann['segmentation'] = rle | |
ann['area'] = float(maskUtils.area(ann['segmentation'])) | |
box = maskUtils.toBbox(ann['segmentation']).astype(int) | |
ann['bbox'] = list(map(int, box)) | |
new_shape = maskUtils.decode(ann['segmentation']).shape | |
assert new_shape == tuple(ann['segmentation']['size']), f"{new_shape} != {ann['segmentation']['size']}" | |
if 'dp_x' in ann: | |
assert 'dp_y' in ann | |
x = np.array(ann['dp_x'], dtype=float) | |
y = np.array(ann['dp_y'], dtype=float) | |
coords = np.stack([x, y], axis=1) | |
assert coords.shape[1] == 2, coords.shape | |
coords = rotate90_dp_coordinated(angle, coords) | |
x = coords[:, 0] | |
y = coords[:, 1] | |
ann['dp_x'] = x.tolist() | |
ann['dp_y'] = y.tolist() | |
assert 'dp_segm' not in ann, 'TODO: implement dp_segm rotation' | |
return ann | |
def resize_annotation(ann, width, height): | |
ann = ann.copy() | |
mask = maskUtils.decode(ann['segmentation']) | |
mask = cv2.resize(mask, (width, height), interpolation=cv2.INTER_NEAREST) | |
rle = maskUtils.encode(np.asfortranarray(mask)) | |
# decode str from bytes | |
rle['counts'] = rle['counts'].decode("utf-8") | |
ann['segmentation'] = rle | |
ann['area'] = float(maskUtils.area(ann['segmentation'])) | |
# x1 = np.nonzero(mask.any(axis=0))[0][0] | |
# x2 = np.nonzero(mask.any(axis=0))[0][-1] | |
# y1 = np.nonzero(mask.any(axis=1))[0][0] | |
# y2 = np.nonzero(mask.any(axis=1))[0][-1] | |
# box = np.array([x1, y1, x2 - x1, y2 - y1], dtype=int) | |
# ann['bbox'] = list(map(str, box)) | |
box = maskUtils.toBbox(ann['segmentation']).astype(int) | |
ann['bbox'] = list(map(int, box)) | |
new_shape = maskUtils.decode(ann['segmentation']).shape | |
assert new_shape == tuple(ann['segmentation']['size']), f"{new_shape} != {ann['segmentation']['size']}" | |
return ann | |
def resize_shortest_edge(img, shortest_edge_length, max_size): | |
size = shortest_edge_length | |
h, w = img.shape[:2] | |
scale = size * 1.0 / min(h, w) | |
if h < w: | |
newh, neww = size, scale * w | |
else: | |
newh, neww = scale * h, size | |
if max(newh, neww) > max_size: | |
scale = max_size * 1.0 / max(newh, neww) | |
newh = newh * scale | |
neww = neww * scale | |
neww = int(neww + 0.5) | |
newh = int(newh + 0.5) | |
img = cv2.resize(img, (neww, newh)) | |
return img | |
def resize_image(image_ann, shortest_edge_length=800, max_size=1333, src_images_dir=None): | |
assert src_images_dir is not None | |
image_ann = image_ann.copy() | |
img = cv2.imread(join(src_images_dir, image_ann['file_name'])) | |
img = resize_shortest_edge(img, shortest_edge_length, max_size) | |
h, w = img.shape[:2] | |
image_ann['height'] = h | |
image_ann['width'] = w | |
return image_ann, img | |
def find_image_info(images, image_id): | |
result = [x for x in images if x['id'] == image_id] | |
assert len(result) == 1, result | |
return result[0] | |
def get_exif_orientation(img_path): | |
pil_img = Image.open(img_path) | |
exif_data = pil_img._getexif() | |
exif_data = {} if exif_data is None else exif_data | |
exif = { | |
PIL.ExifTags.TAGS[k]: v | |
for k, v in exif_data.items() | |
if k in PIL.ExifTags.TAGS | |
} | |
if 'Orientation' in exif: | |
if exif['Orientation'] == 6: | |
return 'rotate_minus90' | |
elif exif['Orientation'] == 8: | |
return 'rotate_90' | |
else: | |
return 'no_rotate' | |
else: | |
return 'no_rotate' |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import os | |
from os.path import join | |
import json | |
import re | |
import cv2 | |
from pathlib import Path | |
import numpy as np | |
from tqdm import tqdm | |
import json | |
from pycocotools import mask as maskUtils | |
from scipy import ndimage | |
from PIL import Image | |
import PIL.ExifTags | |
from tools.annotation_common_tools import ( | |
rotate_annotation, | |
resize_annotation, | |
resize_shortest_edge, | |
find_image_info, | |
get_exif_orientation | |
) | |
""" | |
Fix chimpandsee pseudo label annotations | |
v1: | |
- Added all COCO categories to the 'caetgories' field | |
- add 'coco_url' if not exists | |
- Computed mask area for every instance. | |
- Convert 'image_id' and 'id' to int | |
- fix 'iscrowd' | |
- make 'file_name' conrtain only the base name of the file | |
- Rotate wrongly oriented masks. (no wrongly oriented found) | |
- Check that mask exactly match the images | |
""" | |
###################################### | |
################################### | |
# Init and make annotations compatible to V2 | |
################################### | |
CONVERT_TO_JPG = False | |
coco_dir = '/private/home/asanakoy/workspace/detectron2/datasets/coco' | |
subset_dir_name = 'from_detections_exp_000_dpR50FPN_prob_anmls_n_prs_1x_dp_v2_probmask_dpbinmask_10000' | |
dataset_root_dir = join('/checkpoint/asanakoy/datasets/detectron/chimpandsee/enhanced_databases/', subset_dir_name) | |
original_json_path = join(dataset_root_dir, 'pseudo_labels_10000_adj.json') | |
output_dir = dataset_root_dir | |
output_json_path = join(output_dir, f'pseudo_labels_10000_adj_v1.json') | |
src_images_dir = join(dataset_root_dir, 'images') | |
dst_images_dir = None # do not modify images | |
#if not os.path.exists(dst_images_dir): | |
# os.makedirs(dst_images_dir) | |
src_images_dir_name = os.path.basename(src_images_dir).strip('/') | |
#dst_images_dir_name = os.path.basename(dst_images_dir).strip('/') | |
dst_images_dir_name = src_images_dir_name | |
with open(join(coco_dir, 'annotations/instances_minival2014.json')) as json_file: | |
coco_minival = json.load(json_file) | |
if not os.path.exists(output_dir): | |
os.makedirs(output_dir) | |
with open(original_json_path) as json_file: | |
f = json.load(json_file) | |
print("chimpandsee categories ", f['categories']) | |
# make it contain all COCO categories, not just person | |
f['categories'] = coco_minival['categories'] | |
for x in f['images']: | |
m = re.match(f'{src_images_dir_name}/(.+)$', x['file_name']) | |
if m is not None: | |
x['file_name'] = m.group(1) | |
x['coco_url'] = f'{src_images_dir_name}/' + x['file_name'] | |
x['id'] = int(x['id']) | |
# Caculate mask areas | |
for ann in f['annotations']: | |
# now only support compressed RLE format as segmentation results | |
assert 'segmentation' in ann | |
if 'is_crowd' in ann: | |
ann['iscrowd'] = ann['is_crowd'] | |
del ann['is_crowd'] | |
elif not 'iscrowd' in ann: | |
ann['iscrowd'] = 0 | |
ann['area'] = float(maskUtils.area(ann['segmentation'])) | |
ann['id'] = int(ann['id']) | |
ann['image_id'] = int(ann['image_id']) | |
# rename to avoid confusion. This is dense densepose annotations. For every point. | |
ann['densepose_tensor'] = ann['densepose'] | |
del ann['densepose'] | |
################################### | |
# Resize images and annotations | |
################################### | |
# Fix image sizes | |
image_ids_to_rotate = {} | |
for image_ann in tqdm(f['images']): | |
image_path = join(src_images_dir, image_ann['file_name']) | |
img = cv2.imread(image_path) | |
#orientation = get_exif_orientation(image_path) | |
orientation = 'no_rotate' | |
if orientation != 'no_rotate': | |
image_ids_to_rotate[image_ann['id']] = orientation | |
print(f'rewrite rotated image w/o exif data: {image_ann["file_name"]}') | |
if CONVERT_TO_JPG: | |
image_ann['file_name'] = str(Path(image_ann['file_name']).with_suffix('.jpg')) | |
#image_save_path = join(dst_images_dir, image_ann['file_name']) | |
#cv2.imwrite(image_save_path, img) | |
h, w = img.shape[:2] | |
image_ann['height'] = h | |
image_ann['width'] = w | |
image_ann['coco_url'] = join(f'{dst_images_dir_name}/', image_ann['file_name']) | |
assert not len(image_ids_to_rotate), image_ids_to_rotate | |
# Rotate annotations and resize if necessary to match image size | |
for i, ann in enumerate(tqdm(f['annotations'])): | |
image_info = find_image_info(f['images'], ann['image_id']) | |
img_shape_ann = (image_info['height'], image_info['width']) | |
mask_shape_ann = tuple(ann['segmentation']['size']) | |
if ann['image_id'] in image_ids_to_rotate: | |
assert mask_shape_ann != img_shape_ann and np.all((np.array(mask_shape_ann) - np.array(img_shape_ann)[::-1]) <= 1 ), (mask_shape_ann, img_shape_ann, ann) | |
# The mask is rotated | |
if image_ids_to_rotate[ann['image_id']] == 'rotate_90': | |
angle = 90 | |
else: | |
assert image_ids_to_rotate[ann['image_id']] == 'rotate_minus90' | |
angle = -90 | |
print(f'{i} : H and W inverted (angle={angle}! img_shape={img_shape_ann}', image_info['file_name']) | |
rotated_ann = rotate_annotation(ann, angle) | |
f['annotations'][i] = rotated_ann | |
resized_ann = resize_annotation(ann, | |
width=image_info['width'], | |
height=image_info['height']) | |
f['annotations'][i] = resized_ann | |
############################################### | |
# Save result | |
############################################### | |
new_dataset = f | |
print('=============') | |
for key, val in new_dataset.items(): | |
print(f'New dataset {key}: num = {len(val)}') | |
with open(output_json_path, 'w') as outfile: | |
json.dump(new_dataset, outfile) | |
print('Saved to', output_json_path) | |
print('\n') | |
print('=============') | |
print(f['images'][:1]) | |
print('=============') | |
print('Annotations [0]:') | |
print(f['annotations'][:1]) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment