"""Define classes and functions for interfacing with SageMaker Ground
Truth object detection.

"""

import os
import imageio
import matplotlib.pyplot as plt
import numpy as np


class BoundingBox:
    """Bounding box for an object in an image."""

    def __init__(self, image_id=None, boxdata=None):
        self.image_id = image_id
        if boxdata:
            for datum in boxdata:
                setattr(self, datum, boxdata[datum])

    def __repr__(self):
        return "Box for image {}".format(self.image_id)

    def compute_bb_data(self):
        """Compute the parameters used for IoU."""
        image = self.image
        self.xmin = self.left / image.width
        self.xmax = (self.left + self.width) / image.width
        self.ymin = self.top / image.height
        self.ymax = (self.top + self.height) / image.height


class BoxedImage:
    """Image with bounding boxes."""

    def __init__(
        self,
        id=None,
        consolidated_boxes=None,
        worker_boxes=None,
        gt_boxes=None,
        uri=None,
        size=None,
    ):
        self.id = id
        self.uri = uri
        if uri:
            self.filename = uri.split("/")[-1]
            self.oid_id = self.filename.split(".")[0]
        else:
            self.filename = None
            self.oid_id = None
        self.local = None
        self.im = None
        if size:
            self.width = size["width"]
            self.depth = size["depth"]
            self.height = size["height"]
            self.shape = self.width, self.height, self.depth
        if consolidated_boxes:
            self.consolidated_boxes = consolidated_boxes
        else:
            self.consolidated_boxes = []
        if worker_boxes:
            self.worker_boxes = worker_boxes
        else:
            self.worker_boxes = []
        if gt_boxes:
            self.gt_boxes = gt_boxes
        else:
            self.gt_boxes = []

    def __repr__(self):
        return "Image{}".format(self.id)

    def n_consolidated_boxes(self):
        """Count the number of consolidated boxes."""
        return len(self.consolidated_boxes)

    def n_worker_boxes(self):
        return len(self.worker_boxes)

    def download(self, directory):
        target_fname = os.path.join(directory, self.uri.split("/")[-1])
        if not os.path.isfile(target_fname):
            os.system(f"aws s3 cp {self.uri} {target_fname}")
            print("downloading image to {}".format(target_fname))
        self.local = target_fname

    def imread(self):
        """Cache the image reading process."""
        try:
            return imageio.imread(self.local)#, exifrotate=False)
        except OSError:
            print(
                "You need to download this image {} first. "
                "Use this_image.download(local_directory).".format(self.local)
            )
            raise

    def plot_bbs(self, ax, bbs, img_kwargs, box_kwargs, **kwargs):
        """Master function for plotting images with bounding boxes."""
        img = self.imread()
        ax.imshow(img, **img_kwargs)
        imh, imw, *_ = img.shape
        box_kwargs["fill"] = None
        class_colors = {
            0: "C0",
            1: "C1",
            2: "C2",
            3: "C3",
            4: "C4",
            5: "C5",
            6: "C6",
            7: "C7",
            8: "C8",
            9: "C9",
            10: "C10",
            11: "C11",
            12: "C12",
            13: "C13",
            14: "C14",
            15: "C15",
            16: "C16",
        }
        for bb in bbs:
            class_id = bb.class_id
            if class_id not in class_colors:
                class_colors[class_id] = "C" + str(class_id)
            rec = plt.Rectangle(
                (bb.left, bb.top),
                bb.width,
                bb.height,
                linewidth=1,
                edgecolor=class_colors[class_id],
                **box_kwargs,
            )
            if bb.class_id == 0:
                annotation = "car-" + str(bb.confidence)
            else:
                annotation = str(bb.class_id) + "-" + str(bb.confidence)
            ax.text(
                bb.left,
                (bb.top + bb.height),
                annotation,
                fontsize=6,
                verticalalignment="center",
                bbox=dict(facecolor=class_colors[class_id], alpha=0.5),
            )
            ax.add_patch(rec)
        ax.axis("off")

    def plot_consolidated_bbs(self, ax, img_kwargs={}, box_kwargs={"lw": 1}):
        """Plot the consolidated boxes."""
        self.plot_bbs(
            ax, self.consolidated_boxes, img_kwargs=img_kwargs, box_kwargs=box_kwargs
        )

    def crop_bbs(self, ax, bbs, uri, img_kwargs, box_kwargs, **kwargs):
        """Master function for cropping images around bounding boxes."""
        img = self.imread()
        uri = uri.split("/")[4]
        bb_count = 0
        for bb in bbs:
            y = int(bb.top)
            x = int(bb.left)
            h = int(bb.height)
            w = int(bb.width)
            cropped = img[y : y + h, x : x + w]
            imageio.imwrite(
                "./cropped-images/cropped-{}_{}.png".format(uri, bb_count),
                cropped,
            )
            bb_count += 1

    def plot_crop_consolidated_bbs(self, ax, img_kwargs={}, box_kwargs={"lw": 1}):
        """Plot the consolidated boxes."""
        self.plot_bbs(
            ax, self.consolidated_boxes, img_kwargs=img_kwargs, box_kwargs=box_kwargs
        )
        self.crop_bbs(
            ax,
            self.consolidated_boxes,
            self.uri,
            img_kwargs=img_kwargs,
            box_kwargs=box_kwargs,
        )

    def compute_img_confidence(self):
        """ Compute the mean bb confidence. """
        if len(self.consolidated_boxes) > 0:
            return np.mean([box.confidence for box in self.consolidated_boxes])
        else:
            return 0

"""Define classes and functions for interfacing with SageMaker Ground
Truth object detection.

"""

import os
import imageio
import matplotlib.pyplot as plt
import numpy as np


class BoundingBox:
    """Bounding box for an object in an image."""

    def __init__(self, image_id=None, boxdata=None):
        self.image_id = image_id
        if boxdata:
            for datum in boxdata:
                setattr(self, datum, boxdata[datum])

    def __repr__(self):
        return "Box for image {}".format(self.image_id)

    def compute_bb_data(self):
        """Compute the parameters used for IoU."""
        image = self.image
        self.xmin = self.left / image.width
        self.xmax = (self.left + self.width) / image.width
        self.ymin = self.top / image.height
        self.ymax = (self.top + self.height) / image.height


class BoxedImage:
    """Image with bounding boxes."""

    def __init__(
        self,
        id=None,
        consolidated_boxes=None,
        worker_boxes=None,
        gt_boxes=None,
        uri=None,
        size=None,
    ):
        self.id = id
        self.uri = uri
        if uri:
            self.filename = uri.split("/")[-1]
            self.oid_id = self.filename.split(".")[0]
        else:
            self.filename = None
            self.oid_id = None
        self.local = None
        self.im = None
        if size:
            self.width = size["width"]
            self.depth = size["depth"]
            self.height = size["height"]
            self.shape = self.width, self.height, self.depth
        if consolidated_boxes:
            self.consolidated_boxes = consolidated_boxes
        else:
            self.consolidated_boxes = []
        if worker_boxes:
            self.worker_boxes = worker_boxes
        else:
            self.worker_boxes = []
        if gt_boxes:
            self.gt_boxes = gt_boxes
        else:
            self.gt_boxes = []

    def __repr__(self):
        return "Image{}".format(self.id)

    def n_consolidated_boxes(self):
        """Count the number of consolidated boxes."""
        return len(self.consolidated_boxes)

    def n_worker_boxes(self):
        return len(self.worker_boxes)

    def download(self, directory):
        target_fname = os.path.join(directory, self.uri.split("/")[-1])
        if not os.path.isfile(target_fname):
            os.system(f"aws s3 cp {self.uri} {target_fname}")
            print("downloading image to {}".format(target_fname))
        self.local = target_fname

    def imread(self):
        """Cache the image reading process."""
        try:
            return imageio.imread(self.local)#, exifrotate=False)
        except OSError:
            print(
                "You need to download this image {} first. "
                "Use this_image.download(local_directory).".format(self.local)
            )
            raise

    def plot_bbs(self, ax, bbs, img_kwargs, box_kwargs, **kwargs):
        """Master function for plotting images with bounding boxes."""
        img = self.imread()
        ax.imshow(img, **img_kwargs)
        imh, imw, *_ = img.shape
        box_kwargs["fill"] = None
        class_colors = {
            0: "C0",
            1: "C1",
            2: "C2",
            3: "C3",
            4: "C4",
            5: "C5",
            6: "C6",
            7: "C7",
            8: "C8",
            9: "C9",
            10: "C10",
            11: "C11",
            12: "C12",
            13: "C13",
            14: "C14",
            15: "C15",
            16: "C16",
        }
        for bb in bbs:
            class_id = bb.class_id
            if class_id not in class_colors:
                class_colors[class_id] = "C" + str(class_id)
            rec = plt.Rectangle(
                (bb.left, bb.top),
                bb.width,
                bb.height,
                linewidth=1,
                edgecolor=class_colors[class_id],
                **box_kwargs,
            )
            if bb.class_id == 0:
                annotation = "car-" + str(bb.confidence)
            else:
                annotation = str(bb.class_id) + "-" + str(bb.confidence)
            ax.text(
                bb.left,
                (bb.top + bb.height),
                annotation,
                fontsize=6,
                verticalalignment="center",
                bbox=dict(facecolor=class_colors[class_id], alpha=0.5),
            )
            ax.add_patch(rec)
        ax.axis("off")

    def plot_consolidated_bbs(self, ax, img_kwargs={}, box_kwargs={"lw": 1}):
        """Plot the consolidated boxes."""
        self.plot_bbs(
            ax, self.consolidated_boxes, img_kwargs=img_kwargs, box_kwargs=box_kwargs
        )

    def crop_bbs(self, ax, bbs, uri, img_kwargs, box_kwargs, **kwargs):
        """Master function for cropping images around bounding boxes."""
        img = self.imread()
        uri = uri.split("/")[4]
        bb_count = 0
        for bb in bbs:
            y = int(bb.top)
            x = int(bb.left)
            h = int(bb.height)
            w = int(bb.width)
            cropped = img[y : y + h, x : x + w]
            imageio.imwrite(
                "./cropped-images/cropped-{}_{}.png".format(uri, bb_count),
                cropped,
            )
            bb_count += 1

    def plot_crop_consolidated_bbs(self, ax, img_kwargs={}, box_kwargs={"lw": 1}):
        """Plot the consolidated boxes."""
        self.plot_bbs(
            ax, self.consolidated_boxes, img_kwargs=img_kwargs, box_kwargs=box_kwargs
        )
        self.crop_bbs(
            ax,
            self.consolidated_boxes,
            self.uri,
            img_kwargs=img_kwargs,
            box_kwargs=box_kwargs,
        )

    def compute_img_confidence(self):
        """ Compute the mean bb confidence. """
        if len(self.consolidated_boxes) > 0:
            return np.mean([box.confidence for box in self.consolidated_boxes])
        else:
            return 0

Define classes and functions for interfacing with SageMaker Ground Truth object detection. Bounding box for an object in an image. Compute the parameters used for IoU. Image with bounding boxes. Count the number of consolidated boxes. Cache the image reading process. #, exifrotate=False) Master function for plotting images with bounding boxes. Plot the consolidated boxes. Master function for cropping images around bounding boxes. Plot the consolidated boxes. Compute the mean bb confidence.