Spaces:

Alex1-ai
/

CataractScan-AI

Running on Zero

added cataract image to the vision cataract model

2ecf98b 11 days ago

10.4 kB

import os
import json
import warnings

import gradio as gr
import gdown
import numpy as np
import tensorflow as tf
import spaces

from tensorflow.keras.applications.mobilenet_v2 import preprocess_input


# ============================================================
# CONFIGURATION
# ============================================================

MODEL_FILE_ID = "10axUT2jw9rAEo8FI4i9_ZGTv83UWqMeF"
CLASS_NAMES_FILE_ID = "1ztvyL4h0sQlLeyf97E9RoZZ60yzrpmRi"
MODEL_PATH = "cataract_model.keras"
CLASS_NAMES_PATH = "class_names.json"

IMG_HEIGHT = 180
IMG_WIDTH = 180


# ============================================================
# ZEROGPU STARTUP PLACEHOLDER
# ============================================================
# ZeroGPU refuses to start a Space unless it finds at least one
# @spaces.GPU function. TensorFlow does not work well inside the
# separate GPU worker process, so this dummy function only exists
# to satisfy the startup check. The real prediction function below
# is deliberately NOT decorated and runs in the main process.

@spaces.GPU
def _zerogpu_placeholder():
    return None


# ============================================================
# GOOGLE DRIVE DOWNLOAD LOGIC
# ============================================================

def download_if_missing(file_id, dest_path, label):
    """Download a Google Drive file if it is not already present."""

    if os.path.exists(dest_path) and os.path.getsize(dest_path) > 0:
        print(f"{label} already exists. Skipping download.")
        return

    if not file_id or file_id.startswith("YOUR_"):
        raise ValueError(f"Please configure the Google Drive ID for {label}.")

    print(f"Downloading {label} from Google Drive...")

    url = f"https://drive.google.com/uc?id={file_id}"
    downloaded_path = gdown.download(url=url, output=dest_path, quiet=False)

    if (
        not downloaded_path
        or not os.path.exists(dest_path)
        or os.path.getsize(dest_path) == 0
    ):
        raise RuntimeError(
            f"Failed to download {label}. Check that the Google Drive "
            "file is shared as 'Anyone with the link'."
        )

    print(f"{label} downloaded successfully.")


download_if_missing(MODEL_FILE_ID, MODEL_PATH, "TensorFlow model")
download_if_missing(CLASS_NAMES_FILE_ID, CLASS_NAMES_PATH, "class names")


# ============================================================
# LOAD CLASS NAMES
# ============================================================

with open(CLASS_NAMES_PATH, "r", encoding="utf-8") as file:
    class_names = json.load(file)

if isinstance(class_names, dict):
    categories = [
        class_names[key]
        for key in sorted(class_names, key=lambda item: int(item))
    ]
else:
    categories = list(class_names)

if len(categories) != 2:
    raise ValueError(f"Expected two class names, but received: {categories}")

print("Class index order:", categories)


# ============================================================
# LOAD TENSORFLOW MODEL
# ============================================================

with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    model = tf.keras.models.load_model(MODEL_PATH, compile=False)

print("TensorFlow version:", tf.__version__)
print("Model loaded successfully.")
print("Model input shape:", model.input_shape)


# ============================================================
# IMAGE PREPROCESSING
# ============================================================

def preprocess_image(img):
    """Prepare an uploaded PIL image for MobileNetV2."""

    img = img.convert("RGB")
    img = img.resize((IMG_WIDTH, IMG_HEIGHT))

    image_array = np.asarray(img, dtype=np.float32)
    image_array = np.expand_dims(image_array, axis=0)

    # Must match the preprocessing used during model training.
    # If your model already contains a Rescaling / preprocess_input
    # layer, remove this line to avoid applying it twice.
    return preprocess_input(image_array)


# ============================================================
# PREDICTION (runs in the main process, no @spaces.GPU)
# ============================================================

def classify_image(img):
    """Return class probabilities and a human-readable result."""

    if img is None:
        return {}, "### Upload an eye image to begin.", "—"

    image_array = preprocess_image(img)

    # Calling the model directly is faster and more reliable than
    # model.predict() for a single image.
    prediction = np.asarray(
        model(image_array, training=False)
    ).reshape(-1)

    # Supports either a single sigmoid output or a two-class output.
    if len(prediction) == 1:
        probability_class_1 = float(np.clip(prediction[0], 0.0, 1.0))
        probabilities = {
            categories[0]: 1.0 - probability_class_1,
            categories[1]: probability_class_1,
        }

    elif len(prediction) == 2:
        values = np.clip(prediction.astype(float), 0.0, None)
        total = float(values.sum())

        if total == 0:
            return {}, "### Unable to calculate a prediction.", "—"

        values = values / total
        probabilities = {categories[i]: float(values[i]) for i in range(2)}

    else:
        raise ValueError(f"Unexpected model output shape: {prediction.shape}")

    predicted_class = max(probabilities, key=probabilities.get)
    confidence = probabilities[predicted_class] * 100

    if predicted_class.strip().lower() == "cataract":
        icon = "🔎"
        description = "The model detected features associated with cataracts."
    else:
        icon = "✓"
        description = "The model classified this image as normal."

    result_markdown = f"""
## {icon} {predicted_class}
{description}
**Model confidence: {confidence:.1f}%**
*This is an AI-generated classification, not a medical diagnosis.*
"""

    return probabilities, result_markdown, f"{confidence:.1f}%"


# ============================================================
# GRADIO INTERFACE
# ============================================================

custom_css = """
.gradio-container {
    max-width: 1120px !important;
    margin: auto !important;
}
#hero {
    text-align: center;
    padding: 32px 20px 26px;
    border-radius: 22px;
    background: linear-gradient(135deg, #102b46, #145c67);
    color: white;
    margin-bottom: 18px;
}
#hero h1 {
    color: white !important;
    font-size: 2.4rem !important;
    margin-bottom: 10px;
}
#hero p {
    color: #d5eeee !important;
    font-size: 1.05rem;
}
.panel {
    border: 1px solid var(--border-color-primary);
    border-radius: 18px;
    padding: 18px;
}
#predict-button {
    background: linear-gradient(90deg, #0d9488, #0f766e);
    color: white;
    border: none;
    font-weight: 700;
    border-radius: 12px;
}
#predict-button:hover {
    filter: brightness(1.08);
}
#disclaimer {
    border-left: 4px solid #d97706;
    padding: 12px 16px;
    border-radius: 8px;
    background: var(--block-background-fill);
}
"""

theme = gr.themes.Soft(
    primary_hue="teal",
    secondary_hue="cyan",
    neutral_hue="slate",
)

with gr.Blocks(title="CataractScan AI") as demo:

    gr.HTML("""
    <div id="hero">
        <h1>👁️ CataractScan AI</h1>
        <p>
            AI-powered eye image classification using TensorFlow
            and MobileNetV2
        </p>
    </div>
    """)

    gr.Markdown(
        """
        Upload an eye image to see how the model classifies it.
        Your result includes the predicted class and the model's
        estimated probabilities.
        """
    )

    with gr.Row(equal_height=True):

        with gr.Column(scale=1, elem_classes=["panel"]):
            gr.Markdown("### 📤 Upload an eye image")

            image_input = gr.Image(
                type="pil",
                label="Eye image",
                height=340,
            )

            with gr.Row():
                predict_button = gr.Button(
                    "Analyze Image",
                    variant="primary",
                    elem_id="predict-button",
                    scale=3,
                )
                clear_button = gr.Button("Clear", scale=1)

            gr.Markdown(
                """
                **Tips:** Use a clear, well-lit eye image.
                The model may perform poorly on images that differ
                from its training data.
                """
            )

        with gr.Column(scale=1, elem_classes=["panel"]):
            gr.Markdown("### 🧠 Classification results")

            prediction_output = gr.Label(
                label="Class probabilities",
                num_top_classes=2,
            )

            result_output = gr.Markdown("### Your result will appear here.")

            confidence_output = gr.Textbox(
                label="Confidence of predicted class",
                value="—",
                interactive=False,
            )

    gr.Markdown("### 🖼️ Try an example")

    example_images = [
        [name] for name in ["cataract2.jpg","cataract.jpg", "normal.png", "mine.jpg"] if os.path.exists(name)
    ]

    if example_images:
        gr.Examples(
            examples=example_images,
            inputs=image_input,
            label="Sample eye images",
        )

    gr.Markdown(
        """
        <div id="disclaimer">
        **⚠️ Medical disclaimer**
        CataractScan AI is an educational AI project, not a
        diagnostic medical device. Its output must not be used
        to diagnose or rule out cataracts. Image quality,
        training data, and model limitations can affect results.
        Consult a qualified eye-care professional for medical
        assessment.
        </div>
        """
    )

    outputs = [prediction_output, result_output, confidence_output]

    # Only the button triggers prediction, so each image is analysed once
    # (the old .change() handler made every upload run twice).
    predict_button.click(
        fn=classify_image,
        inputs=image_input,
        outputs=outputs,
    )

    def clear_all():
        return None, {}, "### Your result will appear here.", "—"

    clear_button.click(
        fn=clear_all,
        inputs=None,
        outputs=[image_input] + outputs,
    )

if __name__ == "__main__":
    # In Gradio 6, theme and css are passed to launch(), not Blocks().
    demo.launch(theme=theme, css=custom_css)

	import os
	import json
	import warnings

	import gradio as gr
	import gdown
	import numpy as np
	import tensorflow as tf
	import spaces

	from tensorflow.keras.applications.mobilenet_v2 import preprocess_input


	# ============================================================
	# CONFIGURATION
	# ============================================================

	MODEL_FILE_ID = "10axUT2jw9rAEo8FI4i9_ZGTv83UWqMeF"
	CLASS_NAMES_FILE_ID = "1ztvyL4h0sQlLeyf97E9RoZZ60yzrpmRi"
	MODEL_PATH = "cataract_model.keras"
	CLASS_NAMES_PATH = "class_names.json"

	IMG_HEIGHT = 180
	IMG_WIDTH = 180


	# ============================================================
	# ZEROGPU STARTUP PLACEHOLDER
	# ============================================================
	# ZeroGPU refuses to start a Space unless it finds at least one
	# @spaces.GPU function. TensorFlow does not work well inside the
	# separate GPU worker process, so this dummy function only exists
	# to satisfy the startup check. The real prediction function below
	# is deliberately NOT decorated and runs in the main process.

	@spaces.GPU
	def _zerogpu_placeholder():
	return None


	# ============================================================
	# GOOGLE DRIVE DOWNLOAD LOGIC
	# ============================================================

	def download_if_missing(file_id, dest_path, label):
	"""Download a Google Drive file if it is not already present."""

	if os.path.exists(dest_path) and os.path.getsize(dest_path) > 0:
	print(f"{label} already exists. Skipping download.")
	return

	if not file_id or file_id.startswith("YOUR_"):
	raise ValueError(f"Please configure the Google Drive ID for {label}.")

	print(f"Downloading {label} from Google Drive...")

	url = f"https://drive.google.com/uc?id={file_id}"
	downloaded_path = gdown.download(url=url, output=dest_path, quiet=False)

	if (
	not downloaded_path
	or not os.path.exists(dest_path)
	or os.path.getsize(dest_path) == 0
	):
	raise RuntimeError(
	f"Failed to download {label}. Check that the Google Drive "
	"file is shared as 'Anyone with the link'."
	)

	print(f"{label} downloaded successfully.")


	download_if_missing(MODEL_FILE_ID, MODEL_PATH, "TensorFlow model")
	download_if_missing(CLASS_NAMES_FILE_ID, CLASS_NAMES_PATH, "class names")


	# ============================================================
	# LOAD CLASS NAMES
	# ============================================================

	with open(CLASS_NAMES_PATH, "r", encoding="utf-8") as file:
	class_names = json.load(file)

	if isinstance(class_names, dict):
	categories = [
	class_names[key]
	for key in sorted(class_names, key=lambda item: int(item))
	]
	else:
	categories = list(class_names)

	if len(categories) != 2:
	raise ValueError(f"Expected two class names, but received: {categories}")

	print("Class index order:", categories)


	# ============================================================
	# LOAD TENSORFLOW MODEL
	# ============================================================

	with warnings.catch_warnings():
	warnings.simplefilter("ignore")
	model = tf.keras.models.load_model(MODEL_PATH, compile=False)

	print("TensorFlow version:", tf.__version__)
	print("Model loaded successfully.")
	print("Model input shape:", model.input_shape)


	# ============================================================
	# IMAGE PREPROCESSING
	# ============================================================

	def preprocess_image(img):
	"""Prepare an uploaded PIL image for MobileNetV2."""

	img = img.convert("RGB")
	img = img.resize((IMG_WIDTH, IMG_HEIGHT))

	image_array = np.asarray(img, dtype=np.float32)
	image_array = np.expand_dims(image_array, axis=0)

	# Must match the preprocessing used during model training.
	# If your model already contains a Rescaling / preprocess_input
	# layer, remove this line to avoid applying it twice.
	return preprocess_input(image_array)


	# ============================================================
	# PREDICTION (runs in the main process, no @spaces.GPU)
	# ============================================================

	def classify_image(img):
	"""Return class probabilities and a human-readable result."""

	if img is None:
	return {}, "### Upload an eye image to begin.", "—"

	image_array = preprocess_image(img)

	# Calling the model directly is faster and more reliable than
	# model.predict() for a single image.
	prediction = np.asarray(
	model(image_array, training=False)
	).reshape(-1)

	# Supports either a single sigmoid output or a two-class output.
	if len(prediction) == 1:
	probability_class_1 = float(np.clip(prediction[0], 0.0, 1.0))
	probabilities = {
	categories[0]: 1.0 - probability_class_1,
	categories[1]: probability_class_1,
	}

	elif len(prediction) == 2:
	values = np.clip(prediction.astype(float), 0.0, None)
	total = float(values.sum())

	if total == 0:
	return {}, "### Unable to calculate a prediction.", "—"

	values = values / total
	probabilities = {categories[i]: float(values[i]) for i in range(2)}

	else:
	raise ValueError(f"Unexpected model output shape: {prediction.shape}")

	predicted_class = max(probabilities, key=probabilities.get)
	confidence = probabilities[predicted_class] * 100

	if predicted_class.strip().lower() == "cataract":
	icon = "🔎"
	description = "The model detected features associated with cataracts."
	else:
	icon = "✓"
	description = "The model classified this image as normal."

	result_markdown = f"""
	## {icon} {predicted_class}

	{description}

	Model confidence: {confidence:.1f}%

	This is an AI-generated classification, not a medical diagnosis.
	"""

	return probabilities, result_markdown, f"{confidence:.1f}%"


	# ============================================================
	# GRADIO INTERFACE
	# ============================================================

	custom_css = """
	.gradio-container {
	max-width: 1120px !important;
	margin: auto !important;
	}

	#hero {
	text-align: center;
	padding: 32px 20px 26px;
	border-radius: 22px;
	background: linear-gradient(135deg, #102b46, #145c67);
	color: white;
	margin-bottom: 18px;
	}

	#hero h1 {
	color: white !important;
	font-size: 2.4rem !important;
	margin-bottom: 10px;
	}

	#hero p {
	color: #d5eeee !important;
	font-size: 1.05rem;
	}

	.panel {
	border: 1px solid var(--border-color-primary);
	border-radius: 18px;
	padding: 18px;
	}

	#predict-button {
	background: linear-gradient(90deg, #0d9488, #0f766e);
	color: white;
	border: none;
	font-weight: 700;
	border-radius: 12px;
	}

	#predict-button:hover {
	filter: brightness(1.08);
	}

	#disclaimer {
	border-left: 4px solid #d97706;
	padding: 12px 16px;
	border-radius: 8px;
	background: var(--block-background-fill);
	}
	"""

	theme = gr.themes.Soft(
	primary_hue="teal",
	secondary_hue="cyan",
	neutral_hue="slate",
	)

	with gr.Blocks(title="CataractScan AI") as demo:

	gr.HTML("""
	<div id="hero">
	<h1>👁️ CataractScan AI</h1>
	<p>
	AI-powered eye image classification using TensorFlow
	and MobileNetV2
	</p>
	</div>
	""")

	gr.Markdown(
	"""
	Upload an eye image to see how the model classifies it.
	Your result includes the predicted class and the model's
	estimated probabilities.
	"""
	)

	with gr.Row(equal_height=True):

	with gr.Column(scale=1, elem_classes=["panel"]):
	gr.Markdown("### 📤 Upload an eye image")

	image_input = gr.Image(
	type="pil",
	label="Eye image",
	height=340,
	)

	with gr.Row():
	predict_button = gr.Button(
	"Analyze Image",
	variant="primary",
	elem_id="predict-button",
	scale=3,
	)
	clear_button = gr.Button("Clear", scale=1)

	gr.Markdown(
	"""
	Tips: Use a clear, well-lit eye image.
	The model may perform poorly on images that differ
	from its training data.
	"""
	)

	with gr.Column(scale=1, elem_classes=["panel"]):
	gr.Markdown("### 🧠 Classification results")

	prediction_output = gr.Label(
	label="Class probabilities",
	num_top_classes=2,
	)

	result_output = gr.Markdown("### Your result will appear here.")

	confidence_output = gr.Textbox(
	label="Confidence of predicted class",
	value="—",
	interactive=False,
	)

	gr.Markdown("### 🖼️ Try an example")

	example_images = [
	[name] for name in ["cataract2.jpg","cataract.jpg", "normal.png", "mine.jpg"] if os.path.exists(name)
	]

	if example_images:
	gr.Examples(
	examples=example_images,
	inputs=image_input,
	label="Sample eye images",
	)

	gr.Markdown(
	"""
	<div id="disclaimer">

	⚠️ Medical disclaimer

	CataractScan AI is an educational AI project, not a
	diagnostic medical device. Its output must not be used
	to diagnose or rule out cataracts. Image quality,
	training data, and model limitations can affect results.
	Consult a qualified eye-care professional for medical
	assessment.

	</div>
	"""
	)

	outputs = [prediction_output, result_output, confidence_output]

	# Only the button triggers prediction, so each image is analysed once
	# (the old .change() handler made every upload run twice).
	predict_button.click(
	fn=classify_image,
	inputs=image_input,
	outputs=outputs,
	)

	def clear_all():
	return None, {}, "### Your result will appear here.", "—"

	clear_button.click(
	fn=clear_all,
	inputs=None,
	outputs=[image_input] + outputs,
	)

	if __name__ == "__main__":
	# In Gradio 6, theme and css are passed to launch(), not Blocks().
	demo.launch(theme=theme, css=custom_css)