Upload folder using huggingface_hub

cgi_classification_app.py

@@ -6,57 +6,6 @@ Automatically generated by Colab.
 Original file is located at
     https://colab.research.google.com/drive/1ckzOtXUiFW_NqlIandwoH07lnsLGKTLB
 """
-from scipy.spatial import distance
-import numpy as np
-
-
-class MeanClassifier:
-    def fit(self, X, y):
-        self.mean_0 = np.mean(X[y == 0], axis=0) if np.any(y == 0) else None
-        self.mean_1 = np.mean(X[y == 1], axis=0) if np.any(y == 1) else None
-
-    def predict(self, X):
-        preds = []
-        for x in X:
-            dist_0 = (
-                distance.euclidean(x, self.mean_0)
-                if self.mean_0 is not None
-                else np.inf
-            )
-            dist_1 = (
-                distance.euclidean(x, self.mean_1)
-                if self.mean_1 is not None
-                else np.inf
-            )
-            preds.append(1 if dist_1 < dist_0 else 0)
-        return np.array(preds)
-
-    def predict_proba(self, X):
-        # An implementation of probability prediction which uses a softmax function to determine the probability of each class based on the distance to the mean for each prototype
-        preds = []
-        for x in X:
-            dist_0 = (
-                distance.euclidean(x, self.mean_0) if self.mean_0 is not None else np
-            )
-            dist_1 = (
-                distance.euclidean(x, self.mean_1)
-                if self.mean_1 is not None
-                else np.inf
-            )
-            prob_0 = np.exp(-dist_0) / (np.exp(-dist_0) + np.exp(-dist_1))
-            prob_1 = np.exp(-dist_1) / (np.exp(-dist_0) + np.exp(-dist_1))
-            preds.append([prob_0, prob_1])
-        return np.array(preds)
-
-    def mean_distance(self, x):
-        dist_mean_0 = (
-            distance.euclidean(x, self.mean_0) if self.mean_0 is not None else np.inf
-        )
-        dist_mean_1 = (
-            distance.euclidean(x, self.mean_1) if self.mean_1 is not None else np.inf
-        )
-        return dist_mean_0, dist_mean_1
-
 
 import gradio as gr
 from PIL import Image
@@ -64,11 +13,10 @@ import numpy as np
 from PIL import Image
 from scipy.fftpack import fft2
 from tensorflow.keras.models import load_model, Model
-import pickle
+from xgboost import XGBClassifier
 
-mean_clf = None
-with open("mean_clf.pkl", "rb") as f:
-    mean_clf = pickle.load(f)
+xgb_clf = XGBClassifier()
+xgb_clf.load_model("xgb_cgi_classifier.json")
 
 
 # Function to apply Fourier transform
@@ -113,7 +61,7 @@ def calculate_embeddings(image, model_path="embedding_modelv2.keras"):
 def classify_image(image):
     embeddings = calculate_embeddings(image)
     # Convert to 2D array for model input
-    probabilities = mean_clf.predict_proba(embeddings)[0]
+    probabilities = xgb_clf.predict_proba(embeddings)[0]
     labels = ["Photo", "CGI"]
     return {f"{labels[i]}": prob for i, prob in enumerate(probabilities)}