# Path: core/Extractions.py
# EXTRACTIONS = "extractions"
# CLASS_POSTERIORS = "class_posteriors"
# SEGMENTATION_POSTERIORS = "segmentation_posteriors"
# SEGMENTATION_MASK_ORIGINAL_SIZE = "segmentation_mask_original_size"
# SEGMENTATION_MASK_INPUT_SIZE = "segmentation_mask_input_size"
# RECURRENT_STATE = "recurrent_state"
# REID_FEATURES = "reid_features"
# DET_BOXES = "det_boxes"
# DET_PROBS = "det_probs"
# DET_LABELS = "det_labels"
# DET_MASKS = "det_masks"
# IMAGE_ID = "img_id"
# CUSTOM_STRING = "custom_string"
# PRED_SPACE_TIME_OFFSETS = "pred_space_time_offsets"
# PIXEL_POSITIONS = "pixel_positions"
# FEATUREMAP_SIZE = "featuremap_size"
# def accumulate_extractions(extractions_accumulator, *new_extractions):
# def extract_batch_size_1(extractions, key):
# 
# Path: datasets/DataKeys.py
# SEGMENTATION_LABELS = "segmentation_labels"
# SEGMENTATION_LABELS_ORIGINAL_SIZE = "segmentation_labels_original_size"
# BBOX_GUIDANCE = "bbox_guidance"
# UNSIGNED_DISTANCE_TRANSFORM_GUIDANCE = "unsigned_distance_transform_guidance"
# SIGNED_DISTANCE_TRANSFORM_GUIDANCE = "signed_distance_transform_guidance"
# LASER_GUIDANCE = "laser_guidance"
# IMAGES = "images"
# RAW_IMAGES = "raw_images"
# INPUTS = "inputs"
# IMAGE_FILENAMES = "image_filenames"
# CLASSES = "classes"
# IDS = "ids"  # bounding box number 1,2,3...  or 0 if no bounding box
# DT_METHOD = "dt_method"
# RAW_SEGMENTATION_LABELS = "raw_segmentation_labels"
# CLICK_MAPS = "clicks_maps"
# NEG_CLICKS = "neg_clicks"
# POS_CLICKS = "pos_clicks"
# OBJ_TAGS = "obj_tags"
# FEATUREMAP_LABELS = "featuremap_labels"
# FEATUREMAP_BOXES = "featuremap_boxes"
# IMAGE_ID = "image_id"
# IS_CROWD = "is_crowd"
# SEGMENTATION_MASK = "segmentation_mask"
# SKIP_EXAMPLE = "skip_example"  # If true, filter the corresponding example from the datastream. Useful, when an example
# RAW_IMAGE_SIZES = "raw_image_sizes"  # h, w
# TIME_OFFSETS = "time_offsets"
# SPACE_OFFSETS = "space_offsets"
# SEGMENTATION_INSTANCE_LABELS = "segmentation_instance_labels"
# N_VALID_IDS = "n_valid_ids"
# 
# Path: forwarding/Forwarder.py
# class Forwarder(ABC):
#   def __init__(self, engine):
#     self.engine = engine
#     self.config = engine.config
#     self.session = engine.session
#     self.val_data = self.engine.valid_data
#     self.train_data = self.engine.train_data
#     self.trainer = self.engine.trainer
#     self.saver = self.engine.saver
# 
#   @abstractmethod
#   def forward(self):
#     pass
# 
# Path: forwarding/tracking/TrackingForwarder_util.py
# def export_detections_for_sequence(tag, boxes, scores, reids, classes, masks, model_str, epoch, add_masks, out_folder=""):
#   if out_folder == "":
#     out_folder = "forwarded/" + model_str + "/detections/" + str(epoch)
#   os.makedirs(out_folder, exist_ok=True)
#   out_filename = out_folder + "/" + tag + ".txt"
#   with open(out_filename, "w") as f:
#     t = 0
#     for boxes_t, scores_t, reids_t, classes_t, masks_t in zip(boxes, scores, reids, classes, masks):
#       for box, score, reid, class_, mask in zip(boxes_t, scores_t, reids_t, classes_t, masks_t):
#         if add_masks:
#           print(t, *box, score, class_, *mask['size'], mask['counts'].decode(encoding='UTF-8'), *reid, file=f)
#         else:
#           print(t, *box, score, class_, *reid, file=f)
#       t = t + 1
# 
# Path: datasets/KITTI/segtrack/KITTI_MOTS_info.py
# SEQ_IDS_VAL = ["%04d" % idx for idx in [2, 6, 7, 8, 10, 13, 14, 16, 18]]


import pycocotools.mask as cocomask
import numpy as np
import tensorflow as tf
from core import Extractions
from datasets import DataKeys
from forwarding.Forwarder import Forwarder
from forwarding.tracking.TrackingForwarder_util import export_detections_for_sequence
from datasets.KITTI.segtrack.KITTI_MOTS_info import SEQ_IDS_VAL



class BboxDetectionRefinementForwarder(Forwarder):
  def __init__(self, engine):
    super().__init__(engine)
    self.model_name = self.config.string("model")

  def forward(self):
    out_folder = "forwarded/" + self.model_name + "/detection_bbox_refined/"
    tf.gfile.MakeDirs(out_folder)
    data = self.val_data
    n_examples_per_epoch = data.n_examples_per_epoch()
    extraction_keys = [Extractions.DET_MASKS, DataKeys.IMAGE_FILENAMES, DataKeys.IDS]
    refined_masks = data.masks
    for n in range(n_examples_per_epoch):
      res = self.trainer.validation_step(extraction_keys=extraction_keys)
      filename = res[Extractions.EXTRACTIONS][DataKeys.IMAGE_FILENAMES][0][0].decode("utf-8")
      sp = filename.split("/")
      seq = sp[-2]
      t = int(sp[-1].replace(".png", "").replace(".jpg", ""))
      ids = res[Extractions.EXTRACTIONS][DataKeys.IDS][0][0]

# Path: core/Extractions.py
# EXTRACTIONS = "extractions"
# CLASS_POSTERIORS = "class_posteriors"
# SEGMENTATION_POSTERIORS = "segmentation_posteriors"
# SEGMENTATION_MASK_ORIGINAL_SIZE = "segmentation_mask_original_size"
# SEGMENTATION_MASK_INPUT_SIZE = "segmentation_mask_input_size"
# RECURRENT_STATE = "recurrent_state"
# REID_FEATURES = "reid_features"
# DET_BOXES = "det_boxes"
# DET_PROBS = "det_probs"
# DET_LABELS = "det_labels"
# DET_MASKS = "det_masks"
# IMAGE_ID = "img_id"
# CUSTOM_STRING = "custom_string"
# PRED_SPACE_TIME_OFFSETS = "pred_space_time_offsets"
# PIXEL_POSITIONS = "pixel_positions"
# FEATUREMAP_SIZE = "featuremap_size"
# def accumulate_extractions(extractions_accumulator, *new_extractions):
# def extract_batch_size_1(extractions, key):
# 
# Path: datasets/DataKeys.py
# SEGMENTATION_LABELS = "segmentation_labels"
# SEGMENTATION_LABELS_ORIGINAL_SIZE = "segmentation_labels_original_size"
# BBOX_GUIDANCE = "bbox_guidance"
# UNSIGNED_DISTANCE_TRANSFORM_GUIDANCE = "unsigned_distance_transform_guidance"
# SIGNED_DISTANCE_TRANSFORM_GUIDANCE = "signed_distance_transform_guidance"
# LASER_GUIDANCE = "laser_guidance"
# IMAGES = "images"
# RAW_IMAGES = "raw_images"
# INPUTS = "inputs"
# IMAGE_FILENAMES = "image_filenames"
# CLASSES = "classes"
# IDS = "ids"  # bounding box number 1,2,3...  or 0 if no bounding box
# DT_METHOD = "dt_method"
# RAW_SEGMENTATION_LABELS = "raw_segmentation_labels"
# CLICK_MAPS = "clicks_maps"
# NEG_CLICKS = "neg_clicks"
# POS_CLICKS = "pos_clicks"
# OBJ_TAGS = "obj_tags"
# FEATUREMAP_LABELS = "featuremap_labels"
# FEATUREMAP_BOXES = "featuremap_boxes"
# IMAGE_ID = "image_id"
# IS_CROWD = "is_crowd"
# SEGMENTATION_MASK = "segmentation_mask"
# SKIP_EXAMPLE = "skip_example"  # If true, filter the corresponding example from the datastream. Useful, when an example
# RAW_IMAGE_SIZES = "raw_image_sizes"  # h, w
# TIME_OFFSETS = "time_offsets"
# SPACE_OFFSETS = "space_offsets"
# SEGMENTATION_INSTANCE_LABELS = "segmentation_instance_labels"
# N_VALID_IDS = "n_valid_ids"
# 
# Path: forwarding/Forwarder.py
# class Forwarder(ABC):
#   def __init__(self, engine):
#     self.engine = engine
#     self.config = engine.config
#     self.session = engine.session
#     self.val_data = self.engine.valid_data
#     self.train_data = self.engine.train_data
#     self.trainer = self.engine.trainer
#     self.saver = self.engine.saver
# 
#   @abstractmethod
#   def forward(self):
#     pass
# 
# Path: forwarding/tracking/TrackingForwarder_util.py
# def export_detections_for_sequence(tag, boxes, scores, reids, classes, masks, model_str, epoch, add_masks, out_folder=""):
#   if out_folder == "":
#     out_folder = "forwarded/" + model_str + "/detections/" + str(epoch)
#   os.makedirs(out_folder, exist_ok=True)
#   out_filename = out_folder + "/" + tag + ".txt"
#   with open(out_filename, "w") as f:
#     t = 0
#     for boxes_t, scores_t, reids_t, classes_t, masks_t in zip(boxes, scores, reids, classes, masks):
#       for box, score, reid, class_, mask in zip(boxes_t, scores_t, reids_t, classes_t, masks_t):
#         if add_masks:
#           print(t, *box, score, class_, *mask['size'], mask['counts'].decode(encoding='UTF-8'), *reid, file=f)
#         else:
#           print(t, *box, score, class_, *reid, file=f)
#       t = t + 1
# 
# Path: datasets/KITTI/segtrack/KITTI_MOTS_info.py
# SEQ_IDS_VAL = ["%04d" % idx for idx in [2, 6, 7, 8, 10, 13, 14, 16, 18]]


# Path: forwarding/BboxDetectionRefinementForwarder.py
import pycocotools.mask as cocomask
import numpy as np
import tensorflow as tf
from core import Extractions
from datasets import DataKeys
from forwarding.Forwarder import Forwarder
from forwarding.tracking.TrackingForwarder_util import export_detections_for_sequence
from datasets.KITTI.segtrack.KITTI_MOTS_info import SEQ_IDS_VAL




class BboxDetectionRefinementForwarder(Forwarder):
  def __init__(self, engine):
    super().__init__(engine)
    self.model_name = self.config.string("model")

  def forward(self):
    out_folder = "forwarded/" + self.model_name + "/detection_bbox_refined/"
    tf.gfile.MakeDirs(out_folder)
    data = self.val_data
    n_examples_per_epoch = data.n_examples_per_epoch()
    extraction_keys = [Extractions.DET_MASKS, DataKeys.IMAGE_FILENAMES, DataKeys.IDS]
    refined_masks = data.masks
    for n in range(n_examples_per_epoch):
      res = self.trainer.validation_step(extraction_keys=extraction_keys)
      filename = res[Extractions.EXTRACTIONS][DataKeys.IMAGE_FILENAMES][0][0].decode("utf-8")
      sp = filename.split("/")
      seq = sp[-2]
      t = int(sp[-1].replace(".png", "").replace(".jpg", ""))
      ids = res[Extractions.EXTRACTIONS][DataKeys.IDS][0][0]

      masks = res[Extractions.EXTRACTIONS][Extractions.DET_MASKS][0][0]