Update app2.py

app2.py

@@ -226,6 +226,63 @@ def model_infer(img_name):
             break
     
     return image_vis, gt_mask, pr_mask
+
+def model_infer_new(img_name):
+    
+    model = smp.UnetPlusPlus(
+        encoder_name=ENCODER, 
+        encoder_weights=ENCODER_WEIGHTS, 
+        encoder_depth=5,
+        decoder_channels=(256, 128, 64, 32, 16),
+        classes=len(CLASSES), 
+        activation=ACTIVATION,
+        decoder_attention_type=None,
+    )
+
+
+    model.load_state_dict(torch.load('best.pth', map_location=torch.device('cpu'))['model_state_dict'])
+    model.eval()
+
+    test_dataset = Dataset(
+        img_name,
+        img_name,
+        augmentation=get_validation_augmentation(),
+        preprocessing=get_preprocessing(preprocessing_fn),
+        classes=CLASSES,
+        single_file=True
+    )
+
+    test_dataloader = DataLoader(test_dataset)
+
+    loaders = {"infer": test_dataloader}
+
+    runner = SupervisedRunner()
+    logits = []
+    f = 0
+    for prediction in runner.predict_loader(model=model, loader=loaders['infer'],cpu=True):
+        if f < 3:
+            logits.append(prediction['logits'])
+            f = f + 1
+        else:
+            break
+
+    threshold = 0.5
+    break_at = 1
+
+    for i, (input, output) in enumerate(zip(
+            test_dataset, logits)):
+        image, mask = input
+
+        image_vis = image.transpose(1, 2, 0)
+        pr_mask = (output[0].numpy() > threshold).astype('uint8')[0]
+        i = i + 1
+        if i >= break_at:
+            break
+    
+    return image_vis, pr_mask
+
+
+
 PAGE_TITLE = "Polyp Segmentation" 
 
 def file_selector(folder_path='.'):
@@ -242,33 +299,55 @@ def file_selector_ui():
     return filename
 
 def file_upload(folder_path='.'):
-    filenames = os.listdir(folder_path)
+    #    filenames = os.listdir(folder_path)
     folder_path = './test/test/images'
-    uploaded_file = st.file_uploader("Choose a file")
-    filename =     os.path.join(folder_path, uploaded_file.name)
-    printname = list(filename)
-    printname[filename.rfind('\\')] = '/'
-    st.write('You selected`%s`' % ''.join(printname))
-    return filename
+    uploaded_file = st.file_uploader("Choose a file")    
+    if uploaded_file is not None:
+        filename = os.path.join(folder_path, uploaded_file.name)
+        printname = list(filename)
+        printname[filename.rfind('\\')] = '/'
+        st.write('You selected`%s`' % ''.join(printname))
+        return filename
+
 
 
 def main():
     st.set_page_config(page_title=PAGE_TITLE, layout="wide")
     st.title(PAGE_TITLE)
-    image_path = file_selector_ui()
-#    image_path = file_upload()
-    image_path = os.path.abspath(image_path)
-    to_infer = image_path[image_path.rfind("\\") + 1:]
-
-    if os.path.isfile(image_path) is True:
-        _, file_extension = os.path.splitext(image_path)
-        if file_extension == ".jpg":
-            image_vis, gt_mask, pr_mask = model_infer(to_infer)
-            visualize(
-                image=image_vis, 
-                #ground_truth_mask=gt_mask, 
-                predicted_mask=pr_mask
-            )            
+    choice = st.radio(
+        "Upload your own image or infer on a pre-existing image?",
+        ('Pre-existing', 'Own'))
+
+    
+    if choice == 'Pre-existing':
+        image_path = file_selector_ui()
+        image_path = os.path.abspath(image_path)
+        to_infer = image_path[image_path.rfind("\\") + 1:]
+
+        if os.path.isfile(image_path) is True:
+            _, file_extension = os.path.splitext(image_path)
+            if file_extension == ".jpg":
+                image_vis, gt_mask, pr_mask = model_infer(to_infer)
+                visualize(
+                    image=image_vis, 
+                    ground_truth_mask=gt_mask, 
+                    predicted_mask=pr_mask
+                )       
+
+    if choice == 'Own':
+        image_path = file_upload()
+        if image_path is not None:
+            image_path = os.path.abspath(image_path)
+            to_infer = image_path[image_path.rfind("\\") + 1:]
+
+            if os.path.isfile(image_path) is True:
+                _, file_extension = os.path.splitext(image_path)
+                if file_extension == ".jpg":
+                    image_vis, pr_mask = model_infer_new(to_infer)
+                    visualize(
+                        image=image_vis,
+                        predicted_mask=pr_mask
+                    )       
 
 if __name__ == "__main__":
     main()