cgi_classification_app.py

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

Automatically generated by Colab.

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

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
from xgboost import XGBClassifier

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

    return embeddings


def classify_image(image):
    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)}


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

interface.launch(share=True)