Upload folder using huggingface_hub

cgi_classification_app.py

@@ -8,63 +8,197 @@ 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))
 
-# 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)
+    @staticmethod
+    def load_model(filename):
+        return joblib.load(filename)
 
-    # Preprocess the image
-    preprocessed_image = preprocess_image(image)
-    # Calculate embeddings
-    embeddings = embedding_model.predict(preprocessed_image)
 
-    return embeddings
+model = WaveletClassifier.load_model("wavelet_classifier.pkl")
 
 
 def classify_image(image):
-    embeddings = calculate_embeddings(image)
+    image = Image.fromarray(image)
+    probabilities = model.predict_proba([image.resize((512, 512))])
     # 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)}
+    return {f"{labels[i]}": prob for i, prob in enumerate(probabilities[0])}
 
 
 interface = gr.Interface(

wavelet_classifier.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:db91c7f1d1841b5820ef7d9a043dfebf4ecf89568279fd1f79c490a13a206c0d
+size 465142385