import os
import re
import sys
from abc import abstractmethod
from reproducibility.utils.cacher import cache_hit_or_miss, cache_numpy_object
from reproducibility.embedders.internal_datasets import CLIPImageDataset
from torch.utils.data import DataLoader
from torch.utils import model_zoo
from torchvision.models.resnet import ResNet, model_urls as resnet_urls, BasicBlock, Bottleneck
from torchvision.models.densenet import DenseNet, model_urls as densenet_urls
import torch.nn.functional as F
import numpy as np
from torch import nn
class FeaturesInterface(object):
@abstractmethod
def n_features(self):
pass
import torch
from torch.hub import download_url_to_file
try:
from requests.utils import urlparse
from requests import get as urlopen
requests_available = True
except ImportError:
requests_available = False
from urllib.request import urlopen
from urllib.parse import urlparse
try:
from tqdm import tqdm
except ImportError:
tqdm = None # defined below
def _remove_prefix(s, prefix):
if s.startswith(prefix):
s = s[len(prefix):]
return s
def clean_state_dict(state_dict, prefix, filter=None):
if filter is None:
filter = lambda *args: True
return {_remove_prefix(k, prefix): v for k, v in state_dict.items() if filter(k)}
def load_dox_url(url, filename, model_dir=None, map_location=None, progress=True):
r"""Adapt to fit format file of mtdp pre-trained models
"""
if model_dir is None:
torch_home = os.path.expanduser(os.getenv('TORCH_HOME', '~/.torch'))
model_dir = os.getenv('TORCH_MODEL_ZOO', os.path.join(torch_home, 'models'))
if not os.path.exists(model_dir):
os.makedirs(model_dir)
cached_file = os.path.join(model_dir, filename)
if not os.path.exists(cached_file):
sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file))
sys.stderr.flush()
download_url_to_file(url, cached_file, None, progress=progress)
return torch.load(cached_file, map_location=map_location)
MTDRN_URLS = {
"resnet50": ("https://dox.uliege.be/index.php/s/kvABLtVuMxW8iJy/download", "resnet50-mh-best-191205-141200.pth")
}
class NoHeadResNet(ResNet, FeaturesInterface):
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
return x
def n_features(self):
return [b for b in list(self.layer4[-1].children()) if hasattr(b, 'num_features')][-1].num_features
def build_resnet(download_dir, pretrained=None, arch="resnet50", model_class=NoHeadResNet, **kwargs):
"""Constructs a ResNet-18 model.
Args:
arch (str): Type of densenet (among: resnet18, resnet34, resnet50, resnet101 and resnet152)
pretrained (str|None): If "imagenet", returns a model pre-trained on ImageNet. If "mtdp" returns a model
pre-trained in multi-task on digital pathology data. Otherwise (None), random weights.
model_class (nn.Module): Actual resnet module class
"""
params = {
"resnet18": [BasicBlock, [2, 2, 2, 2]],
"resnet34": [BasicBlock, [3, 4, 6, 3]],
"resnet50": [Bottleneck, [3, 4, 6, 3]],
"resnet101": [Bottleneck, [3, 4, 23, 3]],
"resnet152": [Bottleneck, [3, 8, 36, 3]]
}
model = model_class(*params[arch], **kwargs)
if isinstance(pretrained, str):
if pretrained == "imagenet":
url = resnet_urls[arch] # default imagenet
state_dict = model_zoo.load_url(url)
elif pretrained == "mtdp":
if arch not in MTDRN_URLS:
raise ValueError("No pretrained weights for multi task pretraining with architecture '{}'".format(arch))
url, filename = MTDRN_URLS[arch]
state_dict = load_dox_url(url, filename, model_dir=download_dir, map_location="cpu")
state_dict = clean_state_dict(state_dict, prefix="features.", filter=lambda k: not k.startswith("heads."))
else:
raise ValueError("Unknown pre-training source")
model.load_state_dict(state_dict)
return model
MTDP_URLS = {
"densenet121": ("https://dox.uliege.be/index.php/s/G72InP4xmJvOrVp/download", "densenet121-mh-best-191205-141200.pth")
}
class NoHeadDenseNet(DenseNet, FeaturesInterface):
def forward(self, x):
return F.adaptive_avg_pool2d(self.features(x), (1, 1))
def n_features(self):
return self.features[-1].num_features
def build_densenet(download_dir, pretrained=False, arch="densenet121", model_class=NoHeadDenseNet, **kwargs):
r"""Densenet-XXX model from
`"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_
Args:
arch (str): Type of densenet (among: densenet121, densenet169, densenet201 and densenet161)
pretrained (str|None): If "imagenet", returns a model pre-trained on ImageNet. If "mtdp" returns a model pre-trained
in multi-task on digital pathology data. Otherwise (None), random weights.
model_class (nn.Module): Actual densenet module class
"""
params = {
"densenet121": {"num_init_features": 64, "growth_rate": 32, "block_config": (6, 12, 24, 16)},
"densenet169": {"num_init_features": 64, "growth_rate": 32, "block_config": (6, 12, 32, 32)},
"densenet201": {"num_init_features": 64, "growth_rate": 32, "block_config": (6, 12, 48, 32)},
"densenet161": {"num_init_features": 96, "growth_rate": 48, "block_config": (6, 12, 36, 24)}
}
model = model_class(**(params[arch]), **kwargs)
if isinstance(pretrained, str):
# '.'s are no longer allowed in module names, but pervious _DenseLayer
# has keys 'norm.1', 'relu.1', 'conv.1', 'norm.2', 'relu.2', 'conv.2'.
# They are also in the checkpoints in model_urls. This pattern is used
# to find such keys.
if pretrained == "imagenet":
pattern = re.compile(
r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$')
state_dict = model_zoo.load_url(densenet_urls[arch])
for key in list(state_dict.keys()):
res = pattern.match(key)
if res:
new_key = res.group(1) + res.group(2)
state_dict[new_key] = state_dict[key]
del state_dict[key]
elif pretrained == "mtdp":
if arch not in MTDP_URLS:
raise ValueError("No pretrained weights for multi task pretraining with architecture '{}'".format(arch))
url, filename = MTDP_URLS[arch]
state_dict = load_dox_url(url, filename, model_dir=download_dir, map_location="cpu")
state_dict = clean_state_dict(state_dict, prefix="features.", filter=lambda k: not k.startswith("heads."))
else:
raise ValueError("Unknown pre-training source")
model.load_state_dict(state_dict)
return model
class ResNetBottom(nn.Module):
def __init__(self, original_model):
super(ResNetBottom, self).__init__()
self.features = nn.Sequential(*list(original_model.children())[:-1])
def forward(self, x):
x = self.features(x)
x = torch.flatten(x, 1)
return x
class DenseNetEmbedder:
def __init__(self, model, preprocess, name, backbone):
self.model = model
self.preprocess = preprocess
self.name = name
self.backbone = backbone
def image_embedder(self, list_of_images, device="cuda", num_workers=1, batch_size=32, additional_cache_name=""):
# additional_cache_name: name of the validation dataset (e.g., Kather_7K)
hit_or_miss = cache_hit_or_miss(self.name + "img" + additional_cache_name, self.backbone)
if hit_or_miss is not None:
return hit_or_miss
else:
hit = self.embed_images(list_of_images, device=device, num_workers=num_workers, batch_size=batch_size)
cache_numpy_object(hit, self.name + "img" + additional_cache_name, self.backbone)
return hit
def embed_images(self, list_of_images, device="cuda", num_workers=1, batch_size=32):
dataset = CLIPImageDataset(list_of_images, self.preprocess)
dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers)
all_embs = []
for batch_X in tqdm(dataloader):
batch_X = batch_X.to(device)
embeddings = self.model(batch_X).detach().float().squeeze()
embeddings = embeddings.detach().cpu().numpy()
all_embs.append(embeddings)
return np.concatenate(all_embs)
Datasets
Models
