|
import streamlit as st |
|
import os |
|
import random |
|
from PIL import Image |
|
from transformers import pipeline |
|
import pandas as pd |
|
import matplotlib.pyplot as plt |
|
|
|
|
|
st.sidebar.header('About the App') |
|
st.sidebar.write('This Image Classifier app can classify images as Real or Fake.') |
|
st.sidebar.write('This AI is trained on a dataset of real and deepfake images. It uses a pre-trained model to classify images. You can upload an image or use a random image from the dataset to test the classifier.') |
|
st.sidebar.header('Author') |
|
st.sidebar.write('Jan Mikolon') |
|
st.sidebar.write('📧 Contact: jan.mikolon@ibm.com') |
|
st.sidebar.markdown('[![LinkedIn](https://img.shields.io/badge/LinkedIn-Profile-blue)](https://www.linkedin.com/in/jan-mikolon/)', unsafe_allow_html=True) |
|
|
|
|
|
def load_random_image(folder_path): |
|
images = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if os.path.isfile(os.path.join(folder_path, f))] |
|
random_image_path = random.choice(images) |
|
return Image.open(random_image_path) |
|
|
|
|
|
folder_path = 'data/' |
|
|
|
|
|
st.title('Image Classifier - Real or Fake') |
|
|
|
|
|
uploaded_image = st.file_uploader("Upload an image for classification", type=["png", "jpg", "jpeg"]) |
|
|
|
|
|
col1, col2 = st.columns(2) |
|
|
|
|
|
if uploaded_image is not None: |
|
image = Image.open(uploaded_image) |
|
col1.image(image, caption='Uploaded Image', use_column_width=True) |
|
else: |
|
|
|
if 'image_path' not in st.session_state or st.button('Load Random Image'): |
|
st.session_state.image_path = load_random_image(folder_path) |
|
col1.image(st.session_state.image_path, caption='Random Image', use_column_width=True) |
|
|
|
|
|
if st.button('Classify'): |
|
|
|
|
|
pipe = pipeline("image-classification", model="dima806/deepfake_vs_real_image_detection") |
|
|
|
if uploaded_image is not None: |
|
classification_results = pipe(image) |
|
else: |
|
classification_results = pipe(st.session_state.image_path) |
|
|
|
|
|
|
|
df_results = pd.DataFrame(classification_results) |
|
|
|
|
|
fig, ax = plt.subplots() |
|
ax.bar(df_results['label'], df_results['score'], color=['blue', 'orange']) |
|
ax.set_ylabel('Scores') |
|
ax.set_title('Classification Scores') |
|
plt.tight_layout() |
|
|
|
|
|
col2.pyplot(fig) |
|
|
|
|
|
if uploaded_image is None: |
|
st.session_state.image_path = load_random_image(folder_path) |
|
|