!mkdir -p /root/.cache/torch/hub/checkpoints
!cp -r ../input/landmark-additional-packages/rwightman_gen-efficientnet-pytorch_master/rwightman_gen-efficientnet-pytorch_master /root/.cache/torch/hub
!cp ../input/landmark-additional-packages/tf_efficientnet_b3_aa-84b4657e.pth /root/.cache/torch/hub/checkpoints/
!cp ../input/landmark-additional-packages/tf_efficientnet_b5_ra-9a3e5369.pth /root/.cache/torch/hub/checkpoints/
!cp ../input/landmark-additional-packages/se_resnext50_32x4d-a260b3a4.pth /root/.cache/torch/hub/checkpoints/
!cp ../input/landmark-additional-packages/resnet50d_ra2-464e36ba.pth /root/.cache/torch/hub/checkpoints/
!pip install -q ../input/landmark-additional-packages/timm-0.3.4-py3-none-any.whl
!pip install -q ../input/landmark-additional-packages/geffnet-1.0.0-py3-none-any.whl
!pip install -q ../input/landmark-additional-packages/EfficientNet-PyTorch/EfficientNet-PyTorch-master
!pip install -q ../input/landmark-additional-packages/pycocotools-2.0.2/dist/pycocotools-2.0.2.tar
!pip install -q ../input/landmark-additional-packages/pretrainedmodels-0.7.4/pretrainedmodels-0.7.4
!pip install "/kaggle/input/hpamisc/pytorch_zoo-master"
!pip install "/kaggle/input/hpamisc/pycocotools-2.0-cp37-cp37m-linux_x86_64.whl"
!pip install "/kaggle/input/hpamisc/faiss_gpu-1.7.0-cp37-cp37m-manylinux2014_x86_64.whl"
! python ../input/maozi-no-arcface/maozi_no_arcface.py
import sys
sys.path.append('../input/hpa-singlecell-e050f56/hpa_singlecell-double_level_valid_all/')

from torch import nn
import torch
import torch.nn.functional as F
import torchvision
import timm
from torch.nn.parameter import Parameter
import albumentations as A

from utils import parse_args, prepare_for_result
from torch.utils.data import DataLoader, Dataset
from losses import get_loss, get_class_balanced_weighted
from dataloaders import get_dataloader
from utils import load_matched_state
from configs import Config
from models import get_model
from dataloaders.transform_loader import get_tfms

tensor_tfms = torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406, 0.406], std=[0.229, 0.224, 0.225, 0.225]),
        ])

tta_tfms = A.Compose([
    A.Resize(always_apply=False, p=1, height=256, width=256, interpolation=1),
    A.HorizontalFlip(always_apply=False, p=0.5),
    A.ShiftScaleRotate(always_apply=False, p=0.7, shift_limit_x=(-0.06, 0.06), shift_limit_y=(-0.06, 0.06), scale_limit=(-0.3, 0.3), rotate_limit=(-22.5, 22.5), interpolation=1, border_mode=2, value=None, mask_value=None),
    A.RandomBrightnessContrast(always_apply=False, p=0.5, brightness_limit=(-0.2, 0.2), contrast_limit=(-0.2, 0.2), brightness_by_max=True),
])


import base64
import zlib
from pycocotools import _mask as coco_mask
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import cv2
import tqdm
import seaborn as sns
def binary_mask_to_ascii(mask, mask_val=1):
    """Converts a binary mask into OID challenge encoding ascii text."""
    mask = np.where(mask==mask_val, 1, 0).astype(np.bool)
    
    # check input mask --
    if mask.dtype != np.bool:
        raise ValueError(f"encode_binary_mask expects a binary mask, received dtype == {mask.dtype}")

    mask = np.squeeze(mask)
    if len(mask.shape) != 2:
        raise ValueError(f"encode_binary_mask expects a 2d mask, received shape == {mask.shape}")

    # convert input mask to expected COCO API input --
    mask_to_encode = mask.reshape(mask.shape[0], mask.shape[1], 1)
    mask_to_encode = mask_to_encode.astype(np.uint8)
    mask_to_encode = np.asfortranarray(mask_to_encode)

    # RLE encode mask --
    encoded_mask = coco_mask.encode(mask_to_encode)[0]["counts"]

    # compress and base64 encoding --
    binary_str = zlib.compress(encoded_mask, zlib.Z_BEST_COMPRESSION)
    base64_str = base64.b64encode(binary_str)
    return base64_str.decode()

def process(x):
    iid, msk, img, sz = x
    img = cv2.resize(img, (2048, 2048))
    enc_msk = cv2.resize(msk, (sz, sz))
    cell_mask = msk
    subs = {}
    results = []
    for i in range(1, cell_mask.max() + 1):
        enc = binary_mask_to_ascii(enc_msk, i)
        sub = cv2.resize((cell_mask == i).astype(np.float), (2048, 2048), cv2.INTER_LINEAR)
        xr, yr = np.where(sub == 1)
        xmin, xmax, ymin, ymax = xr.min(), xr.max(), yr.min(), yr.max()
        subs[i] = (img * np.repeat((sub == 1).astype(np.int)[:, :, np.newaxis], 4, 2))[xmin:xmax, ymin: ymax]
#         imsave(f'./seg_png_fix_test/{iid}_{i}.png', (255 * subs[i]).astype(np.uint8))
        results.append(((255 * subs[i]).astype(np.uint8), enc, sz, sz))
    return results

def squarify(M,val):
    (a,b,c)=M.shape
    if a>b:
        padding=((0,0),((a-b)//2,a-b-(a-b)//2),(0, 0))
    else:
        padding=(((b-a)//2,b-a-(b-a)//2),(0,0),(0, 0))
    return np.pad(M,padding,mode='constant',constant_values=val)
"""
## Loading models
* b3
* b5
* r50d
* r200d
* se50
"""
ckpt = {
    0: 13, 1: 12, 2: 12, 3: 11, 4: 14
}

models = []
for i in range(5):
    cfg = Config.load_json('../input/hpa-single-cell-b3-philandrare-5f/5f_double_sin_exp5_rare.yaml/config.json')
    model = get_model(cfg).cuda()
    load_matched_state(model, torch.load(
        f'../input/hpa-single-cell-b3-philandrare-5f/5f_double_sin_exp5_rare.yaml/f{i}_epoch-{ckpt[i]}.pth'))
    _ = model.eval()
    models.append(model)
ckpt = {
    0: 18, 1: 14, 2: 14, 3: 15, 4: 15
}

# models = []
for i in range(5):
    if i in [2, 3, 4]: continue
    cfg = Config.load_json('../input/hpa-b5-final-model/b5_final_hpa_0504/config.json')
    model = get_model(cfg).cuda()
    load_matched_state(model, torch.load(
        f'../input/hpa-b5-final-model/b5_final_hpa_0504/checkpoints/f{i}_epoch-{ckpt[i]}.pth'))
    _ = model.eval()
    models.append(model)
ckpt = {
    0: 19, 1: 19, 2: 17, 3: 17, 4: 18
}

for i in range(5):
    if i in [0, 3, 4]: continue
    cfg = Config.load_json('../input/hpa-resnet50d-0508/resnet50d_final/config.json')
    model = get_model(cfg).cuda()
    load_matched_state(model, torch.load(
        f'../input/hpa-resnet50d-0508/resnet50d_final/checkpoints/f{i}_epoch-{ckpt[i]}.pth'))
    _ = model.eval()
    models.append(model)
!cp ../input/landmark-additional-packages/resnet200d_ra2-bdba9bf9.pth /root/.cache/torch/hub/checkpoints/
ckpt = {
    0: 15, 1: 15, 2: 13, 3: 13
}

for i in range(4):
    if i in [0, 1, 4]: continue
    cfg = Config.load_json('../input/hpa-jakiro-resnet200d/double_sin_exp5_r200d_rarex2_upload/config.json')
    model = get_model(cfg).cuda()
    load_matched_state(model, torch.load(
        f'../input/hpa-jakiro-resnet200d/double_sin_exp5_r200d_rarex2_upload/f{i}_epoch-{ckpt[i]}.pth'))
    _ = model.eval()
    models.append(model)
ckpt = {
    0: 19, 1: 16, 2: 16, 3: 17, 4:19
}

for i in range(5):
    if i in [0, 1, 2]: continue
    print(i)
    cfg = Config.load_json('../input/hpa-se50-final-0509/se50_final/config.json')
    model = get_model(cfg).cuda()
    load_matched_state(model, torch.load(
        f'../input/hpa-se50-final-0509/se50_final/checkpoints/f{i}_epoch-{ckpt[i]}.pth'))
    _ = model.eval()
    models.append(model)
len(models)
"""
## If we read from a csv
"""
df = pd.read_csv('submission.csv')

imgs = []
for i, x in df.iterrows():
    label = x.PredictionString.split(' ')[0::3]
    prob = x.PredictionString.split(' ')[1::3]
    encodes = x.PredictionString.split(' ')[2::3]
    for idx, enc in enumerate(list(set(encodes))):
        imgs.append({
            'image_id': x.ID,
            'cell_id': idx+1,
            'enc': enc,
            'fname': f'{x.ID}_{idx+1}',
        })

tm = pd.DataFrame(imgs)
probs = []
for i, x in df.iterrows():
    label = x.PredictionString.split(' ')[0::3]
    prob = x.PredictionString.split(' ')[1::3]
    encodes = x.PredictionString.split(' ')[2::3]
    for idx, enc in enumerate(encodes):
        probs.append({
            'enc': enc,
            'predict': int(label[idx]),
            'prob': float(prob[idx])
        })

prob = pd.DataFrame(probs)
tm_pred = prob.groupby(['enc', 'predict']).mean().unstack()['prob']
tm_pred.columns.name = ''
team = tm[['enc', 'fname']].merge(tm_pred.reset_index(), on='enc', how='inner').drop('enc', 1)
sample_submission = pd.read_csv('../input/hpa-single-cell-image-classification/sample_submission.csv', index_col=0)
team_pred = team.set_index('fname')
class SliceInferenceDataset(torch.utils.data.Dataset):
    def __init__(self, df, tta=16, cfg=None, tfms=None):
        self.df = df
        self.iids = self.df.image_id.unique()
        self.tta = tta
        
    def __len__(self):
        return len(self.iids)

    def __getitem__(self, idx):
        iid = self.iids[idx]
        mt = f'../input/hpa-single-cell-image-classification/test/{iid}_red.png'
        er = f'../input/hpa-single-cell-image-classification/test/{iid}_yellow.png'
        nu = f'../input/hpa-single-cell-image-classification/test/{iid}_blue.png'
        pr = f'../input/hpa-single-cell-image-classification/test/{iid}_green.png'
        r = cv2.imread(mt, 0).astype(np.float) / 255.0
        g = cv2.imread(pr, 0).astype(np.float) / 255.0
        b = cv2.imread(nu, 0).astype(np.float) / 255.0
        a = cv2.imread(er, 0).astype(np.float) / 255.0
        sz = r.shape[0]
        img = np.stack([r, g, b, a], -1)
        sli = []
        for i, x in self.df[self.df.image_id == iid].iterrows():
            bd = base64.b64decode(x.enc)
            zd = zlib.decompress(bd)
            encoded = [{'counts': zd, 'size': (sz, sz)}]
            ded = coco_mask.decode(encoded)[:, :, 0]

            xr, yr = np.where(ded == 1)
            sub = img[xr.min(): xr.max(), yr.min(): yr.max()]
            crop_sub_mask = ded[xr.min(): xr.max(), yr.min(): yr.max()]
            crop_sub_mask = np.repeat(crop_sub_mask[:, :, np.newaxis], 4, axis=2)
            r = sub * crop_sub_mask
            sli.append((cv2.resize(squarify(r, 0), (256, 256)).astype(np.float32), x.fname))
        BS, tta=len(sli) + 1, self.tta
        ipts = []
        raw_ipt = [e[0] for e in sli]
        for tt in range(tta):
            ipts.append(torch.stack([tensor_tfms(tta_tfms(image=x)['image']) for x in raw_ipt]).float())
        return ipts, BS, len(sli), tta, iid, [x[1] for x in sli]
# tm
sid = SliceInferenceDataset(tm, tta=8)
dl = torch.utils.data.DataLoader(sid, batch_size=1, num_workers=2)
pdfs = []
whole_dfs = []
for ipts, BS, lsli, tta, iid, fnames_raw in tqdm.tqdm(dl):
    BS, tta, iid, fnames, lsli = BS.item(), tta.item(), iid[0], [e[0] for e in fnames_raw], lsli.item()
    predicted_ps = []
    exp_ps = []
    for i in range(0, lsli, BS):
    #   ipt = torch.stack([tensor_tfms(cv2.resize(squarify(s[0], 0), (256, 256))) for s in ress[i: BS+i]]).cuda()
        with torch.no_grad():
            res = []
            exp = []
            for tt in range(tta):
                ipt = ipts[tt][0].cuda()
                for model in models:
                    with torch.cuda.amp.autocast():
                        ifr = model(ipt, len(ipt))
                    res.append(ifr[0].float())
                    exp.append(ifr[1].float())
        predict_p = [torch.sigmoid(r.cpu()) for r in res]
        exp_p = [torch.sigmoid(r.cpu()) for r in exp]
        predict_p = np.stack(predict_p).mean(0)
        exp_p = np.stack(exp_p).mean(0)
        predicted_ps.append(predict_p)
        exp_ps.append(exp_p)
    p = np.concatenate(predicted_ps)
    image_df = pd.DataFrame(p, index=fnames)
    whole_df = pd.DataFrame(np.concatenate(exp_ps).mean(0).reshape(1, 19), index=[iid])
    whole_dfs.append(whole_df)
    pdfs.append(image_df) 
# tm = tm.reset_index('fname')
image_level = pd.concat(whole_dfs)
image_pred = image_level.reset_index().merge(
    tm[['image_id', 'fname']], left_on='index', right_on='image_id', how='left'
).set_index('fname').drop(['index', 'image_id'], 1)
pub_pred = pd.concat(pdfs)
merge_pred = pub_pred * image_pred.loc[pub_pred.index]
"""
## If any ensemble
"""
merge_pred
ensem = merge_pred + team_pred.loc[merge_pred.index]
merge_pred = ensem
"""
## Save prediction
"""
df = df.set_index('ID')
merge_pred.index.name = 'fname'
merge_pred = merge_pred.reset_index()
tm = tm.set_index('fname')

merge_pred['ID'] = merge_pred['fname'].str.split('_', expand=True)[0]
j_pred = []
for iid in merge_pred.ID.unique():
    enc = ''
    sub_df = merge_pred[merge_pred.ID == iid]
    for idx, row in sub_df.iterrows():
        for i in range(19):
            enc += f'{i} {row[i]} {tm.loc[row.fname].enc} '
    j_pred.append({
        'ID': iid,
        'ImageWidth': df.loc[iid].ImageWidth,
        'ImageHeight': df.loc[iid].ImageHeight,
        'PredictionString': enc[:-1]
    })
fast_sub = pd.DataFrame(j_pred)
fast_sub.to_csv('pub.csv')
fast_sub = fast_sub.set_index('ID')
fast_sub.head(2)
"""
## save
"""
sub2 = pd.concat([sample_submission.drop(fast_sub.index), fast_sub], 0)
sub2 = sub2.loc[sample_submission.index]
sub2.to_csv('submission.csv')