feat(model):

- add support for new model (model_7) for image classification task

- extended CLASS_NAMES to support new model class labels

- added resource loading for the new model

- accessing the new model in the predict_image function

- updated the combined_results and combined_outputs to include new model output

- added a new tile in the HTML results for the new model

♻️ style(frontend):
- increase w-24 to w-30 in HTML CSS snippet

Note: Preferred writing "feat" for adding new model, even if it is just adding an URLs, and "chore" for any small tutorials added.

app.py

@@ -27,7 +27,8 @@ MODEL_PATHS = {
     "model_4": "cmckinle/sdxl-flux-detector",
     "model_5": "prithivMLmods/Deep-Fake-Detector-v2-Model",
     "model_5b": "prithivMLmods/Deepfake-Detection-Exp-02-22",
-    "model_6": "ideepankarsharma2003/AI_ImageClassification_MidjourneyV6_SDXL"
+    "model_6": "ideepankarsharma2003/AI_ImageClassification_MidjourneyV6_SDXL",
+    "model_7": "date3k2/vit-real-fake-classification-v4"
 }
 
 CLASS_NAMES = {
@@ -38,6 +39,7 @@ CLASS_NAMES = {
     "model_5": ['Realism', 'Deepfake'],
     "model_5b": ['Real', 'Deepfake'],
     "model_6": ['ai_gen', 'human'],
+    "model_7": ['Fake', 'Real'],
 
 }
 
@@ -63,9 +65,13 @@ def load_models():
     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)
 
-    return clf_1, clf_2, feature_extractor_3, model_3, feature_extractor_4, model_4, clf_5, clf_5b, clf_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)
 
-clf_1, clf_2, feature_extractor_3, model_3, feature_extractor_4, model_4, clf_5, clf_5b, clf_6 = load_models()
+    return clf_1, clf_2, feature_extractor_3, model_3, feature_extractor_4, model_4, clf_5, clf_5b, clf_6, model_7, clf_7
+
+clf_1, clf_2, feature_extractor_3, model_3, feature_extractor_4, model_4, clf_5, clf_5b, clf_6, model_7, clf_7 = load_models()
 
 @spaces.GPU(duration=10)
 def predict_with_model(img_pil, clf, class_names, confidence_threshold, model_name, model_id, feature_extractor=None):
@@ -118,6 +124,7 @@ def predict_image(img, confidence_threshold):
     label_5, result_5output = predict_with_model(img_pilvits, clf_5, CLASS_NAMES["model_5"], confidence_threshold, "ViT-base Newcomer", 5)
     label_5b, result_5boutput = predict_with_model(img_pilvits, clf_5b, CLASS_NAMES["model_5b"], confidence_threshold, "ViT-base Newcomer", 6)
     label_6, result_6output = predict_with_model(img_pilvits, clf_6, CLASS_NAMES["model_6"], confidence_threshold, "Swin Midjourney/SDXL", 7)
+    label_7, result_7output = predict_with_model(img_pilvits, clf_7, CLASS_NAMES["model_7"], confidence_threshold, "Vit", 7)
 
     combined_results = {
         "SwinV2/detect": label_1,
@@ -126,11 +133,12 @@ def predict_image(img, confidence_threshold):
         "Swin/SDXL-FLUX": label_4,
         "prithivMLmods": label_5,
         "prithivMLmods-2-22": label_5b,
-        "SwinMidSDXL": label_6
+        "SwinMidSDXL": label_6,
+        "Vit": label_7
     }
     print(combined_results)
 
-    combined_outputs = [result_1output, result_2output, result_3output, result_4output, result_5output, result_5boutput, result_6output]
+    combined_outputs = [result_1output, result_2output, result_3output, result_4output, result_5output, result_5boutput, result_6output, result_7output]
     return img_pil, combined_outputs
 # Define a function to generate the HTML content
 
@@ -159,7 +167,7 @@ def generate_results_html(results):
                 class="-m-4 h-24 {header_colors[0]} rounded-sm rounded-b-none transition border group-hover:border-gray-100 group-hover:shadow-lg group-hover:{header_colors[4]}">
                 <span class="text-gray-300 font-mono tracking-widest p-4 pb-3 block text-xs text-center">MODEL {index + 1}:</span>
                 <span
-                    class="flex w-24 mx-auto tracking-wide items-center justify-center rounded-full {header_colors[2]} px-1 py-0.5 {header_colors[3]}"
+                    class="flex w-30 mx-auto tracking-wide items-center justify-center rounded-full {header_colors[2]} px-1 py-0.5 {header_colors[3]}"
                 >
                     <svg xmlns="http://www.w3.org/2000/svg" fill="none" viewBox="0 0 24 24" stroke-width="3" stroke="currentColor" class="w-4 h-4 mr-2 -ml-3 group-hover:animate group-hover:animate-pulse">
                         {'<path stroke-linecap="round" stroke-linejoin="round" d="M9 12.75 11.25 15 15 9.75M21 12a9 9 0 1 1-18 0 9 9 0 0 1 18 0Z" />' if label == 'REAL' else '<path stroke-linecap="round" stroke-linejoin="round" d="m9.75 9.75 4.5 4.5m0-4.5-4.5 4.5M21 12a9 9 0 1 1-18 0 9 9 0 0 1 18 0Z" />'}
@@ -207,7 +215,8 @@ def generate_results_html(results):
             {generate_tile_html(3, results[3], "SDXL + FLUX", "cmckinle", MODEL_PATHS["model_4"])}
             {generate_tile_html(4, results[4], "Vit Based", "prithivMLmods", MODEL_PATHS["model_5"])}
             {generate_tile_html(5, results[5], "Vit Based, Newer Dataset", "prithivMLmods", MODEL_PATHS["model_5b"])}
-            {generate_tile_html(6, results[6], "Swin, Midjourney+SDXL", "ideepankarsharma2003", MODEL_PATHS["model_6"])}
+            {generate_tile_html(6, results[6], "Swin, Midj + SDXL", "ideepankarsharma2003", MODEL_PATHS["model_6"])}
+            {generate_tile_html(7, results[7], "ViT", "temp", MODEL_PATHS["model_7"])}
         </div>
     </div>
     """