♻️ refactor(app): code cleanup

- removed unused import: warnings
- added logging configuration
-【docs】modified logging level
-【refactor】renamed old logging config to current logging module
- optimized importance of models by introducing model hook configuration modes
-【docs】map classes of various models
-【feat】move and refactor models into reusable function
- modified logic containment
- reorganized and refactored prediction logic
- moved all augmentation related function to util
- clarified all parameters to be consistent
-【refactor】modified function name consistent
▪️️ feat(app): new model prediction logic with gpu decorator
-【refactor】added model index to output [[model id, model name, class a confidence, class b confidence, label] recommended output]

---·

+ refactor(file management): code cleanup

- remove previous unused paths
- moved function into utils

app.py

@@ -5,284 +5,116 @@ from torchvision import transforms
 import torch
 from PIL import Image
 import numpy as np
-# from utils.goat import call_inference / announcement soon
 import io
-import warnings
+import logging
+from utils.utils import softmax, augment_image, convert_pil_to_bytes
+
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# Suppress warnings
-warnings.filterwarnings("ignore", category=UserWarning, message="Using a slow image processor as `use_fast` is unset")
 
 # Ensure using GPU if available
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# Load the first model and processor
-image_processor_1 = AutoImageProcessor.from_pretrained("haywoodsloan/ai-image-detector-deploy", use_fast=True)
-model_1 = Swinv2ForImageClassification.from_pretrained("haywoodsloan/ai-image-detector-deploy")
-model_1 = model_1.to(device)
-clf_1 = pipeline(model=model_1, task="image-classification", image_processor=image_processor_1, device=device)
 
-# Load the second model
-model_2_path = "Heem2/AI-vs-Real-Image-Detection"
-clf_2 = pipeline("image-classification", model=model_2_path, device=device)
+# 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",
+    "model_5": "prithivMLmods/Deep-Fake-Detector-v2-Model",
+    "model_5b": "prithivMLmods/Deepfake-Detection-Exp-02-22"
+}
 
-# Load additional models
-models = ["Organika/sdxl-detector", "cmckinle/sdxl-flux-detector"]
-feature_extractor_3 = AutoFeatureExtractor.from_pretrained(models[0], device=device)
-model_3 = AutoModelForImageClassification.from_pretrained(models[0]).to(device)
-feature_extractor_4 = AutoFeatureExtractor.from_pretrained(models[1], device=device)
-model_4 = AutoModelForImageClassification.from_pretrained(models[1]).to(device)
+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']
+}
 
-# Load the second model
-model_5_path = "prithivMLmods/Deep-Fake-Detector-v2-Model"
-clf_5 = pipeline("image-classification", model=model_5_path, device=device)
+# Load models and processors
+def load_models():
+    image_processor_1 = AutoImageProcessor.from_pretrained(MODEL_PATHS["model_1"], use_fast=True)
+    model_1 = Swinv2ForImageClassification.from_pretrained(MODEL_PATHS["model_1"])
+    model_1 = model_1.to(device)
+    clf_1 = pipeline(model=model_1, task="image-classification", image_processor=image_processor_1, device=device)
 
-model_5b_path = "prithivMLmods/Deepfake-Detection-Exp-02-22"
-clf_5b = pipeline("image-classification", model=model_5b_path, device=device)
+    clf_2 = pipeline("image-classification", model=MODEL_PATHS["model_2"], device=device)
 
-# Define class names for all models
-class_names_1 = ['artificial', 'real']
-class_names_2 = ['AI Image', 'Real Image']
-labels_3 = ['AI', 'Real']
-labels_4 = ['AI', 'Real']
-class_names_5 = ['Realism', 'Deepfake']
-class_names_5b = ['Real', 'Deepfake']
+    feature_extractor_3 = AutoFeatureExtractor.from_pretrained(MODEL_PATHS["model_3"], device=device)
+    model_3 = AutoModelForImageClassification.from_pretrained(MODEL_PATHS["model_3"]).to(device)
 
-def softmax(vector):
-    e = np.exp(vector - np.max(vector))  # for numerical stability
-    return e / e.sum()
+    feature_extractor_4 = AutoFeatureExtractor.from_pretrained(MODEL_PATHS["model_4"], device=device)
+    model_4 = AutoModelForImageClassification.from_pretrained(MODEL_PATHS["model_4"]).to(device)
 
-def augment_image(img_pil):
-    # Example augmentation: horizontal flip
-    transform_flip = transforms.Compose([
-        transforms.RandomHorizontalFlip(p=1.0)  # Flip the image horizontally with probability 1.0
-    ])
-    
-    # Example augmentation: rotation
-    transform_rotate = transforms.Compose([
-        transforms.RandomRotation(degrees=(90, 90))  # Rotate the image by 90 degrees
-    ])
-    
-    augmented_img_flip = transform_flip(img_pil)
-    augmented_img_rotate = transform_rotate(img_pil)
-    
-    return augmented_img_flip, augmented_img_rotate
+    clf_5 = pipeline("image-classification", model=MODEL_PATHS["model_5"], device=device)
+    clf_5b = pipeline("image-classification", model=MODEL_PATHS["model_5b"], device=device)
+
+    return clf_1, clf_2, feature_extractor_3, model_3, feature_extractor_4, model_4, clf_5, clf_5b
 
-# def convert_pil_to_bytes(img_pil):
-#     img_byte_arr = io.BytesIO()
-#     img_pil.save(img_byte_arr, format='PNG')
-#     img_byte_arr = img_byte_arr.getvalue()
-#     return img_byte_arr
+clf_1, clf_2, feature_extractor_3, model_3, feature_extractor_4, model_4, clf_5, clf_5b = load_models()
 
-def convert_pil_to_bytes(image, format='JPEG'):
-    img_byte_arr = io.BytesIO()
-    image.save(img_byte_arr, format=format)
-    img_byte_arr = img_byte_arr.getvalue()
-    return img_byte_arr
+@spaces.GPU(duration=10)
+def predict_with_model(img_pil, clf, class_names, confidence_threshold, model_name, model_id):
+    try:
+        prediction = clf(img_pil)
+        result = {pred['label']: pred['score'] for pred in prediction}
+        result_output = [model_id, model_name, result.get(class_names[1], 0.0), result.get(class_names[0], 0.0)]
+        logger.info(result_output)
+        for class_name in class_names:
+            if class_name not in result:
+                result[class_name] = 0.0
+        if result[class_names[0]] >= confidence_threshold:
+            label = f"AI, Confidence: {result[class_names[0]]:.4f}"
+            result_output.append('AI')
+        elif result[class_names[1]] >= confidence_threshold:
+            label = f"Real, Confidence: {result[class_names[1]]:.4f}"
+            result_output.append('REAL')
+        else:
+            label = "Uncertain Classification"
+            result_output.append('UNCERTAIN')
+    except Exception as e:
+        label = f"Error: {str(e)}"
+    return label, result_output
 
 @spaces.GPU(duration=10)
 def predict_image(img, confidence_threshold):
-    # Ensure the image is a PIL Image
     if not isinstance(img, Image.Image):
         raise ValueError(f"Expected a PIL Image, but got {type(img)}")
-    
-    # Convert the image to RGB if not already
     if img.mode != 'RGB':
         img_pil = img.convert('RGB')
     else:
         img_pil = img
-
-    # Resize the image
     img_pil = transforms.Resize((256, 256))(img_pil)
-    # Size 224 for vits models
     img_pilvits = transforms.Resize((224, 224))(img_pil)
-    
-    # Predict using the first model
-    try:
-        prediction_1 = clf_1(img_pil)
-        result_1 = {pred['label']: pred['score'] for pred in prediction_1}
-        result_1output = [1, 'SwinV2-base', result_1['real'], result_1['artificial']]
-        print(result_1output)
-        # Ensure the result dictionary contains all class names
-        for class_name in class_names_1:
-            if class_name not in result_1:
-                result_1[class_name] = 0.0
-        # Check if either class meets the confidence threshold
-        if result_1['artificial'] >= confidence_threshold:
-            label_1 = f"AI, Confidence: {result_1['artificial']:.4f}"
-            result_1output += ['AI']
-        elif result_1['real'] >= confidence_threshold:
-            label_1 = f"Real, Confidence: {result_1['real']:.4f}"
-            result_1output += ['REAL']
-        else:
-            label_1 = "Uncertain Classification"
-            result_1output += ['UNCERTAIN']
 
-    except Exception as e:
-        label_1 = f"Error: {str(e)}"
-    print(result_1output)
-    # Predict using the second model
-    try:
-        prediction_2 = clf_2(img_pilvits)
-        result_2 = {pred['label']: pred['score'] for pred in prediction_2}
-        result_2output = [2, 'ViT-base Classifer', result_2['Real Image'], result_2['AI Image']]
-        print(result_2output)
-        # Ensure the result dictionary contains all class names
-        for class_name in class_names_2:
-            if class_name not in result_2:
-                result_2[class_name] = 0.0
-        # Check if either class meets the confidence threshold
-        if result_2['AI Image'] >= confidence_threshold:
-            label_2 = f"AI, Confidence: {result_2['AI Image']:.4f}"
-            result_2output += ['AI']
-        elif result_2['Real Image'] >= confidence_threshold:
-            label_2 = f"Real, Confidence: {result_2['Real Image']:.4f}"
-            result_2output += ['REAL']
-        else:
-            label_2 = "Uncertain Classification"
-            result_2output += ['UNCERTAIN']
-    except Exception as e:
-        label_2 = f"Error: {str(e)}"
-    
-    # Predict using the third model with softmax
-    try:
-        inputs_3 = feature_extractor_3(img_pil, return_tensors="pt").to(device)
-        with torch.no_grad():
-            outputs_3 = model_3(**inputs_3)
-            logits_3 = outputs_3.logits
-            probabilities_3 = softmax(logits_3.cpu().numpy()[0])
-        result_3 = {
-            labels_3[1]: float(probabilities_3[1]),   # Real
-            labels_3[0]: float(probabilities_3[0])  # AI
-        }
-        result_3output = [3, 'SDXL-Trained', float(probabilities_3[1]), float(probabilities_3[0])]
-        print(result_3output)
-        # Ensure the result dictionary contains all class names
-        for class_name in labels_3:
-            if class_name not in result_3:
-                result_3[class_name] = 0.0
-        # Check if either class meets the confidence threshold
-        if result_3['AI'] >= confidence_threshold:
-            label_3 = f"AI, Confidence: {result_3['AI']:.4f}"
-            result_3output += ['AI']
-        elif result_3['Real'] >= confidence_threshold:
-            label_3 = f"Real, Confidence: {result_3['Real']:.4f}"
-            result_3output += ['REAL']
-        else:
-            label_3 = "Uncertain Classification"
-            result_3output += ['UNCERTAIN']
-    except Exception as e:
-        label_3 = f"Error: {str(e)}"
-    
-    # Predict using the fourth model with softmax
-    try:
-        inputs_4 = feature_extractor_4(img_pil, return_tensors="pt").to(device)
-        with torch.no_grad():
-            outputs_4 = model_4(**inputs_4)
-            logits_4 = outputs_4.logits
-            probabilities_4 = softmax(logits_4.cpu().numpy()[0])
-        result_4 = {
-            labels_4[1]: float(probabilities_4[1]),   # Real
-            labels_4[0]: float(probabilities_4[0])  # AI
-        }
-        result_4output = [4, 'SDXL + FLUX', float(probabilities_4[1]), float(probabilities_4[0])]
-        print(result_4)
-        # Ensure the result dictionary contains all class names
-        for class_name in labels_4:
-            if class_name not in result_4:
-                result_4[class_name] = 0.0
-        # Check if either class meets the confidence threshold
-        if result_4['AI'] >= confidence_threshold:
-            label_4 = f"AI, Confidence: {result_4['AI']:.4f}"
-            result_4output += ['AI']
-        elif result_4['Real'] >= confidence_threshold:
-            label_4 = f"Real, Confidence: {result_4['Real']:.4f}"
-            result_4output += ['REAL']
-        else:
-            label_4 = "Uncertain Classification"
-            result_4output += ['UNCERTAIN']
-    except Exception as e:
-        label_4 = f"Error: {str(e)}"
-    
-    try:
-        prediction_5 = clf_5(img_pilvits)
-        result_5 = {pred['label']: pred['score'] for pred in prediction_5}
-        result_5output = [5, 'ViT-base Newcomer', result_5['Realism'], result_5['Deepfake']]
-        
-        # Ensure the result dictionary contains all class names
-        for class_name in class_names_5:
-            if class_name not in result_5:
-                result_5[class_name] = 0.0
-        # Check if either class meets the confidence threshold
-        if result_5['Deepfake'] >= confidence_threshold:
-            label_5 = f"AI, Confidence: {result_5['Deepfake']:.4f}"
-            result_5output += ['AI']
-        elif result_5['Real Image'] >= confidence_threshold:
-            label_5 = f"Real, Confidence: {result_5['Realism']:.4f}"
-            result_5output += ['REAL']
-        else:
-            label_5 = "Uncertain Classification"
-            result_5output += ['UNCERTAIN']
-    except Exception as e:
-        label_5 = f"Error: {str(e)}"
-        
-    print(result_5output)
-    
-    try:
-        prediction_5b = clf_5b(img_pilvits)
-        result_5b = {pred['label']: pred['score'] for pred in prediction_5b}
-        result_5boutput = [6, 'ViT-base Newcomer', result_5b['Real'], result_5b['Deepfake']]
-        
-        # Ensure the result dictionary contains all class names
-        for class_name in class_names_5b:
-            if class_name not in result_5b:
-                result_5b[class_name] = 0.0
-        # Check if either class meets the confidence threshold
-        if result_5b['Deepfake'] >= confidence_threshold:
-            label_5b = f"AI, Confidence: {result_5b['Deepfake']:.4f}"
-            result_5boutput += ['AI']
-        elif result_5b['Real Image'] >= confidence_threshold:
-            label_5b = f"Real, Confidence: {result_5b['Real']:.4f}"
-            result_5boutput += ['REAL']
-        else:
-            label_5b = "Uncertain Classification"
-            result_5boutput += ['UNCERTAIN']
-    except Exception as e:
-        label_5b = f"Error: {str(e)}"
-        
-    print(result_5boutput)
-    
-    
-    # try:
-    #     result_5output = [5, 'TBA', 0.0, 0.0, 'MAINTENANCE']
-    #     img_bytes = convert_pil_to_bytes(img_pil)
-    #     # print(img)
-    #     # print(img_bytes)
-    #     response5_raw = call_inference(img)
-    #     print(response5_raw)
-    #     response5 = response5_raw
-    #     print(response5)
-    #     label_5 = f"Result: {response5}"
-        
-    # except Exception as e:
-    #     label_5 = f"Error: {str(e)}"
-    
+    label_1, result_1output = predict_with_model(img_pil, clf_1, CLASS_NAMES["model_1"], confidence_threshold, "SwinV2-base", 1)
+    label_2, result_2output = predict_with_model(img_pilvits, clf_2, CLASS_NAMES["model_2"], confidence_threshold, "ViT-base Classifer", 2)
+    label_3, result_3output = predict_with_model(img_pil, feature_extractor_3, model_3, CLASS_NAMES["model_3"], confidence_threshold, "SDXL-Trained", 3)
+    label_4, result_4output = predict_with_model(img_pil, feature_extractor_4, model_4, CLASS_NAMES["model_4"], confidence_threshold, "SDXL + FLUX", 4)
+    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)
 
-    # Combine results
     combined_results = {
         "SwinV2/detect": label_1,
         "ViT/AI-vs-Real": label_2,
         "Swin/SDXL": label_3,
         "Swin/SDXL-FLUX": label_4,
         "prithivMLmods": label_5,
-        "prithivMLmods-2-22": label_5b 
+        "prithivMLmods-2-22": label_5b
     }
-        # Generate HTML content
-    
-    combined_outputs = [ result_1output, result_2output, result_3output, result_4output, result_5output, result_5boutput ]
-    # html_content = generate_results_html(combined_outputs)
+    print(combined_results)
 
+    combined_outputs = [result_1output, result_2output, result_3output, result_4output, result_5output, result_5boutput]
     return img_pil, combined_outputs
 
+
 # Define a function to generate the HTML content
 # Define a function to generate the HTML content
 def generate_results_html(results):

utils/utils.py

@@ -0,0 +1,25 @@
+import numpy as np
+import io
+from PIL import Image
+from torchvision import transforms
+
+def softmax(vector):
+    e = np.exp(vector - np.max(vector))  # for numerical stability
+    return e / e.sum()
+
+def augment_image(img_pil):
+    transform_flip = transforms.Compose([
+        transforms.RandomHorizontalFlip(p=1.0)
+    ])
+    transform_rotate = transforms.Compose([
+        transforms.RandomRotation(degrees=(90, 90))
+    ])
+    augmented_img_flip = transform_flip(img_pil)
+    augmented_img_rotate = transform_rotate(img_pil)
+    return augmented_img_flip, augmented_img_rotate
+
+def convert_pil_to_bytes(image, format='JPEG'):
+    img_byte_arr = io.BytesIO()
+    image.save(img_byte_arr, format=format)
+    img_byte_arr = img_byte_arr.getvalue()
+    return img_byte_arr