Spaces:

catolim3
/

BN4601A_IndividualAssignment

Paused

Update app.py

cea7f00 verified about 1 month ago

7.19 kB

from __future__ import annotations

import logging
import os
from pathlib import Path

import gradio as gr
import torch
from PIL import Image
from torch import nn
from torchvision import transforms
from torchvision.models import resnet18


logging.basicConfig(level=logging.INFO)
LOGGER = logging.getLogger(__name__)


# This order must match the order used when the checkpoint was trained.
CLASS_NAMES = [
    "Normal",
    "Mild",
    "Moderate",
    "Severe",
    "Proliferative",
]

IMAGE_SIZE = 224
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]
MODEL_PATH = Path(
    os.getenv("MODEL_PATH", Path(__file__).with_name("best_retinamnist_resnet18.pth"))
)
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def _torch_load(path: Path):
    """Load a tensor-only checkpoint on CPU before moving the model to DEVICE."""
    try:
        # The supplied checkpoint contains a state_dict plus simple metadata.
        return torch.load(path, map_location="cpu", weights_only=True)
    except TypeError:
        # Compatibility with older PyTorch versions without weights_only.
        return torch.load(path, map_location="cpu")


def _get_state_dict(checkpoint):
    if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
        state_dict = checkpoint["model_state_dict"]
    elif isinstance(checkpoint, dict):
        # Also accept a raw state_dict for easier future checkpoint replacement.
        state_dict = checkpoint
    else:
        raise RuntimeError("The checkpoint is not a PyTorch state_dict or checkpoint dictionary.")

    if not isinstance(state_dict, dict) or not state_dict:
        raise RuntimeError("The checkpoint does not contain a usable model_state_dict.")

    # Make checkpoints saved through DataParallel usable as well.
    return {
        key.removeprefix("module."): value for key, value in state_dict.items()
    }


def build_model() -> nn.Module:
    if not MODEL_PATH.exists():
        raise FileNotFoundError(
            f"Checkpoint not found at {MODEL_PATH}. Upload "
            "best_retinamnist_resnet18.pth to the Space repository."
        )

    checkpoint = _torch_load(MODEL_PATH)
    metadata = checkpoint if isinstance(checkpoint, dict) else {}

    saved_model_name = str(metadata.get("model_name", "resnet18")).lower()
    if saved_model_name != "resnet18":
        raise RuntimeError(
            f"This app expects a ResNet-18 checkpoint, but the file declares "
            f"{saved_model_name!r}."
        )

    saved_image_size = int(metadata.get("image_size", IMAGE_SIZE))
    if saved_image_size != IMAGE_SIZE:
        raise RuntimeError(
            f"This app expects image_size={IMAGE_SIZE}, but the checkpoint declares "
            f"image_size={saved_image_size}."
        )

    model = resnet18(weights=None)
    model.fc = nn.Linear(model.fc.in_features, len(CLASS_NAMES))

    # strict=True catches an incorrect architecture instead of silently producing bad output.
    model.load_state_dict(_get_state_dict(checkpoint), strict=True)
    model.to(DEVICE)
    model.eval()
    LOGGER.info("Loaded %s on %s", MODEL_PATH.name, DEVICE)
    return model


MODEL = build_model()

# Match the validation/test transform in the training notebook exactly:
# resize to 224x224, convert to tensor, then ImageNet-normalize.
PREPROCESS = transforms.Compose(
    [
        transforms.Resize((IMAGE_SIZE, IMAGE_SIZE), antialias=True),
        transforms.ToTensor(),
        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
    ]
)


def _as_rgb_pil(image) -> Image.Image:
    if isinstance(image, Image.Image):
        return image.convert("RGB")
    # Defensive support for callers that pass a NumPy array instead of a PIL image.
    return Image.fromarray(image).convert("RGB")


@torch.inference_mode()
def predict(image):
    """Return five-class probabilities and a user-facing interpretation."""
    if image is None:
        return {}, "Upload or capture a fundus image, then select **Assess image**."

    try:
        pil_image = _as_rgb_pil(image)
        input_tensor = PREPROCESS(pil_image).unsqueeze(0).to(DEVICE)
        logits = MODEL(input_tensor)
        probabilities = torch.softmax(logits, dim=1)[0].cpu().tolist()
        predicted_index = int(torch.argmax(logits, dim=1).item())
        confidence = float(probabilities[predicted_index])

        scores = {
            name: float(probability)
            for name, probability in zip(CLASS_NAMES, probabilities)
        }
        predicted_class = CLASS_NAMES[predicted_index]

        result = (
            f"## Predicted category: {predicted_class}\n\n"
            f"**Model confidence:** {confidence:.1%}\n\n"
            "The confidence value is the softmax score for this image; it is "
            "not a calibrated probability of disease.\n\n"
            "> **Important:** This research demonstration is not a medical diagnosis. "
            "It was trained on RetinaMNIST images and may perform poorly on images "
            "from different cameras, image qualities, fields of view, or populations. "
            "Do not use it to make treatment or screening decisions."
        )
        return scores, result
    except Exception:
        LOGGER.exception("Prediction failed")
        return {}, (
            "**Could not assess this image.** Please upload a readable JPG/PNG fundus "
            "image and try again. Check the Space logs for the technical error."
        )


with gr.Blocks(theme=gr.themes.Soft(), title="RetinaMNIST DR Grader") as demo:
    gr.Markdown(
        """
        # RetinaMNIST diabetic-retinopathy grader
        Upload a fundus photograph or use your webcam. The ResNet-18 model returns
        the five RetinaMNIST/International Clinical DR severity categories.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(
                label="Fundus image",
                type="pil",
                sources=["upload", "webcam"],
                height=420,
            )
            with gr.Row():
                assess_button = gr.Button("Assess image", variant="primary")
                clear_button = gr.Button("Clear")

        with gr.Column(scale=1):
            score_output = gr.Label(
                label="Class probabilities",
                num_top_classes=5,
            )
            result_output = gr.Markdown(
                "Upload an image to see the assessment.",
            )

    gr.Markdown(
        "**Image guidance:** Use a centred, well-lit retinal fundus photograph with "
        "the optic disc and macula visible where possible. Non-fundus images are "
        "accepted by the interface but are out of scope for this model."
    )

    assess_button.click(
        fn=predict,
        inputs=input_image,
        outputs=[score_output, result_output],
        api_name="predict",
    )
    clear_button.click(
        fn=lambda: (
            None,
            {},
            "Upload an image to see the assessment.",
        ),
        outputs=[input_image, score_output, result_output],
    )


if __name__ == "__main__":
    demo.launch()

	from __future__ import annotations

	import logging
	import os
	from pathlib import Path

	import gradio as gr
	import torch
	from PIL import Image
	from torch import nn
	from torchvision import transforms
	from torchvision.models import resnet18


	logging.basicConfig(level=logging.INFO)
	LOGGER = logging.getLogger(__name__)


	# This order must match the order used when the checkpoint was trained.
	CLASS_NAMES = [
	"Normal",
	"Mild",
	"Moderate",
	"Severe",
	"Proliferative",
	]

	IMAGE_SIZE = 224
	IMAGENET_MEAN = [0.485, 0.456, 0.406]
	IMAGENET_STD = [0.229, 0.224, 0.225]
	MODEL_PATH = Path(
	os.getenv("MODEL_PATH", Path(__file__).with_name("best_retinamnist_resnet18.pth"))
	)
	DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


	def _torch_load(path: Path):
	"""Load a tensor-only checkpoint on CPU before moving the model to DEVICE."""
	try:
	# The supplied checkpoint contains a state_dict plus simple metadata.
	return torch.load(path, map_location="cpu", weights_only=True)
	except TypeError:
	# Compatibility with older PyTorch versions without weights_only.
	return torch.load(path, map_location="cpu")


	def _get_state_dict(checkpoint):
	if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
	state_dict = checkpoint["model_state_dict"]
	elif isinstance(checkpoint, dict):
	# Also accept a raw state_dict for easier future checkpoint replacement.
	state_dict = checkpoint
	else:
	raise RuntimeError("The checkpoint is not a PyTorch state_dict or checkpoint dictionary.")

	if not isinstance(state_dict, dict) or not state_dict:
	raise RuntimeError("The checkpoint does not contain a usable model_state_dict.")

	# Make checkpoints saved through DataParallel usable as well.
	return {
	key.removeprefix("module."): value for key, value in state_dict.items()
	}


	def build_model() -> nn.Module:
	if not MODEL_PATH.exists():
	raise FileNotFoundError(
	f"Checkpoint not found at {MODEL_PATH}. Upload "
	"best_retinamnist_resnet18.pth to the Space repository."
	)

	checkpoint = _torch_load(MODEL_PATH)
	metadata = checkpoint if isinstance(checkpoint, dict) else {}

	saved_model_name = str(metadata.get("model_name", "resnet18")).lower()
	if saved_model_name != "resnet18":
	raise RuntimeError(
	f"This app expects a ResNet-18 checkpoint, but the file declares "
	f"{saved_model_name!r}."
	)

	saved_image_size = int(metadata.get("image_size", IMAGE_SIZE))
	if saved_image_size != IMAGE_SIZE:
	raise RuntimeError(
	f"This app expects image_size={IMAGE_SIZE}, but the checkpoint declares "
	f"image_size={saved_image_size}."
	)

	model = resnet18(weights=None)
	model.fc = nn.Linear(model.fc.in_features, len(CLASS_NAMES))

	# strict=True catches an incorrect architecture instead of silently producing bad output.
	model.load_state_dict(_get_state_dict(checkpoint), strict=True)
	model.to(DEVICE)
	model.eval()
	LOGGER.info("Loaded %s on %s", MODEL_PATH.name, DEVICE)
	return model


	MODEL = build_model()

	# Match the validation/test transform in the training notebook exactly:
	# resize to 224x224, convert to tensor, then ImageNet-normalize.
	PREPROCESS = transforms.Compose(
	[
	transforms.Resize((IMAGE_SIZE, IMAGE_SIZE), antialias=True),
	transforms.ToTensor(),
	transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
	]
	)


	def _as_rgb_pil(image) -> Image.Image:
	if isinstance(image, Image.Image):
	return image.convert("RGB")
	# Defensive support for callers that pass a NumPy array instead of a PIL image.
	return Image.fromarray(image).convert("RGB")


	@torch.inference_mode()
	def predict(image):
	"""Return five-class probabilities and a user-facing interpretation."""
	if image is None:
	return {}, "Upload or capture a fundus image, then select Assess image."

	try:
	pil_image = _as_rgb_pil(image)
	input_tensor = PREPROCESS(pil_image).unsqueeze(0).to(DEVICE)
	logits = MODEL(input_tensor)
	probabilities = torch.softmax(logits, dim=1)[0].cpu().tolist()
	predicted_index = int(torch.argmax(logits, dim=1).item())
	confidence = float(probabilities[predicted_index])

	scores = {
	name: float(probability)
	for name, probability in zip(CLASS_NAMES, probabilities)
	}
	predicted_class = CLASS_NAMES[predicted_index]

	result = (
	f"## Predicted category: {predicted_class}\n\n"
	f"Model confidence: {confidence:.1%}\n\n"
	"The confidence value is the softmax score for this image; it is "
	"not a calibrated probability of disease.\n\n"
	"> Important: This research demonstration is not a medical diagnosis. "
	"It was trained on RetinaMNIST images and may perform poorly on images "
	"from different cameras, image qualities, fields of view, or populations. "
	"Do not use it to make treatment or screening decisions."
	)
	return scores, result
	except Exception:
	LOGGER.exception("Prediction failed")
	return {}, (
	"Could not assess this image. Please upload a readable JPG/PNG fundus "
	"image and try again. Check the Space logs for the technical error."
	)


	with gr.Blocks(theme=gr.themes.Soft(), title="RetinaMNIST DR Grader") as demo:
	gr.Markdown(
	"""
	# RetinaMNIST diabetic-retinopathy grader

	Upload a fundus photograph or use your webcam. The ResNet-18 model returns
	the five RetinaMNIST/International Clinical DR severity categories.
	"""
	)

	with gr.Row():
	with gr.Column(scale=1):
	input_image = gr.Image(
	label="Fundus image",
	type="pil",
	sources=["upload", "webcam"],
	height=420,
	)
	with gr.Row():
	assess_button = gr.Button("Assess image", variant="primary")
	clear_button = gr.Button("Clear")

	with gr.Column(scale=1):
	score_output = gr.Label(
	label="Class probabilities",
	num_top_classes=5,
	)
	result_output = gr.Markdown(
	"Upload an image to see the assessment.",
	)

	gr.Markdown(
	"Image guidance: Use a centred, well-lit retinal fundus photograph with "
	"the optic disc and macula visible where possible. Non-fundus images are "
	"accepted by the interface but are out of scope for this model."
	)

	assess_button.click(
	fn=predict,
	inputs=input_image,
	outputs=[score_output, result_output],
	api_name="predict",
	)
	clear_button.click(
	fn=lambda: (
	None,
	{},
	"Upload an image to see the assessment.",
	),
	outputs=[input_image, score_output, result_output],
	)


	if __name__ == "__main__":
	demo.launch()