File size: 6,383 Bytes
e93bfbd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | """
Browser-based UI for MedicalAI - Light Weight.
Auto-installs gradio if missing.
Auto-generates default models if none are trained.
"""
import importlib
import os
import subprocess
import sys
import tempfile
from PIL import Image
from optimize import (
clear_memory,
get_available_models,
get_device,
infer_blip,
infer_fusion,
infer_fusion_onnx,
set_cpu_threads,
)
CHECKPOINT_PATH = "./checkpoints/fusion_model.pth"
ONNX_FULL_DIR = "./checkpoints/onnx_full"
DEFAULT_MODEL_DIR = "./models/default"
def _ensure_gradio():
try:
return importlib.import_module("gradio")
except ImportError:
from rich.console import Console
console = Console()
console.print("[yellow]'gradio' is required for the web UI.[/yellow]")
import questionary
install = questionary.confirm("Install gradio now?", default=True).ask()
if not install:
console.print("[red]gradio is required. Exiting.[/red]")
sys.exit(1)
console.print("[cyan]Installing gradio...[/cyan]")
subprocess.check_call([sys.executable, "-m", "pip", "install", "gradio"])
return importlib.import_module("gradio")
def _ensure_default_models():
"""Generate default ONNX models if none exist at all."""
if os.path.exists(os.path.join(DEFAULT_MODEL_DIR, "fusion_classifier.onnx")):
return
if os.path.exists(CHECKPOINT_PATH):
return
print("No models found. Generating default models...")
subprocess.check_call(
[sys.executable, "setup_default.py"],
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
)
print("Default models ready.")
def analyze_vision(image):
if image is None:
return "Please upload an X-ray image."
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
path = f.name
Image.fromarray(image).save(path)
try:
caption = infer_blip(path, use_onnx=False)
return caption
except Exception as e:
return f"Error: {e}"
finally:
os.unlink(path)
clear_memory()
def analyze_symptom(image, symptoms, use_onnx):
if image is None:
return "Please upload an X-ray image.", ""
if not symptoms:
symptoms = "No symptoms provided"
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
path = f.name
Image.fromarray(image).save(path)
try:
if use_onnx:
diagnosis, confidence = infer_fusion_onnx(path, symptoms)
else:
diagnosis, confidence = infer_fusion(path, symptoms)
if diagnosis is None:
return confidence, "No result"
return diagnosis, f"{confidence:.1%}"
except Exception as e:
return f"Error: {e}", ""
finally:
os.unlink(path)
clear_memory()
def main():
set_cpu_threads()
_ensure_default_models()
gr = _ensure_gradio()
models = get_available_models()
has_trained = models["trained_pytorch"]
has_default = models["default_classifier"]
has_onnx_full = models["onnx_full_pipeline"]
model_source = "ONNX (full pipeline)" if has_onnx_full else \
"Trained PyTorch" if has_trained else \
"Default (random weights)" if has_default else \
"NOT AVAILABLE"
print(f"Device: {get_device().upper()}")
print(f"Model: {model_source}")
print()
print("Launching web UI... open the URL below in your browser.")
with gr.Blocks(title="MedicalAI - Light Weight", theme=gr.themes.Soft()) as demo:
gr.Markdown(
"""
# MedicalAI - Light Weight
Upload a chest X-ray for AI-assisted analysis.
"""
)
with gr.Tab("Vision - Generate Report"):
gr.Markdown("Upload an X-ray and get an AI-generated radiology caption.")
with gr.Row():
img_in = gr.Image(label="X-ray Image", type="numpy")
with gr.Row():
btn_vision = gr.Button("Generate Report", variant="primary")
with gr.Row():
caption_out = gr.Textbox(label="Radiology Caption", lines=6)
btn_vision.click(fn=analyze_vision, inputs=img_in, outputs=caption_out)
with gr.Tab("Symptom Check - Diagnosis"):
gr.Markdown("Upload an X-ray, enter symptoms, and get a diagnosis.")
with gr.Row():
img_in2 = gr.Image(label="X-ray Image", type="numpy")
with gr.Row():
symptoms_in = gr.Textbox(
label="Patient Symptoms / Clinical Indication",
placeholder="e.g., shortness of breath, cough, fever...",
lines=3,
)
with gr.Row():
onnx_checkbox = gr.Checkbox(
label="Use ONNX (no PyTorch backend)",
value=has_onnx_full,
interactive=has_onnx_full,
)
with gr.Row():
btn_diag = gr.Button("Diagnose", variant="primary")
with gr.Row():
with gr.Column():
diag_out = gr.Textbox(label="Diagnosis", lines=4)
conf_out = gr.Textbox(label="Confidence")
btn_diag.click(
fn=analyze_symptom,
inputs=[img_in2, symptoms_in, onnx_checkbox],
outputs=[diag_out, conf_out],
)
with gr.Tab("System Info"):
gr.Markdown(f"""
**Device:** `{get_device().upper()}`
**Model:** `{model_source}`
**Trained checkpoint:** {'Yes' if has_trained else 'No'}
**ONNX full pipeline:** {'Yes' if has_onnx_full else 'No'}
**Default model:** {'Yes' if has_default else 'No'}
**Dataset:** `{os.path.abspath('./data/dataset.csv')}`
### 3 ways to improve accuracy
1. **Train** — `python training.py --mode prepare-data && python training.py --mode train`
2. **Update** — `python update.py --models` (downloads pre-trained model from Hugging Face)
3. **Export to ONNX** — `python quantization.py --mode export-full`
""")
demo.launch(server_name="127.0.0.1", server_port=7860, share=False)
if __name__ == "__main__":
main()
|