feat: enhance model registration with metadata

- Introduced a new function `register_model_with_metadata` to include display name, contributor, and model path in model entries.
- Updated model registration calls for all models to utilize the new function, enhancing the metadata associated with each model.
- Modified `ModelEntry` class in `registry.py` to accommodate additional metadata fields.

app_mcp.py

@@ -16,7 +16,7 @@ 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 forensics.registry import register_model, MODEL_REGISTRY
+from forensics.registry import register_model, MODEL_REGISTRY, ModelEntry
 
 
 # Configure logging
@@ -107,25 +107,30 @@ def postprocess_logits(outputs, class_names):
     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(
-    "model_1",
-    clf_1,
-    preprocess_resize_256,
-    postprocess_pipeline,
-    CLASS_NAMES["model_1"]
+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(
-    "model_2",
-    clf_2,
-    preprocess_resize_224,
-    postprocess_pipeline,
-    CLASS_NAMES["model_2"]
+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
@@ -144,12 +149,9 @@ def model3_infer(image):
     with torch.no_grad():
         outputs = model_3(**inputs)
     return outputs
-register_model(
-    "model_3",
-    model3_infer,
-    preprocess_256,
-    postprocess_logits_model3,
-    CLASS_NAMES["model_3"]
+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)
@@ -163,52 +165,37 @@ 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(
-    "model_4",
-    model4_infer,
-    preprocess_256,
-    postprocess_logits_model4,
-    CLASS_NAMES["model_4"]
+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(
-    "model_5",
-    clf_5,
-    preprocess_resize_224,
-    postprocess_pipeline,
-    CLASS_NAMES["model_5"]
+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(
-    "model_5b",
-    clf_5b,
-    preprocess_resize_224,
-    postprocess_pipeline,
-    CLASS_NAMES["model_5b"]
+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(
-    "model_6",
-    clf_6,
-    preprocess_resize_224,
-    postprocess_pipeline,
-    CLASS_NAMES["model_6"]
+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(
-    "model_7",
-    clf_7,
-    preprocess_resize_224,
-    postprocess_pipeline,
-    CLASS_NAMES["model_7"]
+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
@@ -218,11 +205,28 @@ def infer(image: Image.Image, model_id: str, confidence_threshold: float = 0.75)
     img = entry.preprocess(image)
     try:
         result = entry.model(img)
-        result = entry.postprocess(result, entry.class_names)
-        # Add confidence threshold logic if needed
-        return result
+        scores = entry.postprocess(result, entry.class_names)
+        # Flatten output for Dataframe: include metadata and both class scores
+        ai_score = scores.get(entry.class_names[0], 0.0)
+        real_score = 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 {"error": str(e)}
+        return {
+            "Model": entry.display_name,
+            "Contributor": entry.contributor,
+            "HF Model Path": entry.model_path,
+            "AI Score": None,
+            "Real Score": None,
+            "Label": f"Error: {str(e)}"
+        }
 
 # Update predict_image to use all registered models in order
 
@@ -287,8 +291,12 @@ with gr.Blocks(css="#post-gallery { overflow: hidden !important;} .grid-wrap{ ov
 
 
                     with gr.Column(scale=2):
-                        # Use Gradio-native Dataframe to display results
-                        results_table = gr.Dataframe(label="Model Predictions", headers=None, datatype="auto")
+                        # Use Gradio-native Dataframe to display results with headers
+                        results_table = gr.Dataframe(
+                            label="Model Predictions",
+                            headers=["Model", "Contributor", "HF Model Path", "AI Score", "Real Score", "Label"],
+                            datatype=["str", "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")
 
                         outputs = [image_output, forensics_gallery, results_table]

forensics/registry.py

@@ -1,11 +1,15 @@
-from typing import Callable, Dict, Any, List
+from typing import Callable, Dict, Any, List, Optional
 
 class ModelEntry:
-    def __init__(self, model: Any, preprocess: Callable, postprocess: Callable, class_names: List[str]):
+    def __init__(self, model: Any, preprocess: Callable, postprocess: Callable, class_names: List[str],
+                 display_name: Optional[str] = None, contributor: Optional[str] = None, model_path: Optional[str] = None):
         self.model = model
         self.preprocess = preprocess
         self.postprocess = postprocess
         self.class_names = class_names
+        self.display_name = display_name
+        self.contributor = contributor
+        self.model_path = model_path
 
 MODEL_REGISTRY: Dict[str, ModelEntry] = {}