# 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"])


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

    demo = gr.Interface(
        fn=predict,
        inputs=[
            gr.Image(type="pil", label="Input Image"),
            gr.Textbox(
                label="Prompts",
                placeholder="Enter prompts separated by commas (e.g., neoplastic cells, inflammatory cells)",
            ),
        ],
        outputs=gr.Image(type="pil", label="Prediction"),
        title="BiomedParse Demo",
        description=description,
        allow_flagging="never",
        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"],
        ],
    )
    return demo


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/).
"""


demo = run()

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