import os import numpy as np import tensorflow as tf import gradio as gr from PIL import Image # ========================= # CONFIG # ========================= MODEL_PATH = "best_model.keras" PATCH_SIZE = 256 INFER_STRIDE = 192 BATCH_SIZE = 8 USE_GRAYSCALE = False CHANNELS = 3 TITLE = "Research Demo: Urine Sediment Virtual Phase Contrast Conversion" DESCRIPTION = ( "This Space demonstrates a research prototype that converts a brightfield urine " "sediment microscopy image into a virtual phase contrast-style image. The output " "is intended only to illustrate the model concept and should not be interpreted " "as a clinical result." ) DISCLAIMER = ( "⚠️ Demonstration only — not for clinical use. " "This model is a research prototype and is not validated as a medical device. " "It must not be used for diagnosis, patient care, treatment decisions, laboratory reporting, " "or any clinical decision-making. The generated image may be inaccurate, may suppress subtle " "features, or may fail to represent clinically important urinary sediment findings, especially " "findings outside the training/evaluation data. Upload only de-identified, non-patient-care images." ) FOOTER = ( "For research and educational demonstration only. Always rely on validated laboratory methods " "and qualified professional review for clinical interpretation." ) # ========================= # LOAD MODEL # ========================= if not os.path.exists(MODEL_PATH): raise FileNotFoundError(f"Model file not found: {MODEL_PATH}") model = tf.keras.models.load_model(MODEL_PATH, compile=False) # ========================= # IMAGE CONVERSION # ========================= def pil_to_model_input(img: Image.Image, grayscale=False): if grayscale: img = img.convert("L") arr = np.array(img, dtype=np.float32) / 255.0 arr = arr[..., None] else: img = img.convert("RGB") arr = np.array(img, dtype=np.float32) / 255.0 return arr def array_to_pil(arr, grayscale=False): arr = np.clip(arr, 0.0, 1.0) if grayscale: if arr.ndim == 3 and arr.shape[-1] == 1: arr = arr[..., 0] arr_uint8 = (arr * 255).astype(np.uint8) return Image.fromarray(arr_uint8, mode="L") if arr.ndim == 3 and arr.shape[-1] == 1: arr = np.repeat(arr, 3, axis=-1) arr_uint8 = (arr * 255).astype(np.uint8) return Image.fromarray(arr_uint8, mode="RGB") def make_side_by_side(left_pil: Image.Image, right_pil: Image.Image): left_rgb = left_pil.convert("RGB") right_rgb = right_pil.convert("RGB") w1, h1 = left_rgb.size w2, h2 = right_rgb.size canvas = Image.new("RGB", (w1 + w2, max(h1, h2)), color=(255, 255, 255)) canvas.paste(left_rgb, (0, 0)) canvas.paste(right_rgb, (w1, 0)) return canvas # ========================= # TILING / STITCHING # ========================= def pad_to_min_size(img, min_h, min_w): h, w = img.shape[:2] pad_h = max(0, min_h - h) pad_w = max(0, min_w - w) if pad_h == 0 and pad_w == 0: return img return np.pad(img, ((0, pad_h), (0, pad_w), (0, 0)), mode="reflect") def compute_start_positions(length, patch_size, stride): if length <= patch_size: return [0] starts = list(range(0, length - patch_size + 1, stride)) if starts[-1] != length - patch_size: starts.append(length - patch_size) return starts def make_blend_weight(patch_size): w = np.hanning(patch_size).astype(np.float32) w = np.outer(w, w).astype(np.float32) w = np.maximum(w, 1e-3) return w[..., None] def predict_full_image(model, bright_img, patch_size=256, stride=192, channels=3, batch_size=8): orig_h, orig_w = bright_img.shape[:2] bright_pad = pad_to_min_size(bright_img, patch_size, patch_size) H, W = bright_pad.shape[:2] ys = compute_start_positions(H, patch_size, stride) xs = compute_start_positions(W, patch_size, stride) weight = make_blend_weight(patch_size) pred_sum = np.zeros((H, W, channels), dtype=np.float32) weight_sum = np.zeros((H, W, 1), dtype=np.float32) patches = [] coords = [] for y in ys: for x in xs: patch = bright_pad[y:y + patch_size, x:x + patch_size, :] patches.append(patch) coords.append((y, x)) patches = np.asarray(patches, dtype=np.float32) preds = model.predict(patches, batch_size=batch_size, verbose=0) for pred, (y, x) in zip(preds, coords): pred_sum[y:y + patch_size, x:x + patch_size, :] += pred * weight weight_sum[y:y + patch_size, x:x + patch_size, :] += weight stitched = pred_sum / np.clip(weight_sum, 1e-6, None) stitched = stitched[:orig_h, :orig_w, :] stitched = np.clip(stitched, 0.0, 1.0) return stitched # ========================= # INFERENCE FUNCTION # ========================= def run_inference(input_image: Image.Image): if input_image is None: raise gr.Error("Please upload a de-identified brightfield microscopy image for demonstration.") original_display = input_image.convert("RGB") x = pil_to_model_input(input_image, grayscale=USE_GRAYSCALE) y_pred = predict_full_image( model=model, bright_img=x, patch_size=PATCH_SIZE, stride=INFER_STRIDE, channels=CHANNELS, batch_size=BATCH_SIZE, ) pred_pil = array_to_pil(y_pred, grayscale=USE_GRAYSCALE) compare_pil = make_side_by_side(original_display, pred_pil) return pred_pil, compare_pil # ========================= # UI # ========================= with gr.Blocks(title=TITLE) as demo: gr.Markdown(f"# {TITLE}") gr.Markdown(DESCRIPTION) gr.Markdown( f"""
{DISCLAIMER}
""" ) with gr.Row(): input_image = gr.Image(type="pil", label="Input brightfield image for demonstration") pred_image = gr.Image(type="pil", label="Demonstration output: virtual phase contrast-style image") compare_image = gr.Image(type="pil", label="Side-by-side demonstration comparison") with gr.Row(): run_button = gr.Button("Generate demonstration output") clear_button = gr.ClearButton([input_image, pred_image, compare_image]) gr.Markdown( f"""
{FOOTER}
""" ) run_button.click( fn=run_inference, inputs=input_image, outputs=[pred_image, compare_image], ) if __name__ == "__main__": demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))