app.py

import gradio as gr
import torch
from tqdm import tqdm
from monai.utils import set_determinism
from torch.cuda.amp import autocast
# from generative.inferers import DiffusionInferer
from generative.networks.nets import DiffusionModelUNet,AutoencoderKL
from generative.networks.schedulers import DDPMScheduler
from generative.networks.schedulers.ddim import DDIMScheduler
import cv2
from lib_image_processing.contrast_brightness_lib import controller
from lib_image_processing.removebg_lib import get_mask
import matplotlib.pyplot as plt
import numpy as np
set_determinism(42)
torch.cuda.empty_cache()

## Load autoencoder

#device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device = torch.device('cpu')
autoencoderkl = AutoencoderKL(
    spatial_dims=2,
    in_channels=1,
    out_channels=1,
    num_channels=(128, 128, 256),
    latent_channels=3,
    num_res_blocks=2,
    attention_levels=(False, False, False),
    with_encoder_nonlocal_attn=False,
    with_decoder_nonlocal_attn=False,
)
root_dir = "models"
PATH_auto = f'{root_dir}/auto_encoder_model.pt'

autoencoderkl.load_state_dict(torch.load(PATH_auto,map_location=device))
autoencoderkl = autoencoderkl.to(device)

#### Load unet and embedings

embedding_dimension = 64
unet = DiffusionModelUNet(
    spatial_dims=2,
    in_channels=3,
    out_channels=3,
    num_res_blocks=2,
    num_channels=(128, 256, 512),
    attention_levels=(False, True, True),
    num_head_channels=(0, 256, 512),
    with_conditioning=True,
    cross_attention_dim=embedding_dimension
)

embed = torch.nn.Embedding(num_embeddings=6, embedding_dim=embedding_dimension, padding_idx=0)

#### Load the Model here ##########################################################
# PATH_check_point = 'checkpoints/275.pth'
# checkpoint = torch.load(PATH_check_point)

PATH_unet_condition = f'{root_dir}/unet_latent_space_model_condition.pt'
PATH_embed_condition = f'{root_dir}/embed_latent_space_model_condition.pt'

unet.load_state_dict(torch.load(PATH_unet_condition,map_location=device))
embed.load_state_dict(torch.load(PATH_embed_condition,map_location=device))

# unet.load_state_dict(checkpoint['model_state_dict'])
# embed.load_state_dict(checkpoint['embed_state_dict'])
####################################################################

unet.to(device)
embed.to(device)


###---------------> Global variables for anomaly detection <------------------##

input_unhealthy = ''
output_healthy = ''

### ------------------------> Anomaly detection <-----------------------###########

scheduler_ddims = DDIMScheduler(num_train_timesteps=1000,schedule="linear_beta", beta_start=0.0015, beta_end=0.0195)

def get_healthy(un_img): # un_img is in range 0-255 but model takes in range 0-1. conversion is needed.
    global input_unhealthy
    global output_healthy
    
    img = cv2.resize(un_img,(112,112)) # resizing here
    gray_image = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    input_unhealthy = gray_image.copy()
    gray_image.resize(112,112,1)
    img_tensor = torch.from_numpy(gray_image*1.0)
    img_tensor = img_tensor.permute(2,0,1)
    img_tensor /= 255.
    img_tensor = img_tensor.float()
    input = img_tensor.reshape((1,1,112,112))
    z_mu, z_sigma = autoencoderkl.encode(input.to(device))
    z = autoencoderkl.sampling(z_mu, z_sigma)
    
    unet.eval()
    guidance_scale = 3.0
    total_timesteps = 1000
    latent_space_depth = int(total_timesteps * 0.5) 
    current_img = z
    current_img = current_img.float()
    scheduler_ddims.set_timesteps(num_inference_steps=total_timesteps)
    ## Ecodings
    scheduler_ddims.clip_sample = False
    class_embedding = embed(torch.zeros(1).long().to(device)).unsqueeze(1)
    progress_bar = tqdm(range(30))
    for i in progress_bar:  # go through the noising process
        t = i
        with torch.no_grad():
            model_output = unet(current_img, timesteps=torch.Tensor((t,)).to(current_img.device), context=class_embedding)
        current_img, _ = scheduler_ddims.reversed_step(model_output, t, current_img)
        progress_bar.set_postfix({"timestep input": t})
    
    latent_img = current_img
    ## Decoding
    conditioning = torch.cat([torch.zeros(1).long(), torch.ones(1).long()], dim=0).to(device)
    class_embedding = embed(conditioning).unsqueeze(1)
    
    progress_bar = tqdm(range(500))
    for i in progress_bar:  # go through the denoising process
        t = latent_space_depth - i
        current_img_double = torch.cat([current_img] * 2)
        with torch.no_grad():
            model_output = unet(
                current_img_double, timesteps=torch.Tensor([t, t]).to(current_img.device), context=class_embedding
            )
        noise_pred_uncond, noise_pred_text = model_output.chunk(2)
        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
        current_img, _ = scheduler_ddims.step(noise_pred, t, current_img)
        progress_bar.set_postfix({"timestep input": t})
    # torch.cuda.empty_cache()
    current_img_decode = autoencoderkl.decode(current_img)
    
    out_image = current_img_decode[0][0].to('cpu').detach().numpy()
    out_image = 255*out_image
    out_image = (out_image).astype('uint8')
    output_healthy = out_image.copy()
    return cv2.resize(out_image,(448,448))

##------------------> Anomaly detection , contrast and background removal <-------------------##

def update(brightness,contrast): ##def update(brightness,contrast,thr1,thr2):
    unhealthy_c = controller(input_unhealthy,brightness,contrast)
    healthy_c = controller(output_healthy,brightness,contrast)
    # unhealthy_remove_bg =  get_mask(unhealthy_c,thr1,thr2)
    # healthy_remove_bg = get_mask(healthy_c,thr1,thr2)
    # diff_img = unhealthy_remove_bg - healthy_remove_bg
    diff_img = unhealthy_c - healthy_c
    cmap = plt.get_cmap('inferno')
    diff_img_a = cmap(diff_img)
    diff_img = np.delete(diff_img_a, 3, 2)
    return cv2.resize(healthy_c,(448,448)),cv2.resize(diff_img,(448,448))



### --------------> Image generation <----------------------------##############



scheduler = DDPMScheduler(num_train_timesteps=1000, schedule="linear_beta", beta_start=0.0015, beta_end=0.0195)
# scale_factor = 0.943597137928009
# inferer = LatentDiffusionInferer(scheduler, scale_factor=scale_factor)


def get_value(grad):
    info_dict = {"Normal":1, "Level_1":2, "Level_2":3,"Level_3":4,"Worse":5}
    return info_dict[grad]

def generate_condition_bone_images(grad=0):
    grad_value = get_value(grad)
    unet.eval()
    scheduler.clip_sample = True
    guidance_scale = 3
    conditioning = torch.cat([torch.zeros(1).long(), grad_value * torch.ones(1).long()], dim=0).to(
        device
    )  # 2*torch.ones(1).long() is the class label for the UNHEALTHY (tumor) class
    class_embedding = embed(conditioning).unsqueeze(
        1
    )  # cross attention expects shape [batch size, sequence length, channels]
    scheduler.set_timesteps(num_inference_steps=1000)
    noise = torch.randn((1, 3, 28, 28))
    noise = noise.to(device)
    
    progress_bar = tqdm(scheduler.timesteps)
    for t in progress_bar:
        with autocast(enabled=True):
            with torch.no_grad():
                noise_input = torch.cat([noise] * 2)
                model_output = unet(noise_input, timesteps=torch.Tensor((t,)).to(noise.device), context=class_embedding,)
                noise_pred_uncond, noise_pred_text = model_output.chunk(2)
                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
    
        noise, _ = scheduler.step(noise_pred, t, noise)
    with torch.no_grad():
        noise = autoencoderkl.decode(noise)
        img = (noise[0][0].to('cpu')).numpy()
    return cv2.resize(img,(448,448))



##--------------------------------> UI <-----------------------------##

my_theme = 'YenLai/Superhuman'


with gr.Blocks(theme=my_theme,title="Knee Predict") as demo:
    gr.Markdown(""" # Knee Predict 
    ## Generative AI for Anomaly Detection and Analysis for Bone Diseases - Knee Osteoarthritis """  )
    
    with gr.Tab("Generate Image on conditions"):
        gr.Markdown("#### Generate Knee X-ray images with condition. You can select the level of Osteoarthritis and click on generate . Then the AI will generate Knee X-ray image of the given condition.")
        with gr.Row():
            output =gr.Image(height=450,width=450)
            gr.Image(value="images/doc_bone.png",label="AI-Assisted Healthcare")
        # output= gr.Textbox(label="Output Box")
        gr.Markdown(" ### Select the level of disease severity you want to generate !!")
        input = gr.Radio(["Normal", "Level_1", "Level_2","Level_3","Worse"], label="Knee Osteoarthritis Disease Severity Levels",scale=1)
        with gr.Row():
            greet_btn = gr.Button("Generate",size="lg",scale=1,interactive=True)
            gr.Markdown()
            gr.Markdown()
        


    with gr.Tab("Anomaly Detection"):
        gr.Markdown("### From a given unhealthy x-ray image generate a healthy image keeping the size and other important features")
        with gr.Row():
            image_input = gr.Image(height=450,width=450,label="Upload your knee x-ray here")
            img_out_heal = gr.Image(height=450,width=450,label="Healthy image")
        with gr.Row():
            gr.Markdown()
            generate_healthy_button = gr.Button("Generate",size="lg")
            gr.Markdown()

        gr.Markdown("""### Generate Anomaly by comparing the healthy and unhealthy Knee x-rays
                    #### Click the update button to update the anomaly after changing the contrast and brightness. 
                     """)
        with gr.Row():
            # image_input = gr.Image()
            image_output = [gr.Image(height=450,width=450,label="Contrasted"),gr.Image(height=450,width=450,label="Anomaly map")] # contrast and anomaly
        with gr.Row():
            gr.Markdown()
            update_anomaly_button = gr.Button("Update",size="lg")
            gr.Markdown()
        inputs_vlaues = [gr.Slider(0, 510, value=284, label="Brightness", info="Choose between 0 and 510"),
                        gr.Slider(0, 254, value=234, label="Contrast", info="Choose between 0 and 254"),
                        # gr.Slider(0, 50, value=7, label="Canny Threshold 1", info="Choose between 0 and 50"),
                        # gr.Slider(0, 50, value=20, label="Canny Threshold 2", info="Choose between 0 and 50"),
                ]
        
        # inputs_vlaues.append(image_input)
        gr.Examples(examples='examples' , fn=get_healthy, cache_examples=True, inputs=image_input, outputs=img_out_heal)
    greet_btn.click(fn=generate_condition_bone_images, inputs=input,outputs=output, api_name="generate_bone")
    generate_healthy_button.click(get_healthy,inputs=image_input,outputs=img_out_heal)
    update_anomaly_button.click(update, inputs=inputs_vlaues, outputs=image_output)



if __name__ == "__main__":
    
    demo.launch(share=True,server_name='0.0.0.0')