# Script for creating json with annotations in the COCO format out of the Cell Tracking Challenge format of annotations
# Author: Kristyna Janku

import json
import cv2
import glob
import numpy as np
from os import listdir, path

dataset_name = 'Fluo-N2DH-SIM+'
videos_folder = '/home/kika/data/Fluo-N2DH-SIM+/train/data'
videos = ['01', '02']


def get_image_info(video):
    images = []

    img_id = 0
    video_path = path.join(videos_folder, video)
    for filename in sorted(listdir(video_path)):
        img = cv2.imread(path.join(video_path, filename))
        height, width, _ = img.shape
        images.append(
            {
                'file_name': filename,
                'height': height,
                'id': img_id,
                'width': width
            }
        )
        img_id += 1

    return images


def get_annotations(video):
    annotations = []
    annotation_path = path.join(videos_folder, f'{video}_GT')
    segmentation_folder_path = path.join(annotation_path, 'SEG')
    tracking_file_path = path.join(annotation_path, 'TRA', 'man_track.txt')

    # Load numbers of all frames from filenames
    frames = sorted([''.join(filter(str.isdigit, frame)) for frame in listdir(segmentation_folder_path) if
                     frame.endswith('.tif')])
    id_set = []
    id_frames = {}

    # Load tracking file and get information about objects and frames they appear in
    with open(tracking_file_path, 'r') as track_file:
        for line in track_file:
            parts = line.split(' ')
            obj_id = parts[0]
            id_set.append(int(obj_id))
            id_frames[int(obj_id)] = list(range(int(parts[1]), (int(parts[2]) + 1)))

    for frame in frames:
        # Load segmented image
        seg = glob.glob(path.join(segmentation_folder_path, '*' + frame + '*'))
        seg_image = cv2.imread(seg[0], cv2.IMREAD_ANYDEPTH)

        for obj_id in range(1, seg_image.max() + 1):
            if obj_id in id_set and int(frame) in id_frames[obj_id]:
                # Mask only pixels with current object, multiply by 1 to change True/False to 1/0
                mask = (seg_image == obj_id) * 1
                # Convert to 8-bit image to enable further transformations
                mask = cv2.convertScaleAbs(mask)
                # Checking if tracked object really is in segmented image
                if mask.max() == 0:
                    print(f'Video: {video} Frame: {frame} Missing object id in seg: {str(obj_id)}')
                else:
                    # Get bounding box of object from contour
                    contours = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
                    areas = [cv2.contourArea(c) for c in contours]
                    max_index = np.argmax(areas)
                    max_contour = contours[max_index]
                    bbox = cv2.boundingRect(max_contour)
                    area = bbox[2] * bbox[3]

                    if area > 5:
                        annotations.append(
                            {
                                'area': area,
                                'bbox': bbox,
                                'category_id': 0,
                                'id': obj_id - 1,
                                'ignore': 0,
                                'image_id': int(frame)
                            }
                        )

    return annotations


def main():
    for video in videos:
        # Add category info
        categories = [
                {
                    'id': 0,
                    'name': 'cell',
                    'supercategory': 'cell'
                }
            ]

        coco_dict = {
            'categories': categories,
            'images': get_image_info(video),
            'annotations': get_annotations(video)
        }

        with open(f'Fluo-N2DH-SIM+_video_{video}.json', 'w') as ann_json_file:
            json.dump(coco_dict, ann_json_file, indent=4, sort_keys=True)


if __name__ == '__main__':
    main()

