feat: enhance image handling in predictions and dynamically configure model weights based on MODEL_REGISTRY

agents/ensemble_weights.py

@@ -1,5 +1,6 @@
 import logging
 import torch
+from utils.registry import MODEL_REGISTRY # Import MODEL_REGISTRY
 
 logger = logging.getLogger(__name__)
 
@@ -35,17 +36,27 @@ class ContextualWeightOverrideAgent:
 
 
 class ModelWeightManager:
-    def __init__(self):
-        self.base_weights = {
-            "model_1": 0.15,  # SwinV2 Based
-            "model_2": 0.15,  # ViT Based
-            "model_3": 0.15,  # SDXL Dataset
-            "model_4": 0.15,  # SDXL + FLUX
-            "model_5": 0.15,  # ViT Based
-            "model_5b": 0.10, # ViT Based, Newer Dataset
-            "model_6": 0.10,  # Swin, Midj + SDXL
-            "model_7": 0.05   # ViT
-        }
+    def __init__(self, strongest_model_id: str = None):
+        # Dynamically initialize base_weights from MODEL_REGISTRY
+        num_models = len(MODEL_REGISTRY)
+        if num_models > 0:
+            if strongest_model_id and strongest_model_id in MODEL_REGISTRY:
+                # Assign a high weight to the strongest model (e.g., 50%)
+                strongest_weight_share = 0.5
+                self.base_weights = {strongest_model_id: strongest_weight_share}
+                remaining_models = [mid for mid in MODEL_REGISTRY.keys() if mid != strongest_model_id]
+                if remaining_models:
+                    other_models_weight_share = (1.0 - strongest_weight_share) / len(remaining_models)
+                    for model_id in remaining_models:
+                        self.base_weights[model_id] = other_models_weight_share
+                else: # Only one model, which is the strongest
+                    self.base_weights[strongest_model_id] = 1.0
+            else:
+                initial_weight = 1.0 / num_models
+                self.base_weights = {model_id: initial_weight for model_id in MODEL_REGISTRY.keys()}
+        else:
+            self.base_weights = {} # Handle case with no registered models
+        
         self.situation_weights = {
             "high_confidence": 1.2,    # Boost weights for high confidence predictions
             "low_confidence": 0.8,     # Reduce weights for low confidence

app.py

@@ -5,6 +5,7 @@ import numpy as np
 import os
 import time
 import logging
+import io
 
 # Assuming these are available from your utils and agents directories
 # You might need to adjust paths or copy these functions/classes if they are not directly importable.
@@ -185,8 +186,14 @@ def postprocess_simple_prediction(result, class_names):
 
 def simple_prediction(img):
     client = Client("aiwithoutborders-xyz/OpenSight-Community-Forensics-Preview")
+    
+    # Convert PIL Image to a file-like object in memory
+    img_byte_arr = io.BytesIO()
+    img.save(img_byte_arr, format='PNG') # Using PNG for lossless conversion, can be JPEG if preferred
+    img_byte_arr.seek(0) # Rewind to the beginning of the stream
+
     result = client.predict(
-            input_image=handle_file(img),
+            input_image=handle_file(img_byte_arr),
             api_name="/simple_predict"
     )
     return result
@@ -251,7 +258,7 @@ def ensemble_prediction_stream(img, confidence_threshold, augment_methods, rotat
             raise ValueError("Input image could not be converted to PIL Image.")
 
     monitor_agent = EnsembleMonitorAgent()
-    weight_manager = ModelWeightManager()
+    weight_manager = ModelWeightManager(strongest_model_id="simple_prediction")
     optimization_agent = WeightOptimizationAgent(weight_manager)
     health_agent = SystemHealthAgent()
     context_agent = ContextualIntelligenceAgent()