from ..datasets.Dataset import DatasetInfo
from .output import ModelOutput


def postprocess(result: ModelOutput, info: DatasetInfo):
	return bbox_rescale_back(result, info)


def bbox_rescale_back(result: ModelOutput, info: DatasetInfo):
	if "detections" not in result or len(result["detections"]) == 0:
		return result
	assert len(info.source_data_shape) == 3
	assert len(info.data_shape) == 3

	# OpenCV, PyTorch, TF uses (height, width) formats
	source_height, source_width = info.source_data_shape[:2]
	height, width = (
		info.data_shape[:2] if info.data_shape[2] == 3 else info.data_shape[1:3]
	)

	rescale = min(height / source_height, width / source_width)
	x_padding = (width - source_width * rescale) // 2
	y_padding = (height - source_height * rescale) // 2

	for object in result["detections"]:
		if (
			"bbox_xmin_absolute" not in object
			or "bbox_xmax_absolute" not in object
			or "bbox_ymin_absolute" not in object
			or "bbox_ymax_absolute" not in object
		):
			continue
		object["bbox_xmin_absolute"] = int(
			(object["bbox_xmin_absolute"] - x_padding) / rescale
		)
		object["bbox_xmax_absolute"] = int(
			(object["bbox_xmax_absolute"] - x_padding) / rescale
		)
		object["bbox_ymin_absolute"] = int(
			(object["bbox_ymin_absolute"] - y_padding) / rescale
		)
		object["bbox_ymax_absolute"] = int(
			(object["bbox_ymax_absolute"] - y_padding) / rescale
		)

	return result
