Spaces:
Sleeping
Sleeping
File size: 4,063 Bytes
5371cce | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | import torch
import torch.nn as nn
import torchvision.models as models
from huggingface_hub import hf_hub_download
from transformers import AutoImageProcessor, AutoModelForImageClassification
# Global variables to cache models
_wbc_model = None
_skin_processor = None
_skin_model = None
# Custom layers for WBC Model (from fastai)
class AdaptiveConcatPool2d(nn.Module):
def __init__(self, sz=None):
super().__init__()
self.ap = nn.AdaptiveAvgPool2d(sz or 1)
self.mp = nn.AdaptiveMaxPool2d(sz or 1)
def forward(self, x):
return torch.cat([self.ap(x), self.mp(x)], 1)
def get_wbc_model_architecture():
# Instantiate standard resnet18 architecture
backbone = models.resnet18()
backbone.avgpool = nn.Identity()
backbone.fc = nn.Identity()
# Sequential layout to match the fastai sequential wrappers
backbone_seq = nn.Sequential(
backbone.conv1,
backbone.bn1,
backbone.relu,
backbone.maxpool,
backbone.layer1,
backbone.layer2,
backbone.layer3,
backbone.layer4
)
# Custom head structure
head_seq = nn.Sequential(
AdaptiveConcatPool2d(),
nn.Flatten(),
nn.BatchNorm1d(1024),
nn.Dropout(0.25),
nn.Linear(1024, 512, bias=False),
nn.ReLU(inplace=True),
nn.BatchNorm1d(512),
nn.Dropout(0.5),
nn.Linear(512, 8, bias=True) # Modified below if state_dict does not contain bias
)
model = nn.Sequential(backbone_seq, head_seq)
return model
def load_wbc_model():
global _wbc_model
if _wbc_model is not None:
return _wbc_model
# Check for custom local keras model
import os
local_keras_path = os.path.join(os.path.dirname(__file__), "model", "my_model.keras")
if os.path.exists(local_keras_path):
print(f"Loading custom Keras WBC model from {local_keras_path}...")
import tensorflow as tf
model = tf.keras.models.load_model(local_keras_path)
_wbc_model = {"framework": "tensorflow", "model": model}
print("Custom Keras WBC Model loaded successfully.")
return _wbc_model
print("Loading fallback PyTorch WBC Blood Cell Model (ResNet-18)...")
model = get_wbc_model_architecture()
# Download weights
weights_path = hf_hub_download(repo_id="esab/pbc-cell-classifier", filename="cell_classifier_weights.pth")
state_dict = torch.load(weights_path, map_location="cpu")
# Adjust output layer bias if not present in weights
if "1.8.bias" not in state_dict:
model[1][8] = nn.Linear(512, 8, bias=False)
model.load_state_dict(state_dict)
model.eval()
_wbc_model = {"framework": "pytorch", "model": model}
print("WBC Model loaded successfully.")
return _wbc_model
def load_skin_model():
global _skin_processor, _skin_model
if _skin_model is not None and _skin_processor is not None:
return _skin_processor, _skin_model
print("Loading Skin Cancer Model (Vision Transformer)...")
model_name = "Anwarkh1/Skin_Cancer-Image_Classification"
try:
_skin_processor = AutoImageProcessor.from_pretrained(model_name)
_skin_model = AutoModelForImageClassification.from_pretrained(model_name, attn_implementation="eager")
except Exception as e:
print(f"Network error preloading ViT ({str(e)}). Attempting to load from local cache...")
try:
_skin_processor = AutoImageProcessor.from_pretrained(model_name, local_files_only=True)
_skin_model = AutoModelForImageClassification.from_pretrained(model_name, attn_implementation="eager", local_files_only=True)
except Exception as e_local:
print(f"Failed to load from local Hugging Face cache: {str(e_local)}")
raise e_local
_skin_model.eval()
print("Skin Cancer Model loaded successfully.")
return _skin_processor, _skin_model
def preload_all_models():
load_wbc_model()
load_skin_model()
|