Spaces:
Sleeping
Sleeping
# -*- 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 | |
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) | |