Spaces:

ruanhelansantos
/

EndtoEnd_AI_Solution_Implementation

Sleeping

EndtoEnd_AI_Solution_Implementation / utils.py

Upload utils.py

7461183 verified 8 months ago

3.1 kB

# utils.py
# ECG Condition Classifier - Image Preprocessing

from PIL import Image, ImageFilter, ImageEnhance, ImageOps
import numpy as np

TARGET_SIZE = (224, 224)
MIN_SIZE = 50


def load_and_preprocess(image):
    """
    Preprocesses an ECG image for ViT inference.
    Key steps:
    1. Convert to RGB
    2. Auto-crop whitespace/borders
    3. Convert to grayscale then back to RGB (removes color noise from
       pink/beige ECG paper, watermarks, etc.)
    4. Enhance contrast to make waveform lines pop
    5. Resize to 224x224 for ViT
    """
    if image is None:
        raise ValueError("No image provided.")

    # Step 1: Ensure RGB
    img = image.convert("RGB")

    # Step 2: Auto-crop borders/whitespace
    img = _autocrop(img)

    # Step 3: Grayscale normalization
    # Converts pink/beige/colored ECG paper to clean black-on-white.
    # This removes color distribution shift between training data and
    # stock photos, dramatically improving model confidence.
    gray = img.convert("L")

    # Step 4: Enhance contrast — make ECG traces sharper
    enhancer = ImageEnhance.Contrast(gray)
    gray = enhancer.enhance(2.5)

    # Step 5: Sharpen edges (waveform lines)
    gray = gray.filter(ImageFilter.SHARPEN)

    # Step 6: Back to RGB (model expects 3-channel input)
    img = gray.convert("RGB")

    # Step 7: Resize to ViT input size
    img = img.resize(TARGET_SIZE, Image.LANCZOS)

    return img


def _autocrop(img, threshold=240, padding=10):
    """
    Crops away white/light borders around the ECG tracing area.
    Helps the model focus on the waveform rather than margins.
    """
    try:
        gray = img.convert("L")
        arr = np.array(gray)

        # Find rows/cols that contain dark pixels (the ECG trace)
        mask = arr < threshold
        rows = np.any(mask, axis=1)
        cols = np.any(mask, axis=0)

        if not rows.any() or not cols.any():
            return img  # Nothing to crop

        rmin, rmax = np.where(rows)[0][[0, -1]]
        cmin, cmax = np.where(cols)[0][[0, -1]]

        # Add small padding back
        h, w = arr.shape
        rmin = max(0, rmin - padding)
        rmax = min(h, rmax + padding)
        cmin = max(0, cmin - padding)
        cmax = min(w, cmax + padding)

        return img.crop((cmin, rmin, cmax, rmax))
    except Exception:
        return img  # Fall back to original if crop fails


def validate_image(image):
    """
    Validates that the uploaded image is usable.
    Returns (is_valid: bool, message: str)
    """
    if image is None:
        return False, "No image uploaded. Please upload an ECG image."

    try:
        w, h = image.size
        if w < MIN_SIZE or h < MIN_SIZE:
            return False, f"Image too small ({w}x{h}). Please upload a higher resolution ECG."
    except Exception as e:
        return False, f"Invalid image: {str(e)}"

    return True, "Valid"


def get_image_info(image):
    """Returns a summary string of image properties (for debugging)."""
    if image is None:
        return "No image"
    return f"Size: {image.size}, Mode: {image.mode}"

	# utils.py
	# ECG Condition Classifier - Image Preprocessing

	from PIL import Image, ImageFilter, ImageEnhance, ImageOps
	import numpy as np

	TARGET_SIZE = (224, 224)
	MIN_SIZE = 50


	def load_and_preprocess(image):
	"""
	Preprocesses an ECG image for ViT inference.

	Key steps:
	1. Convert to RGB
	2. Auto-crop whitespace/borders
	3. Convert to grayscale then back to RGB (removes color noise from
	pink/beige ECG paper, watermarks, etc.)
	4. Enhance contrast to make waveform lines pop
	5. Resize to 224x224 for ViT
	"""
	if image is None:
	raise ValueError("No image provided.")

	# Step 1: Ensure RGB
	img = image.convert("RGB")

	# Step 2: Auto-crop borders/whitespace
	img = _autocrop(img)

	# Step 3: Grayscale normalization
	# Converts pink/beige/colored ECG paper to clean black-on-white.
	# This removes color distribution shift between training data and
	# stock photos, dramatically improving model confidence.
	gray = img.convert("L")

	# Step 4: Enhance contrast — make ECG traces sharper
	enhancer = ImageEnhance.Contrast(gray)
	gray = enhancer.enhance(2.5)

	# Step 5: Sharpen edges (waveform lines)
	gray = gray.filter(ImageFilter.SHARPEN)

	# Step 6: Back to RGB (model expects 3-channel input)
	img = gray.convert("RGB")

	# Step 7: Resize to ViT input size
	img = img.resize(TARGET_SIZE, Image.LANCZOS)

	return img


	def _autocrop(img, threshold=240, padding=10):
	"""
	Crops away white/light borders around the ECG tracing area.
	Helps the model focus on the waveform rather than margins.
	"""
	try:
	gray = img.convert("L")
	arr = np.array(gray)

	# Find rows/cols that contain dark pixels (the ECG trace)
	mask = arr < threshold
	rows = np.any(mask, axis=1)
	cols = np.any(mask, axis=0)

	if not rows.any() or not cols.any():
	return img # Nothing to crop

	rmin, rmax = np.where(rows)[0][[0, -1]]
	cmin, cmax = np.where(cols)[0][[0, -1]]

	# Add small padding back
	h, w = arr.shape
	rmin = max(0, rmin - padding)
	rmax = min(h, rmax + padding)
	cmin = max(0, cmin - padding)
	cmax = min(w, cmax + padding)

	return img.crop((cmin, rmin, cmax, rmax))
	except Exception:
	return img # Fall back to original if crop fails


	def validate_image(image):
	"""
	Validates that the uploaded image is usable.
	Returns (is_valid: bool, message: str)
	"""
	if image is None:
	return False, "No image uploaded. Please upload an ECG image."

	try:
	w, h = image.size
	if w < MIN_SIZE or h < MIN_SIZE:
	return False, f"Image too small ({w}x{h}). Please upload a higher resolution ECG."
	except Exception as e:
	return False, f"Invalid image: {str(e)}"

	return True, "Valid"


	def get_image_info(image):
	"""Returns a summary string of image properties (for debugging)."""
	if image is None:
	return "No image"
	return f"Size: {image.size}, Mode: {image.mode}"