jkottu/image-captioning-chest-xrays

1import gradio as gr2from PIL import Image3import clipGPT4import vitGPT5import skimage.io as io6import PIL.Image7import difflib8import ViTCoAtt9import cnnrnn10from build_vocab import Vocabulary11import pickle12 13# Caption generation functions14def generate_caption_clipgpt(image, max_tokens, temperature):15    caption = clipGPT.generate_caption_clipgpt(image, max_tokens, temperature)16    return caption17 18def generate_caption_vitgpt(image, max_tokens, temperature):19    caption = vitGPT.generate_caption(image, max_tokens, temperature)20    return caption21 22def generate_caption_vitCoAtt(image):23    caption = ViTCoAtt.CaptionSampler.main(image)24    return caption25 26def generate_caption_cnnrnn(image):27    with open('Image_features_ecoder_decoder.pickle', 'rb') as f:28        Xnet_features = pickle.load(f)29        image = Xnet_features[image]30        caption = cnnrnn.get_result(image)31    return caption32 33 34with gr.Row():35    image = gr.Image(label="Upload Chest X-ray", type="pil", height='50',width='50')   36                     37 38with gr.Row():39    with gr.Column(): # Column for dropdowns and model choice40        max_tokens = gr.Dropdown(list(range(50, 101)), label="Max Tokens", value=75)41        temperature = gr.Slider(0.5, 0.9, step=0.1, label="Temperature", value=0.9)42        imgID = gr.Dropdown(["1","2","3","4","5","6"], label="Choose the ID of the image selected")43 44    model_choice = gr.Radio(["CLIP-GPT2", "ViT-GPT2", "ViT-CoAttention", "Baseline Model CNN-RNN"], label="Select Model") 45    generate_button = gr.Button("Generate Caption") 46    47 48caption = gr.Textbox(label="Generated Caption")49real_caption = gr.Textbox(label="Actual Caption")50 51def getCaption(imgID):   52    real_captions = {"1" : "No acute cardiopulmonary abnormality. 2. Stable bilateral emphysematous and lower lobe fibrotic changes. Bilateral emphysematous again noted and lower lobe fibrotic changes. Postsurgical changes of the chest including CABG procedure, stable. Stable valve artifact. There are no focal areas of consolidation. No large pleural effusions. No evidence of pneumothorax. Degenerative changes noted of the visualized thoracic spine. Nodular right lower lobe opacity, XXXX nipple XXXX. Contour abnormality of the posterior aspect of the right 7th rib again noted, stable.", 53                     "2":"Hypoinflation with bibasilar focal atelectasis. Lung volumes are XXXX. XXXX opacities are present in both lung bases. A hiatal hernia is present. Heart and pulmonary XXXX are normal.",54                     "3":"No evidence of acute cardiopulmonary process. The XXXX examination consists of frontal and lateral radiographs of the chest. External monitor leads XXXX the thorax. The cardiomediastinal contours are within normal limits. Pulmonary vascularity is within normal limits. No focal consolidation, pleural effusion, or pneumothorax identified. The visualized osseous structures and upper abdomen are unremarkable.",55                     "4":"Negative chest . The lungs are clear. The cardiomediastinal silhouette is within normal limits. No pneumothorax or pleural effusion.",56                    "5": "No Actual Caption",57                    "6": "No Actual Caption"}58 59    return real_captions[imgID]60    61def getImageID(imgID):62    imgIDs = {"1":"/content/drive/MyDrive/cnn-rnn/NLMCXR_png/CXR412_IM-2056_0",63          "2":"/content/drive/MyDrive/cnn-rnn/NLMCXR_png/CXR545_IM-2149_0",64          "3":"/content/drive/MyDrive/cnn-rnn/NLMCXR_png/CXR3044_IM-1418_0",65          "4":"/content/drive/MyDrive/cnn-rnn/NLMCXR_png/CXR3587_IM-1765_0"}66    return imgIDs[imgID]67    68def predict(img, model_name, max_tokens, temperature, imgID):69    if model_name == "CLIP-GPT2":70        return generate_caption_clipgpt(img, max_tokens, temperature), getCaption(imgID)71    elif model_name == "ViT-GPT2":72        return generate_caption_vitgpt(img, max_tokens, temperature), getCaption(imgID)73    elif model_name == "ViT-CoAttention":74        return generate_caption_vitCoAtt(img),  getCaption(imgID)75    elif model_name == "Baseline Model CNN-RNN":76        img = getImageID(imgID)77        return generate_caption_cnnrnn(img), getCaption(imgID)78    else:79        return "select a model","select an image"   80 81 82 83#main call84        85examples = [[f"example{i}.jpg"] for i in range(1,7)]86 87description= "You can generate captions by uploading an X-Ray and selecting a model of your choice below. Please select the number of Max Tokens and Temperature setting, if you are testing CLIP GPT2 and VIT GPT2 Models"88title = "A Vision Transformer-Driven Method for Generating Medical Reports based on X-ray radiology 🏥🤖"89 90interface = gr.Interface(91        fn=predict,92        inputs = [image, model_choice, max_tokens, temperature, imgID],93        theme="sudeepshouche/minimalist",94        outputs=[caption,real_caption],95        examples = examples,96        title = title,97        description = description98    )99 100 101interface.launch(debug=True)102    103 104