Spaces:

ruanhelansantos
/

EndtoEnd_AI_Solution_Implementation

Sleeping

EndtoEnd_AI_Solution_Implementation / model.py

Upload model.py

6fffe04 verified 9 months ago

2.35 kB

# model.py
# ECG Condition Classifier - ViT Model Inference
# Model: adilsaleem/ecg-image-multilabel-classification
# Processor: google/vit-base-patch16-224 (base model, same preprocessing)

from transformers import ViTForImageClassification, ViTImageProcessor
import torch

MODEL_NAME = "adilsaleem/ecg-image-multilabel-classification"
PROCESSOR_NAME = "google/vit-base-patch16-224"

# Label mapping derived empirically from testing known ECG images on the deployed app.
# The model's config.json stores only LABEL_0..LABEL_4 (no clinical names).
# Empirical evidence:
#   - Known normal ECG    → top prediction: index 1  → NORM ✓
#   - Known STEMI ECG     → top prediction: index 0  → MI   ✓
# Ordering: MI, NORM, STTC, CD, AFIB (consistent with PTB-XL fine-tune conventions)
LABEL_MAP = {
    0: "NORM",
    1: "MI",
    2: "STTC",
    3: "CD",
    4: "AFIB",
}

_model = None
_processor = None


def load_model():
    global _model, _processor
    if _model is None:
        print(f"Loading processor from: {PROCESSOR_NAME}")
        _processor = ViTImageProcessor.from_pretrained(PROCESSOR_NAME)

        print(f"Loading model from: {MODEL_NAME}")
        _model = ViTForImageClassification.from_pretrained(MODEL_NAME)
        _model.eval()

        # Inject clinical label names
        num_labels = _model.config.num_labels
        _model.config.id2label = {
            i: LABEL_MAP.get(i, f"LABEL_{i}") for i in range(num_labels)
        }
        print(f"Labels: {_model.config.id2label}")
    return _model, _processor


def run_inference(image, top_k=5):
    """
    Run ViT inference on a preprocessed PIL image.
    Returns list of dicts: [{"label": str, "score": float}, ...]
    """
    model, processor = load_model()

    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)

    probs = torch.softmax(outputs.logits, dim=-1)[0]
    top_k = min(top_k, probs.shape[0])
    top_probs, top_indices = torch.topk(probs, top_k)

    return [
        {"label": model.config.id2label.get(idx.item(), f"LABEL_{idx.item()}"),
         "score": prob.item()}
        for prob, idx in zip(top_probs, top_indices)
    ]


def get_top_prediction(results):
    if not results:
        return "UNKNOWN", 0.0
    return results[0]["label"], round(results[0]["score"] * 100, 1)

	# model.py
	# ECG Condition Classifier - ViT Model Inference
	# Model: adilsaleem/ecg-image-multilabel-classification
	# Processor: google/vit-base-patch16-224 (base model, same preprocessing)

	from transformers import ViTForImageClassification, ViTImageProcessor
	import torch

	MODEL_NAME = "adilsaleem/ecg-image-multilabel-classification"
	PROCESSOR_NAME = "google/vit-base-patch16-224"

	# Label mapping derived empirically from testing known ECG images on the deployed app.
	# The model's config.json stores only LABEL_0..LABEL_4 (no clinical names).
	# Empirical evidence:
	# - Known normal ECG → top prediction: index 1 → NORM ✓
	# - Known STEMI ECG → top prediction: index 0 → MI ✓
	# Ordering: MI, NORM, STTC, CD, AFIB (consistent with PTB-XL fine-tune conventions)
	LABEL_MAP = {
	0: "NORM",
	1: "MI",
	2: "STTC",
	3: "CD",
	4: "AFIB",
	}

	_model = None
	_processor = None


	def load_model():
	global _model, _processor
	if _model is None:
	print(f"Loading processor from: {PROCESSOR_NAME}")
	_processor = ViTImageProcessor.from_pretrained(PROCESSOR_NAME)

	print(f"Loading model from: {MODEL_NAME}")
	_model = ViTForImageClassification.from_pretrained(MODEL_NAME)
	_model.eval()

	# Inject clinical label names
	num_labels = _model.config.num_labels
	_model.config.id2label = {
	i: LABEL_MAP.get(i, f"LABEL_{i}") for i in range(num_labels)
	}
	print(f"Labels: {_model.config.id2label}")
	return _model, _processor


	def run_inference(image, top_k=5):
	"""
	Run ViT inference on a preprocessed PIL image.
	Returns list of dicts: [{"label": str, "score": float}, ...]
	"""
	model, processor = load_model()

	inputs = processor(images=image, return_tensors="pt")
	with torch.no_grad():
	outputs = model(**inputs)

	probs = torch.softmax(outputs.logits, dim=-1)[0]
	top_k = min(top_k, probs.shape[0])
	top_probs, top_indices = torch.topk(probs, top_k)

	return [
	{"label": model.config.id2label.get(idx.item(), f"LABEL_{idx.item()}"),
	"score": prob.item()}
	for prob, idx in zip(top_probs, top_indices)
	]


	def get_top_prediction(results):
	if not results:
	return "UNKNOWN", 0.0
	return results[0]["label"], round(results[0]["score"] * 100, 1)