# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# # OOD Detection
#
# The idea of this set of experiments is to evaluate a `ResNet50` backbone pretrained on ImageNet with CF data on OOD benchmarks such as `ImageNet Sketch`. If findings are worthy, we can extend it to other datasets.

# %load_ext autoreload
# %autoreload 2

# +
import re
from os.path import join, basename, dirname
from glob import glob
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from natsort import natsorted
from tqdm import tqdm
from PIL import Image
import torch
import torchvision
from torchvision.io import read_image
from torchvision.utils import make_grid
from torchvision import transforms

from gradcam.utils import visualize_cam
from gradcam import GradCAM, GradCAMpp
# -

from experiment_utils import set_env, REPO_PATH, seed_everything
set_env()

from experiments.image_utils import denormalize, show_single_image
from experiments.imagenet_utils import IMModel, AverageEnsembleModel
from experiments.ood_utils import validate as ood_validate

plt.rcParams.update({
    "text.usetex": True,
    "font.family": "serif",
    "font.serif": ["Computer Modern Roman"],
})

# ### Set environment

# +
seed_everything(0)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# -

# ### Load model

from imagenet.models.classifier_ensemble import InvariantEnsemble

model = InvariantEnsemble("resnet50", pretrained=True)

# load weights from a checkpoint
ckpt_path = "imagenet/experiments/classifier_2022_01_19_15_36_sample_run/model_best.pth"
ckpt = torch.load(ckpt_path, map_location="cpu")
ckpt_state_dict = ckpt["state_dict"]
ckpt_state_dict = {k.replace("module.", ""):v for k, v in ckpt_state_dict.items()}

model.load_state_dict(ckpt_state_dict)

model = model.eval()

model = model.to(device)

shape_model = IMModel(base_model=model, mode="shape")
avg_model = AverageEnsembleModel(base_model=model)
pytorch_model = torchvision.models.resnet50(pretrained=True)

# +
# check models
x = torch.randn((1, 3, 224, 224))

y = shape_model(x)
assert y.shape == torch.Size([1, 1000])

y = avg_model(x)
assert y.shape == torch.Size([1, 1000])

y = pytorch_model(x)
assert y.shape == torch.Size([1, 1000])
# -

# ### Download data

# Use `download_imagenet_sketch.sh` to download the data in folder `cgn_framework/imagenet/data/sketch/`

# ### Define the validation loader

valdir = join(join(REPO_PATH, "cgn_framework", "imagenet", "data/sketch"), 'val')

combined_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

val_loader = torch.utils.data.DataLoader(
    torchvision.datasets.ImageFolder(valdir, combined_transform),
    batch_size=36,
    shuffle=False,
    drop_last=False,
    num_workers=4,
    pin_memory=True,
)

len(val_loader)

# ### Benchmark performance

acc1 = ood_validate(val_loader, pytorch_model)




StarcoderdataJupyterScriptsDedupFiltered