force use CPU

app.py

@@ -87,6 +87,7 @@ def load_model():
     cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3
     cfg.MODEL.WEIGHTS = MODEL
     cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = TH
+    cfg.MODEL.DEVICE = "cpu"
 
 
     metadata = Metadata()
@@ -138,7 +139,10 @@ def draw_box(file_name, box, type, model, resize_input=False):
 
 
 def explain(img, model):
+    state.write("Loading features...")
     database = json.load(open(FEATURES_DATABASE))
+
+    state.write("Computing logits...")
     instances, input = forward_model_full(model["model"], model["cfg"], img)
     
     instances.remove("pred_masks")
@@ -150,23 +154,28 @@ def explain(img, model):
     pred = cv2.resize(pred, (800, 800))
     pred = cv2.cvtColor(pred, cv2.COLOR_BGR2RGB)
     
-    tabs = st.tabs(["Detection"] + [f"Lesion #{i}" for i in range(0, len(instances))])
-    lesion_tabs = tabs[1:]
+    tabs = st.tabs(["Result", "Detection"] + [f"Lesion #{i}" for i in range(0, len(instances))])
+    lesion_tabs = tabs[2:]
 
     with tabs[0]:
+        st.header("Image processed")
+        st.success("Use the tabs on the right to see the detected lesions and detailed explanations for each lesion")
+
+    state.write("Populating first tab...")
+    with tabs[1]:
         st.header("Detected lesions")
-        state.text("All done...")
-        tooltip.success("Use the tabs for a detailed explanation of each lesion")
         st.image(pred)
 
     
     for i, (tab, box, type, scores, features) in enumerate(zip(lesion_tabs, instances.pred_boxes, instances.pred_classes, instances.probs, instances.features)):
+        state.write(f"Populating tab for lesion #{i}...")
         healthy_prob = scores[-1].item()
         scores = scores[:-1]
         features = features.tolist()
 
         with tab:
             st.header(f"Lesion #{i}")
+            state.write(f"Populating classes for lesion #{i}...")
             lesion_img = draw_box(img, box.cpu(), type, model)
             lesion_img = cv2.cvtColor(lesion_img, cv2.COLOR_BGR2RGB)
 
@@ -190,16 +199,21 @@ def explain(img, model):
             st.subheader("Feature space")
             col1, col2 = st.columns(2)
 
+            state.write(f"Populating PCA for lesion #{i}...")
             fig = plot_pca_point(point=features, features_database=FEATURES_DATABASE, pca_model=PCA_MODEL, fig_h=800, fig_w=600, fig_dpi=100)
             col1.pyplot(fig)
-
+            
+            state.write(f"Populating histogram for lesion #{i}...")
             fig = plot_histogram_dist(point=features, features_database=FEATURES_DATABASE, fig_h=800, fig_w=600, fig_dpi=100)
             col2.pyplot(fig)
 
+            state.write(f"Populating Gradcam++ for lesion #{i}...")
             st.subheader("Gradcam++")
             fig = plot_gradcam(model=MODEL, file=FILE, instance=i, fig_h=1600, fig_w=1200, fig_dpi=200, th=TH, layer="backbone.bottom_up.res5.2.conv3")
             st.pyplot(fig)
 
+    state.write("All done...")
+
 FILE = "./test.jpg"
 MODEL = "./models/model.pth"
 PCA_MODEL = "./models/pca.pkl"