annotation_common_tools.py

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'

fix_pseudo_label_dataset.py

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])