AHA/app.py at main · DigBioLab/AHA · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
import os
import glob
import requests

import numpy as np
from PIL import Image
import streamlit as st

from venomai import unet, predictor, preprocess

st.set_page_config(
     page_title="AI-assisted Haemorrhage Analysis",
     page_icon="",
    #  layout="wide",
     initial_sidebar_state="expanded"
 )

def download_models():
#    print('Downloading models...')
   for i in range(5):
       filename = f'models/unet_inference_{i}.ckpt'
       if os.path.exists(filename):
#           print(f'Inference model {i} is already downloaded. Skipping...')
          continue
       URL = f'https://github.com/laprade117/venom-ai-haemorrhage-analysis-tool/releases/latest/download/unet_inference_{i}.ckpt'
       response = requests.get(URL)
       os.makedirs(os.path.dirname(filename), exist_ok=True)
       open(f'models/unet_inference_{i}.ckpt', 'wb').write(response.content)
#        print(f'Downloading inference model {i}...')

if __name__ == '__main__':

    download_models()

    st.title('AI-assisted Haemorrhage Analysis')

    # uploaded_file = st.sidebar.file_uploader("hello")
    uploaded_file = st.sidebar.file_uploader("Upload an image with the template and black sheet of paper separating the mice. After uploading, wait a few seconds while the tool computes the severity scores.", type=['.jpg','.png','.tif'], accept_multiple_files=False)

    if uploaded_file is not None:

        image = Image.open(uploaded_file)

        image = np.array(image)
        image = preprocess.preprocess_image(image)

        final_predictions = None
        for i in range(5):
            model = unet.UNet.load_from_checkpoint(f'models/unet_inference_{i}.ckpt')
            predictions, windows = predictor.predict_image(model, image, apply_preprocessing=False)
            if i == 0:
                final_predictions = predictions
            else:
                final_predictions += predictions
        predictions = final_predictions / 5

        haus, real_areas, luminance_values, mean_rgb_values = predictor.compute_haemorrhagic_units(predictions, windows, return_stats=True)
        masks = np.round((predictions > 0.5)[:,:,:,None] * mean_rgb_values[:,None,None,:]).astype('uint8')

        windows = list(np.array(windows, dtype=object))
        masks = list(np.array(masks, dtype=object))

        captions_hau = [f'HaU: {haus[i]:.02f}' for i in range(4)]
        captions = [f'Area: {real_areas[i]:.02f} mm\N{SUPERSCRIPT TWO}\nLuminance: {luminance_values[i]:.02f}' for i in range(4)]

        col1, col2, col3 = st.columns(3)
        col1.metric("HaU", f"{np.nanmean(haus):.02f}")
        col2.metric("Area", f"{np.nanmean(real_areas):.02f} mm\N{SUPERSCRIPT TWO}")
        col3.metric("Luminance", f"{np.nanmean(luminance_values):.02f}")

        st.image(windows, caption=captions_hau, width=174, clamp=[0,255])
        st.image(masks, caption=captions, width=174, clamp=[0,255])
        st.image(image)