Add cache optimization

app.py

@@ -14,11 +14,17 @@ from SingleTreePointCloudLoader import SingleTreePointCloudLoader
 gc.enable()
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
-with st.spinner("Loading PointNet++ model..."):
-    checkpoint = torch.load('checkpoints/best_model.pth', map_location=torch.device(device))
-    classifier = pn2.get_model(num_class=4, normal_channel=False)
-    classifier.load_state_dict(checkpoint['model_state_dict'])
-    classifier.eval()
+
+@st.cache_resource
+def load_model():
+    with st.spinner("Loading PointNet++ model..."):
+        checkpoint = torch.load('checkpoints/best_model.pth', map_location=torch.device(device))
+        classifier = pn2.get_model(num_class=4, normal_channel=False)
+        classifier.load_state_dict(checkpoint['model_state_dict'])
+        classifier.eval()
+    return classifier
+
+classifier = load_model()
 
 side_bg = "static/sidebar.png"
 side_bg_ext = "png"
@@ -160,6 +166,7 @@ if uploaded_file:
         proceed = st.button("Run model")
     except Exception as e:
         st.error(f"An error occured: {str(e)}")
+gc.collect() # Optimize after file is loaded
 
 if proceed:
     try:
@@ -259,6 +266,7 @@ if proceed:
         st.write(f"**Height of tree: {height:.2f}m**")
         st.write(f"**Canopy volume: {canopy_volume:.2f}m\u00b3**")
         st.write(f"**DBH: {dbh:.2f}m**")
+        gc.collect() # Optimize after inference is done
         
     except Exception as e:
         st.error(f"An error occured: {str(e)}")