Update app.py

app.py

@@ -1,15 +1,12 @@
 from huggingface_hub import login
-# import spaces
 import os
 import time
 import gradio as gr
-from typing import *
+from typing import *  # Ensure that the typing module is imported
 from pillow_heif import register_heif_opener
 register_heif_opener()
 import vision_agent as va
-from vision_agent.tools import register_tool
-from vision_agent.tools import load_image, owl_v2, overlay_bounding_boxes, save_image
-
+from vision_agent.tools import register_tool, load_image, owl_v2, overlay_bounding_boxes, save_image
 
 # Perform login using the token
 hf_token = os.getenv("HF_TOKEN")
@@ -18,6 +15,32 @@ login(token=hf_token, add_to_git_credential=True)
 import numpy as np
 from PIL import Image
 
+# Set up logging
+import logging
+import sys
+
+def get_logging_level(default_level=logging.INFO):
+    log_level_str = os.getenv('VISION_AGENT_LOG_LEVEL', '').upper()
+    if log_level_str == 'DEBUG':
+        return logging.DEBUG
+    elif log_level_str == 'INFO':
+        return logging.INFO
+    elif log_level_str == 'WARNING':
+        return logging.WARNING
+    elif log_level_str == 'ERROR':
+        return logging.ERROR
+    elif log_level_str == 'CRITICAL':
+        return logging.CRITICAL
+    else:
+        return default_level
+
+logging_level = get_logging_level()
+logging.basicConfig(stream=sys.stdout, level=logging_level)
+_LOGGER = logging.getLogger(__name__)
+
+# Print the logging level to verify it's set correctly
+print(f"Logging level set to: {logging.getLevelName(logging_level)}")
+
 def detect_brain_tumor(image, seg_input, debug: bool = True):
     """
     Detects a brain tumor in the given image and returns the annotated image.
@@ -31,20 +54,20 @@ def detect_brain_tumor(image, seg_input, debug: bool = True):
         tuple: (numpy array of image, list of (label, (x1, y1, x2, y2)) tuples)
     """
     if debug:
-        print(f"Image received, shape: {image.shape}")
+        _LOGGER.debug(f"Image received, shape: {image.shape}")
 
     # Step 2: Detect brain tumor using owl_v2
     prompt = "detect brain tumor"
     detections = owl_v2(prompt, image)
     if debug:
-        print(f"Raw detections: {detections}")
+        _LOGGER.debug(f"Raw detections: {detections}")
 
     # Step 3: Overlay bounding boxes on the image
     image_with_bboxes = overlay_bounding_boxes(image, detections)
     if debug:
-        print("Bounding boxes overlaid on the image")
+        _LOGGER.debug("Bounding boxes overlaid on the image")
 
-        # Prepare annotations for AnnotatedImage output
+    # Prepare annotations for AnnotatedImage output
     annotations = []
     for detection in detections:
         label = detection['label']
@@ -57,7 +80,7 @@ def detect_brain_tumor(image, seg_input, debug: bool = True):
         annotations.append(((x1, y1, x2, y2), f"{label} {score:.2f}"))
 
     if debug:
-        print(f"Annotations: {annotations}")
+        _LOGGER.debug(f"Annotations: {annotations}")
 
     # Convert image to numpy array if it's not already
     if isinstance(image_with_bboxes, Image.Image):
@@ -106,4 +129,4 @@ with gr.Blocks(css="style.css") as demo:
         )
 
 if __name__ == "__main__":
-    demo.queue(max_size=10).launch(debug=True)
+    demo.queue(max_size=10).launch(debug=True)