feat: enhance model registration by adding architecture and dataset metadata to the ModelEntry class

app.py

@@ -86,11 +86,8 @@ def postprocess_logits(outputs, class_names):
     probabilities = softmax(logits)
     return {class_names[i]: probabilities[i] for i in range(len(class_names))}
 
-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
+def register_model_with_metadata(model_id, model, preprocess, postprocess, class_names, display_name, contributor, model_path, architecture=None, dataset=None):
+    entry = ModelEntry(model, preprocess, postprocess, class_names, display_name=display_name, contributor=contributor, model_path=model_path, architecture=architecture, dataset=dataset)
     MODEL_REGISTRY[model_id] = entry
 
 # Load and register models (copied from app_mcp.py)
@@ -99,13 +96,15 @@ model_1 = Swinv2ForImageClassification.from_pretrained(MODEL_PATHS["model_1"]).t
 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"]
+    display_name="SWIN1", contributor="haywoodsloan", model_path=MODEL_PATHS["model_1"],
+    architecture="SwinV2", dataset="TBA"
 )
 
 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"]
+    display_name="VIT2", contributor="Heem2", model_path=MODEL_PATHS["model_2"],
+    architecture="ViT", dataset="TBA"
 )
 
 feature_extractor_3 = AutoFeatureExtractor.from_pretrained(MODEL_PATHS["model_3"], device=device)
@@ -125,7 +124,8 @@ def model3_infer(image):
     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"]
+    display_name="SDXL3", contributor="Organika", model_path=MODEL_PATHS["model_3"],
+    architecture="VIT", dataset="SDXL"
 )
 
 feature_extractor_4 = AutoFeatureExtractor.from_pretrained(MODEL_PATHS["model_4"], device=device)
@@ -141,13 +141,15 @@ def postprocess_logits_model4(outputs, class_names):
     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"]
+    display_name="XLFLUX4", contributor="cmckinle", model_path=MODEL_PATHS["model_4"],
+    architecture="VIT", dataset="SDXL, FLUX"
 )
 
 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"]
+    display_name="VIT5", contributor="prithivMLmods", model_path=MODEL_PATHS["model_5"],
+    architecture="VIT", dataset="TBA"
 )
 
 image_processor_6 = AutoImageProcessor.from_pretrained(MODEL_PATHS["model_6"], use_fast=True)
@@ -155,7 +157,8 @@ model_6 = SwinForImageClassification.from_pretrained(MODEL_PATHS["model_6"]).to(
 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"]
+    display_name="SWIN6", contributor="ideepankarsharma2003", model_path=MODEL_PATHS["model_6"],
+    architecture="SWINv1", dataset="SDXL, Midjourney"
 )
 
 image_processor_7 = AutoImageProcessor.from_pretrained(MODEL_PATHS["model_7"], use_fast=True)
@@ -163,7 +166,8 @@ model_7 = AutoModelForImageClassification.from_pretrained(MODEL_PATHS["model_7"]
 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"]
+    display_name="VIT7", contributor="date3k2", model_path=MODEL_PATHS["model_7"],
+    architecture="VIT", dataset="TBA"
 )
 
 def preprocess_simple_prediction(image):
@@ -196,7 +200,8 @@ register_model_with_metadata(
     ["AI", "REAL"], 
     display_name="Community Forensics",
     contributor="Jeongsoo Park",
-    model_path="aiwithoutborders-xyz/CommunityForensics-DeepfakeDet-ViT"
+    model_path="aiwithoutborders-xyz/CommunityForensics-DeepfakeDet-ViT",
+    architecture="ViT", dataset="GOAT"
 )
 def infer(image: Image.Image, model_id: str, confidence_threshold: float = 0.75) -> dict:
     """Predict using a specific model.
@@ -449,7 +454,7 @@ detection_model_eval_playground = gr.Interface(
         gr.Gallery(label="Post Processed Images", visible=True, columns=[4], rows=[2], container=False, height="auto", object_fit="contain", elem_id="post-gallery"),
         gr.Dataframe(
             label="Model Predictions",
-            headers=["Model", "Contributor", "AI Score", "Real Score", "Label"],
+            headers=["Arch / Dataset", "By", "AI", "Real", "Label"],
             datatype=["str", "str", "number", "number", "str"]
         ),
         gr.JSON(label="Raw Model Results", visible=False),

utils/registry.py

@@ -2,7 +2,8 @@ from typing import Callable, Dict, Any, List, Optional
 
 class ModelEntry:
     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):
+                 display_name: Optional[str] = None, contributor: Optional[str] = None, model_path: Optional[str] = None,
+                 architecture: Optional[str] = None, dataset: Optional[str] = None):
         self.model = model
         self.preprocess = preprocess
         self.postprocess = postprocess
@@ -10,8 +11,10 @@ class ModelEntry:
         self.display_name = display_name
         self.contributor = contributor
         self.model_path = model_path
+        self.architecture = architecture
+        self.dataset = dataset
 
 MODEL_REGISTRY: Dict[str, ModelEntry] = {}
 
-def register_model(model_id: str, model: Any, preprocess: Callable, postprocess: Callable, class_names: List[str]):
-    MODEL_REGISTRY[model_id] = ModelEntry(model, preprocess, postprocess, class_names) 
+def register_model(model_id: str, model: Any, preprocess: Callable, postprocess: Callable, class_names: List[str], architecture: Optional[str] = None, dataset: Optional[str] = None):
+    MODEL_REGISTRY[model_id] = ModelEntry(model, preprocess, postprocess, class_names, architecture=architecture, dataset=dataset)