GAD-Research-Lab/MedicalAI-Light-Weight

1import importlib2import os3import subprocess4import sys5import tempfile6from datetime import datetime7from pathlib import Path8 9from PIL import Image, ImageOps10 11DATA_DIR = Path("./data")12IMAGES_DIR = DATA_DIR / "images"13 14 15def _ensure_dep(package_name, import_name=None):16    if import_name is None:17        import_name = package_name18    try:19        return importlib.import_module(import_name)20    except ImportError:21        from rich.console import Console22        console = Console()23        console.print(f"[yellow]'{package_name}' is required for this feature.[/yellow]")24        import questionary25        install = questionary.confirm(f"Install {package_name} now?", default=True).ask()26        if not install:27            return None28        console.print(f"[cyan]Installing {package_name}...[/cyan]")29        subprocess.check_call([sys.executable, "-m", "pip", "install", package_name])30        return importlib.import_module(import_name)31 32def _ensure_dirs():33    IMAGES_DIR.mkdir(parents=True, exist_ok=True)34 35def _save_image(pil_image, prefix="capture"):36    _ensure_dirs()37    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")38    filename = f"{prefix}_{timestamp}.jpg"39    path = str(IMAGES_DIR / filename)40    if pil_image.mode != "RGB":41        pil_image = pil_image.convert("RGB")42    pil_image.save(path, quality=95)43    return path44 45# ── Camera ──────────────────────────────────────────────────────46 47def capture_camera():48    cv2 = _ensure_dep("opencv-python", "cv2")49    if cv2 is None:50        return None, "Camera capture requires opencv-python"51 52    cap = cv2.VideoCapture(0)53    if not cap.isOpened():54        return None, "No camera detected (could not open index 0)"55 56    from rich.console import Console57    console = Console()58    console.print("[cyan]Camera opened. Press SPACE to capture, ESC to cancel.[/cyan]")59 60    import questionary61    input("Press Enter when ready for camera preview...")62    ret, frame = cap.read()63    cap.release()64 65    if not ret:66        return None, "Failed to capture frame from camera"67 68    preview_path = tempfile.mktemp(suffix="_preview.jpg")69    cv2.imwrite(preview_path, frame)70    preview = Image.open(preview_path)71    os.unlink(preview_path)72 73    console.print("[cyan]Image captured from camera.[/cyan]")74    path = _save_image(preview, "camera")75    return path, f"Captured from camera -> {path}"76 77# ── File browser ───────────────────────────────────────────────78 79def capture_file():80    try:81        import tkinter as tk82        from tkinter import filedialog83        root = tk.Tk()84        root.withdraw()85        root.attributes("-topmost", True)86        path = filedialog.askopenfilename(87            title="Select an X-ray image",88            filetypes=[89                ("Image files", "*.jpg *.jpeg *.png *.bmp *.tif *.tiff *.dcm"),90                ("All files", "*.*"),91            ],92        )93        root.destroy()94    except Exception as e:95        return None, f"File dialog failed: {e}"96 97    if not path:98        return None, "No file selected"99 100    return _open_and_save(path)101 102# ── Manual path entry ──────────────────────────────────────────103 104def capture_path():105    from rich.console import Console106    console = Console()107    console.print("[cyan]Enter the path to an X-ray image file.[/cyan]")108 109    import questionary110    path = questionary.path("Image path:").ask()111    if not path:112        return None, "No path entered"113    return _open_and_save(path)114 115# ── DICOM ──────────────────────────────────────────────────────116 117def capture_dicom(path=None):118    pydicom = _ensure_dep("pydicom")119    if pydicom is None:120        return None, "DICOM loading requires pydicom"121    np = _ensure_dep("numpy")122    if np is None:123        return None, "DICOM loading requires numpy"124 125    if not path:126        try:127            import tkinter as tk128            from tkinter import filedialog129            root = tk.Tk()130            root.withdraw()131            root.attributes("-topmost", True)132            path = filedialog.askopenfilename(133                title="Select a DICOM file",134                filetypes=[("DICOM files", "*.dcm"), ("All files", "*.*")],135            )136            root.destroy()137        except Exception as e:138            return None, f"File dialog failed: {e}"139 140    if not path:141        return None, "No DICOM file selected"142 143    try:144        ds = pydicom.dcmread(path)145        arr = ds.pixel_array146        arr = arr - arr.min()147        arr = (arr / arr.max() * 255).astype(np.uint8)148        if len(arr.shape) == 2:149            img = Image.fromarray(arr, mode="L")150            img = ImageOps.equalize(img)151        else:152            img = Image.fromarray(arr)153        result_path = _save_image(img, "dicom")154        return result_path, f"DICOM loaded from {path} -> saved as {result_path}"155    except Exception as e:156        return None, f"Failed to read DICOM: {e}"157 158# ── Generic open + save ────────────────────────────────────────159 160def _open_and_save(source_path):161    source_path = str(source_path)162    if source_path.lower().endswith(".dcm"):163        return capture_dicom(source_path)164    try:165        img = Image.open(source_path)166        path = _save_image(img, "import")167        return path, f"Imported from {source_path} -> {path}"168    except Exception as e:169        return None, f"Failed to open image: {e}"170 171# ── Top-level picker ───────────────────────────────────────────172 173def pick_image():174    from rich.console import Console175    import questionary176 177    console = Console()178    method = questionary.select(179        "How do you want to provide the X-ray image?",180        choices=[181            "Browse files on computer",182            "Enter file path manually",183            "Capture from camera",184            "Load DICOM file",185        ],186        pointer=">",187    ).ask()188 189    result = None190    if method == "Browse files on computer":191        result = capture_file()192    elif method == "Enter file path manually":193        result = capture_path()194    elif method == "Capture from camera":195        result = capture_camera()196    elif method == "Load DICOM file":197        result = capture_dicom()198 199    return result200 201 202if __name__ == "__main__":203    path, msg = pick_image()204    print(msg)205