main.py

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os

# If True, then mock init_model() and predict() functions will be used.
DEV_MODE = True if os.getenv("DEV_MODE") else False

import gradio as gr

if DEV_MODE:
    from inference_utils.init_predict_mock import init_model, predict
else:
    from inference_utils.init_predict import init_model, predict


gr.set_static_paths(["assets"])


description = """Upload a biomedical image and enter prompts (separated by commas) to detect specific features.

The model understands these prompts:
![gpt4_ontology_hierarchy.png](file/assets/gpt4_ontology_hierarchy.png)

Above figure is from the [BiomedParse paper](https://arxiv.org/abs/2405.12971).

The model understands these types of biomedical images:

- [Computed Tomography (CT)](https://en.wikipedia.org/wiki/Computed_tomography)
- [Magnetic Resonance Imaging (MRI)](https://en.wikipedia.org/wiki/Magnetic_resonance_imaging)
- [X-ray](https://en.wikipedia.org/wiki/X-ray)
- [Medical Ultrasound](https://en.wikipedia.org/wiki/Medical_ultrasound)
- [Pathology](https://en.wikipedia.org/wiki/Pathology)
- [Fundus (eye)](https://en.wikipedia.org/wiki/Fundus_(eye))
- [Dermoscopy](https://en.wikipedia.org/wiki/Dermoscopy)
- [Endoscopy](https://en.wikipedia.org/wiki/Endoscopy)
- [Optical Coherence Tomography (OCT)](https://en.wikipedia.org/wiki/Optical_coherence_tomography)

This Space is based on the [BiomedParse model](https://microsoft.github.io/BiomedParse/).
"""


examples = [
    ["examples/144DME_as_F.jpeg", "edema"],
    ["examples/C3_EndoCV2021_00462.jpg", "polyp"],
    ["examples/CT-abdomen.png", "liver, pancreas, spleen"],
    ["examples/covid_1585.png", "left lung"],
    ["examples/covid_1585.png", "right lung"],
    ["examples/covid_1585.png", "COVID-19 infection"],
    ["examples/ISIC_0015551.jpg", "lesion"],
    ["examples/LIDC-IDRI-0140_143_280_CT_lung.png", "lung nodule"],
    ["examples/LIDC-IDRI-0140_143_280_CT_lung.png", "COVID-19 infection"],
    ["examples/Part_1_516_pathology_breast.png", "connective tissue cells"],
    ["examples/Part_1_516_pathology_breast.png", "neoplastic cells"],
    [
        "examples/Part_1_516_pathology_breast.png",
        "neoplastic cells, inflammatory cells",
    ],
    ["examples/T0011.jpg", "optic disc"],
    ["examples/T0011.jpg", "optic cup"],
    ["examples/TCGA_HT_7856_19950831_8_MRI-FLAIR_brain.png", "glioma"],
]


def run():
    global model
    model = init_model()

    with gr.Blocks() as demo:
        gr.Markdown("# BiomedParse Demo")
        gr.Markdown(description)

        with gr.Row():
            with gr.Column():
                input_image = gr.Image(type="pil", label="Input Image")
                input_text = gr.Textbox(
                    label="Prompts",
                    placeholder="Enter prompts separated by commas (e.g., neoplastic cells, inflammatory cells)",
                )
            with gr.Column():
                output_image = gr.Image(type="pil", label="Prediction")

        predict_btn = gr.Button("Submit")
        predict_btn.click(
            fn=predict,
            inputs=[input_image, input_text],
            outputs=output_image,
        )

        gr.Examples(
            examples=examples,
            inputs=[input_image, input_text],
            outputs=output_image,
            fn=predict,
            cache_examples=False,
        )

    return demo


demo = run()

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)