major: testing out first implentation of an ensemble team

app_mcp.py

@@ -1,4 +1,5 @@
 import os
+import time
 from typing import Literal
 import spaces
 import gradio as gr
@@ -301,17 +302,232 @@ def get_consensus_label(results):
 
 # Update predict_image_with_json to return consensus label
 
+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
+        }
+        self.situation_weights = {
+            "high_confidence": 1.2,    # Boost weights for high confidence predictions
+            "low_confidence": 0.8,     # Reduce weights for low confidence
+            "conflict": 0.5,          # Reduce weights when models disagree
+            "consensus": 1.5          # Boost weights when models agree
+        }
+    
+    def adjust_weights(self, predictions, confidence_scores):
+        """Dynamically adjust weights based on prediction patterns"""
+        adjusted_weights = self.base_weights.copy()
+        
+        # Check for consensus
+        if self._has_consensus(predictions):
+            for model in adjusted_weights:
+                adjusted_weights[model] *= self.situation_weights["consensus"]
+        
+        # Check for conflicts
+        if self._has_conflicts(predictions):
+            for model in adjusted_weights:
+                adjusted_weights[model] *= self.situation_weights["conflict"]
+        
+        # Adjust based on confidence
+        for model, confidence in confidence_scores.items():
+            if confidence > 0.8:
+                adjusted_weights[model] *= self.situation_weights["high_confidence"]
+            elif confidence < 0.5:
+                adjusted_weights[model] *= self.situation_weights["low_confidence"]
+        
+        return self._normalize_weights(adjusted_weights)
+    
+    def _has_consensus(self, predictions):
+        """Check if models agree on prediction"""
+        return len(set(predictions.values())) == 1
+    
+    def _has_conflicts(self, predictions):
+        """Check if models have conflicting predictions"""
+        return len(set(predictions.values())) > 2
+    
+    def _normalize_weights(self, weights):
+        """Normalize weights to sum to 1"""
+        total = sum(weights.values())
+        return {k: v/total for k, v in weights.items()}
+
+class EnsembleMonitorAgent:
+    def __init__(self):
+        self.performance_metrics = {
+            "model_accuracy": {},
+            "response_times": {},
+            "confidence_distribution": {},
+            "consensus_rate": 0.0
+        }
+        self.alerts = []
+    
+    def monitor_prediction(self, model_id, prediction, confidence, response_time):
+        """Monitor individual model performance"""
+        if model_id not in self.performance_metrics["model_accuracy"]:
+            self.performance_metrics["model_accuracy"][model_id] = []
+            self.performance_metrics["response_times"][model_id] = []
+            self.performance_metrics["confidence_distribution"][model_id] = []
+        
+        self.performance_metrics["response_times"][model_id].append(response_time)
+        self.performance_metrics["confidence_distribution"][model_id].append(confidence)
+        
+        # Check for performance issues
+        self._check_performance_issues(model_id)
+    
+    def _check_performance_issues(self, model_id):
+        """Check for any performance anomalies"""
+        response_times = self.performance_metrics["response_times"][model_id]
+        if len(response_times) > 10:
+            avg_time = sum(response_times[-10:]) / 10
+            if avg_time > 2.0:  # More than 2 seconds
+                self.alerts.append(f"High latency detected for {model_id}: {avg_time:.2f}s")
+
+class WeightOptimizationAgent:
+    def __init__(self, weight_manager):
+        self.weight_manager = weight_manager
+        self.performance_history = []
+        self.optimization_threshold = 0.1  # 10% performance change triggers optimization
+    
+    def analyze_performance(self, predictions, actual_results):
+        """Analyze model performance and suggest weight adjustments"""
+        # Placeholder for actual_results. In a real scenario, this would come from a validation set.
+        # For now, we'll just track predictions.
+        self.performance_history.append(predictions)
+        
+        if self._should_optimize():
+            self._optimize_weights()
+    
+    def _should_optimize(self):
+        """Determine if weights should be optimized"""
+        if len(self.performance_history) < 10:
+            return False
+        
+        # Placeholder for actual performance calculation
+        # For demonstration, let's say we optimize every 10 runs
+        return len(self.performance_history) % 10 == 0
+    
+    def _optimize_weights(self):
+        """Optimize model weights based on performance"""
+        logger.info("Optimizing model weights based on recent performance.")
+        # This is where more sophisticated optimization logic would go.
+        # For example, you could slightly adjust weights of models that consistently predict correctly.
+        pass
+
+class SystemHealthAgent:
+    def __init__(self):
+        self.health_metrics = {
+            "memory_usage": [],
+            "gpu_utilization": [],
+            "model_load_times": {},
+            "error_rates": {}
+        }
+    
+    def monitor_system_health(self):
+        """Monitor overall system health"""
+        self._check_memory_usage()
+        self._check_gpu_utilization()
+        # You might add _check_model_health() here later
+    
+    def _check_memory_usage(self):
+        """Monitor memory usage"""
+        try:
+            import psutil
+            memory = psutil.virtual_memory()
+            self.health_metrics["memory_usage"].append(memory.percent)
+            
+            if memory.percent > 90:
+                logger.warning(f"High memory usage detected: {memory.percent}%")
+        except ImportError:
+            logger.warning("psutil not installed. Cannot monitor memory usage.")
+    
+    def _check_gpu_utilization(self):
+        """Monitor GPU utilization if available"""
+        if torch.cuda.is_available():
+            try:
+                gpu_util = torch.cuda.memory_allocated() / torch.cuda.max_memory_allocated()
+                self.health_metrics["gpu_utilization"].append(gpu_util)
+                
+                if gpu_util > 0.9:
+                    logger.warning(f"High GPU utilization detected: {gpu_util*100:.2f}%")
+            except Exception as e:
+                logger.warning(f"Error monitoring GPU utilization: {e}")
+        else:
+            logger.info("CUDA not available. Skipping GPU utilization monitoring.")
+
 def predict_image_with_json(img, confidence_threshold, augment_methods, rotate_degrees, noise_level, sharpen_strength):
+    # Initialize agents
+    monitor_agent = EnsembleMonitorAgent()
+    weight_manager = ModelWeightManager()
+    optimization_agent = WeightOptimizationAgent(weight_manager)
+    health_agent = SystemHealthAgent()
+    
+    # Monitor system health
+    health_agent.monitor_system_health()
+
     if augment_methods:
         img_pil, _ = augment_image(img, augment_methods, rotate_degrees, noise_level, sharpen_strength)
     else:
         img_pil = img
-    img_pil, results = predict_image(img_pil, confidence_threshold)
-    img_np = np.array(img_pil)  # Convert PIL Image to NumPy array
     img_np_og = np.array(img)  # Convert PIL Image to NumPy array
 
-    gradient_image = gradient_processing(img_np)  # Added gradient processing
-    minmax_image = minmax_preprocess(img_np)  # Added MinMax processing
+    # Get predictions with timing
+    model_predictions = {}
+    confidence_scores = {}
+    results = [] # To store the results for the DataFrame
+
+    for model_id in MODEL_REGISTRY:
+        model_start = time.time()
+        result = infer(img_pil, model_id, confidence_threshold)
+        model_end = time.time()
+        
+        # Monitor individual model performance
+        monitor_agent.monitor_prediction(
+            model_id,
+            result["Label"],
+            max(result.get("AI Score", 0.0), result.get("Real Score", 0.0)),
+            model_end - model_start
+        )
+        
+        model_predictions[model_id] = result["Label"]
+        confidence_scores[model_id] = max(result.get("AI Score", 0.0), result.get("Real Score", 0.0))
+        results.append(result) # Add individual model result to the list
+    
+    # Get adjusted weights
+    adjusted_weights = weight_manager.adjust_weights(model_predictions, confidence_scores)
+    
+    # Optimize weights if needed
+    optimization_agent.analyze_performance(model_predictions, None) # Placeholder for actual results
+    
+    # Calculate weighted consensus
+    weighted_predictions = {
+        "AI": 0.0,
+        "REAL": 0.0,
+        "UNCERTAIN": 0.0
+    }
+    
+    for model_id, prediction in model_predictions.items():
+        # Ensure the prediction label is valid for weighted_predictions
+        if prediction in weighted_predictions:
+            weighted_predictions[prediction] += adjusted_weights[model_id]
+        else:
+            # Handle cases where prediction might be an error or unexpected label
+            logger.warning(f"Unexpected prediction label '{prediction}' from model '{model_id}'. Skipping its weight in consensus.")
+
+    final_prediction_label = "UNCERTAIN"
+    if weighted_predictions["AI"] > weighted_predictions["REAL"] and weighted_predictions["AI"] > weighted_predictions["UNCERTAIN"]:
+        final_prediction_label = "AI"
+    elif weighted_predictions["REAL"] > weighted_predictions["AI"] and weighted_predictions["REAL"] > weighted_predictions["UNCERTAIN"]:
+        final_prediction_label = "REAL"
+    
+    # Rest of your existing code remains the same after this point
+    gradient_image = gradient_processing(img_np_og)  # Added gradient processing
+    minmax_image = minmax_preprocess(img_np_og)  # Added MinMax processing
 
     # First pass - standard analysis
     ela1 = ELA(img_np_og, quality=75, scale=50, contrast=20, linear=False, grayscale=True)
@@ -330,8 +546,11 @@ def predict_image_with_json(img, confidence_threshold, augment_methods, rotate_d
         r.get("Real Score", ""),
         r.get("Label", "")
     ] for r in results]
-    consensus = get_consensus_label(table_rows)
-    return img_pil, forensics_images, table_rows, results, consensus
+    
+    # The get_consensus_label function is now replaced by final_prediction_label from weighted consensus
+    consensus_html = f"<b><span style='color:{'red' if final_prediction_label == 'AI' else ('green' if final_prediction_label == 'REAL' else 'orange')}'>{final_prediction_label}</span></b>"
+    
+    return img_pil, forensics_images, table_rows, results, consensus_html
 
 with gr.Blocks(css="#post-gallery { overflow: hidden !important;} .grid-wrap{ overflow-y: hidden !important;} .ms-gr-ant-welcome-icon{ height:unset !important;} .tabs{margin-top:10px;}") as demo:
     with ms.Application() as app: