import os
import sys
import json
import datetime
import numpy as np
import skimage.draw
import cv2
import random
import matplotlib.pyplot as plt
import matplotlib.patches as patches

# Root directory of the project
ROOT_DIR = os.path.abspath("../")

# Import Mask RCNN
sys.path.append(ROOT_DIR)  # To find local version of the library
from mrcnn.config import Config
from mrcnn import model as modellib, utils
from mrcnn import visualize
from mrcnn.visualize import display_images
from mrcnn.model import log
from workspace import evaluate
from workspace import helper

# Path to trained weights file
COCO_WEIGHTS_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")

# Directory to save logs and model checkpoints, if not provided
# through the command line argument --logs
DEFAULT_LOGS_DIR = os.path.join(ROOT_DIR, "logs")

class CarPlateConfig(Config):

    NAME = 'carplate'

    IMAGES_PER_GPU = 1

    NUM_CLASSES = 1 + 33

    STEPS_PER_EPOCH = 100

    DETECTION_MIN_CONFIDENCE = 0.5

    RPN_ANCHOR_SCALES = (16, 32, 64, 128, 256)

    BACKBONE_STRIDES = [4, 8, 16, 24, 48]
    RPN_ANCHOR_RATIOS = [0.3, 0.6, 1]

    RPN_TRAIN_ANCHORS_PER_IMAGE = 500

    # RPN_NMS_THRESHOLD = 0.6

    IMAGE_MIN_DIM = int(480)
    IMAGE_MAX_DIM = int(640)

    # POST_NMS_ROIS_INFERENCE = 2000

    TRAIN_ROIS_PER_IMAGE = 400

    MEAN_PIXEL = np.array([0.449122045 * 255, 0.449122045 * 255, 0.449122045 * 255 ])

    LEARNING_RATE = 0.005

class InferenceConfig(CarPlateConfig):
            # Set batch size to 1 since we'll be running inference on
            # one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1

class CarPlateDataset(utils.Dataset):

   def load_carplates(self, dataset_dir, subset):

        self.add_class("carplate", 1, "0")
        self.add_class("carplate", 2, "1")
        self.add_class("carplate", 3, "2")
        self.add_class("carplate", 4, "3")
        self.add_class("carplate", 5, "4")
        self.add_class("carplate", 6, "5")
        self.add_class("carplate", 7, "6")
        self.add_class("carplate", 8, "7")
        self.add_class("carplate", 9, "8")
        self.add_class("carplate", 10, "9")
        self.add_class("carplate", 11, "A")
        self.add_class("carplate", 12, "B")
        self.add_class("carplate", 13, "C")
        self.add_class("carplate", 14, "D")
        self.add_class("carplate", 15, "E")
        self.add_class("carplate", 16, "F")
        self.add_class("carplate", 17, "G")
        self.add_class("carplate", 18, "H")
        self.add_class("carplate", 19, "J")
        self.add_class("carplate", 20, "K")
        self.add_class("carplate", 21, "L")
        self.add_class("carplate", 22, "M")
        self.add_class("carplate", 23, "N")
        self.add_class("carplate", 24, "P")
        self.add_class("carplate", 25, "R")
        self.add_class("carplate", 26, "S")
        self.add_class("carplate", 27, "T")
        self.add_class("carplate", 28, "U")
        self.add_class("carplate", 29, "V")
        self.add_class("carplate", 30, "W")
        self.add_class("carplate", 31, "X")
        self.add_class("carplate", 32, "Y")
        self.add_class("carplate", 33, "Z")
        # self.add_class("carplate", 1, "carplate")
        name_dict = {"0": 1,
                         "1": 2,
                         "2": 3,
                         "3": 4,
                         "4": 5,
                         "5": 6,
                         "6": 7,
                         "7": 8,
                         "8": 9,
                         "9": 10,
                         "A": 11,
                         "B": 12,
                         "C": 13,
                         "D": 14,
                         "E": 15,
                         "F": 16,
                         "G": 17,
                         "H": 18,
                         "J": 19,
                         "K": 20,
                         "L": 21,
                         "M": 22,
                         "N": 23,
                         "P": 24,
                         "R": 25,
                         "S": 26,
                         "T": 27,
                         "U": 28,
                         "V": 29,
                         "W": 30,
                         "X": 31,
                         "Y": 32,
                         "Z": 33
                         }
        # name_dict = {
            # "carplate": 1
        # }
        assert subset in ['train', 'val']

        dataset_dir = os.path.join(dataset_dir, subset)
        print("Dataset: ", dataset_dir)

        annotations = json.load(open(os.path.join(dataset_dir, "via_region_data.json")))

        annotations = list(annotations.values())  # don't need the dict keys

        annotations = [a for a in annotations if a['regions']]

        for a in annotations:
            if type(a['regions']) is dict:
                polygons = [r['shape_attributes'] for r in a['regions'].values()]
            else:
                polygons = [r['shape_attributes'] for r in a['regions']]

            class_ids = [r['region_attributes']['type'] for r in a['regions']]

            name_id = [name_dict[n] for n in class_ids]
            image_path = os.path.join(dataset_dir, a['filename'])
            image = skimage.io.imread(image_path)
            height, width = image.shape[:2]
            self.add_image(
                "carplate",
                image_id=a['filename'],  # use file name as a unique image id
                path=image_path,
                width=width, height=height,
                polygons=polygons,
                class_id = name_id)

   def load_mask(self, image_id):

        image_info = self.image_info[image_id]
        if image_info["source"] != "carplate":
            return super(self.__class__, self).load_mask(image_id)
        name_id = image_info["class_id"]
        # Convert polygons to a bitmap mask of shape
        # [height, width, instance_count]
        info = self.image_info[image_id]
        mask = np.zeros([info["height"], info["width"], len(info["polygons"])],
                        dtype=np.uint8)

        class_ids = np.array(name_id, dtype=np.int32)
        for i, p in enumerate(info["polygons"]):
            if 'all_points_x' and 'all_points_y' in p:
            # Get indexes of pixels inside the polygon and set them to 1
                rr, cc = skimage.draw.polygon(p['all_points_y'], p['all_points_x'])
                mask[rr, cc, i] = 1

        # Return mask, and array of class IDs of each instance. Since we have
        # one class ID only, we return an array of 1s
        return (mask.astype(np.bool),class_ids)

   def image_reference(self, image_id):
        """Return the path of the image."""
        info = self.image_info[image_id]
        if info["source"] == "carplate":
            return info["path"]
        else:
            super(self.__class__, self).image_reference(image_id)

def display_dataset(subfolder):
    dataset_train = CarPlateDataset()
    dataset_train.load_carplates(os.path.join(ROOT_DIR, subfolder), "val")
    dataset_train.prepare()

    print("Image Count: {}".format(len(dataset_train.image_ids)))
    print("Class Count: {}".format(dataset_train.num_classes))
    for i, info in enumerate(dataset_train.class_info):
        print("{:3}. {:50}".format(i, info['name']))

    image_ids = np.random.choice(dataset_train.image_ids, 3)
    for image_id in image_ids:
        image = dataset_train.load_image(image_id)
        mask, class_ids = dataset_train.load_mask(image_id)
        visualize.display_top_masks(image, mask, class_ids, dataset_train.class_names, limit=7)

    dataset_test = CarPlateDataset()
    dataset_test.load_carplates(os.path.join(ROOT_DIR, subfolder), "val")
    dataset_test.prepare()

    print("Image Count: {}".format(len(dataset_test.image_ids)))
    print("Class Count: {}".format(dataset_test.num_classes))
    for i, info in enumerate(dataset_test.class_info):
        print("{:3}. {:50}".format(i, info['name']))

    image_ids = np.random.choice(dataset_test.image_ids, 3)
    for image_id in image_ids:
        image = dataset_test.load_image(image_id)
        mask, class_ids = dataset_test.load_mask(image_id)
        visualize.display_top_masks(image, mask, class_ids, dataset_test.class_names, limit=7)

def bounding_boxes(subfolder):
    dataset_train = CarPlateDataset()
    dataset_train.load_carplates(os.path.join(ROOT_DIR, subfolder), "val")
    dataset_train.prepare()

    # Load random image and mask.
    # image_id = random.choice(dataset_train.image_ids)

    for image_id in dataset_train.image_ids:
        image = dataset_train.load_image(image_id)
        mask, class_ids = dataset_train.load_mask(image_id)
        # Compute Bounding box
        bbox = utils.extract_bboxes(mask)

        # Display image and additional stats
        print("image_id ", image_id, dataset_train.image_reference(image_id))

        # Display image and instances
        visualize.display_instances(image, bbox, mask, class_ids, dataset_train.class_names, figsize=(3,3))

def mini_mask(config, subfolder):
    dataset_train = CarPlateDataset()
    dataset_train.load_carplates(os.path.join(ROOT_DIR, subfolder), "val")
    dataset_train.prepare()

    image_id = np.random.choice(dataset_train.image_ids, 1)[0]
    image, image_meta, class_ids, bbox, mask = modellib.load_image_gt(
    dataset_train, config, image_id, augment=True, use_mini_mask=True)

    log("image", image)
    log("image_meta", image_meta)
    log("class_ids", class_ids)
    log("bbox", bbox)
    log("mask", mask)

    display_images([image]+[mask[:,:,i] for i in range(min(mask.shape[-1], 7))])

    visualize.display_instances(image, bbox, mask, class_ids, dataset_train.class_names)

def get_ax(rows=1, cols=1, size=16):
    """Return a Matplotlib Axes array to be used in
    all visualizations in the notebook. Provide a
    central point to control graph sizes.

    Adjust the size attribute to control how big to render images
    """
    _, ax = plt.subplots(rows, cols, figsize=(size*cols, size*rows))
    return ax

def display_anchor(config, subfolder):

    dataset_train = CarPlateDataset()
    dataset_train.load_carplates(os.path.join(ROOT_DIR, subfolder), "val")
    dataset_train.prepare()

    backbone_shapes = modellib.compute_backbone_shapes(config, config.IMAGE_SHAPE)
    anchors = utils.generate_pyramid_anchors(config.RPN_ANCHOR_SCALES,
                                            config.RPN_ANCHOR_RATIOS,
                                            backbone_shapes,
                                            config.BACKBONE_STRIDES,
                                            config.RPN_ANCHOR_STRIDE)

    # Print summary of anchors
    num_levels = len(backbone_shapes)
    anchors_per_cell = len(config.RPN_ANCHOR_RATIOS)
    print("Count: ", anchors.shape[0])
    print("Scales: ", config.RPN_ANCHOR_SCALES)
    print("ratios: ", config.RPN_ANCHOR_RATIOS)
    print("Anchors per Cell: ", anchors_per_cell)
    print("Levels: ", num_levels)
    anchors_per_level = []
    for l in range(num_levels):
        num_cells = backbone_shapes[l][0] * backbone_shapes[l][1]
        anchors_per_level.append(anchors_per_cell * num_cells // config.RPN_ANCHOR_STRIDE**2)
        print("Anchors in Level {}: {}".format(l, anchors_per_level[l]))

    image_id = np.random.choice(dataset_train.image_ids, 1)[0]
    image, image_meta, _, _, _ = modellib.load_image_gt(dataset_train, config, image_id)
    fig, ax = plt.subplots(1, figsize=(10, 10))
    ax.imshow(image)
    levels = len(backbone_shapes)

    for level in range(levels):
        colors = visualize.random_colors(levels)
        # Compute the index of the anchors at the center of the image
        level_start = sum(anchors_per_level[:level]) # sum of anchors of previous levels
        level_anchors = anchors[level_start:level_start+anchors_per_level[level]]
        print("Level {}. Anchors: {:6}  Feature map Shape: {}".format(level, level_anchors.shape[0],
                                                                    backbone_shapes[level]))
        center_cell = backbone_shapes[level] // 2
        center_cell_index = (center_cell[0] * backbone_shapes[level][1] + center_cell[1])
        level_center = center_cell_index * anchors_per_cell
        center_anchor = anchors_per_cell * (
            (center_cell[0] * backbone_shapes[level][1] / config.RPN_ANCHOR_STRIDE**2) \
            + center_cell[1] / config.RPN_ANCHOR_STRIDE)
        level_center = int(center_anchor)

        # Draw anchors. Brightness show the order in the array, dark to bright.
        for i, rect in enumerate(level_anchors[level_center:level_center+anchors_per_cell]):
            y1, x1, y2, x2 = rect
            p = patches.Rectangle((x1, y1), x2-x1, y2-y1, linewidth=2, facecolor='none',
                                edgecolor=(i+1)*np.array(colors[level]) / anchors_per_cell)
            ax.add_patch(p)
    # plt.show()

    random_rois = 2000
    g = modellib.data_generator(
        dataset_train, config, shuffle=True, random_rois=random_rois,
        batch_size=4,
        detection_targets=True)

    # Get Next Image
    if random_rois:
        [normalized_images, image_meta, rpn_match, rpn_bbox, gt_class_ids, gt_boxes, gt_masks, rpn_rois, rois], \
        [mrcnn_class_ids, mrcnn_bbox, mrcnn_mask] = next(g)

        log("rois", rois)
        log("mrcnn_class_ids", mrcnn_class_ids)
        log("mrcnn_bbox", mrcnn_bbox)
        log("mrcnn_mask", mrcnn_mask)
    else:
        [normalized_images, image_meta, rpn_match, rpn_bbox, gt_boxes, gt_masks], _ = next(g)

    log("gt_class_ids", gt_class_ids)
    log("gt_boxes", gt_boxes)
    log("gt_masks", gt_masks)
    log("rpn_match", rpn_match, )
    log("rpn_bbox", rpn_bbox)
    image_id = modellib.parse_image_meta(image_meta)["image_id"][0]
    print("image_id: ", image_id, dataset_train.image_reference(image_id))

    # Remove the last dim in mrcnn_class_ids. It's only added
    # to satisfy Keras restriction on target shape.
    mrcnn_class_ids = mrcnn_class_ids[:,:,0]

    b = 0

    # Restore original image (reverse normalization)
    sample_image = modellib.unmold_image(normalized_images[b], config)

    # Compute anchor shifts.
    indices = np.where(rpn_match[b] == 1)[0]
    refined_anchors = utils.apply_box_deltas(anchors[indices], rpn_bbox[b, :len(indices)] * config.RPN_BBOX_STD_DEV)
    log("anchors", anchors)
    log("refined_anchors", refined_anchors)

    # Get list of positive anchors
    positive_anchor_ids = np.where(rpn_match[b] == 1)[0]
    print("Positive anchors: {}".format(len(positive_anchor_ids)))
    negative_anchor_ids = np.where(rpn_match[b] == -1)[0]
    print("Negative anchors: {}".format(len(negative_anchor_ids)))
    neutral_anchor_ids = np.where(rpn_match[b] == 0)[0]
    print("Neutral anchors: {}".format(len(neutral_anchor_ids)))

    # ROI breakdown by class
    for c, n in zip(dataset_train.class_names, np.bincount(mrcnn_class_ids[b].flatten())):
        if n:
            print("{:23}: {}".format(c[:20], n))

    # Show positive anchors
    # fig, ax = plt.subplots(1, figsize=(16, 16))
    visualize.draw_boxes(sample_image, boxes=anchors[positive_anchor_ids],
                        refined_boxes=refined_anchors, ax=ax)

    visualize.draw_boxes(sample_image, boxes=anchors[negative_anchor_ids])

    if random_rois:
        # Class aware bboxes
        bbox_specific = mrcnn_bbox[b, np.arange(mrcnn_bbox.shape[1]), mrcnn_class_ids[b], :]

        # Refined ROIs
        refined_rois = utils.apply_box_deltas(rois[b].astype(np.float32), bbox_specific[:,:4] * config.BBOX_STD_DEV)

        # Class aware masks
        mask_specific = mrcnn_mask[b, np.arange(mrcnn_mask.shape[1]), :, :, mrcnn_class_ids[b]]

        visualize.draw_rois(sample_image, rois[b], refined_rois, mask_specific, mrcnn_class_ids[b], dataset_train.class_names)

        # Any repeated ROIs?
        rows = np.ascontiguousarray(rois[b]).view(np.dtype((np.void, rois.dtype.itemsize * rois.shape[-1])))
        _, idx = np.unique(rows, return_index=True)
        print("Unique ROIs: {} out of {}".format(len(idx), rois.shape[1]))

    if random_rois:
    # Dispalay ROIs and corresponding masks and bounding boxes
        ids = random.sample(range(rois.shape[1]), 8)

        images = []
        titles = []
        for i in ids:
            image = visualize.draw_box(sample_image.copy(), rois[b,i,:4].astype(np.int32), [255, 0, 0])
            image = visualize.draw_box(image, refined_rois[i].astype(np.int64), [0, 255, 0])
            images.append(image)
            titles.append("ROI {}".format(i))
            images.append(mask_specific[i] * 255)
            titles.append(dataset_train.class_names[mrcnn_class_ids[b,i]][:20])

        display_images(images, titles, cols=4, cmap="Blues", interpolation="none")

    if random_rois:
        limit = 10
        temp_g = modellib.data_generator(
            dataset_train, config, shuffle=True, random_rois=10000,
            batch_size=1, detection_targets=True)
        total = 0
        for i in range(limit):
            _, [ids, _, _] = next(temp_g)
            positive_rois = np.sum(ids[0] > 0)
            total += positive_rois
            print("{:5} {:5.2f}".format(positive_rois, positive_rois/ids.shape[1]))
        print("Average percent: {:.2f}".format(total/(limit*ids.shape[1])))
    plt.show()

def inspect_model(config, model_path):

    dataset_train = CarPlateDataset()
    dataset_train.load_carplates(os.path.join(ROOT_DIR, 'car_plate_data'), "test")
    dataset_train.prepare()

    model = modellib.MaskRCNN(mode='inference', config=config, model_dir=ROOT_DIR)
    model.load_weights(model_path, by_name=True, exclude=[
            "mrcnn_class_logits", "mrcnn_bbox_fc",
            "mrcnn_bbox", "mrcnn_mask"])

    image_id = random.choice(dataset_train.image_ids)
    image, image_meta, gt_class_id, gt_bbox, gt_mask =\
        modellib.load_image_gt(dataset_train, config, image_id, use_mini_mask=False)
    info = dataset_train.image_info[image_id]
    print("image ID: {}.{} ({}) {}".format(info["source"], info["id"], image_id,
                                        dataset_train.image_reference(image_id)))

    # Generate RPN trainig targets
    # target_rpn_match is 1 for positive anchors, -1 for negative anchors
    # and 0 for neutral anchors.
    target_rpn_match, target_rpn_bbox = modellib.build_rpn_targets(
        image.shape, model.get_anchors(image.shape), gt_class_id, gt_bbox, model.config)
    log("target_rpn_match", target_rpn_match)
    log("target_rpn_bbox", target_rpn_bbox)

    positive_anchor_ix = np.where(target_rpn_match[:] == 1)[0]
    negative_anchor_ix = np.where(target_rpn_match[:] == -1)[0]
    neutral_anchor_ix = np.where(target_rpn_match[:] == 0)[0]
    positive_anchors = model.anchors[positive_anchor_ix]
    negative_anchors = model.anchors[negative_anchor_ix]
    neutral_anchors = model.anchors[neutral_anchor_ix]
    log("positive_anchors", positive_anchors)
    log("negative_anchors", negative_anchors)
    log("neutral anchors", neutral_anchors)

    # Apply refinement deltas to positive anchors
    refined_anchors = utils.apply_box_deltas(
        positive_anchors,
        target_rpn_bbox[:positive_anchors.shape[0]] * model.config.RPN_BBOX_STD_DEV)
    log("refined_anchors", refined_anchors, )

    # visualize.draw_boxes(image, boxes=positive_anchors, refined_boxes=refined_anchors, ax=get_ax())

        # Run RPN sub-graph
    pillar = model.keras_model.get_layer("ROI").output  # node to start searching from

    # TF 1.4 and 1.9 introduce new versions of NMS. Search for all names to support TF 1.3~1.10
    nms_node = model.ancestor(pillar, "ROI/rpn_non_max_suppression:0")
    if nms_node is None:
        nms_node = model.ancestor(pillar, "ROI/rpn_non_max_suppression/NonMaxSuppressionV2:0")
    if nms_node is None: #TF 1.9-1.10
        nms_node = model.ancestor(pillar, "ROI/rpn_non_max_suppression/NonMaxSuppressionV3:0")

    rpn = model.run_graph([image], [
        ("rpn_class", model.keras_model.get_layer("rpn_class").output),
        ("pre_nms_anchors", model.ancestor(pillar, "ROI/pre_nms_anchors:0")),
        ("refined_anchors", model.ancestor(pillar, "ROI/refined_anchors:0")),
        ("refined_anchors_clipped", model.ancestor(pillar, "ROI/refined_anchors_clipped:0")),
        ("post_nms_anchor_ix", nms_node),
        ("proposals", model.keras_model.get_layer("ROI").output),
    ])

    limit = 100
    sorted_anchor_ids = np.argsort(rpn['rpn_class'][:,:,1].flatten())[::-1]
    visualize.draw_boxes(image, boxes=model.anchors[sorted_anchor_ids[:limit]], ax=get_ax())

    limit = 50
    ax = get_ax(1, 2)
    pre_nms_anchors = utils.denorm_boxes(rpn["pre_nms_anchors"][0], image.shape[:2])
    refined_anchors = utils.denorm_boxes(rpn["refined_anchors"][0], image.shape[:2])
    refined_anchors_clipped = utils.denorm_boxes(rpn["refined_anchors_clipped"][0], image.shape[:2])
    visualize.draw_boxes(image, boxes=pre_nms_anchors[:limit],
                        refined_boxes=refined_anchors[:limit], ax=ax[0])
    visualize.draw_boxes(image, refined_boxes=refined_anchors_clipped[:limit], ax=ax[1])

    limit = 50
    ixs = rpn["post_nms_anchor_ix"][:limit]
    visualize.draw_boxes(image, refined_boxes=refined_anchors_clipped[ixs], ax=get_ax())

        # Show final proposals
    # These are the same as the previous step (refined anchors
    # after NMS) but with coordinates normalized to [0, 1] range.
    limit = 50
    # Convert back to image coordinates for display
    h, w = config.IMAGE_SHAPE[:2]
    proposals = rpn['proposals'][0, :limit] * np.array([h, w, h, w])
    visualize.draw_boxes(image, refined_boxes=proposals, ax=get_ax())

    mrcnn = model.run_graph([image], [
    ("proposals", model.keras_model.get_layer("ROI").output),
    ("probs", model.keras_model.get_layer("mrcnn_class").output),
    ("deltas", model.keras_model.get_layer("mrcnn_bbox").output),
    ("masks", model.keras_model.get_layer("mrcnn_mask").output),
    ("detections", model.keras_model.get_layer("mrcnn_detection").output),
        ])

    # # Get detection class IDs. Trim zero padding.
    # det_class_ids = mrcnn['detections'][0, :, 4].astype(np.int32)
    # det_count = np.where(np.array(det_class_ids) == 0)
    # print(det_count)
    # det_class_ids = det_class_ids[:det_count]
    # detections = mrcnn['detections'][0, :det_count]

    # print("{} detections: {}".format(
    #     det_count, np.array(dataset_train.class_names)[det_class_ids]))

    # captions = ["{} {:.3f}".format(dataset_train.class_names[int(c)], s) if c > 0 else ""
    #             for c, s in zip(detections[:, 4], detections[:, 5])]
    # visualize.draw_boxes(
    #     image,
    #     refined_boxes=utils.denorm_boxes(detections[:, :4], image.shape[:2]),
    #     visibilities=[2] * len(detections),
    #     captions=captions, title="Detections",
    #     ax=get_ax())

        # Proposals are in normalized coordinates. Scale them
    # to image coordinates.
    h, w = config.IMAGE_SHAPE[:2]
    proposals = np.around(mrcnn["proposals"][0] * np.array([h, w, h, w])).astype(np.int32)

    # Class ID, score, and mask per proposal
    roi_class_ids = np.argmax(mrcnn["probs"][0], axis=1)
    roi_scores = mrcnn["probs"][0, np.arange(roi_class_ids.shape[0]), roi_class_ids]
    roi_class_names = np.array(dataset_train.class_names)[roi_class_ids]
    roi_positive_ixs = np.where(roi_class_ids > 0)[0]

    # How many ROIs vs empty rows?
    print("{} Valid proposals out of {}".format(np.sum(np.any(proposals, axis=1)), proposals.shape[0]))
    print("{} Positive ROIs".format(len(roi_positive_ixs)))

    # Class counts
    print(list(zip(*np.unique(roi_class_names, return_counts=True))))

        # Display a random sample of proposals.
    # Proposals classified as background are dotted, and
    # the rest show their class and confidence score.
    limit = 200
    ixs = np.random.randint(0, proposals.shape[0], limit)
    captions = ["{} {:.3f}".format(dataset_train.class_names[c], s) if c > 0 else ""
                for c, s in zip(roi_class_ids[ixs], roi_scores[ixs])]
    visualize.draw_boxes(image, boxes=proposals[ixs],
                        visibilities=np.where(roi_class_ids[ixs] > 0, 2, 1),
                        captions=captions, title="ROIs Before Refinement",
                        ax=get_ax())
        # Class-specific bounding box shifts.
    roi_bbox_specific = mrcnn["deltas"][0, np.arange(proposals.shape[0]), roi_class_ids]
    log("roi_bbox_specific", roi_bbox_specific)

    # Apply bounding box transformations
    # Shape: [N, (y1, x1, y2, x2)]
    refined_proposals = utils.apply_box_deltas(
        proposals, roi_bbox_specific * config.BBOX_STD_DEV).astype(np.int32)
    log("refined_proposals", refined_proposals)

    # Show positive proposals
    # ids = np.arange(roi_boxes.shape[0])  # Display all
    limit = 5
    ids = np.random.randint(0, len(roi_positive_ixs), limit)  # Display random sample
    captions = ["{} {:.3f}".format(dataset_train.class_names[c], s) if c > 0 else ""
                for c, s in zip(roi_class_ids[roi_positive_ixs][ids], roi_scores[roi_positive_ixs][ids])]
    visualize.draw_boxes(image, boxes=proposals[roi_positive_ixs][ids],
                        refined_boxes=refined_proposals[roi_positive_ixs][ids],
                        visibilities=np.where(roi_class_ids[roi_positive_ixs][ids] > 0, 1, 0),
                        captions=captions, title="ROIs After Refinement",
                        ax=get_ax())

        # Remove boxes classified as background
    keep = np.where(roi_class_ids > 0)[0]
    print("Keep {} detections:\n{}".format(keep.shape[0], keep))

        # Remove low confidence detections
    keep = np.intersect1d(keep, np.where(roi_scores >= config.DETECTION_MIN_CONFIDENCE)[0])
    print("Remove boxes below {} confidence. Keep {}:\n{}".format(
        config.DETECTION_MIN_CONFIDENCE, keep.shape[0], keep))

        # Apply per-class non-max suppression
    pre_nms_boxes = refined_proposals[keep]
    pre_nms_scores = roi_scores[keep]
    pre_nms_class_ids = roi_class_ids[keep]

    nms_keep = []
    for class_id in np.unique(pre_nms_class_ids):
        # Pick detections of this class
        ixs = np.where(pre_nms_class_ids == class_id)[0]
        # Apply NMS
        class_keep = utils.non_max_suppression(pre_nms_boxes[ixs],
                                                pre_nms_scores[ixs],
                                                config.DETECTION_NMS_THRESHOLD)
        # Map indicies
        class_keep = keep[ixs[class_keep]]
        nms_keep = np.union1d(nms_keep, class_keep)
        print("{:22}: {} -> {}".format(dataset_train.class_names[class_id][:20],
                                    keep[ixs], class_keep))

    keep = np.intersect1d(keep, nms_keep).astype(np.int32)
    print("\nKept after per-class NMS: {}\n{}".format(keep.shape[0], keep))

        # Show final detections
    ixs = np.arange(len(keep))  # Display all
    # ixs = np.random.randint(0, len(keep), 10)  # Display random sample
    captions = ["{} {:.3f}".format(dataset_train.class_names[c], s) if c > 0 else ""
                for c, s in zip(roi_class_ids[keep][ixs], roi_scores[keep][ixs])]
    visualize.draw_boxes(
        image, boxes=proposals[keep][ixs],
        refined_boxes=refined_proposals[keep][ixs],
        visibilities=np.where(roi_class_ids[keep][ixs] > 0, 1, 0),
        captions=captions, title="Detections after NMS",
        ax=get_ax())

    plt.show()

def detection(model, image_path=None, video_path=None):
    assert image_path or video_path
    # class_names = ['BG','carplate']
    class_names = ['BG','0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H',
                            'J', 'K', 'L', 'M', 'N', 'P',  'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z']
    # Image or video?
    if image_path:
        # Run model detection and generate the color splash effect
        imgs = os.listdir(image_path)
        for img in imgs:
            print("Running on {}".format(img))
            # Read image
            image = skimage.io.imread(image_path + '/' + img)
            # Detect objects
            # image = skimage.color.gray2rgb(image)
            r = model.detect([image], verbose=1)[0]

            result = helper.get_char_result(image, r['rois'], r['masks'], r['class_ids'],
                                      class_names, r['scores'], show_bbox=True, score_threshold=0.50,
                                      show_mask=False)
            print(result)
            visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'],
                                class_names, r['scores'], figsize=(3,3), show_mask=False)
            cv2.waitKey(0)

def train_model(model, subfolder, mode_train):
    """Train the model."""
    # Training dataset.

    dataset_train = CarPlateDataset()
    dataset_train.load_carplates(os.path.join(ROOT_DIR, subfolder), "train")
    dataset_train.prepare()

    # Validation dataset
    dataset_val = CarPlateDataset()
    dataset_val.load_carplates(os.path.join(ROOT_DIR, subfolder), "val")
    dataset_val.prepare()

    if mode_train == 'all':
        print("Training network all")
        model.train(dataset_train, dataset_val,
                    learning_rate=config.LEARNING_RATE,
                    epochs=1000,
                    layers='all')
    elif mode_train == 'heads':
        print("Training network heads")
        model.train(dataset_train, dataset_val,
                    learning_rate=config.LEARNING_RATE,
                    epochs=1000,
                    layers='all')

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(
       description='Train Mask R-CNN to detect carplates.')

    parser.add_argument('command',
                    metavar='<command>',
                    help="'train, detect, display_data, display_box, mini_mask, display_anchor'")

    parser.add_argument('--dataset', required=False,
                        metavar="/path/to/balloon/dataset/",
                        help='Directory inference dataset')

    parser.add_argument('--dataset_train', required=False,
                        metavar="/path/to/balloon/dataset/",
                        help='Directory training dataset')
    parser.add_argument('--weights', required=False,
                        metavar="/path/to/weights.h5",
                        help="Path to weights .h5 file or 'coco'")
    parser.add_argument('--mode_train', required=False,
                        metavar="/path/to/weights.h5",
                        help="Path to weights .h5 file or 'coco'")

    args = parser.parse_args()

    if args.command == 'train':
        config = CarPlateConfig()
    elif args.command == 'detect':
        config = InferenceConfig()

    #config.display()

    if args.command == 'display_data':
        display_dataset(args.dataset)
    if args.command == 'display_box':
        bounding_boxes(args.dataset)
    if args.command == 'mini_mask':
        mini_mask(config, args.dataset)
    if args.command == 'display_anchor':
        display_anchor(config, args.dataset)
    if args.command == 'inspect_model':
        inspect_model(config, args.weights)
    if args.command == 'eval_carplate':
        model = modellib.MaskRCNN(mode='inference', config=config, model_dir=os.path.join(ROOT_DIR, 'logs'))

        model_path = os.path.join(ROOT_DIR, args.weights)
        print("Loading weights ", model_path)
        model.load_weights(model_path, by_name=True)

        evaluate.evaluate_carplate(args.dataset, model, ROOT_DIR)

    if args.command == 'eval_chars':

        lp_path = args.weights.split(',')[0]
        chars_path = args.weights.split(',')[1]

        lp_config = evaluate.EvalCarPlateConfig()
        lp_config.display()
        chars_config = evaluate.EvalCharConfig()
        chars_config.display

        lp_model = modellib.MaskRCNN(mode='inference', config=lp_config, model_dir=os.path.join(ROOT_DIR, 'logs'))
        model_path = os.path.join(ROOT_DIR, lp_path)
        print("Loading weights ", model_path)
        lp_model.load_weights(model_path, by_name=True)

        chars_model = modellib.MaskRCNN(mode='inference', config=chars_config, model_dir=os.path.join(ROOT_DIR, 'logs'))
        model_path = os.path.join(ROOT_DIR, chars_path)
        print("Loading weights ", model_path)
        chars_model.load_weights(model_path, by_name=True)

        evaluate.evaluate_numbers(args.dataset, lp_model, chars_model, ROOT_DIR)

    if args.command == 'train':
        model = modellib.MaskRCNN(mode='training', config=config, model_dir=os.path.join(ROOT_DIR, 'logs'))
        if args.mode_train == 'all':
            weights_path = os.path.join(ROOT_DIR, 'resnet50_weights.h5')
        elif args.mode_train == 'heads':
            weights_path = args.weights

        print("Logs: ", os.path.join(ROOT_DIR, 'logs'))
        print("Loading weights: ", weights_path)

        # weights_path = 'char_101.h5'
        model.load_weights(weights_path, by_name=True)

        train_model(model, args.dataset_train, args.mode_train)

    if args.command == 'detect':
        model = modellib.MaskRCNN(mode='inference', config=config, model_dir=os.path.join(ROOT_DIR, 'logs'))
        # model_path = model.find_last()
        model_path = os.path.join(ROOT_DIR, args.weights)
        print("Loading weights ", model_path)
        model.load_weights(model_path, by_name=True)

        dataset = os.path.join(ROOT_DIR, args.dataset)
        detection(model, image_path=dataset,
                                video_path=None)


#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
This is a small example which creates a twitch stream to connect with
and changes the color of the video according to the colors provided in
the chat.
"""
from __future__ import print_function
from twitchstream.outputvideo import TwitchBufferedOutputStream
import argparse
import numpy as np

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    required = parser.add_argument_group('required arguments')
    required.add_argument('-s', '--streamkey',
                          help='twitch streamkey',
                          required=True)
    args = parser.parse_args()

    with TwitchBufferedOutputStream(
            twitch_stream_key=args.streamkey,
            width=640,
            height=480,
            fps=30.,
            verbose=True,
            enable_audio=True) as videostream:

        frame = np.zeros((480, 640, 3))

        while True:
            if videostream.get_video_frame_buffer_state() < 30:
                frame = np.random.rand(480, 640, 3)
                videostream.send_video_frame(frame)

            if videostream.get_audio_buffer_state() < 30:
                left_audio = np.random.randn(1470)
                right_audio = np.random.randn(1470)
                videostream.send_audio(left_audio, right_audio)


#!/usr/bin/env python3

"""Unconditional sampling from a diffusion model."""

import argparse
from functools import partial
from pathlib import Path

import jax
import jax.numpy as jnp
from PIL import Image
from tqdm import tqdm, trange

from diffusion import get_model, get_models, load_params, sampling, utils

MODULE_DIR = Path(__file__).resolve().parent

def main():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    p.add_argument('--batch-size', '-bs', type=int, default=1,
                   help='the number of images per batch')
    p.add_argument('--checkpoint', type=str,
                   help='the checkpoint to use')
    p.add_argument('--eta', type=float, default=1.,
                   help='the amount of noise to add during sampling (0-1)')
    p.add_argument('--init', type=str,
                   help='the init image')
    p.add_argument('--model', type=str, choices=get_models(), required=True,
                   help='the model to use')
    p.add_argument('-n', type=int, default=1,
                   help='the number of images to sample')
    p.add_argument('--seed', type=int, default=0,
                   help='the random seed')
    p.add_argument('--starting-timestep', '-st', type=float, default=0.9,
                   help='the timestep to start at (used with init images)')
    p.add_argument('--steps', type=int, default=1000,
                   help='the number of timesteps')
    args = p.parse_args()

    model = get_model(args.model)
    checkpoint = args.checkpoint
    if not checkpoint:
        checkpoint = MODULE_DIR / f'checkpoints/{args.model}.pkl'
    params = load_params(checkpoint)

    if args.init:
        _, y, x = model.shape
        init = Image.open(args.init).convert('RGB').resize((x, y), Image.LANCZOS)
        init = utils.from_pil_image(init)[None]

    key = jax.random.PRNGKey(args.seed)

    def run(key, n):
        tqdm.write('Sampling...')
        key, subkey = jax.random.split(key)
        noise = jax.random.normal(subkey, [n, *model.shape])
        key, subkey = jax.random.split(key)
        sample_step = partial(sampling.jit_sample_step, extra_args={})
        steps = utils.get_ddpm_schedule(jnp.linspace(1, 0, args.steps + 1)[:-1])
        if args.init:
            steps = steps[steps < args.starting_timestep]
            alpha, sigma = utils.t_to_alpha_sigma(steps[0])
            noise = init * alpha + noise * sigma
        return sampling.sample_loop(model, params, subkey, noise, steps, args.eta, sample_step)

    def run_all(key, n, batch_size):
        for i in trange(0, n, batch_size):
            key, subkey = jax.random.split(key)
            cur_batch_size = min(n - i, batch_size)
            outs = run(key, cur_batch_size)
            for j, out in enumerate(outs):
                utils.to_pil_image(out).save(f'out_{i + j:05}.png')

    try:
        run_all(key, args.n, args.batch_size)
    except KeyboardInterrupt:
        pass

if __name__ == '__main__':
    main()


import os, sys
#os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="1"

import numpy as np
import pandas as pd
from tqdm import tqdm
import cv2
from sklearn.utils import shuffle
from ml_stratifiers import MultilabelStratifiedKFold
import albumentations as A
import warnings
warnings.filterwarnings("ignore")

from classification_models.resnet.models import ResNet34
from classification_models.resnet import preprocess_input
MODEL_PATH = 'Christof/models/ResNet34/3/'
SIZE = 256

# Load dataset info
path_to_train = 'Christof/assets/train_rgb_256/'
data = pd.read_csv('Christof/assets/train.csv')

clusters = pd.read_csv('Russ/cluster4_folds.csv')

folds = dict(zip(clusters.Id,clusters.cluster4))
data['fold'] = data['Id'].apply(lambda x: folds[x])

def get_fold_ids(fold_id,data_set_info, shuff = True):
    fold_info = np.array([item['fold'] for item in data_set_info])
    val_ids = np.where(fold_info == fold_id)[0]
    train_ids = np.where(fold_info != fold_id)[0]
    if shuff:
        shuffle(val_ids)
        shuffle(train_ids)
    return train_ids, val_ids

normal_aug = A.Compose([A.OneOf([A.Rotate((-180,180)),
                                 A.Rotate((-180,180),border_mode=cv2.BORDER_CONSTANT)]),
                        A.Flip(p=0.75)
                        ])

class data_generator:

    @staticmethod
    def create_train(dataset_info, batch_size, shape, augument=True, oversample_factor = 0):
        assert shape[2] == 3

        if oversample_factor > 0:

            rare_dataset_info = np.array([item for item in dataset_info if np.isin(item['labels'], rare_classes).any()])
            for i in range(oversample_factor):
                dataset_info = np.append(dataset_info,rare_dataset_info)
        while True:
            dataset_info = shuffle(dataset_info)
            for start in range(0, len(dataset_info), batch_size):
                end = min(start + batch_size, len(dataset_info))
                batch_images = []
                X_train_batch = dataset_info[start:end]
                batch_labels = np.zeros((len(X_train_batch), 28))
                for i in range(len(X_train_batch)):
                    image = data_generator.load_image(X_train_batch[i]['path'])
                    #rare = np.isin(X_train_batch[i]['labels'], rare_classes).any()

                    if augument:
                        image = data_generator.augment(normal_aug,image)

                    batch_images.append(image)
                    batch_labels[i][X_train_batch[i]['labels']] = 1
                yield np.array(batch_images, np.float32), batch_labels

    @staticmethod
    def load_image(path):
        image = cv2.imread(path + '.png', cv2.IMREAD_UNCHANGED)
        image = preprocess_input(image)
        return image

    @staticmethod
    def augment(aug,image):
        image_aug = aug(image=image)['image']
        return image_aug

from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential, load_model
from keras.layers import Activation, Dropout, Flatten, Dense, GlobalMaxPooling2D, BatchNormalization, Input, Conv2D
from keras.callbacks import ModelCheckpoint
from keras import metrics
from keras.optimizers import Adam
from keras import backend as K
import keras
from keras.models import Model
from sklearn.metrics import f1_score

def create_model(input_shape, n_out):
    input_tensor = Input(shape=(SIZE, SIZE, 3))

    base_model = ResNet34(include_top=False,
                             weights='imagenet',
                             input_shape=(SIZE, SIZE, 3),input_tensor=input_tensor)

    x = base_model.output
    x = Conv2D(32, kernel_size=(1, 1), activation='relu')(x)
    x = Flatten()(x)
    x = Dropout(0.5)(x)
    x = Dense(1024, activation='relu')(x)
    x = Dropout(0.5)(x)
    output = Dense(n_out, activation='sigmoid')(x)
    model = Model(input_tensor, output)

    return model

model = create_model(
    input_shape=(SIZE, SIZE, 3),
    n_out=28)

train_dataset_info = []
for name, labels in zip(data['Id'], data['Target'].str.split(' ')):
    train_dataset_info.append({
        'path': os.path.join(path_to_train, name),
        'labels': np.array([int(label) for label in labels]),
        'fold':folds[name]})
train_dataset_info = np.array(train_dataset_info)

y_trues = []
val_preds = []

for fold_id in range(4):
    model.load_weights(MODEL_PATH + 'base_100/' + 'model_f{}.h5'.format(fold_id))

    train_indexes, valid_indexes = get_fold_ids(fold_id, train_dataset_info)
    #train_indexes, valid_indexes = next(kf)

    #validation_generator = data_generator.create_train(train_dataset_info[valid_indexes],
    #                                                   1, (SIZE, SIZE, 3), augument=False,oversample_factor=0)

    val_pred = np.zeros(shape=(len(valid_indexes),28))
    y_true= np.zeros(shape=(len(valid_indexes),28))
    for i, info in tqdm(enumerate(train_dataset_info[valid_indexes])):
        image = data_generator.load_image(info['path'])

        val_pred[i] = model.predict(image[np.newaxis])[0]
        y_true[i][info['labels']]=1

    y_trues.append(y_true)
    val_preds.append(val_pred)

y_trues = np.concatenate(y_trues, axis = 0)
val_preds = np.concatenate(val_preds, axis = 0)

f1_res = f1_score(y_trues, np.round(val_preds), average='macro')

thresholds = np.linspace(0, 1, 101)
scores = np.array([f1_score(y_trues, (val_preds > t).astype(int), average='macro') for t in thresholds])

threshold_best_index = np.argmax(scores)
score_best = scores[threshold_best_index]
threshold_best = thresholds[threshold_best_index]
print('')
print('f1_best',score_best)
print('threshold_best',threshold_best)
print('')

test_preds = []

sample_submit = pd.read_csv('Christof/assets/sample_submission.csv')
test_path = 'Christof/assets/test_rgb_256/'

for fold_id in range(4):
    model.load_weights(MODEL_PATH + 'base_100/' + 'model_f{}.h5'.format(fold_id))

    test_pred = np.zeros(shape=(len(sample_submit['Id']),28))
    for i in tqdm(range(len(sample_submit['Id']))):
        image = data_generator.load_image(test_path + sample_submit['Id'][i])

        test_pred[i] = model.predict(image[np.newaxis])[0]

    test_preds.append(test_pred)

from scipy.stats.mstats import gmean
test_preds_avg = np.mean(test_preds,axis = 0)
#test_preds_avg = gmean(test_preds, axis = 0)

binary_preds = (test_preds_avg > threshold_best).astype(bool)

class_preds = [np.arange(28)[binary_pred] for binary_pred in binary_preds]
str_preds = [' '.join(str(l) for l in class_pred) for class_pred in class_preds]

sample_submit['Predicted'] = str_preds

#np.save('best_ckpt_fold0_preds.npy', preds)
sample_submit.to_csv(MODEL_PATH + 'predictions/' + 'submission_4fold_mean_notta.csv', index=False)

"""
submit = pd.read_csv('Christof/assets/sample_submission.csv')
predicted = []
draw_predict = []

for name in tqdm(submit['Id']):
    path = os.path.join('Christof/assets/test_rgb_256/', name)
    image = data_generator.load_image(path)
    score_predict = model.predict(image[np.newaxis])[0]
    draw_predict.append(score_predict)

    #thresh = max(score_predict[np.argsort(score_predict, axis=-1)[-5]], 0.2)
    label_predict = np.arange(28)[score_predict >= 0.5]
    str_predict_label = ' '.join(str(l) for l in label_predict)
    predicted.append(str_predict_label)

submit['Predicted'] = predicted
# np.save('draw_predict_InceptionV3.npy', score_predict)
submit.to_csv(MODEL_PATH + 'sec_submission.csv', index=False)
"""


import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import random

from collections import namedtuple, deque
from maddpgtennis import DDPG, ReplayBuffer

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

class MADDPG:

    '''Multi Agent DDPG'''

    def __init__(self, n_state, n_action, n_agents, seed=0):

        '''Initialize the MADDPG class

        Params:
            n_state     : size of observation (state-space)
            n_action    : size of action
            n_agents    : how many agents are trained
            seed        : seed random number
        '''

        self.n_state = n_state
        self.n_action = n_action
        self.n_agents = n_agents

        self.lr_act = 1e-4
        self.lr_crit = 3e-4
        self.agents_ = []
        for i in range(self.n_agents):
            agent = DDPG(self.n_state, self.n_action, self.n_agents,
                         self.lr_act, self.lr_crit, i)
            self.agents_.append(agent)

        self.buffer_size = int(1e6)
        self.batch_size = 128
        self.t_step = 0
        self.gamma = 0.99
        self.updateTargetNet = 1 # 20
        self.num_learn = 1 # 15
        self.noise_weight = 1
        self.noise_decay = 0.997
        self.noise_min = 0.1
        self.noise_update = 1

    def add_noise(self):

        '''Adding normal distribuion noise sample to be used later for agents'
           actions in learning

        Return:
            noise       : noise in tensor form
        '''

        random_dist = np.random.normal(0, 1, self.n_action)
        noise = self.noise_weight * random_dist
        return torch.from_numpy(noise).float().to(device)

    def action_(self, states):

        '''Action of the agents

        Params:
            states      : current state of the environment

        Return:
            actions     : calculate action for each agent on top of noise
        '''

        agents_observation = torch.from_numpy(states).float().to(device)

        actions = []
        for i in range(self.n_agents):
            agent = self.agents_[i]
            noise = self.add_noise().to(device)
            action = agent.action(agents_observation[i]) + noise
            actions.append(action)

        return actions

    def step(self, state, action, reward, state_, done):

        '''Gather memory for experience replay and check learning conditions
           as well as update the noise parameters

        Params:
            state       : current state of the environment
            action      : action taken by the agents
            reward      : reward given to the agent based on the action
            state_      : new state of the environment after action resolved
            done        : status of the environment
        '''

        for i in range(self.n_agents):
            self.agents_[i].memory.add(state[i], action[i], reward[i],
                                       state_[i], done[i])

        if self.t_step % self.noise_update == 0:
            self.noise_weight = max(self.noise_weight*self.noise_decay,
                                    self.noise_min)
        if len(self.agents_[0].memory) > self.batch_size:
            if self.t_step % self.updateTargetNet == 0:
                for _ in range(self.num_learn):
                    for i in range(self.n_agents):
                        agent_exp = self.agents_[i].memory.sample_replay()
                        self.learn(agent_exp, i)
        self.t_step += 1

    def learn(self, agent_exp, agent_no):

        '''Learning steps for the agents

        Params:
            agent_exp   : random sample of memory to be learned from
            agent_no    : unique id of the agent
        '''

        agent = self.agents_[agent_no]
        states, actions, rewards, states_, dones = agent_exp

        actions_ = agent.act_target(states_)
        Q_targets_ = agent.crit_target(states_, actions_)

        Q_targets = rewards + (self.gamma * Q_targets_ * (1 - dones))
        Q_expect = agent.crit_local(states, actions)

        crit_loss = F.mse_loss(Q_expect, Q_targets)
        agent.crit_opt.zero_grad()
        crit_loss.backward()
        torch.nn.utils.clip_grad_norm_(agent.crit_local.parameters(), 1)
        agent.crit_opt.step()

        # update actor
        actions_pred = agent.act_local(states)
        actor_loss = -agent.crit_local(states, actions_pred).mean()
        agent.act_opt.zero_grad()
        actor_loss.backward()
        agent.act_opt.step()

        # soft update
        self.soft_update(agent.crit_local, agent.crit_target)
        self.soft_update(agent.act_local, agent.act_target)

    def soft_update(self, local, target, TAU=0.01):

        '''Soft update in updating the networks parameter

        Params:
            local       : local network
            target      : target network
            TAU         : constant tau that governs the soft-update
        '''

        for local_param, target_param in zip(local.parameters(),
                                             target.parameters()):
            target_param.data.copy_(TAU * local_param.data + \
                                    (1.0-TAU) * target_param.data)


import torch
import torch.nn as nn
from torchvision import datasets, models, transforms
from torch.utils.data.sampler import SubsetRandomSampler
from torch.optim import lr_scheduler
from pathlib import Path
from matplotlib import pyplot as plt
import numpy as np
import torch.nn.functional as F
from torch.autograd import Variable
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
from torchsummary import summary
from tqdm import tqdm
import torch.optim as optim
from cv2 import cv2

from torch.utils.data import DataLoader
from tools.dataloader_bg import DataLoaderX

if torch.cuda.is_available():
    device = torch.device('cuda')
    torch.backends.cudnn.benchmark = True
    torch.cuda.empty_cache()
else:
    device = torch.device('cpu')

path_train = r'D:\training_dataset\simpsons\train'
path_val = r'D:\training_dataset\simpsons\val'

# path_train = r'/media/blackyyen/新增磁碟區1/Dataset/simpsons/train'
# path_val = r'/media/blackyyen/新增磁碟區1/Dataset/simpsons/val'
# data_dir = r'/simpsons'

# number of subprocesses to use for data loading
num_workers = 0
# how many samples per batch to load
train_batch_size = 20
val_batch_size = 20
# learning rate
lr = 1e-3

# convert data to a normalized torch.FloatTensor
train_transforms = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

valid_transforms = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# choose the training and test datasets
train_data = datasets.ImageFolder(path_train, transform=train_transforms)
valid_data = datasets.ImageFolder(path_val, transform=valid_transforms)

# prepare data loaders (combine dataset and sampler)
train_loader = DataLoaderX(train_data,
                           batch_size=train_batch_size,
                           num_workers=num_workers,
                           shuffle=True,
                           pin_memory=True)
val_loader = DataLoaderX(valid_data,
                         batch_size=val_batch_size,
                         num_workers=num_workers,
                         shuffle=False,
                         pin_memory=True)
# 讀取模型
device = torch.device("cuda")
model = torch.load('./classification/pretrained_model/resnext101_32x8d.pth')
# 查看可用的模型
# print(torch.hub.list('pytorch/vision'))
# print(torch.hub.list('facebookresearch/WSL-Images))
# 將模型儲存起來
# model = torch.hub.load('facebookresearch/WSL-Images',
#                        'resnext101_32x8d_wsl',
#                        force_reload=False,
#                        verbose=0)
# torch.save(model, './classification/resnext101_32x8d.pth')
# 輸出模型可用函式
# print(dir(model))
# print(model.layer1[0].conv1._parameters)

# 提取參數fc的輸入參數
fc_features = model.fc.in_features
# 將最後輸出類別改為20
model.fc = nn.Linear(fc_features, 20)
# for name, param in model.named_parameters():
#     if name == 'layer4.0.conv1.weight':
#         break
#     param.requires_grad = False
# 將所有層設定為不可訓練
# for param in model.parameters():
#     param.requires_grad = False
# 將最後一層fc設定為可以訓練
# for param in model.fc.parameters():
#     param.requires_grad = True
# 輸出模型參數
# from torchsummary import summary
# summary(model.to(device), (3, 224, 224))

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
# scheduler = lr_scheduler.StepLR(optimizer, step_size=500, gamma=0.9)
scheduler = lr_scheduler.CosineAnnealingLR(optimizer,
                                           T_max=50,
                                           eta_min=1e-5,
                                           last_epoch=-1)

model.to(device)
# number of epochs to train the model
n_epochs = 50

val_loss_min = np.Inf  # track change in validation loss

train_losses, val_acc, val_losses, lr = [], [], [], []

for epoch in range(1, n_epochs + 1):
    # keep track of training and validation loss
    train_loss = 0.0
    val_loss = 0.0
    print('running epoch: {}'.format(epoch))
    # 訓練模型
    model.train()
    with tqdm(train_loader) as pbar:
        for data, target in train_loader:
            # move tensors to GPU if CUDA is available
            data, target = data.to(device), target.to(device)
            # clear the gradients of all optimized variables
            optimizer.zero_grad()
            # forward pass: compute predicted outputs by passing inputs to the model
            output = model(data)
            # calculate the batch loss
            loss = criterion(output, target)
            # backward pass: compute gradient of the loss with respect to model parameters
            loss.backward()
            # perform a single optimization step (parameter update)
            optimizer.step()
            # update training loss
            train_loss += loss.item() * data.size(0)
            pbar.update(1)
            pbar.set_description('train')
            pbar.set_postfix(
                **{
                    'loss': loss.item(),
                    'lr': optimizer.state_dict()['param_groups'][0]['lr']
                })
    lr.append(optimizer.state_dict()['param_groups'][0]['lr'])
    scheduler.step()
    # 驗證模型
    model.eval()
    with tqdm(val_loader) as pbar:
        with torch.no_grad():
            accuracy = 0.0
            for data, target in val_loader:
                # move tensors to GPU if CUDA is available
                data, target = data.to(device), target.to(device)
                # forward pass: compute predicted outputs by passing inputs to the model
                output = model(data)
                # calculate the batch loss
                loss = criterion(output, target)
                # update average validation loss
                val_loss += loss.item() * data.size(0)
                # accuracy
                predict_y = torch.max(output, dim=1)[1]
                accuracy = accuracy + (predict_y
                                       == target.to(device)).sum().item()
                pbar.update(1)
                pbar.set_description('val')
                pbar.set_postfix(**{
                    'loss': loss.item(),
                })
    accuracy = accuracy / len(val_loader.dataset)
    val_acc.append(accuracy)

    # calculate average losses
    train_losses.append(train_loss / len(train_loader.dataset))
    val_losses.append(val_loss / len(val_loader.dataset))
    train_loss = train_loss / len(train_loader.dataset)
    val_loss = val_loss / len(val_loader.dataset)

    # print training/validation statistics
    print('\ttrain Loss: {:.6f} \tval Loss: {:.6f}'.format(
        train_loss, val_loss))
    print('\taccuracy: {:.6f}'.format(accuracy))
    # save model if validation loss has decreased
    if val_loss <= val_loss_min:
        print(
            'val loss decreased ({:.6f} --> {:.6f}).  saving model ...'.format(
                val_loss_min, val_loss))
        torch.save(model.state_dict(), './classification/logs/model.pth')
        val_loss_min = val_loss

# 繪製圖
plt.figure()
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.plot(val_acc, label='Val Accuracy')
plt.legend(loc='best')
plt.savefig('./classification/image/acc.jpg')
plt.show()

plt.figure()
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.plot(train_losses, label='Train Loss')
plt.plot(val_losses, label='Val Loss')
plt.legend(loc='best')
plt.savefig('./classification/image/loss.jpg')
plt.show()

plt.figure()
plt.xlabel('Epochs')
plt.ylabel('Learning rate')
plt.plot(lr, label='Learning Rate')
plt.legend(loc='best')
plt.savefig('./classification/image/lr.jpg')
plt.show()


import onnx
import mxnet as mx
from retinaface import RetinaFace
# from retinaface_cov import RetinaFaceCoV
from mxnet.contrib import onnx as onnx_mxnet
import mxnet.contrib.onnx.mx2onnx.export_onnx as mx_op
from mxnet.contrib.onnx.mx2onnx._op_translations import get_inputs, convert_string_to_list
import numpy as np
import argparse
import json

def get_sym_train(symbol):
    score_group = []
    box_group = []
    landmark_group = []
    mask_group = []
    for layers in symbol:
        transpose_layers = mx.symbol.transpose(layers, (0, 2, 3, 1))
        if 'face_rpn_cls_prob_reshape_stride' in layers.name:
            slice_layers = mx.symbol.slice_axis(transpose_layers, axis=3, begin=2, end=4)
            reshape_layers = mx.symbol.Reshape(data=slice_layers,
                                               shape=(0, -1, 1),
                                               name=layers.name + 'transpose')
            score_group.append(reshape_layers)
        if 'face_rpn_bbox_pred_stride' in layers.name:
            reshape_layers = mx.symbol.Reshape(data=transpose_layers,
                                               shape=(0, -1, 4),
                                               name=layers.name + 'transpose')
            box_group.append(reshape_layers)
        if 'face_rpn_landmark_pred_stride' in layers.name:
            reshape_layers = mx.symbol.Reshape(data=transpose_layers,
                                               shape=(0, -1, 10),
                                               name=layers.name + 'transpose')
            landmark_group.append(reshape_layers)
        if 'face_rpn_type_prob_reshape_stride' in layers.name:
            slice_layers = mx.symbol.slice_axis(transpose_layers, axis=3, begin=4, end=6)
            reshape_layers = mx.symbol.Reshape(data=slice_layers,
                                               shape=(0, -1, 1),
                                               name=layers.name + 'transpose')
            mask_group.append(reshape_layers)
    score_concat = mx.sym.concat(*score_group, dim=1, name='score_concat')
    bbox_concat = mx.sym.concat(*box_group, dim=1, name='bbox_concat')
    landmark_concat = mx.sym.concat(*landmark_group, dim=1, name='landmark_concat')
    if len(mask_group) > 0:
        mask_concat = mx.sym.concat(*mask_group, dim=1, name='mask_concat')
        output = mx.sym.concat(*[score_concat, bbox_concat, landmark_concat, mask_concat], dim=2, name='output')
    else:
        output = mx.sym.concat(*[score_concat, bbox_concat, landmark_concat], dim=2, name='output')
    return output

def change_plus(file_prefix):
    file_name = file_prefix + '-symbol.json'
    with open(file_name, 'r') as f:
        model_dict = json.load(f)
    index = 0
    for node in model_dict['nodes']:
        if 'plus' in node['name']:
            node['name'] = '_plus' + str(index)
            index += 1
    with open(file_name, 'w') as f:
        json.dump(model_dict, f, indent=1)

def create_helper_tensor_node(input_vals, output_name, kwargs):
    """create extra tensor node from numpy values"""
    data_type = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[input_vals.dtype]
    tensor_node = onnx.helper.make_tensor_value_info(
        name=output_name,
        elem_type=data_type,
        shape=input_vals.shape
    )
    kwargs["initializer"].append(
        onnx.helper.make_tensor(
            name=output_name,
            data_type=data_type,
            dims=input_vals.shape,
            vals=input_vals.flatten().tolist(),
            raw=False,
        )
    )
    return tensor_node

def create_helper_shape_node(input_node, node_name):
    """create extra transpose node for dot operator"""
    trans_node = onnx.helper.make_node(
        'Shape',
        inputs=[input_node],
        outputs=[node_name],
        name=node_name
    )
    return trans_node

@mx_op.MXNetGraph.register("SoftmaxActivation")
def convert_softmax_activation(node, **kwargs):
    """
    Map MXNet's softmax operator attributes to onnx's Softmax operator and return the created node.
    """
    name, input_nodes, attrs = get_inputs(node, kwargs)
    outputs = []
    mode = attrs.get("mode", "channel")
    axis = -1
    if mode == "channel":
        # transpose nchw -> nhwc  Softmax: axis=3  nhwc -> nchw
        trans_op_name = 'transpose' + str(kwargs["idx"])
        trans_node = onnx.helper.make_node(
            "Transpose",
            input_nodes,
            [trans_op_name],
            name=trans_op_name,
            perm=[0, 2, 3, 1]
        )
        softmax_op_name = 'softmax' + str(kwargs["idx"])
        softmax_node = onnx.helper.make_node(
            "Softmax",
            [trans_op_name],
            [softmax_op_name],
            axis=3,
            name=softmax_op_name
        )
        output_node = onnx.helper.make_node(
            "Transpose",
            [softmax_op_name],
            [name],
            name=name,
            perm=[0, 3, 1, 2]
        )
        outputs.append(trans_node)
        outputs.append(softmax_node)
        outputs.append(output_node)
    else:
        softmax_node = onnx.helper.make_node(
            "Softmax",
            input_nodes,
            [name],
            axis=axis,
            name=name
        )
        outputs.append(softmax_node)
    return outputs

@mx_op.MXNetGraph.register("UpSampling")
def convert_upsample(node, **kwargs):
    """
    Map MXNet's UpSampling operator attributes to onnx's Upsample operator and return the created node.
    """
    name, input_nodes, attrs = get_inputs(node, kwargs)
    sample_type = attrs.get('sample_type', 'nearest')
    sample_type = 'linear' if sample_type == 'bilinear' else sample_type
    scale = convert_string_to_list(attrs.get('scale'))
    scaleh = scalew = float(scale[0])
    if len(scale) > 1:
        scaleh = float(scale[0])
        scalew = float(scale[1])
    scale = np.array([1.0, 1.0, scaleh, scalew], dtype=np.float32)
    scale_node = create_helper_tensor_node(scale, name + 'scales', kwargs)
    input_nodes.append(name + 'scales')
    node = onnx.helper.make_node(
        'Resize',
        input_nodes,
        [name],
        mode=sample_type,
        name=name
    )
    return [scale_node, node]

@mx_op.MXNetGraph.register("Crop")
def convert_crop(node, **kwargs):
    """Map MXNet's crop operator attributes to onnx's Crop operator
    and return the created node.
    """
    name, inputs, attrs = get_inputs(node, kwargs)
    start = np.array([0, 0, 0, 0], dtype=np.int)  # index是int类型
    start_node = create_helper_tensor_node(start, name + '_starts', kwargs)
    shape_node = create_helper_shape_node(inputs[1], inputs[1] + '_shape')
    crop_node = onnx.helper.make_node(
        "Slice",
        inputs=[inputs[0], name + '_starts', inputs[1] + '_shape'],  # data、start、end
        outputs=[name],
        name=name
    )
    return [start_node, shape_node, crop_node]

@mx_op.MXNetGraph.register("BatchNorm")
def convert_batchnorm(node, **kwargs):
    """Map MXNet's BatchNorm operator attributes to onnx's BatchNormalization operator
    and return the created node.
    """
    name, input_nodes, attrs = get_inputs(node, kwargs)

    momentum = float(attrs.get("momentum", 0.9))
    eps = float(attrs.get("eps", 0.001))

    bn_node = onnx.helper.make_node(
        "BatchNormalization",
        input_nodes,
        [name],
        name=name,
        epsilon=eps,
        momentum=momentum
        # MXNet computes mean and variance per channel for batchnorm.
        # Default for onnx is across all spatial features. Relying on default
        # ONNX behavior of spatial=1 for ONNX opset 8 and below. As the spatial
        # attribute is deprecated in opset 9 and above, not explicitly encoding it.
    )
    return [bn_node]

@mx_op.MXNetGraph.register("slice_axis")
def convert_slice_axis(node, **kwargs):
    """Map MXNet's slice_axis operator attributes to onnx's Slice operator
    and return the created node.
    """
    name, input_nodes, attrs = get_inputs(node, kwargs)

    axes = int(attrs.get("axis"))
    starts = int(attrs.get("begin"))
    ends = attrs.get("end", None)
    if not ends:
        raise ValueError("Slice: ONNX doesnt't support 'None' in 'end' attribute")

    export_nodes = []

    starts = np.atleast_1d(np.asarray(starts, dtype=np.int))
    ends = np.atleast_1d(np.asarray(ends, dtype=np.int))
    axes = np.atleast_1d(np.asarray(axes, dtype=np.int))

    starts_node = create_helper_tensor_node(starts, name + '__starts', kwargs)
    export_nodes.append(starts_node)
    starts_node = starts_node.name

    ends_node = create_helper_tensor_node(ends, name + '__ends', kwargs)
    export_nodes.append(ends_node)
    ends_node = ends_node.name

    axes_node = create_helper_tensor_node(axes, name + '__axes', kwargs)
    export_nodes.append(axes_node)
    axes_node = axes_node.name

    input_node = input_nodes[0]
    node = onnx.helper.make_node(
        "Slice",
        [input_node, starts_node, ends_node, axes_node],
        [name],
        name=name,
    )
    export_nodes.extend([node])

    return export_nodes

def main():
    parser = argparse.ArgumentParser(description='convert arcface models to onnx')
    # general
    parser.add_argument('--prefix', default='./model/R50', help='prefix to load model.')
    parser.add_argument('--epoch', default=0, type=int, help='epoch number to load model.')
    parser.add_argument('--gpuid', default=0, type=int, help='ctx_id in model.')
    parser.add_argument('--network', default='net3', type=str, help='network in model.')
    parser.add_argument('--input_shape', nargs='+', default=[1, 3, 640, 640], type=int, help='input shape.')
    args = parser.parse_args()
    converted_onnx_filename = args.prefix + '.onnx'
    model = RetinaFace(args.prefix, args.epoch, args.gpuid, args.network).model
    # model = RetinaFaceCoV(args.prefix, args.epoch, args.gpuid, args.network).model
    sym, arg_params, aux_params = model.symbol, model._arg_params, model._aux_params
    model = get_sym_train(sym)
    mx.model.save_checkpoint(args.prefix + '_transpose', args.epoch, model, arg_params, aux_params)
    change_plus(args.prefix + '_transpose')
    converted_onnx_filename = onnx_mxnet.export_model(args.prefix + '_transpose-symbol.json',
                                                      f'{args.prefix}_transpose-{args.epoch:04d}.params',
                                                      [args.input_shape], np.float32, converted_onnx_filename)

if __name__ == '__main__':
    main()


import torch
import os
import numpy as np
from PIL import Image
import Constants
from data import cxr_process as preprocess
import pandas as pd
from torchvision import transforms
import pickle
from pathlib import Path
from torch.utils.data import Dataset, ConcatDataset

def get_dfs(envs = [], split = None, only_frontal = False):
    dfs = []
    for e in envs:
        func = preprocess.get_process_func(e)
        paths = Constants.df_paths[e]

        if split is not None:
            splits = [split]
        else:
            splits = ['train', 'val', 'test']

        dfs += [func(pd.read_csv(paths[i]), only_frontal) for i in splits]

    return pd.concat(dfs, ignore_index = True, sort = False).sample(frac=1) #shuffle

def prepare_df_for_cb(df, positive_env = 'CXP'):
    df2 = df.copy()
    df2 = df2.rename(columns = {'path': 'filename', 'Atelectasis': 'label', 'env': 'conf'})
    df2['conf'] = (df2['conf'] == positive_env).astype(int)
    df2['label'] = (df2['label']).astype(int)
    return df2

def dataset_from_cb_output(orig_df, labels_gen, split, causal_type, data_type, cache = False):
    '''
    massages output from labels_gen (which is only filename, label, conf) into a more informative
        dataframe format to allow for generalized caching in dataloader
    '''
    envs = orig_df.env.unique()
    augmented_dfs, labels_env = {i: {} for i in envs}, []
    temp = orig_df.set_index('path').loc[labels_gen[:, 0], :].reset_index()
    for i in envs:
        # assert(len(np.unique(labels_gen[:, 0])) == len(labels_gen)) #  this should give error for deconf
        # augmented_dfs[i][split] = orig_df[(orig_df.path.isin(labels_gen[:, 0])) & (orig_df.env == i)]
        subset = (temp.env == i)
        augmented_dfs[i][split] = temp[subset]
        labels_env.append(labels_gen[subset, :])

    dataset = get_dataset(envs, split, only_frontal = False, imagenet_norm = True, augment = 1 if split == 'train' else 0,
               cache = cache, subset_label = 'Atelectasis', augmented_dfs = augmented_dfs)

    return CXRWrapper(dataset, np.concatenate(labels_env), causal_type, data_type)

def get_dataset(envs = [], split = None, only_frontal = False, imagenet_norm = True, augment = 0, cache = False, subset_label = None,
               augmented_dfs = None):

    if split in ['val', 'test']:
        assert(augment in [0, -1])

    if augment == 1: # image augmentations
        image_transforms = [transforms.RandomHorizontalFlip(),
                            transforms.RandomRotation(10),
                            transforms.RandomResizedCrop(size = 224, scale = (0.75, 1.0)),
                        transforms.ToTensor()]
    elif augment == 0:
        image_transforms = [transforms.ToTensor()]
    elif augment == -1: # only resize, just return a dataset with PIL images; don't ToTensor()
        image_transforms = []

    if imagenet_norm and augment != -1:
        image_transforms.append(transforms.Normalize(Constants.IMAGENET_MEAN, Constants.IMAGENET_STD))

    datasets = []
    for e in envs:
        func = preprocess.get_process_func(e)
        paths = Constants.df_paths[e]

        if split is not None and split != 'all':
            splits = [split]
        else:
            splits = ['train', 'val', 'test']

        if augmented_dfs is not None: # use provided dataframes instead of loading
            dfs = [augmented_dfs[e][i] for i in splits]
        else:
            dfs = [func(pd.read_csv(paths[i]), only_frontal) for i in splits]

        for c, s in enumerate(splits):
            cache_dir = Path(Constants.cache_dir)/ f'{e}_{s}/'
            if cache:
                cache_dir.mkdir(parents=True, exist_ok=True)
            datasets.append(AllDatasetsShared(dfs[c], transform = transforms.Compose(image_transforms)
                                      , split = split, cache = cache, cache_dir = cache_dir, subset_label = subset_label))

    if len(datasets) == 0:
        return None
    elif len(datasets) == 1:
        ds = datasets[0]
    else:
        ds = ConcatDataset(datasets)
        ds.dataframe = pd.concat([i.dataframe for i in datasets])

    return ds

class CXRWrapper(Dataset):
    '''
    Wraps a AllDatasetsShared object and a generated labels array from cb to return (x, y, extras)
      where the size of extras depends on the confounding and data types
    '''
    def __init__(self, ds, labels, causal_type, data_type):
        super().__init__()
        self.ds = ds
        self.labels = labels
        self.causal_type = causal_type
        self.data_type = data_type

    def get_img_id(self, path, env):
        path = Path(path)
        if env in ['PAD', 'NIH']:
            return path.stem
        elif env in ['MIMIC', 'CXP']:
            return  '_'.join(path.parts[-3:])

    def check_id(self, p1, p2, env):
        return self.get_img_id(p1, env) == self.get_img_id(p2, env)

    def __len__(self):
        return len(self.ds)

    def __getitem__(self, idx):
        x, y, meta = self.ds[idx]
        lst_info = self.labels[idx]
        assert self.check_id(lst_info[0], meta['path'], meta['env'])
        if self.data_type == 'IF':
            if self.causal_type == 'back':
                extras = np.array([float(lst_info[2])])
            else:
                extras = np.array([float(lst_info[2]), float(lst_info[3])])
        else:
            extras = np.array([])

        return x, y, extras

class AllDatasetsShared(Dataset):
    def __init__(self, dataframe, transform=None, split = None, cache = False, cache_dir = '', subset_label = None):
        super().__init__()
        self.dataframe = dataframe
        self.dataset_size = self.dataframe.shape[0]
        self.transform = transform
        self.split = split
        self.cache = cache
        self.cache_dir = Path(cache_dir)
        self.subset_label = subset_label # (str) select one label instead of returning all Constants.take_labels

    def get_cache_path(self, cache_dir, meta):
        path = Path(meta['path'])
        if meta['env'] in ['PAD', 'NIH']:
            return cache_dir / (path.stem + '.pkl')
        elif meta['env'] in ['MIMIC', 'CXP']:
            return (cache_dir / '_'.join(path.parts[-3:])).with_suffix('.pkl')

    def __getitem__(self, idx):
        item = self.dataframe.iloc[idx]
        cache_path = self.get_cache_path(self.cache_dir, item)

        if self.cache and cache_path.is_file():
            img, label, meta = pickle.load(cache_path.open('rb'))
            meta = item.to_dict() # override
        else:
            img = np.array(Image.open(item["path"]))

            if img.dtype == 'int32':
                img = np.uint8(img/(2**16)*255)
            elif img.dtype == 'bool':
                img = np.uint8(img)
            else: #uint8
                pass

            if len(img.shape) == 2:
                img = img[:, :, np.newaxis]
                img = np.concatenate([img, img, img], axis=2)
            elif len(img.shape)>2:
                # print('different shape', img.shape, item)
                img = img[:,:,0]
                img = img[:, :, np.newaxis]
                img = np.concatenate([img, img, img], axis=2)

            img = Image.fromarray(img)
            resize_transform = transforms.Resize(size = [224, 224])
            img = transforms.Compose([resize_transform])(img)

            label = torch.FloatTensor(np.zeros(len(Constants.take_labels), dtype=float))
            for i in range(0, len(Constants.take_labels)):
                if (self.dataframe[Constants.take_labels[i].strip()].iloc[idx].astype('float') > 0):
                    label[i] = self.dataframe[Constants.take_labels[i].strip()].iloc[idx].astype('float')

            meta = item.to_dict()

            if self.cache:
                pickle.dump((img, label, meta), cache_path.open('wb'))

        if self.transform is not None: # apply image augmentations after caching
            img = self.transform(img)

        if self.subset_label:
            if self.subset_label in Constants.take_labels:
                label = int(label[Constants.take_labels.index(self.subset_label)])
            else:
                raise NotImplementedError

        return img, label, meta

    def __len__(self):
        return self.dataset_size