updated st.cache

appTest.py

@@ -4,6 +4,17 @@ from modInference import main
 import numpy as np
 import math
 
+import numpy as np
+from skimage import transform
+import os
+from keras.models import Model
+from keras.applications.vgg16 import VGG16, preprocess_input
+from keras.layers import Dense, Dropout, Flatten
+import numpy as np
+
+import torch
+from models.create_fasterrcnn_model import create_model
+
 st.set_page_config(layout="wide") 
 st.markdown("")
 
@@ -18,6 +29,56 @@ st.subheader("Enter an image of a thin blood smear. Preview the image and run th
 
 input, outputIm = st.columns(2)
 
+@st.cache_resource
+def load_model():
+   NUM_CLASSES = 2
+   CLASSES = ['__background__', 'Cell']
+
+   DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+   print(DEVICE)
+   checkpoint = torch.load("best_model.pth", map_location=DEVICE)
+
+   NUM_CLASSES = checkpoint['config']['NC']
+   CLASSES = checkpoint['config']['CLASSES']
+   build_model = create_model[checkpoint['model_name']]
+
+   model = build_model(num_classes=NUM_CLASSES, coco_model=False)
+   model.load_state_dict(checkpoint['model_state_dict'])
+   model.to(DEVICE).eval()
+
+   COLORS = np.random.uniform(0, 255, size=(len(CLASSES), 3))
+
+   conv_base = VGG16(include_top=False,
+               weights='imagenet', 
+               input_shape=(200,200,3))
+
+   if 2 > 0:
+      for layer in conv_base.layers[:-2]:
+         layer.trainable = False
+   else:
+      for layer in conv_base.layers:
+         layer.trainable = False
+
+
+   top_model = conv_base.output
+   top_model = Flatten(name="flatten")(top_model)
+   top_model = Dense(4096, activation='relu')(top_model)
+   top_model = Dense(1048, activation='relu')(top_model)
+   top_model = Dense(256, activation='relu')(top_model)
+   top_model = Dense(128, activation='relu')(top_model)
+   top_model = Dense(64, activation='relu')(top_model)
+   top_model = Dropout(0.2)(top_model)
+   output_layer = Dense(5, activation='softmax')(top_model)
+
+   CNN = Model(inputs=conv_base.input, outputs=output_layer)
+   CNN.load_weights("CNN.hdf5")
+
+   return CNN, model
+
+CNN, model = load_model()
+
+
+
 with input:
    st.header("Input")
    imageInput = st.file_uploader("Enter an image of a thin blood smear.")
@@ -26,7 +87,7 @@ with input:
             image = Image.open(imageInput)
             ogInp = image
             img_array = np.array(image)
-            output, cellImgs = main("best_model.pth", img_array)
+            output, cellImgs = main(CNN, model, img_array)
             showImg = Image.fromarray(output)
    if st.button("Preview"):
       if imageInput is not None: