# -*- coding: utf-8 -*-
"""CGI Classification App.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1ckzOtXUiFW_NqlIandwoH07lnsLGKTLB
"""

!pip install gradio

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
import numpy as np
from PIL import Image
from scipy.fftpack import fft2
from tensorflow.keras.models import load_model, Model
import pickle

mean_clf = None
with open('mean_clf.pkl', 'rb') as f:
  mean_clf = pickle.load(f)


# 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)

    return embeddings


def classify_image(image):
    embeddings = calculate_embeddings(image)
      # Convert to 2D array for model input
    probabilities = mean_clf.predict_proba(embeddings)[0]
    labels = ["Photo", "CGI"]
    return {f"{labels[i]}": prob for i, prob in enumerate(probabilities)}

interface = gr.Interface(
    fn=classify_image,
    inputs=["image"],
    outputs=gr.Label(num_top_classes=2)
)

interface.launch(share=True)