Upload folder using huggingface_hub

cgi_classification_app.py

@@ -8,197 +8,63 @@ Original file is located at
 """
 
 import gradio as gr
-
-from sklearn.model_selection import train_test_split
-from sklearn.metrics import (
-    accuracy_score,
-    f1_score,
-    confusion_matrix,
-    ConfusionMatrixDisplay,
-)
-from sklearn.preprocessing import StandardScaler
-from sklearn.decomposition import PCA
-import umap
-import pywt
-import os
 from PIL import Image
-import matplotlib.pyplot as plt
 import numpy as np
+from PIL import Image
+from scipy.fftpack import fft2
+from tensorflow.keras.models import load_model, Model
 from xgboost import XGBClassifier
-from sklearn.model_selection import cross_val_score, KFold
-from sklearn.dummy import DummyClassifier
-from sklearn.ensemble import RandomForestClassifier
-from sklearn.svm import SVC
-from sklearn.neighbors import KNeighborsClassifier
-from sklearn.model_selection import train_test_split
-from sklearn.metrics import classification_report
-import plotly.express as px
-import pandas as pd
-import joblib
-from tqdm import tqdm
-import lzma
-
-
-class WaveletClassifier:
-    def __init__(
-        self,
-        wavelets=["db4", "db10"],
-        umap_n_neighbors=16,
-        umap_n_components=32,
-        random_state=42,
-    ):
-        self.wavelets = wavelets
-        self.umap_n_neighbors = umap_n_neighbors
-        self.umap_n_components = umap_n_components
-        self.random_state = random_state
-        self.reducer = umap.UMAP(
-            n_neighbors=self.umap_n_neighbors,
-            n_components=self.umap_n_components,
-            random_state=self.random_state,
-        )
-        self.classifier = KNeighborsClassifier(n_neighbors=7)  # Default classifier
-
-    def load_images_from_folder(self, folder):
-        images = []
-        labels = []
-        print(f"Loading images from {folder}")
-        for filename in tqdm(os.listdir(folder)):
-            if not (
-                filename.endswith(".jpg")
-                or filename.endswith(".png")
-                or filename.endswith("jpeg")
-                or filename.endswith("webp")
-            ):
-                continue
-            img = Image.open(os.path.join(folder, filename))
-            img = img.resize((512, 512))
-            if img is not None:
-                images.append(img)
-                labels.append(
-                    1 if "CGI" in folder else 0
-                )  # Assuming folder names contain "AI" or not
-        return images, labels
-
-    def extract_wavelet_features(self, images):
-        all_features = []
-        for img in images:
-            img_gray = img.convert("L")
-            img_array = np.array(img_gray)
-            features = []
-            for wavelet in self.wavelets:
-                cA, cD = pywt.dwt(img_array, wavelet)
-                features.extend(cD.flatten())
-            all_features.append(features)
-        return np.array(all_features)
-
-    def fit(self, train_folder1, train_folder2):
-        # Load images and extract features
-        images1, labels1 = self.load_images_from_folder(train_folder1)
-        images2, labels2 = self.load_images_from_folder(train_folder2)
-
-        min_length = min(len(images1), len(images2))
-        images1 = images1[:min_length]
-        images2 = images2[:min_length]
-        labels1 = labels1[:min_length]
-        labels2 = labels2[:min_length]
-
-        images = images1 + images2
-        labels = labels1 + labels2
-        features = self.extract_wavelet_features(images)
-
-        # Apply UMAP dimensionality reduction
-        embeddings = self.reducer.fit_transform(features)
-        X_train, X_test, y_train, y_test = train_test_split(
-            embeddings, labels, test_size=0.2, random_state=42
-        )
-
-        # Train the classifier
-        self.classifier.fit(X_train, y_train)
-
-        acc = self.classifier.score(X_test, y_test)
-        y_pred = self.classifier.predict(X_test)
-        print(f"Classifier accuracy = {acc}")
-
-        f1 = f1_score(y_test, y_pred)
-        print(f"Classifier F1 = {f1}")
-        print(classification_report(y_test, y_pred))
-
-    def predict(self, images):
-        # Load images and extract features
-        features = self.extract_wavelet_features(images)
-
-        # Apply UMAP dimensionality reduction
-        embeddings = self.reducer.transform(features)
-
-        # Make predictions
-        return self.classifier.predict(embeddings)
-
-    def predict_proba(self, images):
-        # Load images and extract features
-        features = self.extract_wavelet_features(images)
-
-        # Apply UMAP dimensionality reduction
-        embeddings = self.reducer.transform(features)
-
-        # Make predictions
-        return self.classifier.predict_proba(embeddings)
-
-    def score(self, test_folder):
-        # Load images and extract features
-        images, labels = self.load_images_from_folder(test_folder)
-        features = self.extract_wavelet_features(images)
-
-        # Apply UMAP dimensionality reduction
-        embeddings = self.reducer.transform(features)
-
-        # Evaluate the classifier
-        return self.classifier.score(embeddings, labels)
-
-    def cross_val_score(self, folder1, folder2, n_splits=5):
-        # Load images and extract features
-        # Load images and extract features
-        images1, labels1 = self.load_images_from_folder(folder1)
-        images2, labels2 = self.load_images_from_folder(folder2)
-
-        min_length = min(len(images1), len(images2))
-        images1 = images1[:min_length]
-        images2 = images2[:min_length]
-        labels1 = labels1[:min_length]
-        labels2 = labels2[:min_length]
-
-        images = images1 + images2
-        labels = labels1 + labels2
-        features = self.extract_wavelet_features(images)
-
-        # Apply UMAP dimensionality reduction
-        embeddings = self.reducer.fit_transform(features)
-        # Perform four-fold cross-validation
-        kfold = KFold(n_splits=n_splits, shuffle=True, random_state=42)
-        scores = cross_val_score(
-            self.classifier, embeddings, labels, cv=kfold, scoring="accuracy"
-        )
-
-        # Print the cross-validation scores
-        print("Cross-validation scores:", scores)
-        print("Average cross-validation score:", scores.mean())
-
-    def save_model(self, filename):
-        joblib.dump(self, filename, compress=("lzma", 9))
 
-    @staticmethod
-    def load_model(filename):
-        return joblib.load(filename)
+# classifier
+xgb_clf = XGBClassifier()
+xgb_clf.load_model("xgb_cgi_classifier.json")
+
+
+# Function to apply Fourier transform
+def apply_fourier_transform(image):
+    image = np.array(image)
+    fft_image = fft2(image)
+    return np.abs(fft_image)
+
+
+def preprocess_image(image):
+    try:
+        image = Image.fromarray(image)
+        image = image.convert("L")
+        image = image.resize((256, 256))
+        image = apply_fourier_transform(image)
+        image = np.expand_dims(
+            image, axis=-1
+        )  # Expand dimensions to match model input shape
+        image = np.expand_dims(image, axis=0)  # Expand to add batch dimension
+        return image
+    except Exception as e:
+        print(f"Error processing image: {e}")
+        return None
+
+
+# Function to load embedding model and calculate embeddings
+def calculate_embeddings(image, model_path="embedding_modelv2.keras"):
+    # Load the trained model
+    model = load_model(model_path)
+
+    # Remove the final classification layer to get embeddings
+    embedding_model = Model(inputs=model.input, outputs=model.output)
 
+    # Preprocess the image
+    preprocessed_image = preprocess_image(image)
+    # Calculate embeddings
+    embeddings = embedding_model.predict(preprocessed_image)
 
-model = WaveletClassifier.load_model("wavelet_classifier.pkl")
+    return embeddings
 
 
 def classify_image(image):
-    image = Image.fromarray(image)
-    probabilities = model.predict_proba([image.resize((512, 512))])
+    embeddings = calculate_embeddings(image)
     # Convert to 2D array for model input
+    probabilities = xgb_clf.predict_proba(embeddings)[0]
     labels = ["Photo", "CGI"]
-    return {f"{labels[i]}": prob for i, prob in enumerate(probabilities[0])}
+    return {f"{labels[i]}": prob for i, prob in enumerate(probabilities)}
 
 
 interface = gr.Interface(