File size: 29,474 Bytes
932e7b4
19d34ca
932e7b4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c6683ee
16601a0
1af0cb5
ae1c2a2
c1d03da
 
c6683ee
ac9c2b2
c1d03da
73783f9
febce11
5967d4d
932e7b4
 
c1d03da
932e7b4
73783f9
932e7b4
5967d4d
489bc21
73783f9
c1d03da
73783f9
0f427f3
 
 
 
 
 
 
932e7b4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ac9c2b2
 
 
 
 
 
 
 
 
 
 
932e7b4
 
 
 
ac9c2b2
 
 
932e7b4
 
 
ac9c2b2
 
 
932e7b4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ac9c2b2
 
 
932e7b4
 
 
 
 
 
 
 
 
 
 
 
 
ac9c2b2
 
 
932e7b4
 
 
ac9c2b2
 
 
932e7b4
 
 
ac9c2b2
 
 
932e7b4
 
 
 
 
ac9c2b2
 
 
932e7b4
 
 
 
 
ac9c2b2
 
 
932e7b4
 
 
 
 
 
 
 
 
ac9c2b2
 
0f427f3
 
ac9c2b2
 
 
 
 
 
 
 
 
932e7b4
ac9c2b2
 
 
 
febce11
 
ac9c2b2
 
932e7b4
 
 
 
 
 
 
 
 
 
82e63cc
 
 
 
 
 
 
 
 
932e7b4
 
d20c076
 
 
 
 
 
 
 
 
 
19d34ca
 
 
 
 
1af0cb5
 
 
19d34ca
 
 
 
932e7b4
 
 
 
 
 
1af0cb5
 
19d34ca
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d7aaf8c
19d34ca
 
 
1af0cb5
 
 
 
 
 
 
 
 
 
 
 
19d34ca
1af0cb5
 
 
 
19d34ca
 
 
 
 
 
1af0cb5
19d34ca
91479f7
 
 
19d34ca
 
91479f7
19d34ca
 
 
 
 
 
1af0cb5
 
 
19d34ca
1af0cb5
19d34ca
 
932e7b4
 
 
 
 
 
 
 
 
1af0cb5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e8835e7
1af0cb5
 
88d2155
 
 
 
febce11
 
 
88d2155
19d34ca
febce11
 
 
 
19d34ca
 
1af0cb5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
febce11
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0f427f3
 
 
 
 
 
 
 
febce11
0f427f3
febce11
 
932e7b4
e2101d4
932e7b4
 
 
88d2155
932e7b4
88d2155
932e7b4
4274392
 
932e7b4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ac9c2b2
 
 
88d2155
 
ac9c2b2
932e7b4
aa96f77
 
82e63cc
932e7b4
82e63cc
932e7b4
 
 
 
 
 
 
 
 
2f1900e
392abb7
932e7b4
 
 
 
 
 
 
 
 
 
 
2f1900e
392abb7
932e7b4
c9d5b11
1af0cb5
c9d5b11
932e7b4
febce11
 
932e7b4
 
 
ae1c2a2
c9d5b11
 
 
 
9fe5b90
2f1900e
392abb7
c9d5b11
9fe5b90
 
ae1c2a2
0457419
 
 
 
 
 
 
 
c9d5b11
0457419
c9d5b11
 
 
 
 
 
 
 
9fe5b90
2f1900e
392abb7
c9d5b11
9fe5b90
 
 
 
 
 
 
932e7b4
0af824e
 
5967d4d
 
 
 
 
 
 
1cbfb0b
 
5967d4d
6432eef
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
import os
import time
from typing import Literal
import spaces
import gradio as gr
import modelscope_studio.components.antd as antd
import modelscope_studio.components.antdx as antdx
import modelscope_studio.components.base as ms
from transformers import pipeline, AutoImageProcessor, SwinForImageClassification, Swinv2ForImageClassification, AutoFeatureExtractor, AutoModelForImageClassification
from torchvision import transforms
import torch
from PIL import Image
import numpy as np
import io
import logging
from utils.utils import softmax, augment_image, convert_pil_to_bytes
from utils.gradient import gradient_processing
from utils.minmax import preprocess as minmax_preprocess
from utils.ela import genELA as ELA
from utils.wavelet import wavelet_blocking_noise_estimation
from utils.bitplane import bit_plane_extractor
from utils.hf_logger import log_inference_data
from utils.text_content import QUICK_INTRO, IMPLEMENTATION
from agents.monitoring_agents import EnsembleMonitorAgent, WeightOptimizationAgent, SystemHealthAgent
from agents.smart_agents import ContextualIntelligenceAgent, ForensicAnomalyDetectionAgent

from forensics.registry import register_model, MODEL_REGISTRY, ModelEntry
from agents.weight_management import ModelWeightManager
from dotenv import load_dotenv
import json
from huggingface_hub import CommitScheduler

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
os.environ['HF_HUB_CACHE'] = './models'

LOCAL_LOG_DIR = "./hf_inference_logs"
HF_DATASET_NAME="degentic_rd0"
load_dotenv()
# print(os.getenv("HF_HUB_CACHE"))

# Custom JSON Encoder to handle numpy types
class NumpyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.float32):
            return float(obj)
        return json.JSONEncoder.default(self, obj)

# Ensure using GPU if available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

header_style = {
    "textAlign": 'center',
    "color": '#fff',
    "height": 64,
    "paddingInline": 48,
    "lineHeight": '64px',
    "backgroundColor": '#4096ff',
}

content_style = {
    "textAlign": 'center',
    "minHeight": 120,
    "lineHeight": '120px',
    "color": '#fff',
    "backgroundColor": '#0958d9',
}

sider_style = {
    "textAlign": 'center',
    "lineHeight": '120px',
    "color": '#fff',
    "backgroundColor": '#1677ff',
}

footer_style = {
    "textAlign": 'center',
    "color": '#fff',
    "backgroundColor": '#4096ff',
}

layout_style = {
    "borderRadius": 8,
    "overflow": 'hidden',
    "width": 'calc(100% - 8px)',
    "maxWidth": 'calc(100% - 8px)',
}
# Model paths and class names
MODEL_PATHS = {
    "model_1": "haywoodsloan/ai-image-detector-deploy",
    "model_2": "Heem2/AI-vs-Real-Image-Detection",
    "model_3": "Organika/sdxl-detector",
    "model_4": "cmckinle/sdxl-flux-detector_v1.1",
    "model_5": "prithivMLmods/Deep-Fake-Detector-v2-Model",
    "model_5b": "prithivMLmods/Deepfake-Detection-Exp-02-22",
    "model_6": "ideepankarsharma2003/AI_ImageClassification_MidjourneyV6_SDXL",
    "model_7": "date3k2/vit-real-fake-classification-v4"
}

CLASS_NAMES = {
    "model_1": ['artificial', 'real'],
    "model_2": ['AI Image', 'Real Image'],
    "model_3": ['AI', 'Real'],
    "model_4": ['AI', 'Real'],
    "model_5": ['Realism', 'Deepfake'],
    "model_5b": ['Real', 'Deepfake'],
    "model_6": ['ai_gen', 'human'],
    "model_7": ['Fake', 'Real'],

}

def preprocess_resize_256(image):
    if image.mode != 'RGB':
        image = image.convert('RGB')
    return transforms.Resize((256, 256))(image)

def preprocess_resize_224(image):
    if image.mode != 'RGB':
        image = image.convert('RGB')
    return transforms.Resize((224, 224))(image)

def postprocess_pipeline(prediction, class_names):
    # Assumes HuggingFace pipeline output
    return {pred['label']: pred['score'] for pred in prediction}

def postprocess_logits(outputs, class_names):
    # Assumes model output with logits
    logits = outputs.logits.cpu().numpy()[0]
    probabilities = softmax(logits)
    return {class_names[i]: probabilities[i] for i in range(len(class_names))}

# Expand ModelEntry to include metadata
# (Assume ModelEntry is updated in registry.py to accept display_name, contributor, model_path)
# If not, we will update registry.py accordingly after this.

def register_model_with_metadata(model_id, model, preprocess, postprocess, class_names, display_name, contributor, model_path):
    entry = ModelEntry(model, preprocess, postprocess, class_names)
    entry.display_name = display_name
    entry.contributor = contributor
    entry.model_path = model_path
    MODEL_REGISTRY[model_id] = entry

# Load and register models (example for two models)
image_processor_1 = AutoImageProcessor.from_pretrained(MODEL_PATHS["model_1"], use_fast=True)
model_1 = Swinv2ForImageClassification.from_pretrained(MODEL_PATHS["model_1"]).to(device)
clf_1 = pipeline(model=model_1, task="image-classification", image_processor=image_processor_1, device=device)
register_model_with_metadata(
    "model_1", clf_1, preprocess_resize_256, postprocess_pipeline, CLASS_NAMES["model_1"],
    display_name="SwinV2 Based", contributor="haywoodsloan", model_path=MODEL_PATHS["model_1"]
)

clf_2 = pipeline("image-classification", model=MODEL_PATHS["model_2"], device=device)
register_model_with_metadata(
    "model_2", clf_2, preprocess_resize_224, postprocess_pipeline, CLASS_NAMES["model_2"],
    display_name="ViT Based", contributor="Heem2", model_path=MODEL_PATHS["model_2"]
)

# Register remaining models
feature_extractor_3 = AutoFeatureExtractor.from_pretrained(MODEL_PATHS["model_3"], device=device)
model_3 = AutoModelForImageClassification.from_pretrained(MODEL_PATHS["model_3"]).to(device)
def preprocess_256(image):
    if image.mode != 'RGB':
        image = image.convert('RGB')
    return transforms.Resize((256, 256))(image)
def postprocess_logits_model3(outputs, class_names):
    logits = outputs.logits.cpu().numpy()[0]
    probabilities = softmax(logits)
    return {class_names[i]: probabilities[i] for i in range(len(class_names))}
def model3_infer(image):
    inputs = feature_extractor_3(image, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model_3(**inputs)
    return outputs
register_model_with_metadata(
    "model_3", model3_infer, preprocess_256, postprocess_logits_model3, CLASS_NAMES["model_3"],
    display_name="SDXL Dataset", contributor="Organika", model_path=MODEL_PATHS["model_3"]
)

feature_extractor_4 = AutoFeatureExtractor.from_pretrained(MODEL_PATHS["model_4"], device=device)
model_4 = AutoModelForImageClassification.from_pretrained(MODEL_PATHS["model_4"]).to(device)
def model4_infer(image):
    inputs = feature_extractor_4(image, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model_4(**inputs)
    return outputs
def postprocess_logits_model4(outputs, class_names):
    logits = outputs.logits.cpu().numpy()[0]
    probabilities = softmax(logits)
    return {class_names[i]: probabilities[i] for i in range(len(class_names))}
register_model_with_metadata(
    "model_4", model4_infer, preprocess_256, postprocess_logits_model4, CLASS_NAMES["model_4"],
    display_name="SDXL + FLUX", contributor="cmckinle", model_path=MODEL_PATHS["model_4"]
)

clf_5 = pipeline("image-classification", model=MODEL_PATHS["model_5"], device=device)
register_model_with_metadata(
    "model_5", clf_5, preprocess_resize_224, postprocess_pipeline, CLASS_NAMES["model_5"],
    display_name="Vit Based", contributor="prithivMLmods", model_path=MODEL_PATHS["model_5"]
)

clf_5b = pipeline("image-classification", model=MODEL_PATHS["model_5b"], device=device)
register_model_with_metadata(
    "model_5b", clf_5b, preprocess_resize_224, postprocess_pipeline, CLASS_NAMES["model_5b"],
    display_name="Vit Based, Newer Dataset", contributor="prithivMLmods", model_path=MODEL_PATHS["model_5b"]
)

image_processor_6 = AutoImageProcessor.from_pretrained(MODEL_PATHS["model_6"], use_fast=True)
model_6 = SwinForImageClassification.from_pretrained(MODEL_PATHS["model_6"]).to(device)
clf_6 = pipeline(model=model_6, task="image-classification", image_processor=image_processor_6, device=device)
register_model_with_metadata(
    "model_6", clf_6, preprocess_resize_224, postprocess_pipeline, CLASS_NAMES["model_6"],
    display_name="Swin, Midj + SDXL", contributor="ideepankarsharma2003", model_path=MODEL_PATHS["model_6"]
)

image_processor_7 = AutoImageProcessor.from_pretrained(MODEL_PATHS["model_7"], use_fast=True)
model_7 = AutoModelForImageClassification.from_pretrained(MODEL_PATHS["model_7"]).to(device)
clf_7 = pipeline(model=model_7, task="image-classification", image_processor=image_processor_7, device=device)
register_model_with_metadata(
    "model_7", clf_7, preprocess_resize_224, postprocess_pipeline, CLASS_NAMES["model_7"],
    display_name="ViT", contributor="temp", model_path=MODEL_PATHS["model_7"]
)

# Generic inference function

def infer(image: Image.Image, model_id: str, confidence_threshold: float = 0.75) -> dict:
    entry = MODEL_REGISTRY[model_id]
    img = entry.preprocess(image)
    try:
        result = entry.model(img)
        scores = entry.postprocess(result, entry.class_names)
        # Flatten output for Dataframe: include metadata and both class scores
        ai_score = float(scores.get(entry.class_names[0], 0.0))
        real_score = float(scores.get(entry.class_names[1], 0.0))
        label = "AI" if ai_score >= confidence_threshold else ("REAL" if real_score >= confidence_threshold else "UNCERTAIN")
        return {
            "Model": entry.display_name,
            "Contributor": entry.contributor,
            "HF Model Path": entry.model_path,
            "AI Score": ai_score,
            "Real Score": real_score,
            "Label": label
        }
    except Exception as e:
        return {
            "Model": entry.display_name,
            "Contributor": entry.contributor,
            "HF Model Path": entry.model_path,
            "AI Score": 0.0,  # Ensure it's a float even on error
            "Real Score": 0.0, # Ensure it's a float even on error
            "Label": f"Error: {str(e)}"
        }

# Update predict_image to use all registered models in order

def predict_image(img, confidence_threshold):
    model_ids = [
        "model_1", "model_2", "model_3", "model_4", "model_5", "model_5b", "model_6", "model_7"
    ]
    results = [infer(img, model_id, confidence_threshold) for model_id in model_ids]
    return img, results

def get_consensus_label(results):
    labels = [r[4] for r in results if len(r) > 4]
    if not labels:
        return "No results"
    consensus = max(set(labels), key=labels.count)
    color = {"AI": "red", "REAL": "green", "UNCERTAIN": "orange"}.get(consensus, "gray")
    return f"<b><span style='color:{color}'>{consensus}</span></b>"

# Update predict_image_with_json to return consensus label

def predict_image_with_json(img, confidence_threshold, augment_methods, rotate_degrees, noise_level, sharpen_strength):
    # Ensure img is a PIL Image (if it's not already)
    if not isinstance(img, Image.Image):
        try:
            # If it's a numpy array, convert it
            img = Image.fromarray(img)
        except Exception as e:
            logger.error(f"Error converting input image to PIL: {e}")
            # If conversion fails, it's a critical error for the whole process
            raise ValueError("Input image could not be converted to PIL Image.")

    # Initialize agents
    monitor_agent = EnsembleMonitorAgent()
    weight_manager = ModelWeightManager()
    optimization_agent = WeightOptimizationAgent(weight_manager)
    health_agent = SystemHealthAgent()
    # New smart agents
    context_agent = ContextualIntelligenceAgent()
    anomaly_agent = ForensicAnomalyDetectionAgent()
    
    # 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_np_og = np.array(img)  # Convert PIL Image to NumPy array

    # 1. Get initial predictions from all models
    model_predictions_raw = {}
    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_raw[model_id] = result # Store the full result dictionary
        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
    
    # 2. Infer context tags using ContextualIntelligenceAgent
    image_data_for_context = {
        "width": img.width,
        "height": img.height,
        "mode": img.mode,
        # Add more features like EXIF data if exif_full_dump is used
    }
    detected_context_tags = context_agent.infer_context_tags(image_data_for_context, model_predictions_raw)
    logger.info(f"Detected context tags: {detected_context_tags}")

    # 3. Get adjusted weights, passing context tags
    adjusted_weights = weight_manager.adjust_weights(model_predictions_raw, confidence_scores, context_tags=detected_context_tags)
    
    # 4. Optimize weights if needed
    # `final_prediction_label` is determined AFTER weighted consensus, so analyze_performance will be called later

    # 5. Calculate weighted consensus
    weighted_predictions = {
        "AI": 0.0,
        "REAL": 0.0,
        "UNCERTAIN": 0.0
    }
    
    for model_id, prediction in model_predictions_raw.items(): # Use raw predictions for weighting
        # Ensure the prediction label is valid for weighted_predictions
        prediction_label = prediction.get("Label") # Extract the label
        if prediction_label in weighted_predictions:
            weighted_predictions[prediction_label] += adjusted_weights[model_id]
        else:
            # Handle cases where prediction might be an error or unexpected label
            logger.warning(f"Unexpected prediction label '{prediction_label}' 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"

    # Call analyze_performance after final_prediction_label is known
    optimization_agent.analyze_performance(final_prediction_label, None)
    
    # 6. Perform forensic processing
    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)

    # Second pass - enhanced visibility
    ela2 = ELA(img_np_og, quality=75, scale=75, contrast=25, linear=False, grayscale=True)
    ela3 = ELA(img_np_og, quality=75, scale=75, contrast=25, linear=False, grayscale=False)
    
    forensics_images = [img_pil, ela1, ela2, ela3, gradient_image, minmax_image]

    # 7. Generate boilerplate descriptions for forensic outputs for anomaly agent
    forensic_output_descriptions = [
        f"Original augmented image (PIL): {img_pil.width}x{img_pil.height}",
        "ELA analysis (Pass 1): Grayscale error map, quality 75.",
        "ELA analysis (Pass 2): Grayscale error map, quality 75, enhanced contrast.",
        "ELA analysis (Pass 3): Color error map, quality 75, enhanced contrast.",
        "Gradient processing: Highlights edges and transitions.",
        "MinMax processing: Deviations in local pixel values."
    ]
    # You could also add descriptions for Wavelet and Bit Plane if they were dynamic outputs
    # For instance, if wavelet_blocking_noise_estimation had parameters that changed and you wanted to describe them.

    # 8. Analyze forensic outputs for anomalies using ForensicAnomalyDetectionAgent
    anomaly_detection_results = anomaly_agent.analyze_forensic_outputs(forensic_output_descriptions)
    logger.info(f"Forensic anomaly detection: {anomaly_detection_results['summary']}")


    # Prepare table rows for Dataframe (exclude model path)
    table_rows = [[
        r.get("Model", ""),
        r.get("Contributor", ""),
        r.get("AI Score", 0.0) if r.get("AI Score") is not None else 0.0,
        r.get("Real Score", 0.0) if r.get("Real Score") is not None else 0.0,
        r.get("Label", "Error")
    ] for r in results]
    
    logger.info(f"Type of table_rows: {type(table_rows)}")
    for i, row in enumerate(table_rows):
        logger.info(f"Row {i} types: {[type(item) for item in row]}")

    # 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>"

    # Prepare data for logging to Hugging Face dataset
    inference_params = {
        "confidence_threshold": confidence_threshold,
        "augment_methods": augment_methods,
        "rotate_degrees": rotate_degrees,
        "noise_level": noise_level,
        "sharpen_strength": sharpen_strength,
        "detected_context_tags": detected_context_tags
    }

    ensemble_output_data = {
        "final_prediction_label": final_prediction_label,
        "weighted_predictions": weighted_predictions,
        "adjusted_weights": adjusted_weights
    }

    # Collect agent monitoring data
    agent_monitoring_data_log = {
        "ensemble_monitor": {
            "alerts": monitor_agent.alerts,
            "performance_metrics": monitor_agent.performance_metrics
        },
        "weight_optimization": {
            "prediction_history_length": len(optimization_agent.prediction_history),
            # You might add a summary of recent accuracy here if _calculate_accuracy is exposed
        },
        "system_health": {
            "memory_usage": health_agent.health_metrics["memory_usage"],
            "gpu_utilization": health_agent.health_metrics["gpu_utilization"]
        },
        "context_intelligence": {
            "detected_context_tags": detected_context_tags
        },
        "forensic_anomaly_detection": anomaly_detection_results
    }

    # Log the inference data
    log_inference_data(
        original_image=img, # Use the original uploaded image
        inference_params=inference_params,
        model_predictions=results, # This already contains detailed results for each model
        ensemble_output=ensemble_output_data,
        forensic_images=forensics_images, # This is the list of PIL images generated by forensic tools
        agent_monitoring_data=agent_monitoring_data_log,
        human_feedback=None # This can be populated later with human review data
    )

    # Final type safety check for forensic_images before returning
    cleaned_forensics_images = []
    for f_img in forensics_images:
        if isinstance(f_img, Image.Image):
            cleaned_forensics_images.append(f_img)
        elif isinstance(f_img, np.ndarray):
            try:
                cleaned_forensics_images.append(Image.fromarray(f_img))
            except Exception as e:
                logger.warning(f"Could not convert numpy array to PIL Image for gallery: {e}")
                # Optionally, append a placeholder or skip
        else:
            logger.warning(f"Unexpected type in forensic_images: {type(f_img)}. Skipping.")

    logger.info(f"Cleaned forensic images types: {[type(img) for img in cleaned_forensics_images]}")

    # Ensure numerical values in results are standard Python floats before JSON serialization
    for i, res_dict in enumerate(results):
        for key in ["AI Score", "Real Score"]:
            value = res_dict.get(key)
            if isinstance(value, np.float32):
                res_dict[key] = float(value)
                logger.info(f"Converted {key} for result {i} from numpy.float32 to float.")

    # Return raw model results as JSON string for debug_json component
    json_results = json.dumps(results, cls=NumpyEncoder)

    return img_pil, cleaned_forensics_images, table_rows, json_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:
        with antd.ConfigProvider():
            antdx.Welcome(
                icon="https://cdn-avatars.huggingface.co/v1/production/uploads/639daf827270667011153fbc/WpeSFhuB81DY-1TjNUmV_.png",
                title="Welcome to Project OpenSight",
                description="The OpenSight aims to be an open-source SOTA generated image detection model. This HF Space is not only an introduction but a educational playground for the public to evaluate and challenge current open source models.  **Space will be upgraded shortly; inference on all 6 models should take about 1.2~ seconds.** "
            )
            with gr.Tab("👀 Rethinking Detection Models: Multi-Model, Multi-Strategy Ensemble Team and Agentic Pipelines"):
                gr.Markdown("# Open Source Detection Models Found on the Hub\n\n - **IMPORTANT UPDATE REGARDING YOUR DATA AND PRIVACY: [PLEASE REFER TO THE MCP SERVER HACKATHON SUBMISSION FOR CRUCIAL DETAILS](https://huggingface.co/spaces/Agents-MCP-Hackathon/mcp-deepfake-forensics).** ")
                
                with gr.Row():
                    with gr.Column(scale=1):
                        image_input = gr.Image(label="Upload Image to Analyze", sources=['upload', 'webcam'], type='pil')
                        with gr.Accordion("Settings (Optional)", open=False, elem_id="settings_accordion"):
                            augment_checkboxgroup = gr.CheckboxGroup(["rotate", "add_noise", "sharpen"], label="Augmentation Methods")
                            rotate_slider = gr.Slider(0, 45, value=2, step=1, label="Rotate Degrees", visible=False)
                            noise_slider = gr.Slider(0, 50, value=4, step=1, label="Noise Level", visible=False)
                            sharpen_slider = gr.Slider(0, 50, value=11, step=1, label="Sharpen Strength", visible=False)
                            confidence_slider = gr.Slider(0.0, 1.0, value=0.75, step=0.05, label="Confidence Threshold")
                        inputs = [image_input, confidence_slider, augment_checkboxgroup, rotate_slider, noise_slider, sharpen_slider]
                        predict_button = gr.Button("Predict")
                        augment_button = gr.Button("Augment & Predict")
                        image_output = gr.Image(label="Processed Image", visible=False)


                    with gr.Column(scale=2):
                        # Use Gradio-native Dataframe to display results with headers
                        results_table = gr.Dataframe(
                            label="Model Predictions",
                            headers=["Model", "Contributor", "AI Score", "Real Score", "Label"],
                            datatype=["str", "str", "number", "number", "str"]
                        )
                        forensics_gallery = gr.Gallery(label="Post Processed Images", visible=True, columns=[4], rows=[2], container=False, height="auto", object_fit="contain", elem_id="post-gallery")
                        with gr.Accordion("Debug Output (Raw JSON)", open=False):
                            debug_json = gr.JSON(label="Raw Model Results")
                        consensus_md = gr.Markdown(label="Consensus", value="")

                        outputs = [image_output, forensics_gallery, results_table, debug_json, consensus_md]
                
                # Show/hide rotate slider based on selected augmentation method
                augment_checkboxgroup.change(lambda methods: gr.update(visible="rotate" in methods), inputs=[augment_checkboxgroup], outputs=[rotate_slider])
                augment_checkboxgroup.change(lambda methods: gr.update(visible="add_noise" in methods), inputs=[augment_checkboxgroup], outputs=[noise_slider])
                augment_checkboxgroup.change(lambda methods: gr.update(visible="sharpen" in methods), inputs=[augment_checkboxgroup], outputs=[sharpen_slider])
                
                predict_button.click(
                    fn=predict_image_with_json, 
                    inputs=inputs, 
                    outputs=outputs,
                    api_name="/predict"
                )
                augment_button.click(  # Connect Augment button to the function
                    fn=predict_image_with_json, 
                    inputs=[
                        image_input, 
                        confidence_slider, 
                        gr.CheckboxGroup(["rotate", "add_noise", "sharpen"], value=["rotate", "add_noise", "sharpen"], visible=False),  # Default values
                        rotate_slider, 
                        noise_slider, 
                        sharpen_slider
                    ], 
                    outputs=outputs,
                    api_name="/augment"
                )
            with gr.Tab("🙈 Project Introduction"):
                gr.Markdown(QUICK_INTRO)
                
            with gr.Tab("👑 Community Forensics Preview"):
                # temp_space = gr.load("aiwithoutborders-xyz/OpenSight-Community-Forensics-Preview", src="spaces")
                gr.Markdown("Community Forensics Preview coming soon!") # Placeholder for now
            with gr.Tab("🥇 Leaderboard"):
                gr.Markdown("# AI Generated / Deepfake Detection Models Leaderboard: Soon™")
                
            with gr.Tab("Wavelet Blocking Noise Estimation", visible=False):
                gr.Interface(
                    fn=wavelet_blocking_noise_estimation,
                    inputs=[gr.Image(type="pil"), gr.Slider(1, 32, value=8, step=1, label="Block Size")],
                    outputs=gr.Image(type="pil"),
                    title="Wavelet-Based Noise Analysis",
                    description="Analyzes image noise patterns using wavelet decomposition. This tool helps detect compression artifacts and artificial noise patterns that may indicate image manipulation. Higher noise levels in specific regions can reveal areas of potential tampering.",
                    api_name="/tool_waveletnoise"
                )


            with gr.Tab("Bit Plane Values", visible=False):
                """Forensics Tool: Bit Plane Extractor

                Args:
                    image: PIL Image to analyze
                    channel: Color channel to extract bit plane from ("Luminance", "Red", "Green", "Blue", "RGB Norm")
                    bit_plane: Bit plane index to extract (0-7)
                    filter_type: Filter to apply ("Disabled", "Median", "Gaussian")
                """
                gr.Interface(

                    fn=bit_plane_extractor,
                    inputs=[
                        gr.Image(type="pil"),
                        gr.Dropdown(["Luminance", "Red", "Green", "Blue", "RGB Norm"], label="Channel", value="Luminance"),
                        gr.Slider(0, 7, value=0, step=1, label="Bit Plane"),
                        gr.Dropdown(["Disabled", "Median", "Gaussian"], label="Filter", value="Disabled")
                    ],
                    outputs=gr.Image(type="pil"),
                    title="Bit Plane Analysis",
                    description="Extracts and visualizes individual bit planes from different color channels. This forensic tool helps identify hidden patterns and artifacts in image data that may indicate manipulation. Different bit planes can reveal inconsistencies in image processing or editing.",
                    api_name="/tool_bitplane"
                )
            # with gr.Tab("EXIF Full Dump"):
            #     gr.Interface(
            #         fn=exif_full_dump,
            #         inputs=gr.Image(type="pil"),
            #         outputs=gr.JSON(),
            #         description="Extract all EXIF metadata from the uploaded image."
            #     )

# --- MCP-Ready Launch ---
if __name__ == "__main__":
    # Initialize CommitScheduler
    # The scheduler will monitor LOCAL_LOG_DIR and push changes to HF_DATASET_NAME
    with CommitScheduler(
        repo_id=HF_DATASET_NAME, # Your Hugging Face dataset repository ID
        repo_type="dataset",
        folder_path=LOCAL_LOG_DIR,
        every=5, # Commit every 5 minutes
        private=False, # Keep your dataset private
        token=os.getenv("HF_TOKEN") # Uncomment and set if token is not saved globally
    ) as scheduler:
        demo.launch(mcp_server=True)