less load: inference shouldn't take longer than 3 seconds...

utils/hf_logger.py

@@ -5,13 +5,13 @@ import io
 import datetime
 from PIL import Image
 import logging
-from huggingface_hub import HfApi, CommitScheduler # Keep HfApi for repo creation, but remove CommitOperationAdd for direct upload
+from huggingface_hub import HfApi, CommitScheduler
 import numpy as np
 
 logger = logging.getLogger(__name__)
 
 HF_DATASET_NAME = "aiwithoutborders-xyz/degentic_rd0"
-LOCAL_LOG_DIR = "./hf_inference_logs" # Define a local directory to store logs
+LOCAL_LOG_DIR = "./hf_inference_logs"
 
 # Custom JSON Encoder to handle numpy types
 class NumpyEncoder(json.JSONEncoder):
@@ -20,18 +20,21 @@ class NumpyEncoder(json.JSONEncoder):
             return float(obj)
         return json.JSONEncoder.default(self, obj)
 
-def _pil_to_base64(image: Image.Image) -> str:
-    """Converts a PIL Image to a base64 string."""
-    # Explicitly check if the input is a PIL Image
+def _save_pil_image_to_file(image: Image.Image, directory: str, prefix: str) -> str:
+    """Saves a PIL Image to a file and returns its filename."""
     if not isinstance(image, Image.Image):
         raise TypeError(f"Expected a PIL Image, but received type: {type(image)}")
 
-    buffered = io.BytesIO()
-    # Ensure image is in RGB mode before saving as JPEG
+    os.makedirs(directory, exist_ok=True)
+    timestamp_str = datetime.datetime.now().strftime("%Y%m%d%H%M%S%f")
+    filename = f"{prefix}_{timestamp_str}.png"
+    file_path = os.path.join(directory, filename)
+
     if image.mode != 'RGB':
         image = image.convert('RGB')
-    image.save(buffered, format="JPEG", quality=85)
-    return base64.b64encode(buffered.getvalue()).decode('utf-8')
+    image.save(file_path, format="PNG")
+    logger.info(f"Saved image to: {file_path}")
+    return filename
 
 # The initialize_dataset function will change significantly or be removed/simplified
 # as we are no longer appending to a datasets.Dataset object directly in memory
@@ -59,26 +62,26 @@ def log_inference_data(
     try:
         api = initialize_dataset_repo() # Get or create the repository
 
-        original_image_b64 = _pil_to_base64(original_image)
+        original_image_filename = _save_pil_image_to_file(original_image, LOCAL_LOG_DIR, "original")
 
-        forensic_images_b64 = []
+        forensic_images_filenames = []
         for img_item in forensic_images:
             if img_item is not None:
                 if not isinstance(img_item, Image.Image):
                     try:
                         img_item = Image.fromarray(img_item)
                     except Exception as e:
-                        logger.error(f"Error converting forensic image to PIL for base64 encoding: {e}")
+                        logger.error(f"Error converting forensic image to PIL for saving: {e}")
                         continue
-                forensic_images_b64.append(_pil_to_base64(img_item))
+                forensic_images_filenames.append(_save_pil_image_to_file(img_item, LOCAL_LOG_DIR, "forensic"))
 
         new_entry = {
             "timestamp": datetime.datetime.now().isoformat(),
-            "image": original_image_b64,
+            "image": original_image_filename,
             "inference_request": inference_params,
             "model_predictions": model_predictions,
             "ensemble_output": ensemble_output,
-            "forensic_outputs": forensic_images_b64,
+            "forensic_outputs": forensic_images_filenames,
             "agent_monitoring_data": agent_monitoring_data,
             "human_feedback": human_feedback if human_feedback is not None else {}
         }
@@ -102,7 +105,8 @@ def log_inference_data(
             every=10  # Commit every 10 files
         )
         
-        logger.info(f"Inference data logged successfully to local file: {log_file_path}")
+        with scheduler:
+            logger.info(f"Inference data logged successfully to local file: {log_file_path}")
 
     except Exception as e:
         logger.error(f"Failed to log inference data to local file: {e}")