# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/03_bbox_east_training.ipynb (unless otherwise specified).

__all__ = ['east_config', 'ObjectCategoryProcessor', 'ObjectCategoryList', 'ObjectItemList', 'bb_pad_collate',
           'gaussian_blur', 'UnetTypeConcat', 'EastModel', 'intersection', 'tlbr2hw', 'normalized_intersection',
           'tlbr2cthw', 'IoU', 'tlbr2hw', 'hw2area', 'match_anchors', 'create_grid', 'anchors', 'cthw2tlbr',
           'show_anchors_on_images', 'tlbr2cthw', 'bbox_to_target', 'bbox_to_target', 'tlbr2cc', 'prepare_bboxes',
           'prepare_labels', 'BBLossMetrics', 'target_to_bbox', 'nms', 'check_overlap', 'calc_precision_recall', 'F1',
           'iou_loss', 'dice_loss', 'encode_class', 'SigmaL1SmoothLoss', 'RetinaNetFocalLoss', 'threshold_images',
           'prepare_masked_bboxes', 'PreciseFocalLoss', 'MSEloss', 'MSE', 'mse']

# Cell
from fastai import *
from fastai.vision import *
import pandas as pd
import numpy as np
import cv2
from tqdm.notebook import tqdm
import xml.etree.ElementTree as ET
from torch.autograd import Variable
from fastai.callbacks import LossMetrics

# Cell
import matplotlib.pyplot as plt

# Cell
class east_config:
    HEIGHT = 512
    WIDTH = 384
    out_shape = (4,4) # 4 times smaller than input

east_config.input_shape = (east_config.HEIGHT, east_config.WIDTH)

# Cell
class ObjectCategoryProcessor(MultiCategoryProcessor):
    "`PreProcessor` for labelled bounding boxes."
    def __init__(self, ds:ItemList, pad_idx:int=0):
        super().__init__(ds)
        self.pad_idx = pad_idx
        self.state_attrs.append('pad_idx')

    def process(self, ds:ItemList):
        ds.pad_idx = self.pad_idx
        super().process(ds)

    def process_one(self,item): return [item[0], [self.c2i.get(o,None) for o in item[1]]]

    def generate_classes(self, items):
        "Generate classes from unique `items` and add `background`."
        classes = super().generate_classes([o[1] for o in items])
        classes = ['background'] + list(classes)
        return classes

def _get_size(xs,i):
    size = xs.sizes.get(i,None)
    if size is None:
        # Image hasn't been accessed yet, so we don't know its size
        _ = xs[i]
        size = xs.sizes[i]
    return size

# Cell
class ObjectCategoryList(MultiCategoryList):
    "`ItemList` for labelled bounding boxes."
    _processor = ObjectCategoryProcessor

    def get(self, i):
        return ImageBBox.create(*_get_size(self.x,i), *self.items[i], classes=self.classes, pad_idx=self.pad_idx)

    def analyze_pred(self, pred): return pred

    def reconstruct(self, t, x):
        (bboxes, labels, *_) = t # CHANGED
        if len((labels - self.pad_idx).nonzero()) == 0: return
        i = (labels - self.pad_idx).nonzero().min()
        bboxes,labels = bboxes[i:],labels[i:]
        return ImageBBox.create(*x.size, bboxes, labels=labels, classes=self.classes, scale=False)

# Cell
class ObjectItemList(ImageList):
    "`ItemList` suitable for object detection."
    _label_cls,_square_show_res = ObjectCategoryList,False

# Cell
def bb_pad_collate(samples:BatchSamples, pad_idx:int=0) -> Tuple[FloatTensor, Tuple[LongTensor, LongTensor]]:
    "Function that collect `samples` of labelled bboxes and adds padding with `pad_idx`."
    if isinstance(samples[0][1], int): return data_collate(samples)
    max_len = max([len(s[1].data[1]) for s in samples])
    bboxes = torch.zeros(len(samples), max_len, 4)
    labels = torch.zeros(len(samples), max_len).long() + pad_idx
    imgs = []
    for i, (image, bbox) in enumerate(samples):
        bbs, lbls = bbox.data
        imgs.append(image.data[None])
        assert len(bbs) == len(lbls)

        if bbs.nelement() != 0:
            bboxes[i,-len(bbs):] = bbs
            labels[i,-len(lbls):] = tensor(lbls)
    images = torch.cat(imgs,0)
    return images, (bboxes,labels,images)

# Cell
def _gaussian_blur(x, size:uniform_int):
    blurred = cv2.blur(image2np(x), (size,size)) # np.arr
#     blurred = cv2.GaussianBlur(image2np(x), (size,size), 0)
    return tensor(blurred).permute(2,0,1)

def gaussian_blur(size, p=1.0):
    return RandTransform(tfm=TfmPixel(_gaussian_blur), kwargs={'size':size}, p=p, resolved={}, do_run=True, is_random=True, use_on_y=False)

# Cell
class UnetTypeConcat(nn.Module):
    def __init__(self, unpool_layer, conv_layer):
        super().__init__()
        # conv_layer: has to ...
        # unpool_layer: has to increase [h,w] -> [hx2, wx2]
        self.unpool = unpool_layer
        self.conv = conv_layer

    def forward(self, e, d): # e/d - encoder/decoder output
        '''
        @param:  :encoder output
        @param:  :decoder output
        '''
        o = self.unpool(d)
        o = torch.cat((o, e), 1)
        return self.conv(o)

# Cell

# To work w/ ResNet18 change conv_layers in unet_concat:
# 3072 -> 768
# 640 -> 256
# 320 ->128
# (also pixelshuffle)

class EastModel(nn.Module):
    def __init__( self, pixelshuffle=False, conv_out=(listify(east_config.out_shape) == listify((8,8))) ):
        super().__init__()
        self.conv_out = conv_out

        self.means = tensor(imagenet_stats[0]).reshape(1,3,1,1) # [R,G,B]
        self.stds = tensor(imagenet_stats[1]).reshape(1,3,1,1) # [R,G,B]

        if pixelshuffle:
            unpool1 = nn.Sequential(nn.PixelShuffle(2), nn.Conv2d(512, 512*4, kernel_size=1))
            unpool2 = nn.Sequential(nn.PixelShuffle(2), nn.Conv2d(32,  32*4,  kernel_size=1))
            unpool3 = nn.Sequential(nn.PixelShuffle(2), nn.Conv2d(16,  16*4,  kernel_size=1))
        else:
            unpool1 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)
            unpool2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)
            unpool3 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)

#         self.resnet = models.resnet50(pretrained=True)
#         self.unet_block1 = UnetTypeConcat(unpool1, conv_layer(ni=3072, nf=128, ks=1))
#         self.unet_block2 = UnetTypeConcat(unpool2, conv_layer(ni=640,  nf=64,  ks=1))
#         self.unet_block3 = UnetTypeConcat(unpool3, conv_layer(ni=320,  nf=64,  ks=1))

        self.resnet = models.resnet18(pretrained=False) # (if using pixelshuffle u have to change those params too)
        self.unet_block1 = UnetTypeConcat(unpool1, conv_layer(ni=768, nf=128, ks=1))
        self.unet_block2 = UnetTypeConcat(unpool2, conv_layer(ni=256,  nf=64,  ks=1))
        self.unet_block3 = UnetTypeConcat(unpool3, conv_layer(ni=128,  nf=64,  ks=1))

        self.conv_bonus1 = conv_layer(ni=128,  nf=128, ks=3)
        self.conv_bonus2 = conv_layer(ni=64,   nf=64,  ks=3)
        self.conv1 = conv_layer(ni=64,   nf=32,  ks=3)
        self.conv2 = conv_layer(ni=32,   nf=32,  ks=3)
        self.conv_score = conv_layer(ni=32,   nf=1,   ks=1, norm_type=None, use_activ=False)
        self.conv_geo   = conv_layer(ni=32,   nf=4,   ks=1, norm_type=None, use_activ=False)
#         self.conv10 = conv_layer(ni=32,  nf=1,   ks=1, norm_type=None, use_activ=False)

        if self.conv_out:
            self.conv_out = conv_layer(ni=32, nf=32, ks=3, stride=2, norm_type=None)

    def forward(self, x):
        x = (x - self.means.to(x.device)) / self.stds.to(x.device) # normalize

        f = []
        for layer in list(self.resnet.children())[:-1]:
            x = layer(x)
            if type(layer) == torch.nn.Sequential: f.append(x)

        # [b,128,h/16,w/16] = unet_block1([b,1024,h/16,w/16], [b,2048,h/32,w/32])
        o = self.unet_block1(f[2], f[3])
        o = self.conv_bonus1(o)
        # [b,64,h/8,w/8] = unet_block1([b,512,h/8,w/8], [b,128,h/16,w/16])
        o = self.unet_block2(f[1], o)
        o = self.conv_bonus2(o)
        # [b,64,h/4,w/4] = unet_block1([b,256,h/4,w/4], [b,64,h/8,w/8])
        o = self.unet_block3(f[0], o)

        o = self.conv1(o)  # [b,32,h/4,w/4]
        o = self.conv2(o)  # [b,32,h/4,w/4]

        if self.conv_out: o = self.conv_out(o) # [b,32,h/2,w/2]

        score = self.conv_score(o) # [b,1,h/4,w/4]
        geo_map = self.conv_geo(o) # [b,4,h/4,w/4]

        score = nn.Sigmoid()(score)
        geo_map = nn.Sigmoid()(geo_map)*2

#         angle_map = self.conv10(g)
#         angle_map = self.sigmoid(angle_map)
#         angle_map = (angle_map - 0.5) * math.pi / 2
#         geo = torch.cat((geo_map, angle_map), 1)  # bs 5 w/4 w/4
        geo = geo_map

        return score, geo

# Cell
# tlbr - top, left, bottom, right
# bb - bounding box
tlbr2hw = lambda bbs: bbs[:,2:] - bbs[:,:2] # [N,4] -> [N,2]

def intersection(bbs1, bbs2, normalize=False):
    ''' Calculates intersection area of bounding boxes
    @param: [N1,4]  :N1 tlbr coordinates
    @param: [N2,4]  :N2 tlbr coordinates
    @return [N1,N2] :area sizes for each coord combination
    '''
    N1, N2 = bbs1.size(0), bbs2.size(0)
    bbs1, bbs2 = bbs1.unsqueeze(1).expand(N1,N2,4), bbs2.unsqueeze(0).expand(N1,N2,4)
    intersection_tl_corner = torch.max(bbs1[...,:2], bbs2[...,:2]) # [N1,N2,2]
    intersection_br_corner = torch.min(bbs1[...,2:], bbs2[...,2:]) # [N1,N2,2]
    sizes = torch.clamp(intersection_br_corner - intersection_tl_corner, min=0) # [N1,N2,2]

    if normalize:
        hw = tlbr2hw(bbs1[:,0,:]) # [N1,2]
        sizes[...,:] /= hw.unsqueeze(1) # [N1,N2,2] * [N1,1,2]

    return sizes[...,0] * sizes[...,1] # [N1,N2]

# Cell
normalized_intersection = partial(intersection, normalize=True)

# Cell
def tlbr2cthw(bbs):
    ''' Convert top/left bottom/right format `boxes` to center/size corners
    @param: [N,4] :tlbr coordinates
    @return [N,4] :cthw coordinates
    '''
    center = (bbs[:,:2] + bbs[:,2:])/2 # [N,2]
    sizes = bbs[:,2:] - bbs[:,:2] # [N,2]
    return torch.cat([center, sizes], 1) #

# Cell
tlbr2hw = lambda bbs: bbs[:,2:] - bbs[:,:2] # [N,4] -> [N,2]
hw2area = lambda hw: hw[:,0] * hw[:,1] # [N,2] -> [N]

def IoU(bbs1, bbs2):
    ''' Calculates intersection area of bounding boxes
    @param: [N1,4]  :N1 tlbr coordinates
    @param: [N2,4]  :N2 tlbr coordinates
    @return [N1,N2] :area sizes for each coord combination
    '''
    inter = intersection(bbs1, bbs2) # [N1,N2]
    hw1, hw2 = tlbr2hw(bbs1), tlbr2hw(bbs2) # [N,2]
    bbs1_areas = hw2area(hw1).unsqueeze(1) # [N1,1]
    bbs2_areas = hw2area(hw2).unsqueeze(0) # [1,N2]
    union = bbs1_areas + bbs2_areas - inter
    return inter/(union+1e-8)

# Cell
# a bit problematic w/ int input and by_iou=False
def match_anchors(anchors, targets, match_thr=0.01, bg_thr=0.01, by_iou=False):
    ''' Match `anchors` to targets. -1 is match to background, -2 is ignore
    @param: [N1,4]  :N1 tlbr coordinates
    @param: [N2,4]  :N2 tlbr coordinates
    @param: int     :if this < IoU -> this match is a match (int index of match in targets)
    @param: int     :if this > IoU -> this match is a background (int -1)
    @param: bool    :to
    @return [N1]    :every match from anchors to targets index (if doesnt exist it is -2/-1)
    '''
    f = IoU if by_iou else normalized_intersection
    ious = f(anchors, targets) # [N1,N2]
#     print(anchors)
#     print(targets)
#     print('ious: \n', ious)
    matches = anchors.new(anchors.size(0)).zero_().long() - 2
#     matches = torch.zeros(anchors.size(0)).zero_().long() - 2

    if ious.shape[1] == 0:
        print('ious.shape[1] == 0')
        return matches

    iou, idxs = torch.max(ious, dim=1) # [N1]
    matches[iou < bg_thr] = -1
    matches[iou > match_thr] = idxs[iou > match_thr]

    return matches

# Cell
def create_grid(size):
    ''' Creates a x,y grid of size `size`, coords start from -1 to 1
    @param: tuple(H,W) :tuple of 2 ints H(height), W(width)
    @return [H,W,4]    :tlbr coords for each (h,w) cell
    '''
    H, W = size if is_tuple(size) else (size,size)
    grid = FloatTensor(H, W, 4)

    # not precise, but good enough
    linear_points_left = torch.linspace(0, 1-1/W, W) # [0, 1]
    linear_points_right = linear_points_left + 1/W
    linear_points_top = torch.linspace(0, 1-1/H, H) # [0, 1]
    linear_points_btm = linear_points_top + 1/H

    grid[:, :, 0] = linear_points_top.unsqueeze(1).expand(H,W)
    grid[:, :, 1] = linear_points_left.unsqueeze(0).expand(H,W)
    grid[:, :, 2] = linear_points_btm.unsqueeze(1).expand(H,W)
    grid[:, :, 3] = linear_points_right.unsqueeze(0).expand(H,W)
    return grid*2-1

# Cell
anchors = create_grid((east_config.HEIGHT//4, east_config.WIDTH//4)).view(-1,4)

# Cell
def _draw_outline(o:Patch, lw:int):
    "Outline bounding box onto image `Patch`."
    o.set_path_effects([patheffects.Stroke(
        linewidth=lw, foreground='black'), patheffects.Normal()])

def _draw_rect(ax:plt.Axes, b:Collection[int], color:str='white', text=None, text_size=14):
    "Draw bounding box on `ax`."
    patch = ax.add_patch(patches.Rectangle(b[:2], *b[-2:], fill=False, edgecolor=color, lw=2))
    _draw_outline(patch, 4)
    if text is not None:
        patch = ax.text(*b[:2], text, verticalalignment='top', color=color, fontsize=text_size, weight='bold')
        _draw_outline(patch,1)

# Cell
def cthw2tlbr(boxes):
    "Convert center/size format `boxes` to top/left bottom/right corners."
    top_left = boxes[:,:2] - boxes[:,2:]/2
    bot_right = boxes[:,:2] + boxes[:,2:]/2
    return torch.cat([top_left, bot_right], 1)

def show_anchors_on_images(data, anchors, figsize=(15,15)):
    all_boxes = []
    all_labels = []
    x, y = data.one_batch(DatasetType.Train, True, True)
    for image, bboxes, labels in zip(x, y[0], y[1]):
        image = Image(image.float().clamp(min=0, max=1))

        # 0=not found; 1=found; found 2=anchor
        processed_boxes = []
        processed_boxes_gt = []
        processed_boxes_bg = []
        processed_labels = []
        processed_labels_gt = []
        a = 0
        bgs = torch.zeros(len(anchors), dtype=bool) + 1
        for gt_box in bboxes[labels > 0]:
            matches = match_anchors(anchors, gt_box[None, :])
            non_bg = matches >= 0
            a += non_bg.sum()
#             print('1)', non_bg.sum(), (matches >= -1).sum(), (matches >= -2).sum())
            if non_bg.sum() != 0:
                processed_boxes_gt.append(to_np(gt_box))
                processed_labels_gt.append(2)
                for bb in anchors[non_bg]:
                    processed_boxes.append(to_np(bb))
                    processed_labels.append(3)
                bgs[non_bg] = 0
            else:
                processed_boxes.append(to_np(gt_box))
                processed_labels.append(0)
                val, idx = torch.max(IoU(anchors, gt_box[None, :]), 0)
                best_fitting_anchor = anchors[idx][0]
                processed_boxes.append(to_np(best_fitting_anchor))
                processed_labels.append(1)
            if a > 500:
                break
        all_boxes.extend(processed_boxes)
        all_labels.extend(processed_labels)

        processed_boxes = np.array(processed_boxes)
        processed_labels = np.array(processed_labels)

#         print('2)', np.unique(processed_labels, return_counts=True))

        _, ax = plt.subplots(nrows=1, ncols=3, figsize=figsize)
        ax[0].set_title("Matched Anchors")
        ax[1].set_title("Actual GT")
        ax[2].set_title("BG Anchors")

        # Anchors
        imageBB = ImageBBox.create(*image.size, tensor(processed_boxes[processed_labels > 1]),
#                                    labels=processed_labels[processed_labels > 1],
                                   classes=["", "", "", ""], scale=False)
#                                        classes=["", "", "Match", "Anchor"], scale=False)

        image.show(ax=ax[0], y=imageBB)

        # Actual
        imageBB = ImageBBox.create(*image.size, tensor(processed_boxes_gt),
#                                    labels=[processed_labels[0]],
                                   classes=["", "", "", ""], scale=False)
#                                        classes=["", "", "Match", "Anchor"], scale=False)
        image.show(ax=ax[1], y=imageBB)

        # BG Anchors
        imageBB = ImageBBox.create(*image.size, tensor(anchors[bgs == 1][1000:2000]),
                                   classes=["", "", "", ""], scale=False)
        image.show(ax=ax[2], y=imageBB)
        break

#     return np.array(all_boxes), np.array(all_labels)

# Cell
def tlbr2cthw(boxes): # [N,4] -> [N,4]
    center = (boxes[:,:2] + boxes[:,2:])/2
    sizes = boxes[:,2:] - boxes[:,:2]
    return torch.cat([center, sizes], 1)


def bbox_to_target(bboxes, anchors):
    ''' Return the target of the model on `anchors` for the `bboxes`
    @param: [N,4] :N tlbr coordinates (true bboxes), (bboxes must match anchors)
    @param: [N,4] :N tlbr coordinates (selected anchors)
    @return [N,4] :distance of how much bboxes are from anchors
    '''
    bboxes, anchors = tlbr2cthw(bboxes), tlbr2cthw(anchors)
    t_centers = (bboxes[...,:2] - anchors[...,:2]) / anchors[...,2:]
    t_sizes = torch.log(bboxes[...,2:] / anchors[...,2:] + 1e-8)
    return torch.cat([t_centers, t_sizes], -1)#.div_(bboxes.new_tensor([[0.1, 0.1, 0.2, 0.2]]))

# Cell
# tlbr to center coords
tlbr2cc = lambda boxes: (boxes[:,:2] + boxes[:,2:])/2 # [N,4] -> [N,2]

# Override
def bbox_to_target(bboxes, anchors):
    ''' Return the target of the model on `anchors` for the `bboxes`
    @param: [N,4] :N tlbr coordinates (true bboxes), (bboxes must match anchors)
    @param: [N,4] :N tlbr coordinates (selected anchors)
    @return [N,4] :distance of how much bboxes are from anchors
    '''
    a_centers = tlbr2cc(anchors)
    return torch.cat([a_centers - bboxes[...,:2], bboxes[...,2:] - a_centers], -1) # .div_(bboxes.new_tensor([[0.1, 0.1, 0.1, 0.1]]))

# Cell
def prepare_bboxes(anchors, bbox_true, bbox_pred, matches=None):
    ''' Return the target of the model on `anchors` for the `bboxes`
    @param: [N1,4] :tlbr coordinates (all anchors), N1 = H x W
    @param: [N2,4] :tlbr coordinates (true bboxes)
    @param: [N1,4] :tlbr coordinates (nn output bboxes)
    @param: [N3,4] :optional for caching
    @return ([N3,4], [N3,4], [N3,4]) or None :bbox_true, bbox_pred, matches (for caching)
    '''
    if matches is None: matches = match_anchors(anchors, bbox_true) # [N3, 4]
    non_bg = matches >= 0 # [N3, 4] (non background)

    if non_bg.sum() == 0:
#         print('non_bg.sum() == 0') # creates a bit of spam (0-2 per epoch)
        return None, None, None

    bbox_pred = bbox_pred[non_bg] # [N3,4]
    bbox_true = bbox_true[matches[non_bg]] # [N3,4] (matches return bbox_true indexes)
    anchors = anchors[non_bg] # [N3,4]
    return bbox_to_target(bbox_true, anchors), bbox_pred, matches

# Cell
def prepare_labels(y_true, y_pred, matches):
    ''' Prepare the probabilities of target and prediction of the model
    @param: [N1]           :
    @param: [H/4 x W/4, 1] :
    @param: [H/4 x W/4]    :
    @return [N2], [N2, 1]  :
    '''
    if matches is None: matches = torch.zeros_like(y_pred).squeeze()
    non_ignore = matches >= -1
    if non_ignore.sum() != matches.size(0):
        print('non_ignore.sum() != matches.size(0)', non_ignore.sum(), matches.size(0))

    target = y_true.new_zeros(len(matches)).float()
    non_bg = matches >= 0
    target[non_bg] = 1

    return target[non_ignore][:, None], y_pred[non_ignore]

# Cell
class BBLossMetrics(LossMetrics):
    def on_batch_end(self, last_target, train, **kwargs):
        "Update the metrics if not `train`"
        if train: return
        bs = last_target[0].size(0) # CHANGED
        for name in self.names:
            self.metrics[name] += bs * self.learn.loss_func.metrics[name].detach().cpu()
        self.nums += bs

# Cell
def target_to_bbox(output, anchors):
#     output.mul_(output.new_tensor([[0.1, 0.1, 0.1, 0.1]]))
    a_centers = tlbr2cc(anchors)
    return torch.cat([a_centers - output[...,:2], a_centers + output[...,2:]], -1)

# Cell
def nms(boxes, scores, threshold=0.2):
    ''' non maximum suppression
    @param: [N,4]
    @param: [N]
    @return [N1]
    '''
    sorted_idx = scores.argsort(descending=True) # [N]
    boxes, scores = boxes[sorted_idx], scores[sorted_idx]
    to_keep, indexes = [], torch.LongTensor(range(len(scores))) # [], [N]
    while len(scores) > 0:
        # put in highest scored index
        to_keep.append(sorted_idx[indexes[0]])
        # calculate IoU of highest_scored_box vs all_boxes
        ious = IoU(boxes, boxes[0].view(1,4)).squeeze() # [N]
        other_boxes = ious <= threshold # [N] (not the same boxes)
        if other_boxes.sum() == 0: break
        if other_boxes.sum() == len(boxes):
            if len(indexes) < 3: break
            boxes, scores, indexes = boxes[1:], scores[1:], indexes[1:]
            continue
        boxes, scores, indexes = boxes[other_boxes], scores[other_boxes], indexes[other_boxes]

    return LongTensor(to_keep)

# Cell
def _unpad(bbox_true, y_true, pad_idx=0):
    ''' [N_max,4], [N_max] -> [N1,4], [N1] '''
    i = torch.min(torch.nonzero(y_true - pad_idx)) if sum(y_true)>0 else 0
    return bbox_true[i:], y_true[i:] + pad_idx

def check_overlap(i_true, _list, is_overlaping):
    for i_pred in _list:
        if is_overlaping[i_pred, i_true]:
            return i_pred
    return None

def calc_precision_recall(bbox_pred, bbox_true, iou_thresh=0.5):
    ''' [P,4], [T,4] (Preds, Trues)'''
    num_preds, num_trues = len(bbox_pred), len(bbox_true)

    ious = IoU(bbox_pred, bbox_true)
    is_overlaping = ious > iou_thresh

    correct = 0
    pred_idxes =list(range(len(bbox_pred)))
    for i_true in range(len(bbox_true)):
        i_overlap = check_overlap(i_true, pred_idxes, is_overlaping)
        if i_overlap is not None:
            pred_idxes.remove(i_overlap)
            correct += 1

    precision = correct / num_preds
    recall = correct / num_trues
    # print(precision, recall, correct, num_preds)
    return precision, recall

def F1(last_output, bbox_true_b, y_true_b, images, detect_thresh=0.8, iou_thresh=0.5, **kwargs):
    y_pred_b, bbox_pred_b = last_output
    if len(y_true_b) == 0: return tensor(0.)

    f1 = 0
    for y_pred, bbox_pred, y_true, bbox_true in zip(y_pred_b, bbox_pred_b, y_true_b, bbox_true_b):

        # 1.get true bbox
        bbox_true, y_true = _unpad(bbox_true, y_true) # [N1,4]

        # 2. get pred bbox
        bbox_pred = bbox_pred.permute(1,2,0).view(-1,4)
        y_pred = y_pred.view(-1,1)

        bbox_pred = target_to_bbox(bbox_pred, anchors.to(y_pred.device))
    #     bbox_pred = bbox_pred.to('cpu')

        detect_mask = y_pred.squeeze() > detect_thresh # [H x W]
        if detect_mask.sum() == 0:
#             print('iou_metric: detect_mask.sum() == 0') # (too much spam when network start learning)
            return tensor(0)

        keep_idxes = nms(bbox_pred[detect_mask], y_pred[detect_mask].view(-1), threshold=0.2)
        bbox_pred = bbox_pred[detect_mask][keep_idxes] # [N2,4]

        # 3. calc result
        precision, recall = calc_precision_recall(bbox_pred, bbox_true, iou_thresh=iou_thresh)
        if precision + recall < 0.001: f1 += 0.
        else: f1 += 2 * (precision * recall) / (precision + recall)
    return tensor(f1 / len(y_true_b))

# Cell
def iou_loss(y_pred_geo, y_true_geo, eps=1.0):
    '''
    @param: [N,4] :
    @param: [N,4] :
    '''
    # d1 -> top, d2->left, d3->bottom, d4->right
    d1_true, d2_true, d3_true, d4_true = torch.split(y_true_geo, 1, 1) # [N,1] x 4
    d1_pred, d2_pred, d3_pred, d4_pred = torch.split(y_pred_geo, 1, 1) # [N,1] x 4
    area_true = (d1_true + d3_true) * (d2_true + d4_true)
    area_pred = (d1_pred + d3_pred) * (d2_pred + d4_pred)
    inter_width = torch.min(d2_true, d2_pred) + torch.min(d4_true, d4_pred)
    inter_height = torch.min(d1_true, d1_pred) + torch.min(d3_true, d3_pred)
    area_intersect = inter_width * inter_height
    area_union = area_true + area_pred - area_intersect
    L_AABB = -torch.log((area_intersect + eps) / (area_union + eps))
    return L_AABB

def dice_loss(y_pred, y_true):
    inter = torch.sum(y_true * y_pred)
    union = torch.sum(y_true) + torch.sum(y_pred) + 1e-5
    return 1. - (2 * inter / union)

# Cell
def encode_class(idxs, n_classes):
    target = idxs.new_zeros(len(idxs), n_classes).float()
    mask = idxs != 0
    i1s = LongTensor(list(range(len(idxs))))
    target[i1s[mask],idxs[mask]-1] = 1
    return target

class SigmaL1SmoothLoss(nn.Module):
    def forward(self, output, target):
        reg_diff = torch.abs(target - output)
        reg_loss = torch.where(torch.le(reg_diff, 1/9), 4.5 * torch.pow(reg_diff, 2), reg_diff - 1/18)
        return reg_loss.mean()

class RetinaNetFocalLoss(nn.Module):
    def __init__(self, anchors: Collection[float], gamma=1.0, pad_idx=0, class_loss_weight=1.0):
        super().__init__()
        self.class_weight = class_loss_weight
        self.gamma, self.pad_idx = gamma, pad_idx
        self.anchors, self.metric_names = anchors, ['bbox_loss', 'class_loss']

    def _unpad(self, bbox_true, y_true):
        ''' [N_max,4], [N_max] -> [N1,4], [N1] '''
        i = torch.min(torch.nonzero(y_true - self.pad_idx)) if sum(y_true)>0 else 0
        return bbox_true[i:], y_true[i:] + self.pad_idx

    def _cls_loss(self, y_pred, y_true, matches):
        ''' [N,1], [N,1] '''
#         return dice_loss(y_pred, y_true)
#         return F.binary_cross_entropy(y_pred, y_true)

        # CE = -log(pt), where pt is prob of correct class
        CE = F.binary_cross_entropy(y_pred, y_true, reduction='none')
        pt = torch.exp(-CE)
        return torch.mean( (1-pt)**self.gamma * CE )

    def _reg_loss(self, bbox_pred, bbox_true, y_pred):
        loss = iou_loss(bbox_pred*256, torch.clamp(bbox_true, min=0)*256)
        weights = -torch.log( torch.clamp( 1-y_pred, min=1e-3, max=0.7) )
#         weights = y_pred + 0.3
        return torch.mean(loss * weights)

    def _one_loss(self, y_pred, bbox_pred, y_true, bbox_true):
        ''' [N,C], [N,4], [N_max], [N_max,4] -> float, float '''
        bbox_true, y_true = self._unpad(bbox_true, y_true) # [N1,4], [N1]
        bbox_true, bbox_pred, matches = prepare_bboxes(self.anchors, bbox_true, bbox_pred) # [N2,4], [N2,4]
        if matches is None or matches is None: return tensor(0.)
        y_true, y_pred = prepare_labels(y_true, y_pred, matches) # [N2], [N2]
        bb_loss = self._reg_loss(bbox_pred, bbox_true, y_pred[matches[matches >= 0]])
        y_loss = self._cls_loss(y_pred, y_true, matches)
        self.metrics += tensor([bb_loss, y_loss])
        return (1-self.class_weight) * bb_loss + self.class_weight * y_loss

    def forward(self, nn_output, bbox_true, y_true, images):
        y_pred, bbox_pred = nn_output
        '''
        @param: [B, 1, H/4, W/4] :probability for each h,w cell
        @param: [B, 4, H/4, W/4] :bboxes for each h,w cell
        @param: [B, N_max, 4]    :true bboxes
        @param: [B, N_max]       :true label (1/0)
        @param: [B, H, W]
        '''
        if bbox_true.device != self.anchors.device:
            self.anchors = self.anchors.to(y_pred.device)

        loss = torch.tensor(0, dtype=torch.float32).to(y_pred.device)
        self.metrics = tensor([0., 0.])
        for yp, bp, yt, bt in zip(y_pred, bbox_pred, y_true, bbox_true):
            loss += self._one_loss(yp.view(-1,1), bp.permute(1,2,0).view(-1,4), yt, bt)

        bs = y_true.size(0)
        self.metrics = dict(zip(self.metric_names, [m/bs for m in self.metrics]))
        return loss/bs

# Cell
def threshold_images(imgs): # torch tensor [b,3,h,w] -> torch tensor [b,1,h,w] w/ 0 or 1
    # denormalize to 0-255
    imgs = (imgs - imgs.min())
    imgs = (imgs / imgs.max()) * 255.
    assert imgs.min() >= 0 and imgs.max() <= 255
    # threshold
    out_imgs = []
    for im in imgs:
        im = image2np(im).astype(np.uint8) # [h,w,3]
        im = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
        _,th = cv2.threshold(im,0,1,cv2.THRESH_BINARY_INV+cv2.THRESH_OTSU)
        out_imgs.append( tensor(th)[None] )
    return torch.cat(out_imgs, 0)

def prepare_masked_bboxes(anchors, bbox_true, bbox_pred, mask):
    matches = match_anchors(anchors, bbox_true) # [N3, 4]
    assert (tensor(matches.shape) == tensor(mask.shape)).all() == True
    matches[mask] = -1
    return prepare_bboxes(anchors, bbox_true, bbox_pred, matches)

class PreciseFocalLoss(RetinaNetFocalLoss):
    def forward(self, nn_output, bbox_true, y_true, images):
        y_pred, bbox_pred = nn_output

        # CHANGED
        b,c,h,w = y_pred.shape
        images = F.interpolate(images, size=(h,w), mode='bilinear', align_corners=True)
        text_mask = threshold_images(images).bool().to(device=y_pred.device)

        if bbox_true.device != self.anchors.device:
            self.anchors = self.anchors.to(y_pred.device)

        loss = torch.tensor(0, dtype=torch.float32).to(y_pred.device)
        self.metrics = tensor([0., 0.])
        for yp, bp, yt, bt, m in zip(y_pred, bbox_pred, y_true, bbox_true, text_mask): # CHANGED
            loss += self._one_loss(yp.view(-1,1), bp.permute(1,2,0).view(-1,4), yt, bt, m.view(-1)) # CHANGED

        bs = y_true.size(0)
        self.metrics = dict(zip(self.metric_names, [m/bs for m in self.metrics]))
        return loss/bs

    def _one_loss(self, y_pred, bbox_pred, y_true, bbox_true, mask):
        ''' [N,C], [N,4], [N_max], [N_max,4], [N_max] -> float, float '''
        bbox_true, y_true = self._unpad(bbox_true, y_true) # [N1,4], [N1]
        bbox_true, bbox_pred, matches = prepare_masked_bboxes(self.anchors, bbox_true, bbox_pred, mask) # [N2,4], [N2,4]
        if matches is None or matches is None: return tensor(0.)
        y_true, y_pred = prepare_labels(y_true, y_pred, matches) # CHANGED
        bb_loss = self._reg_loss(bbox_pred, bbox_true, y_pred[matches[matches >= 0]])
        y_loss = self._cls_loss(y_pred, y_true, matches)
        self.metrics += tensor([bb_loss, y_loss])
        return (1-self.class_weight) * bb_loss + self.class_weight * y_loss

# Cell
class MSEloss(RetinaNetFocalLoss):
    def _reg_loss(self, bbox_pred, bbox_true, y_pred):
        return nn.MSELoss()(bbox_pred, bbox_true)

    def _one_loss(self, y_pred, bbox_pred, y_true, bbox_true):
        ''' [N,C], [N,4], [N_max], [N_max,4] -> float, float '''
        bbox_true, y_true = self._unpad(bbox_true, y_true) # [N1,4], [N1]
        bbox_true, bbox_pred, matches = prepare_bboxes(self.anchors, bbox_true, bbox_pred) # [N2,4], [N2,4]
        if matches is None or matches is None: return tensor(0.)
        y_true, y_pred = prepare_labels(y_true, y_pred, matches) # [N2], [N2]
        bb_loss = self._reg_loss(bbox_pred, bbox_true, y_pred[matches[matches >= 0]])
        return bb_loss

# Cell
mse = MSEloss(anchors)
def MSE(last_output, bbox_true, y_true, **kwargs):
    return mse(last_output, bbox_true, y_true)