Update prediction.py

prediction.py

@@ -1,13 +1,19 @@
+import logging
 from pathlib import Path
 import numpy as np
 import torch
 from transformers import SamModel, SamConfig, SamProcessor
 
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 # Constants
 MODEL_DIR = Path(__file__).parent
 my_SAM_model = MODEL_DIR / "www/sidewalkSAM.pth"
 
 def load_model_and_processor(model_path, config_path):
+    logger.info("Loading model and processor...")
     model = SamModel(config=SamConfig.from_pretrained(config_path))
     processor = SamProcessor.from_pretrained(config_path)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -17,13 +23,14 @@ def load_model_and_processor(model_path, config_path):
         model.load_state_dict(torch.load(model_path, map_location=device))
         model.to(device)
     except Exception as e:
-        print(f"Error loading model: {e}")
+        logger.error(f"Error loading model: {e}")
         exit(1)
 
-    print("Model is loaded.")
+    logger.info("Model and processor loaded.")
     return model, processor, device
 
 def predict_bbox(model, image, processor, device, prediction_threshold=0.5):
+    logger.info("Starting prediction...")
     prompt = [0, 0, image.width-5, image.height-5]
     inputs = processor(image, input_boxes=[[prompt]], return_tensors="pt")
     inputs = {k: v.to(device) for k, v in inputs.items()}
@@ -37,6 +44,7 @@ def predict_bbox(model, image, processor, device, prediction_threshold=0.5):
     medsam_seg_prob = medsam_seg_prob.cpu().numpy().squeeze()
     mask = (medsam_seg_prob > prediction_threshold).astype(np.uint8)
 
+    logger.info("Prediction completed.")
     return mask
 
 def get_sidewalk_prediction(image, model, processor, device, threshold=0.7):